add docstrings and types

examples/lw_example.ipynb

@@ -27,7 +27,7 @@
     "\n",
     "atmosphere_file = \"multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc\"\n",
     "atmosphere_path = os.path.join(input_dir, atmosphere_file)\n",
-    "atmosphere = xr.load_dataset(atmosphere_path)#.sel(expt=0)\n",
+    "atmosphere = xr.load_dataset(atmosphere_path)\n",
     "\n",
     "gas_optics_lw.gas_optics.compute(atmosphere, problem_type=\"absorption\")\n",
     "\n",
@@ -39,8 +39,12 @@
     "rlu = xr.load_dataset(rlu_reference, decode_cf=False)\n",
     "rld = xr.load_dataset(rld_reference, decode_cf=False)\n",
     "\n",
-    "assert np.isclose(fluxes[\"lw_flux_up_broadband\"], rlu[\"rlu\"], atol=ERROR_TOLERANCE).all()\n",
-    "assert np.isclose(fluxes[\"lw_flux_down_broadband\"], rld[\"rld\"], atol=ERROR_TOLERANCE).all()\n"
+    "assert np.isclose(\n",
+    "    fluxes[\"lw_flux_up_broadband\"], rlu[\"rlu\"], atol=ERROR_TOLERANCE\n",
+    ").all()\n",
+    "assert np.isclose(\n",
+    "    fluxes[\"lw_flux_down_broadband\"], rld[\"rld\"], atol=ERROR_TOLERANCE\n",
+    ").all()"
    ]
   }
  ],

pyrte_rrtmgp/config.py

@@ -1,9 +1,16 @@
-from typing import Dict
+"""Default mappings for gas names, dimensions and variables used in RRTMGP.
 
-DEFAULT_GAS_MAPPING: Dict[str, str] = {
+This module contains dictionaries that map standard names to dataset-specific names
+for gases, dimensions and variables used in radiative transfer calculations.
+"""
+
+from typing import Dict, Final
+
+# Mapping of standard gas names to RRTMGP-specific names
+DEFAULT_GAS_MAPPING: Final[Dict[str, str]] = {
     "h2o": "water_vapor",
     "co2": "carbon_dioxide_GM",
-    "o3": "ozone",
+    "o3": "ozone", 
     "n2o": "nitrous_oxide_GM",
     "co": "carbon_monoxide_GM",
     "ch4": "methane_GM",
@@ -21,3 +28,25 @@
     "cf4": "cf4_GM",
     "no2": "no2",
 }
+
+# Mapping of standard dimension names to dataset-specific names
+DEFAULT_DIM_MAPPING: Final[Dict[str, str]] = {
+    "site": "site",
+    "layer": "layer", 
+    "level": "level",
+}
+
+# Mapping of standard variable names to dataset-specific names
+DEFAULT_VAR_MAPPING: Final[Dict[str, str]] = {
+    "pres_layer": "pres_layer",
+    "pres_level": "pres_level",
+    "temp_layer": "temp_layer",
+    "temp_level": "temp_level",
+    "surface_temperature": "surface_temperature",
+    "solar_zenith_angle": "solar_zenith_angle",
+    "surface_albedo": "surface_albedo",
+    "surface_albedo_dir": "surface_albedo_dir",
+    "surface_albedo_dif": "surface_albedo_dif",
+    "surface_emissivity": "surface_emissivity",
+    "surface_emissivity_jacobian": "surface_emissivity_jacobian",
+}

pyrte_rrtmgp/constants.py

@@ -1,9 +1,48 @@
-HELMERT1 = 9.80665
-HELMERT2 = 0.02586
-M_DRY = 0.028964
-M_H2O = 0.018016
-AVOGAD = 6.02214076e23
-SOLAR_CONSTANTS = {
-    "A_OFFSET": 0.1495954,
-    "B_OFFSET": 0.00066696,
+"""Physical and mathematical constants used in radiative transfer calculations.
+
+This module contains various physical and mathematical constants needed for
+radiative transfer calculations, including gravitational parameters, molecular
+masses, and Gaussian quadrature weights and points.
+"""
+
+from typing import Dict, Final
+import numpy as np
+from numpy.typing import NDArray
+
+# Gravitational parameters from Helmert's equation (m/s^2)
+HELMERT1: Final[float] = 9.80665  # Standard gravity at sea level
+HELMERT2: Final[float] = 0.02586  # Gravity variation with latitude
+
+# Molecular masses (kg/mol)
+M_DRY: Final[float] = 0.028964  # Dry air
+M_H2O: Final[float] = 0.018016  # Water vapor
+
+# Avogadro's number (molecules/mol)
+AVOGAD: Final[float] = 6.02214076e23
+
+# Solar constants for orbit calculations
+SOLAR_CONSTANTS: Final[Dict[str, float]] = {
+    "A_OFFSET": 0.1495954,  # Semi-major axis offset (AU)
+    "B_OFFSET": 0.00066696,  # Orbital eccentricity factor
 }
+
+# Gaussian quadrature constants for radiative transfer
+GAUSS_DS: NDArray[np.float64] = np.reciprocal(
+    np.array(
+        [
+            [0.6096748751, np.inf, np.inf, np.inf],
+            [0.2509907356, 0.7908473988, np.inf, np.inf],
+            [0.1024922169, 0.4417960320, 0.8633751621, np.inf],
+            [0.0454586727, 0.2322334416, 0.5740198775, 0.9030775973],
+        ]
+    )
+)
+
+GAUSS_WTS: NDArray[np.float64] = np.array(
+    [
+        [1.0, 0.0, 0.0, 0.0],
+        [0.2300253764, 0.7699746236, 0.0, 0.0],
+        [0.0437820218, 0.3875796738, 0.5686383044, 0.0],
+        [0.0092068785, 0.1285704278, 0.4323381850, 0.4298845087],
+    ]
+)

pyrte_rrtmgp/data_types.py

@@ -1,17 +1,40 @@
-from enum import Enum
+from enum import Enum, StrEnum
 
 
-class GasOpticsFiles(Enum):
-    """Enumeration of default RRTMGP gas optics data files."""
+class GasOpticsFiles(StrEnum):
+    """Enumeration of default RRTMGP gas optics data files.
+    
+    This enum defines the available pre-configured gas optics data files that can be used
+    with RRTMGP. The files contain absorption coefficients and other optical properties
+    needed for radiative transfer calculations.
+
+    Attributes:
+        LW_G128: Longwave gas optics file with 128 g-points
+        LW_G256: Longwave gas optics file with 256 g-points
+        SW_G112: Shortwave gas optics file with 112 g-points
+        SW_G224: Shortwave gas optics file with 224 g-points
+    """
 
     LW_G128 = "rrtmgp-gas-lw-g128.nc"
     LW_G256 = "rrtmgp-gas-lw-g256.nc"
     SW_G112 = "rrtmgp-gas-sw-g112.nc"
     SW_G224 = "rrtmgp-gas-sw-g224.nc"
 
 
-class ProblemTypes(Enum):
+class ProblemTypes(StrEnum):
+    """Enumeration of available radiation calculation types.
+    
+    This enum defines the different types of radiation calculations that can be performed,
+    including both longwave and shortwave calculations with different solution methods.
+
+    Attributes:
+        LW_ABSORPTION: Longwave absorption-only calculation
+        LW_2STREAM: Longwave two-stream approximation calculation
+        SW_DIRECT: Shortwave direct beam calculation
+        SW_2STREAM: Shortwave two-stream approximation calculation
+    """
+
     LW_ABSORPTION = "Longwave absorption"
-    LW_2STREAM = "Longwave 2-stream"
+    LW_2STREAM = "Longwave 2-stream" 
     SW_DIRECT = "Shortwave direct"
     SW_2STREAM = "Shortwave 2-stream"

pyrte_rrtmgp/data_validation.py

@@ -1,28 +1,54 @@
 from dataclasses import asdict, dataclass
-from typing import Dict, Optional, Set
+from typing import Dict, Optional, Set, Union
 
 import xarray as xr
 
-from pyrte_rrtmgp.config import DEFAULT_GAS_MAPPING
+from pyrte_rrtmgp.config import (
+    DEFAULT_DIM_MAPPING,
+    DEFAULT_GAS_MAPPING,
+    DEFAULT_VAR_MAPPING,
+)
 
 
 @dataclass
 class GasMapping:
+    """Class for managing gas name mappings between standard and dataset-specific names.
+
+    Attributes:
+        _mapping: Dictionary mapping standard gas names to dataset-specific names
+        _required_gases: Set of required gas names that must be present
+    """
     _mapping: Dict[str, str]
     _required_gases: Set[str]
 
     @classmethod
     def create(
-        cls, gas_names: Set[str], custom_mapping: Dict[str, str] | None = None
+        cls, gas_names: Set[str], custom_mapping: Optional[Dict[str, str]] = None
     ) -> "GasMapping":
+        """Create a new GasMapping instance with default and custom mappings.
+
+        Args:
+            gas_names: Set of required gas names
+            custom_mapping: Optional custom mapping to override defaults
+
+        Returns:
+            New GasMapping instance
+        """
         mapping = DEFAULT_GAS_MAPPING.copy()
         if custom_mapping:
             mapping.update(custom_mapping)
 
         return cls(mapping, gas_names)
 
     def validate(self) -> Dict[str, str]:
-        """Validates and returns the final mapping."""
+        """Validate and return the final gas name mapping.
+
+        Returns:
+            Dictionary mapping standard gas names to dataset-specific names
+
+        Raises:
+            ValueError: If a required gas is not found in any mapping
+        """
         validated_mapping = {}
 
         for gas in self._required_gases:
@@ -38,57 +64,91 @@ def validate(self) -> Dict[str, str]:
 
 @dataclass
 class DatasetMapping:
-    """Container for dimension and variable mappings"""
+    """Container for dimension and variable name mappings.
 
+    Attributes:
+        dim_mapping: Dictionary mapping standard dimension names to dataset-specific names
+        var_mapping: Dictionary mapping standard variable names to dataset-specific names
+    """
     dim_mapping: Dict[str, str]
     var_mapping: Dict[str, str]
 
-    def __post_init__(self):
-        """Validate mappings upon initialization"""
+    def __post_init__(self) -> None:
+        """Validate mappings upon initialization."""
         pass
 
     @classmethod
-    def from_dict(cls, d: Dict) -> "DatasetMapping":
-        """Create mapping from dictionary"""
+    def from_dict(cls, d: Dict[str, Dict[str, str]]) -> "DatasetMapping":
+        """Create mapping from dictionary representation.
+
+        Args:
+            d: Dictionary containing dim_mapping and var_mapping
+
+        Returns:
+            New DatasetMapping instance
+        """
         return cls(dim_mapping=d["dim_mapping"], var_mapping=d["var_mapping"])
 
 
 @xr.register_dataset_accessor("mapping")
 class DatasetMappingAccessor:
-    """
-    An accessor for xarray datasets that provides information about variable mappings.
-    The mapping is stored in the dataset's attributes.
+    """Accessor for xarray datasets that provides variable mapping functionality.
+
+    The mapping is stored in the dataset's attributes to maintain persistence.
     """
 
-    def __init__(self, xarray_obj):
+    def __init__(self, xarray_obj: xr.Dataset) -> None:
         self._obj = xarray_obj
 
     def set_mapping(self, mapping: DatasetMapping) -> None:
-        """Set the mapping in dataset attributes"""
-        # Validate that mapped variables exist in dataset
+        """Set the mapping in dataset attributes.
+
+        Args:
+            mapping: DatasetMapping instance to store
+
+        Raises:
+            ValueError: If mapped dimensions don't exist in dataset
+        """
         missing_dims = set(mapping.dim_mapping.values()) - set(self._obj.dims)
         if missing_dims:
             raise ValueError(f"Dataset missing required dimensions: {missing_dims}")
 
-        # Store mapping in attributes
         self._obj.attrs["dataset_mapping"] = asdict(mapping)
 
     @property
     def mapping(self) -> Optional[DatasetMapping]:
-        """Get the mapping from dataset attributes"""
+        """Get the mapping from dataset attributes.
+
+        Returns:
+            DatasetMapping if exists, None otherwise
+        """
         if "dataset_mapping" not in self._obj.attrs:
             return None
         return DatasetMapping.from_dict(self._obj.attrs["dataset_mapping"])
 
     def get_var(self, standard_name: str) -> Optional[str]:
-        """Get the actual variable name in the dataset for a standard name"""
+        """Get the dataset-specific variable name for a standard name.
+
+        Args:
+            standard_name: Standard variable name
+
+        Returns:
+            Dataset-specific variable name if found, None otherwise
+        """
         mapping = self.mapping
         if mapping is None:
             return None
         return mapping.var_mapping.get(standard_name)
 
     def get_dim(self, standard_name: str) -> Optional[str]:
-        """Get the actual dimension name in the dataset for a standard name"""
+        """Get the dataset-specific dimension name for a standard name.
+
+        Args:
+            standard_name: Standard dimension name
+
+        Returns:
+            Dataset-specific dimension name if found, None otherwise
+        """
         mapping = self.mapping
         if mapping is None:
             return None
@@ -97,10 +157,17 @@ def get_dim(self, standard_name: str) -> Optional[str]:
 
 @dataclass
 class AtmosphericMapping(DatasetMapping):
-    """Specific mapping for atmospheric data"""
+    """Specific mapping for atmospheric data with required dimensions and variables.
 
-    def __post_init__(self):
-        """Validate atmospheric-specific mappings"""
+    Inherits from DatasetMapping and adds validation for required atmospheric fields.
+    """
+
+    def __post_init__(self) -> None:
+        """Validate atmospheric-specific mappings.
+
+        Raises:
+            ValueError: If required dimensions or variables are missing
+        """
         required_dims = {"site", "layer", "level"}
         missing_dims = required_dims - set(self.dim_mapping.keys())
         if missing_dims:
@@ -113,20 +180,12 @@ def __post_init__(self):
 
 
 def create_default_mapping() -> AtmosphericMapping:
-    """Create a default mapping configuration"""
+    """Create a default atmospheric mapping configuration.
+
+    Returns:
+        AtmosphericMapping instance with default dimension and variable mappings
+    """
     return AtmosphericMapping(
-        dim_mapping={"site": "site", "layer": "layer", "level": "level"},
-        var_mapping={
-            "pres_layer": "pres_layer",
-            "pres_level": "pres_level",
-            "temp_layer": "temp_layer",
-            "temp_level": "temp_level",
-            "surface_temperature": "surface_temperature",
-            "solar_zenith_angle": "solar_zenith_angle",
-            "surface_albedo": "surface_albedo",
-            "surface_albedo_dir": "surface_albedo_dir",
-            "surface_albedo_dif": "surface_albedo_dif",
-            "surface_emissivity": "surface_emissivity",
-            "surface_emissivity_jacobian": "surface_emissivity_jacobian",
-        },
+        dim_mapping=DEFAULT_DIM_MAPPING,
+        var_mapping=DEFAULT_VAR_MAPPING,
     )