Add lw and sw solver integration test to CI (#39)

Co-authored-by: tcmetzger <[email protected]>

.github/workflows/conda.yml

@@ -23,6 +23,9 @@ jobs:
 
     runs-on: ${{ matrix.platform }}
 
+    env:  # We need to add the path to the rrtmgp-data directory
+      RRTMGP_DATA: /home/runner/work/pyRTE-RRTMGP/pyRTE-RRTMGP/rrtmgp-data
+
     # The setup-miniconda action needs this to activate miniconda
     defaults:
       run:

docs/source/user_guide/usage.md

@@ -12,9 +12,7 @@ The `pyrte_rrtmgp` package contains the following submodules:
 - `pyrte_rrtmgp.utils`: A module that provides utility functions for working with RTE-RRTMGP data. This module is still under development and will be expanded in future releases. See [](module_ref) for details.
 
 ```{seealso}
-The folder `examples` in the repository contains a Jupyter notebook example that demonstrates how to use the package.
-
-The functions in `pyrte_rrtmgp.rte`, `pyrte_rrtmgp.rrtmgp`, and `pyrte_rrtmgp.utils` contain docstrings that are available in IDEs with features such as IntelliSense (in VSCode).
+The folder `/examples` in the repository contains Jupyter notebook examples that demonstrate how to use the package.
 ```
 
 ## Importing the Package

tests/test_python_frontend/test_lw_solver.py

@@ -0,0 +1,49 @@
+import os
+import numpy as np
+import xarray as xr
+from pyrte_rrtmgp.gas_optics import GasOptics
+from pyrte_rrtmgp.rte import lw_solver_noscat
+
+ERROR_TOLERANCE = 1e-4
+
+rte_rrtmgp_dir = os.environ.get("RRTMGP_DATA", "rrtmgp-data") 
+clear_sky_example_files = f"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/inputs"
+
+rfmip = xr.load_dataset(
+    f"{clear_sky_example_files}/multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc"
+)
+rfmip = rfmip.sel(expt=0)  # only one experiment
+kdist = xr.load_dataset(f"{rte_rrtmgp_dir}/rrtmgp-gas-lw-g256.nc")
+
+rlu = xr.load_dataset(
+    "tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc"
+)
+ref_flux_up = rlu.isel(expt=0)["rlu"].values
+
+rld = xr.load_dataset(
+    "tests/test_python_frontend/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc"
+)
+ref_flux_down = rld.isel(expt=0)["rld"].values
+
+
+def test_lw_solver_noscat():
+    min_index = np.argmin(rfmip["pres_level"].values)
+    rfmip["pres_level"][:, min_index] = 1.0051835744630002
+
+    gas_optics = GasOptics(kdist, rfmip)
+    tau, _, _, layer_src, level_src, sfc_src, sfc_src_jac = gas_optics.gas_optics()
+
+    sfc_emis = rfmip["surface_emissivity"].values
+    sfc_emis = np.stack([sfc_emis] * tau.shape[2]).T
+
+    _, solver_flux_up, solver_flux_down, _, _ = lw_solver_noscat(
+        tau=tau,
+        lay_source=layer_src,
+        lev_source=level_src,
+        sfc_emis=sfc_emis,
+        sfc_src=sfc_src,
+        sfc_src_jac=sfc_src_jac,
+    )
+
+    assert np.isclose(solver_flux_up, ref_flux_up, atol=ERROR_TOLERANCE).all()
+    assert np.isclose(solver_flux_down, ref_flux_down, atol=ERROR_TOLERANCE).all()

tests/test_python_frontend/test_sw_solver.py

@@ -0,0 +1,78 @@
+import os
+import numpy as np
+import xarray as xr
+from pyrte_rrtmgp.gas_optics import GasOptics
+from pyrte_rrtmgp.rte import sw_solver_2stream
+from pyrte_rrtmgp.utils import compute_mu0, get_usecols
+
+ERROR_TOLERANCE = 1e-4
+
+rte_rrtmgp_dir = os.environ.get("RRTMGP_DATA", "rrtmgp-data") 
+clear_sky_example_files = f"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/inputs"
+
+rfmip = xr.load_dataset(
+    f"{clear_sky_example_files}/multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc"
+)
+rfmip = rfmip.sel(expt=0)  # only one experiment
+kdist = xr.load_dataset(f"{rte_rrtmgp_dir}/rrtmgp-gas-sw-g224.nc")
+
+rsu = xr.load_dataset(
+    "tests/test_python_frontend/rsu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc"
+)
+ref_flux_up = rsu.isel(expt=0)["rsu"].values
+
+rsd = xr.load_dataset(
+    "tests/test_python_frontend/rsd_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc"
+)
+ref_flux_down = rsd.isel(expt=0)["rsd"].values
+
+
+def test_lw_solver_noscat():
+    min_index = np.argmin(rfmip["pres_level"].values)
+    rfmip["pres_level"][:, min_index] = 1.0051835744630002
+
+    gas_optics = GasOptics(kdist, rfmip)
+    gas_optics.source_is_internal
+    tau, g, ssa, toa_flux = gas_optics.gas_optics()
+
+    pres_layers = rfmip["pres_layer"]["layer"]
+    top_at_1 = (pres_layers[0] < pres_layers[-1]).values.item()
+
+    # Expand the surface albedo to ngpt
+    ngpt = len(kdist["gpt"])
+    surface_albedo = rfmip["surface_albedo"].values
+    surface_albedo = np.stack([surface_albedo] * ngpt)
+    sfc_alb_dir = surface_albedo.T.copy()
+    sfc_alb_dif = surface_albedo.T.copy()
+
+    nlayer = len(rfmip["layer"])
+    mu0 = compute_mu0(rfmip["solar_zenith_angle"].values, nlayer=nlayer)
+
+    total_solar_irradiance = rfmip["total_solar_irradiance"].values
+    toa_flux = np.stack([toa_flux] * mu0.shape[0])
+    def_tsi = toa_flux.sum(axis=1)
+    toa_flux = (toa_flux.T * (total_solar_irradiance / def_tsi)).T
+
+    _, _, _, solver_flux_up, solver_flux_down, _ = sw_solver_2stream(
+        top_at_1,
+        tau,
+        ssa,
+        g,
+        mu0,
+        sfc_alb_dir,
+        sfc_alb_dif,
+        inc_flux_dir=toa_flux,
+        inc_flux_dif=None,
+        has_dif_bc=False,
+        do_broadband=True,
+    )
+
+    # RTE will fail if passed solar zenith angles greater than 90 degree. We replace any with
+    #   nighttime columns with a default solar zenith angle. We'll mask these out later, of
+    #   course, but this gives us more work and so a better measure of timing.
+    usecol = get_usecols(rfmip["solar_zenith_angle"].values)
+    solver_flux_up = solver_flux_up * usecol[:, np.newaxis]
+    solver_flux_down = solver_flux_down * usecol[:, np.newaxis]
+
+    assert np.isclose(solver_flux_up, ref_flux_up, atol=ERROR_TOLERANCE).all()
+    assert np.isclose(solver_flux_down, ref_flux_down, atol=ERROR_TOLERANCE).all()