Update of the cloud optics kernels pybinds (#80)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

CMakeLists.txt

@@ -33,23 +33,27 @@ endif()
 
 include(ExternalProject REQUIRED)
 
-if(CMAKE_HOST_SYSTEM_NAME STREQUAL "Windows")
-    set(BUILD_COMMAND_STRING "set FC=%CMAKE_Fortran_COMPILER% && cd build && nmake /A")
-else()
-    set(BUILD_COMMAND_STRING "FC=${CMAKE_Fortran_COMPILER} FCFLAGS='-fPIC' make -C build -j ${N}")
-endif()
+set(FFLAGS "-ffree-line-length-none -m64 -std=f2008 -march=native -fbounds-check -fmodule-private -fimplicit-none -finit-real=nan")
 
 ExternalProject_Add(
     rte-rrtmgp
-    GIT_REPOSITORY https://github.com/earth-system-radiation/rte-rrtmgp.git
-    GIT_TAG 06ea8284071b80d393ff8df659a2ef4b4bfb2aa8
-    GIT_CONFIG "advice.detachedHead=false"
-    # GIT_SHALLOW TRUE
+    # For development only
+    GIT_REPOSITORY https://github.com/makepath-alex/rte-rrtmgp.git
+    GIT_TAG better-header-placement-dev
+    GIT_SHALLOW TRUE
     SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp
-    CONFIGURE_COMMAND ""
+    CMAKE_ARGS
+        -DCMAKE_Fortran_COMPILER=${CMAKE_Fortran_COMPILER}
+        -DCMAKE_Fortran_FLAGS=${FFLAGS}
+        -DRTE_ENABLE_SP=OFF
+        -DKERNEL_MODE=default
+        -DBUILD_TESTING=OFF
+        -DFAILURE_THRESHOLD=7.e-4
+        -DCMAKE_BUILD_TYPE=Debug
+        -DCMAKE_INSTALL_PREFIX=${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp/
     BUILD_IN_SOURCE TRUE
-    BUILD_COMMAND eval ${BUILD_COMMAND_STRING}
-    INSTALL_COMMAND ""
+    BUILD_COMMAND ${CMAKE_COMMAND} --build . --target all --parallel
+    INSTALL_COMMAND ${CMAKE_COMMAND} --install .
 )
 
 # Compile C bindings
@@ -63,12 +67,11 @@ pybind11_add_module(${TARGET_NAME} ${SOURCES})
 add_dependencies(${TARGET_NAME} rte-rrtmgp)
 
 target_include_directories(${TARGET_NAME} PUBLIC
-    ${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp/rte-kernels/api/
-    ${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp/rrtmgp-kernels/api/
+    ${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp/include/
 )
 
 target_link_directories(${TARGET_NAME} PUBLIC
-    ${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp/build
+    ${CMAKE_CURRENT_BINARY_DIR}/rte-rrtmgp/lib/
 )
 
 target_link_libraries(${TARGET_NAME} PUBLIC

pyproject.toml

@@ -49,7 +49,7 @@ addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
 xfail_strict = true
 log_cli_level = "INFO"
 filterwarnings = ["error"]
-testpaths = ["tests"]
+testpaths = ["pyrte_rrtmgp/tests"]
 
 [tool.cibuildwheel]
 test-command = "pytest {project}/tests"

pyrte_rrtmgp/kernels/rrtmgp.py

@@ -84,7 +84,7 @@ def interpolation(
     fmajor = np.ndarray([2, 2, 2, ncol, nlay, nflav], dtype=np.float64, order="F")
     fminor = np.ndarray([2, 2, ncol, nlay, nflav], dtype=np.float64, order="F")
     col_mix = np.ndarray([2, ncol, nlay, nflav], dtype=np.float64, order="F")
-    tropo = np.ndarray([ncol, nlay], dtype=np.int32, order="F")
+    tropo = np.ndarray([ncol, nlay], dtype=np.bool_, order="F")
     jeta = np.ndarray([2, ncol, nlay, nflav], dtype=np.int32, order="F")
     jpress = np.ndarray([ncol, nlay], dtype=np.int32, order="F")
 
@@ -177,6 +177,7 @@ def compute_planck_source(
         tropo: Troposphere mask with shape (ncol, nlay)
         jtemp: Temperature interpolation indices with shape (ncol, nlay)
         jpress: Pressure interpolation indices with shape (ncol, nlay)
+        gpoint_bands: TODO: Add information (ngpt)
         band_lims_gpt: Band limits in g-point space with shape (2, nbnd)
         pfracin: Planck fractions with shape (ntemp, neta, npres+1, ngpt)
         temp_ref_min: Minimum reference temperature
@@ -192,7 +193,7 @@ def compute_planck_source(
             - sfc_src_jac: Surface emission Jacobian with shape (ncol, ngpt)
     """
     sfc_lay = nlay if top_at_1 else 1
-    gpoint_bands = []
+    gpoint_bands = np.ndarray((ngpt), dtype=np.int32, order="F")
     totplnk_delta = (temp_ref_max - temp_ref_min) / (nPlanckTemp - 1)
 
     # Initialize output arrays
@@ -425,7 +426,7 @@ def compute_tau_rayleigh(
         krayl: Rayleigh scattering coefficients with shape (ntemp, neta, ngpt, 2)
         idx_h2o: Index of water vapor
         col_dry: Dry air column amounts with shape (ncol, nlay)
-        col_gas: Gas concentrations with shape (ncol, nlay, ngas)
+        col_gas: Gas concentrations with shape (ncol, nlay, ngas + 1)
         fminor: Minor gas interpolation weights
         jeta: Eta interpolation indices
         tropo: Troposphere mask with shape (ncol, nlay)
@@ -437,6 +438,8 @@ def compute_tau_rayleigh(
     # Initialize output array
     tau_rayleigh = np.ndarray((ncol, nlay, ngpt), dtype=np.float64, order="F")
 
+    ngas = ngas - 1  # Fortran uses index 0 here
+
     args = [
         ncol,
         nlay,

pyrte_rrtmgp/rrtmgp_data.py

@@ -8,8 +8,8 @@
 import requests
 
 # URL of the file to download
-TAG = "v1.8.2"
-DATA_URL = f"https://github.com/earth-system-radiation/rrtmgp-data/archive/refs/tags/{TAG}.tar.gz"
+REF = "v1.9"  # Can be a tag (e.g. "v1.8.2") or branch name (e.g. "main")
+DATA_URL = f"https://github.com/earth-system-radiation/rrtmgp-data/archive/refs/{'tags' if REF.startswith('v') else 'heads'}/{REF}.tar.gz"
 
 
 def get_cache_dir() -> str:
@@ -51,10 +51,10 @@ def download_rrtmgp_data() -> str:
     cache_dir = get_cache_dir()
 
     # Path to the downloaded file
-    file_path = os.path.join(cache_dir, f"{TAG}.tar.gz")
+    file_path = os.path.join(cache_dir, f"{REF}.tar.gz")
 
     # Path to the file containing the checksum of the downloaded file
-    checksum_file_path = os.path.join(cache_dir, f"{TAG}.tar.gz.sha256")
+    checksum_file_path = os.path.join(cache_dir, f"{REF}.tar.gz.sha256")
 
     # Download the file if it doesn't exist or if the checksum doesn't match
     if not os.path.exists(file_path) or (
@@ -77,7 +77,11 @@ def download_rrtmgp_data() -> str:
     with tarfile.open(file_path) as tar:
         tar.extractall(path=cache_dir, filter="data")
 
-    return os.path.join(cache_dir, f"rrtmgp-data-{TAG[1:]}")
+    # Handle both tag and branch names in the extracted directory name
+    # For tags like "v1.8.2", remove the "v" prefix
+    # For branches like "main", use as-is
+    ref_dirname = REF[1:] if REF.startswith("v") else REF
+    return os.path.join(cache_dir, f"rrtmgp-data-{ref_dirname}")
 
 
 def _get_file_checksum(filepath: Union[str, Path], mode: str = "r") -> str:

pyrte_rrtmgp/rrtmgp_gas_optics.py

@@ -98,8 +98,8 @@ def __init__(
         self.is_internal = is_internal
 
         # Get the gas names from the dataset
-        self._gas_names: tuple[str, ...] = self.extract_names(
-            self._dataset["gas_names"].values
+        self._gas_names: tuple[str, ...] = tuple(
+            self.extract_names(self._dataset["gas_names"].values)
         )
 
         if selected_gases is not None:
@@ -135,9 +135,17 @@ def _initialize_pressure_levels(
         """
         pres_level_var = atmosphere.mapping.get_var("pres_level")
 
+        min_press = self._dataset["press_ref"].min().item()
+
         min_index = np.argmin(atmosphere[pres_level_var].data)
         min_press = self._dataset["press_ref"].min().item() + sys.float_info.epsilon
-        atmosphere[pres_level_var][:, min_index] = min_press
+
+        # Replace values smaller than min_press with min_press at min_index
+        atmosphere[pres_level_var][:, min_index] = xr.where(
+            atmosphere[pres_level_var][:, min_index] < min_press,
+            min_press,
+            atmosphere[pres_level_var][:, min_index],
+        )
 
         if not inplace:
             return atmosphere
@@ -360,19 +368,47 @@ def tau_absorption(
         ngas = gas_interpolation_data["gas"].size
         nflav = self.flavors_sets["flavor"].size
 
-        nminorlower = self._dataset["minor_scales_with_density_lower"].size
-        nminorupper = self._dataset["minor_scales_with_density_upper"].size
-        nminorklower = self._dataset["contributors_lower"].size
-        nminorkupper = self._dataset["contributors_upper"].size
-
         minor_gases_lower = self.extract_names(self._dataset["minor_gases_lower"].data)
         minor_gases_upper = self.extract_names(self._dataset["minor_gases_upper"].data)
+
+        lower_gases_mask = np.isin(minor_gases_lower, self._gas_names)
+        upper_gases_mask = np.isin(minor_gases_upper, self._gas_names)
+
+        # TODO: Hardcoded 16, but shouldn't it be nbnd?
+        upper_gases_mask_expanded = np.repeat(upper_gases_mask, 16)
+        lower_gases_mask_expanded = np.repeat(lower_gases_mask, 16)
+
+        reduced_dataset = self._dataset.isel(
+            contributors_lower=lower_gases_mask_expanded
+        )
+        reduced_dataset = reduced_dataset.isel(
+            contributors_upper=upper_gases_mask_expanded
+        )
+        reduced_dataset = reduced_dataset.isel(
+            minor_absorber_intervals_lower=lower_gases_mask
+        )
+        reduced_dataset = reduced_dataset.isel(
+            minor_absorber_intervals_upper=upper_gases_mask
+        )
+
+        minor_gases_lower_reduced = minor_gases_lower[lower_gases_mask]
+        minor_gases_upper_reduced = minor_gases_upper[upper_gases_mask]
+
+        nminorlower = reduced_dataset.sizes["minor_absorber_intervals_lower"]
+        nminorupper = reduced_dataset.sizes["minor_absorber_intervals_upper"]
+        nminorklower = reduced_dataset.sizes["contributors_lower"]
+        nminorkupper = reduced_dataset.sizes["contributors_upper"]
+
         # check if the index is correct
-        idx_minor_lower = self.get_idx_minor(minor_gases_lower)
-        idx_minor_upper = self.get_idx_minor(minor_gases_upper)
+        idx_minor_lower = self.get_idx_minor(minor_gases_lower_reduced)
+        idx_minor_upper = self.get_idx_minor(minor_gases_upper_reduced)
 
-        scaling_gas_lower = self.extract_names(self._dataset["scaling_gas_lower"].data)
-        scaling_gas_upper = self.extract_names(self._dataset["scaling_gas_upper"].data)
+        scaling_gas_lower = self.extract_names(
+            reduced_dataset["scaling_gas_lower"].data
+        )
+        scaling_gas_upper = self.extract_names(
+            reduced_dataset["scaling_gas_upper"].data
+        )
 
         idx_minor_scaling_lower = self.get_idx_minor(scaling_gas_lower)
         idx_minor_scaling_upper = self.get_idx_minor(scaling_gas_upper)
@@ -397,22 +433,22 @@ def tau_absorption(
             nminorkupper,
             self._selected_gas_names_ext.index("h2o"),
             self.gpoint_flavor,
-            self._dataset["bnd_limits_gpt"],
-            self._dataset["kmajor"],
-            self._dataset["kminor_lower"],
-            self._dataset["kminor_upper"],
-            self._dataset["minor_limits_gpt_lower"],
-            self._dataset["minor_limits_gpt_upper"],
-            self._dataset["minor_scales_with_density_lower"],
-            self._dataset["minor_scales_with_density_upper"],
-            self._dataset["scale_by_complement_lower"],
-            self._dataset["scale_by_complement_upper"],
+            reduced_dataset["bnd_limits_gpt"],
+            reduced_dataset["kmajor"],
+            reduced_dataset["kminor_lower"],
+            reduced_dataset["kminor_upper"],
+            reduced_dataset["minor_limits_gpt_lower"],
+            reduced_dataset["minor_limits_gpt_upper"],
+            reduced_dataset["minor_scales_with_density_lower"],
+            reduced_dataset["minor_scales_with_density_upper"],
+            reduced_dataset["scale_by_complement_lower"],
+            reduced_dataset["scale_by_complement_upper"],
             idx_minor_lower,
             idx_minor_upper,
             idx_minor_scaling_lower,
             idx_minor_scaling_upper,
-            self._dataset["kminor_start_lower"],
-            self._dataset["kminor_start_upper"],
+            reduced_dataset["kminor_start_lower"],
+            reduced_dataset["kminor_start_upper"],
             gas_interpolation_data["tropopause_mask"],
             gas_interpolation_data["column_mix"],
             gas_interpolation_data["fmajor"],
@@ -588,7 +624,9 @@ def extract_names(names: npt.NDArray) -> tuple[str, ...]:
         Returns:
             Tuple of decoded and cleaned names
         """
-        output = tuple(gas.tobytes().decode().strip().split("_")[0] for gas in names)
+        output = np.array(
+            [gas.tobytes().decode().strip().split("_")[0] for gas in names]
+        )
         return output
 
     @staticmethod
@@ -662,21 +700,34 @@ def compute(
         """
         # Create and validate gas mapping
         gas_mapping = GasMapping.create(self._gas_names, gas_name_map).validate()
+        gas_mapping = {
+            k: v for k, v in gas_mapping.items() if v in list(atmosphere.data_vars)
+        }
+        self._gas_names = [
+            k for k, v in gas_mapping.items() if v in list(atmosphere.data_vars)
+        ]
 
         if variable_mapping is None:
             variable_mapping = create_default_mapping()
         # Set mapping in accessor
         atmosphere.mapping.set_mapping(variable_mapping)
 
+        pres_layer_var = atmosphere.mapping.get_var("pres_layer")
+        top_at_1 = (
+            atmosphere[pres_layer_var].values[0, 0]
+            < atmosphere[pres_layer_var].values[0, -1]
+        )
+
         # Modify pressure levels to avoid division by zero, runs inplace
         self._initialize_pressure_levels(atmosphere)
 
         gas_interpolation_data = self.interpolate(atmosphere, gas_mapping)
         problem = self.compute_problem(atmosphere, gas_interpolation_data)
         sources = self.compute_sources(atmosphere, gas_interpolation_data)
         boundary_conditions = self.compute_boundary_conditions(atmosphere)
+        gas_data = self._dataset["bnd_limits_gpt"].to_dataset()
 
-        gas_optics = xr.merge([sources, problem, boundary_conditions])
+        gas_optics = xr.merge([sources, problem, boundary_conditions, gas_data])
 
         # Add problem type to dataset attributes
         if problem_type == "absorption" and self.is_internal:
@@ -695,9 +746,11 @@ def compute(
         if add_to_input:
             atmosphere.update(gas_optics)
             atmosphere.attrs["problem_type"] = problem_type
+            atmosphere.attrs["top_at_1"] = top_at_1
         else:
             output_ds = gas_optics
             output_ds.attrs["problem_type"] = problem_type
+            output_ds.attrs["top_at_1"] = top_at_1
             output_ds.mapping.set_mapping(variable_mapping)
             return output_ds
 
@@ -757,6 +810,7 @@ def compute_boundary_conditions(self, atmosphere: xr.Dataset) -> xr.DataArray:
                 coords={
                     site_dim: atmosphere[site_dim],
                 },
+                name=surface_emissivity_var,
             )
         else:
             return atmosphere[surface_emissivity_var]
@@ -782,7 +836,11 @@ def compute_planck(
         surface_temperature_var = atmosphere.mapping.get_var("surface_temperature")
 
         # Check if the top layer is at the first level
-        top_at_1 = atmosphere[layer_dim][0] < atmosphere[layer_dim][-1]
+        pres_layer_var = atmosphere.mapping.get_var("pres_layer")
+        top_at_1 = (
+            atmosphere[pres_layer_var].values[0, 0]
+            < atmosphere[pres_layer_var].values[0, -1]
+        )
 
         ncol = atmosphere.sizes[site_dim]
         nlay = atmosphere.sizes[layer_dim]