Improve xarray aspect of interface (#41)

docs/source/reference_guide/pyrte_rrtmgp_python_modules.rst

@@ -8,18 +8,18 @@ The documentation below provides details for several of the high-level functions
 For more information about the low-level functions available in the ``pyrte_rrtmgp.pyrte_rrtmgp`` submodule, please go to :ref:`low_level_interface`.
 
 
-pyrte\_rrtmgp.rte module
-------------------------
+pyrte\_rrtmgp.kernels.rte module
+--------------------------------
 
-.. automodule:: pyrte_rrtmgp.rte
+.. automodule:: pyrte_rrtmgp.kernels.rte
    :members:
    :undoc-members:
    :show-inheritance:
 
-pyrte\_rrtmgp.rrtmgp module
----------------------------
+pyrte\_rrtmgp.kernels.rrtmgp module
+-----------------------------------
 
-.. automodule:: pyrte_rrtmgp.rrtmgp
+.. automodule:: pyrte_rrtmgp.kernels.rrtmgp
    :members:
    :undoc-members:
    :show-inheritance:

docs/source/user_guide/usage.md

@@ -7,8 +7,8 @@ This section provides a brief overview of how to use `pyrte_rrtmgp` with Python.
 The `pyrte_rrtmgp` package contains the following submodules:
 
 - `pyrte_rrtmgp.pyrte_rrtmgp`: The main module that provides access to a subset of RTE-RRTMGP's Fortran functions in Python. The functions available in this module mirror the Fortran functions (see below). You can think of this as the low-level implementation that allows you to access the respective Fortran functions directly in Python. See [](low_level_interface) for more information.
-- `pyrte_rrtmgp.rte`: A high-level module that provides a more user-friendly Python interface for select RTE functions. This module is still under development and will be expanded in future releases. See [](module_ref) for details.
-- `pyrte_rrtmgp.rrtmgp`: A high-level module that provides a more user-friendly Python interface for select RRTMGP functions. This module is still under development and will be expanded in future releases. See [](module_ref) for details.
+- `pyrte_rrtmgp.kernels.rte`: A high-level module that provides a more user-friendly Python interface for select RTE functions. This module is still under development and will be expanded in future releases. See [](module_ref) for details.
+- `pyrte_rrtmgp.kernels.rrtmgp`: A high-level module that provides a more user-friendly Python interface for select RRTMGP functions. This module is still under development and will be expanded in future releases. See [](module_ref) for details.
 - `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}

examples/lw_example.ipynb

@@ -6,10 +6,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from pyrte_rrtmgp.gas_optics import GasOptics\n",
-    "import xarray as xr\n",
     "import numpy as np\n",
-    "from pyrte_rrtmgp.rte import lw_solver_noscat\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-4\n",
     "\n",
@@ -22,33 +24,25 @@
     "rfmip = rfmip.sel(expt=0)  # only one experiment\n",
     "\n",
     "kdist = xr.load_dataset(f\"{rte_rrtmgp_dir}/rrtmgp-gas-lw-g256.nc\")\n",
-    "\n",
-    "# RRTMGP won't run with pressure less than its minimum. so we add a small value to the minimum pressure\n",
-    "# There was an issue replicating k_dist%get_press_min() + epsilon(k_dist%get_press_min()) in python, sys.epsilon is not the same\n",
-    "min_index = rfmip[\"pres_level\"].argmin()\n",
-    "rfmip[\"pres_level\"][:, min_index] = 1.0051835744630002\n",
-    "\n",
-    "gas_optics = GasOptics(kdist, rfmip)\n",
-    "gas_optics.source_is_internal\n",
-    "tau, _, _, layer_src, level_src, sfc_src, sfc_src_jac = gas_optics.gas_optics()\n",
-    "\n",
-    "# Expand the surface emissivity to ngpt\n",
-    "sfc_emis = rfmip[\"surface_emissivity\"].values\n",
-    "sfc_emis = np.stack([sfc_emis]*tau.shape[2]).T\n",
+    "rrtmgp_gas_optics = kdist.gas_optics.load_atmosferic_conditions(rfmip)\n",
     "\n",
     "_, solver_flux_up, solver_flux_down, _, _ = lw_solver_noscat(\n",
-    "    tau=tau,\n",
-    "    lay_source=layer_src,\n",
-    "    lev_source=level_src,\n",
-    "    sfc_emis=sfc_emis,\n",
-    "    sfc_src=sfc_src,\n",
-    "    sfc_src_jac=sfc_src_jac\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",
     "\n",
-    "rlu = xr.load_dataset(\"../tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\")\n",
+    "rlu = xr.load_dataset(\n",
+    "    \"../tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
+    ")\n",
     "ref_flux_up = rlu.isel(expt=0)[\"rlu\"].values\n",
     "\n",
-    "rld = xr.load_dataset(\"../tests/test_python_frontend/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\")\n",
+    "rld = xr.load_dataset(\n",
+    "    \"../tests/test_python_frontend/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
+    ")\n",
     "ref_flux_down = rld.isel(expt=0)[\"rld\"].values\n",
     "\n",
     "assert np.isclose(solver_flux_up, ref_flux_up, atol=ERROR_TOLERANCE).all()\n",

examples/sw_example.ipynb

@@ -2,15 +2,16 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from pyrte_rrtmgp.gas_optics import GasOptics\n",
-    "import xarray as xr\n",
     "import numpy as np\n",
-    "from pyrte_rrtmgp.rte import sw_solver_2stream\n",
-    "from pyrte_rrtmgp.utils import compute_mu0, get_usecols\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, compute_toa_flux\n",
     "\n",
     "ERROR_TOLERANCE = 1e-4\n",
     "\n",
@@ -23,42 +24,24 @@
     "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_atmosferic_conditions(rfmip)\n",
     "\n",
-    "# RRTMGP won't run with pressure less than its minimum. so we add a small value to the minimum pressure\n",
-    "# There was an issue replicating k_dist%get_press_min() + epsilon(k_dist%get_press_min()) in python, sys.epsilon is not the same\n",
-    "min_index = rfmip[\"pres_level\"].argmin()\n",
-    "rfmip[\"pres_level\"][:, min_index] = 1.0051835744630002\n",
-    "\n",
-    "gas_optics = GasOptics(kdist, rfmip)\n",
-    "gas_optics.source_is_internal\n",
-    "tau, g, ssa, toa_flux = gas_optics.gas_optics()\n",
-    "\n",
-    "pres_layers = rfmip[\"pres_layer\"][\"layer\"]\n",
-    "top_at_1 = (pres_layers[0] < pres_layers[-1]).values.item()\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",
     "\n",
-    "total_solar_irradiance = rfmip[\"total_solar_irradiance\"].values\n",
-    "toa_flux = np.stack([toa_flux]*mu0.shape[0])\n",
-    "def_tsi = toa_flux.sum(axis=1)\n",
-    "toa_flux = (toa_flux.T * (total_solar_irradiance/def_tsi)).T\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",
-    "    top_at_1,\n",
-    "    tau,\n",
-    "    ssa,\n",
-    "    g,\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",
@@ -73,10 +56,14 @@
     "solver_flux_down = solver_flux_down * usecol[:, np.newaxis]\n",
     "\n",
     "# Compare the results with the reference data\n",
-    "rsu = xr.load_dataset(\"../tests/test_python_frontend/rsu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\")\n",
+    "rsu = xr.load_dataset(\n",
+    "    \"../tests/test_python_frontend/rsu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
+    ")\n",
     "ref_flux_up = rsu.isel(expt=0)[\"rsu\"].values\n",
     "\n",
-    "rsd = xr.load_dataset(\"../tests/test_python_frontend/rsd_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\")\n",
+    "rsd = xr.load_dataset(\n",
+    "    \"../tests/test_python_frontend/rsd_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\"\n",
+    ")\n",
     "ref_flux_down = rsd.isel(expt=0)[\"rsd\"].values\n",
     "\n",
     "assert np.isclose(solver_flux_up, ref_flux_up, atol=ERROR_TOLERANCE).all()\n",

pyrte_rrtmgp/constants.py

@@ -0,0 +1,5 @@
+HELMERT1 = 9.80665
+HELMERT2 = 0.02586
+M_DRY = 0.028964
+M_H2O = 0.018016
+AVOGAD = 6.02214076e23

pyrte_rrtmgp/exceptions.py

@@ -0,0 +1,16 @@
+class NotInternalSourceError(ValueError):
+    pass
+
+
+class NotExternalSourceError(ValueError):
+    pass
+
+
+class MissingAtmosfericConditionsError(AttributeError):
+    message = (
+        "You need to load the atmospheric conditions first."
+        "Use the method load_atmosferic_conditions with an appropriated file."
+    )
+
+    def __init__(self, message=message):
+        super().__init__(message)