final adjustments

examples/lw_example.ipynb

@@ -2,12 +2,10 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sys\n",
-    "\n",
     "import numpy as np\n",
     "import xarray as xr\n",
     "\n",
@@ -28,15 +26,11 @@
     "kdist = xr.load_dataset(f\"{rte_rrtmgp_dir}/rrtmgp-gas-lw-g256.nc\")\n",
     "rrtmgp_gas_optics = kdist.gas_optics.load_atmosferic_conditions(rfmip)\n",
     "\n",
-    "# Expand the surface emissivity to ngpt\n",
-    "sfc_emis = rfmip[\"surface_emissivity\"].values\n",
-    "sfc_emis = np.stack([sfc_emis] * rrtmgp_gas_optics.tau.shape[2]).T\n",
-    "\n",
     "_, solver_flux_up, solver_flux_down, _, _ = lw_solver_noscat(\n",
     "    tau=rrtmgp_gas_optics.tau,\n",
     "    lay_source=rrtmgp_gas_optics.lay_src,\n",
     "    lev_source=rrtmgp_gas_optics.lev_src,\n",
-    "    sfc_emis=sfc_emis,\n",
+    "    sfc_emis=rfmip[\"surface_emissivity\"].data,\n",
     "    sfc_src=rrtmgp_gas_optics.sfc_src,\n",
     "    sfc_src_jac=rrtmgp_gas_optics.sfc_src_jac,\n",
     ")\n",

examples/sw_example.ipynb

@@ -6,14 +6,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sys\n",
-    "\n",
     "import numpy as np\n",
     "import xarray as xr\n",
     "\n",
     "from pyrte_rrtmgp import rrtmgp_gas_optics\n",
     "from pyrte_rrtmgp.kernels.rte import sw_solver_2stream\n",
-    "from pyrte_rrtmgp.utils import compute_mu0, get_usecols\n",
+    "from pyrte_rrtmgp.utils import compute_mu0, get_usecols, compute_toa_flux\n",
     "\n",
     "ERROR_TOLERANCE = 1e-4\n",
     "\n",
@@ -28,30 +26,22 @@
     "kdist = xr.load_dataset(f\"{rte_rrtmgp_dir}/rrtmgp-gas-sw-g224.nc\")\n",
     "gas_optics = kdist.gas_optics.load_atmosferic_conditions(rfmip)\n",
     "\n",
-    "pres_layers = rfmip[\"pres_layer\"][\"layer\"]\n",
-    "\n",
-    "# Expand the surface albedo to ngpt\n",
-    "ngpt = len(kdist[\"gpt\"])\n",
-    "surface_albedo = rfmip[\"surface_albedo\"].values\n",
-    "surface_albedo = np.stack([surface_albedo] * ngpt)\n",
-    "sfc_alb_dir = surface_albedo.T.copy()\n",
-    "sfc_alb_dif = surface_albedo.T.copy()\n",
+    "surface_albedo = rfmip[\"surface_albedo\"].data\n",
+    "total_solar_irradiance = rfmip[\"total_solar_irradiance\"].data\n",
     "\n",
     "nlayer = len(rfmip[\"layer\"])\n",
     "mu0 = compute_mu0(rfmip[\"solar_zenith_angle\"].values, nlayer=nlayer)\n",
-    "total_solar_irradiance = rfmip[\"total_solar_irradiance\"].values\n",
-    "toa_flux = np.stack([gas_optics.solar_source] * mu0.shape[0])\n",
-    "def_tsi = toa_flux.sum(axis=1)\n",
-    "toa_flux = (toa_flux.T * (total_solar_irradiance / def_tsi)).T\n",
+    "\n",
+    "toa_flux = compute_toa_flux(total_solar_irradiance, gas_optics.solar_source)\n",
     "\n",
     "_, _, _, solver_flux_up, solver_flux_down, _ = sw_solver_2stream(\n",
     "    kdist.gas_optics.top_at_1,\n",
     "    gas_optics.tau,\n",
     "    gas_optics.ssa,\n",
     "    gas_optics.g,\n",
     "    mu0,\n",
-    "    sfc_alb_dir,\n",
-    "    sfc_alb_dif,\n",
+    "    sfc_alb_dir=surface_albedo,\n",
+    "    sfc_alb_dif=surface_albedo,\n",
     "    inc_flux_dir=toa_flux,\n",
     "    inc_flux_dif=None,\n",
     "    has_dif_bc=False,\n",

pyrte_rrtmgp/constants.py

@@ -2,4 +2,4 @@
 HELMERT2 = 0.02586
 M_DRY = 0.028964
 M_H2O = 0.018016
-AVOGAD = 6.02214076e23
+AVOGAD = 6.02214076e23

pyrte_rrtmgp/kernels/rte.py

@@ -79,6 +79,9 @@ def lw_solver_noscat(
 
     ncol, nlay, ngpt = tau.shape
 
+    if len(sfc_emis.shape) == 1:
+        sfc_emis = np.stack([sfc_emis] * ngpt).T
+
     # default values
     n_quad_angs = nmus
 
@@ -211,6 +214,11 @@ def sw_solver_2stream(
     """
     ncol, nlay, ngpt = tau.shape
 
+    if len(sfc_alb_dir.shape) == 1:
+        sfc_alb_dir = np.stack([sfc_alb_dir] * ngpt).T
+    if len(sfc_alb_dif.shape) == 1:
+        sfc_alb_dif = np.stack([sfc_alb_dif] * ngpt).T
+
     if inc_flux_dif is None:
         inc_flux_dif = np.zeros((ncol, ngpt), dtype=np.float64)
 

pyrte_rrtmgp/rrtmgp_gas_optics.py

@@ -1,16 +1,17 @@
 import sys
-import xarray as xr
+from dataclasses import dataclass
+from typing import Optional
+
 import numpy as np
 import numpy.typing as npt
-from typing import Optional
-from dataclasses import dataclass
+import xarray as xr
 
+from pyrte_rrtmgp.constants import AVOGAD, HELMERT1, HELMERT2, M_DRY, M_H2O
 from pyrte_rrtmgp.exceptions import (
-    NotInternalSourceError,
-    NotExternalSourceError,
     MissingAtmosfericConditionsError,
+    NotExternalSourceError,
+    NotInternalSourceError,
 )
-from pyrte_rrtmgp.constants import HELMERT1, HELMERT2, M_DRY, M_H2O, AVOGAD
 from pyrte_rrtmgp.kernels.rrtmgp import (
     compute_planck_source,
     compute_tau_absorption,
@@ -124,7 +125,6 @@ def load_atmosferic_conditions(self, atmosferic_conditions: xr.Dataset):
         return self.gas_optics
 
     def get_col_gas(self):
-
         if self._atm_cond is None:
             raise MissingAtmosfericConditionsError()
 
@@ -203,7 +203,6 @@ def gas_mappings(self):
     @property
     def top_at_1(self):
         if self._top_at_1 is None:
-
             if self._atm_cond is None:
                 raise MissingAtmosfericConditionsError()
 
@@ -268,7 +267,6 @@ def compute_planck(self):
         )
 
     def compute_gas_taus(self):
-
         minor_gases_lower = self.extract_names(self._obj["minor_gases_lower"].data)
         minor_gases_upper = self.extract_names(self._obj["minor_gases_upper"].data)
         # check if the index is correct

pyrte_rrtmgp/utils.py

@@ -25,3 +25,19 @@ def compute_mu0(solar_zenith_angle, nlayer=None):
     if nlayer is not None:
         mu0 = np.stack([mu0] * nlayer).T
     return mu0
+
+
+def compute_toa_flux(total_solar_irradiance, solar_source):
+    """Compute the top of atmosphere flux
+
+    Args:
+        total_solar_irradiance (np.ndarray): Total solar irradiance
+        solar_source (np.ndarray): Solar source
+
+    Returns:
+        np.ndarray: Top of atmosphere flux
+    """
+    ncol = total_solar_irradiance.shape[0]
+    toa_flux = np.stack([solar_source] * ncol)
+    def_tsi = toa_flux.sum(axis=1)
+    return (toa_flux.T * (total_solar_irradiance / def_tsi)).T

tests/test_python_frontend/test_lw_solver.py

@@ -1,10 +1,9 @@
 import os
-import sys
 
 import numpy as np
 import pytest
 import xarray as xr
-from pyrte_rrtmgp.rrtmgp_gas_optics import GasOptics
+from pyrte_rrtmgp import rrtmgp_gas_optics
 from pyrte_rrtmgp.kernels.rte import lw_solver_noscat
 
 ERROR_TOLERANCE = 1e-4
@@ -30,23 +29,15 @@
 
 
 def test_lw_solver_noscat():
-    min_index = np.argmin(rfmip["pres_level"].data)
-    min_press = kdist["press_ref"].min().item() + sys.float_info.epsilon
-    rfmip["pres_level"][:, min_index] = min_press
-
-    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
+    rrtmgp_gas_optics = kdist.gas_optics.load_atmosferic_conditions(rfmip)
 
     _, 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,
+        tau=rrtmgp_gas_optics.tau,
+        lay_source=rrtmgp_gas_optics.lay_src,
+        lev_source=rrtmgp_gas_optics.lev_src,
+        sfc_emis=rfmip["surface_emissivity"].data,
+        sfc_src=rrtmgp_gas_optics.sfc_src,
+        sfc_src_jac=rrtmgp_gas_optics.sfc_src_jac,
     )
 
     assert np.isclose(solver_flux_up, ref_flux_up, atol=ERROR_TOLERANCE).all()

tests/test_python_frontend/test_sw_solver.py

@@ -1,11 +1,11 @@
 import os
-import sys
 
 import numpy as np
+import pytest
 import xarray as xr
-from pyrte_rrtmgp.rrtmgp_gas_optics import GasOptics
+from pyrte_rrtmgp import rrtmgp_gas_optics
 from pyrte_rrtmgp.kernels.rte import sw_solver_2stream
-from pyrte_rrtmgp.utils import compute_mu0, get_usecols
+from pyrte_rrtmgp.utils import compute_mu0, compute_toa_flux, get_usecols
 
 ERROR_TOLERANCE = 1e-4
 
@@ -30,40 +30,24 @@
 
 
 def test_lw_solver_noscat():
-    min_index = np.argmin(rfmip["pres_level"].data)
-    min_press = kdist["press_ref"].min().item() + sys.float_info.epsilon
-    rfmip["pres_level"][:, min_index] = min_press
+    gas_optics = kdist.gas_optics.load_atmosferic_conditions(rfmip)
 
-    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()
+    surface_albedo = rfmip["surface_albedo"].data
+    total_solar_irradiance = rfmip["total_solar_irradiance"].data
 
     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
+    toa_flux = compute_toa_flux(total_solar_irradiance, gas_optics.solar_source)
 
     _, _, _, solver_flux_up, solver_flux_down, _ = sw_solver_2stream(
-        top_at_1,
-        tau,
-        ssa,
-        g,
+        kdist.gas_optics.top_at_1,
+        gas_optics.tau,
+        gas_optics.ssa,
+        gas_optics.g,
         mu0,
-        sfc_alb_dir,
-        sfc_alb_dif,
+        sfc_alb_dir=surface_albedo,
+        sfc_alb_dif=surface_albedo,
         inc_flux_dir=toa_flux,
         inc_flux_dif=None,
         has_dif_bc=False,