encapsulating variables for frontend

examples/lw_example.ipynb

@@ -2,15 +2,14 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "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 lw_solver_noscat\n",
     "\n",
     "\n",
     "ERROR_TOLERANCE = 1e-7\n",
@@ -24,21 +23,15 @@
     "rfmip = rfmip.sel(expt=0)  # only one experiment\n",
     "\n",
     "kdist = xr.load_dataset(f\"{rte_rrtmgp_dir}/rrtmgp-gas-lw-g256.nc\")\n",
-    "rrtmgp_gas_optics = kdist.gas_optics.load_atmospheric_conditions(rfmip)\n",
+    "lw_problem = kdist.gas_optics.load_atmospheric_conditions(rfmip)\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=rfmip[\"surface_emissivity\"].data,\n",
-    "    sfc_src=rrtmgp_gas_optics.sfc_src,\n",
-    "    sfc_src_jac=rrtmgp_gas_optics.sfc_src_jac,\n",
-    ")\n",
+    "lw_problem.sfc_emis = rfmip[\"surface_emissivity\"].data\n",
+    "\n",
+    "solver_flux_up, solver_flux_down = lw_problem.rte_solve()\n",
     "\n",
     "rlu_reference = f\"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/reference/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
     "rld_reference = f\"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/reference/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
     "\n",
-    "\n",
     "rlu = xr.load_dataset(rlu_reference, decode_cf=False)\n",
     "ref_flux_up = rlu.isel(expt=0)[\"rlu\"].values\n",
     "\n",

examples/sw_example.ipynb

@@ -10,8 +10,6 @@
     "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, compute_toa_flux\n",
     "\n",
     "ERROR_TOLERANCE = 1e-7\n",
     "\n",
@@ -24,36 +22,13 @@
     "rfmip = rfmip.sel(expt=0)  # only one experiment\n",
     "\n",
     "kdist = xr.load_dataset(f\"{rte_rrtmgp_dir}/rrtmgp-gas-sw-g224.nc\")\n",
-    "gas_optics = kdist.gas_optics.load_atmospheric_conditions(rfmip)\n",
+    "sw_problem = kdist.gas_optics.load_atmospheric_conditions(rfmip)\n",
     "\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",
-    "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=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",
-    "    do_broadband=True,\n",
-    ")\n",
-    "\n",
-    "# RTE will fail if passed solar zenith angles greater than 90 degree. We replace any with\n",
-    "#   nighttime columns with a default solar zenith angle. We'll mask these out later, of\n",
-    "#   course, but this gives us more work and so a better measure of timing.\n",
-    "usecol = get_usecols(rfmip[\"solar_zenith_angle\"].values)\n",
-    "solver_flux_up = solver_flux_up * usecol[:, np.newaxis]\n",
-    "solver_flux_down = solver_flux_down * usecol[:, np.newaxis]\n",
+    "sw_problem.sfc_alb_dir = rfmip[\"surface_albedo\"].data\n",
+    "sw_problem.total_solar_irradiance = rfmip[\"total_solar_irradiance\"].data\n",
+    "sw_problem.solar_zenith_angle = rfmip[\"solar_zenith_angle\"].values\n",
     "\n",
+    "solver_flux_up, solver_flux_down = sw_problem.solve()\n",
     "\n",
     "rsu_reference = f\"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/reference/rsu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
     "rsd_reference = f\"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/reference/rsd_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",

pyrte_rrtmgp/constants.py

@@ -3,3 +3,7 @@
 M_DRY = 0.028964
 M_H2O = 0.018016
 AVOGAD = 6.02214076e23
+SOLAR_CONSTANTS = {
+    'A_OFFSET': 0.1495954,
+    'B_OFFSET': 0.00066696,
+}

pyrte_rrtmgp/exceptions.py

@@ -1,11 +1,3 @@
-class NotInternalSourceError(ValueError):
-    pass
-
-
-class NotExternalSourceError(ValueError):
-    pass
-
-
 class MissingAtmosphericConditionsError(AttributeError):
     message = (
         "You need to load the atmospheric conditions first."

pyrte_rrtmgp/kernels/rte.py

@@ -174,36 +174,37 @@ def sw_solver_noscat(
 
 
 def sw_solver_2stream(
-    top_at_1,
     tau,
     ssa,
     g,
     mu0,
     sfc_alb_dir,
     sfc_alb_dif,
     inc_flux_dir,
+    top_at_1=True,
     inc_flux_dif=None,
     has_dif_bc=False,
-    do_broadband=False,
+    do_broadband=True,
 ):
     """
     Solve the shortwave radiative transfer equation using the 2-stream approximation.
 
     Args:
-        top_at_1 (bool): Flag indicating whether the top of the atmosphere is at level 1.
         tau (ndarray): Array of optical depths with shape (ncol, nlay, ngpt).
         ssa (ndarray): Array of single scattering albedos with shape (ncol, nlay, ngpt).
         g (ndarray): Array of asymmetry parameters with shape (ncol, nlay, ngpt).
         mu0 (ndarray): Array of cosine of solar zenith angles with shape (ncol, ngpt).
         sfc_alb_dir (ndarray): Array of direct surface albedos with shape (ncol, ngpt).
         sfc_alb_dif (ndarray): Array of diffuse surface albedos with shape (ncol, ngpt).
         inc_flux_dir (ndarray): Array of direct incident fluxes with shape (ncol, ngpt).
+        top_at_1 (bool): Flag indicating whether the top of the atmosphere is at level 1.
+            Defaults to True.
         inc_flux_dif (ndarray, optional): Array of diffuse incident fluxes with shape (ncol, ngpt).
             Defaults to None.
         has_dif_bc (bool, optional): Flag indicating whether the boundary condition includes diffuse fluxes.
             Defaults to False.
         do_broadband (bool, optional): Flag indicating whether to compute broadband fluxes.
-            Defaults to False.
+            Defaults to True.
 
     Returns:
         Tuple of ndarrays: Tuple containing the following arrays:

pyrte_rrtmgp/rrtmgp_data.py

@@ -37,7 +37,6 @@ def download_rrtmgp_data():
     # Path to the file containing the checksum of the downloaded file
     checksum_file_path = os.path.join(cache_dir, f"{TAG}.tar.gz.sha256")
 
-    # Download the file if it doesn't exist or if the checksum doesn't match
     # Download the file if it doesn't exist or if the checksum doesn't match
     if not os.path.exists(file_path) or (
         os.path.exists(checksum_file_path)