diff --git a/.github/workflows/test_docs_build.yml b/.github/workflows/test_docs_build.yml deleted file mode 100644 index 21ef041..0000000 --- a/.github/workflows/test_docs_build.yml +++ /dev/null @@ -1,43 +0,0 @@ -name: Build Sphinx Documentation - -on: - push: - branches: - - main - paths: - - 'docs/**/*' - pull_request: - branches: - - main - paths: - - 'docs/**/*' - -jobs: - build-docs: - runs-on: ubuntu-latest - steps: - - name: Check out repository - uses: actions/checkout@v2 - - - name: Set up Python - uses: actions/setup-python@v2 - with: - python-version: '3.11' # Specify the Python version you need for your project - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -r docs/requirements-docs.txt - - - name: Build Sphinx Documentation - run: | - cd docs - sphinx-build -b html source build/html - - # # Uncomment the following lines to check for build warnings - # - name: Check for build errors - # run: | - # if [ -s "build/html/warnings.txt" ]; then - # echo "Build failed with errors." - # exit 1 - # fi diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d6a8fca..3e77369 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -4,12 +4,12 @@ Thanks for considering making a contribution to pyRTE-RRTMGP! ## How to Report a Bug? -Please file a bug report on the [Github page](https://github.com/earth-system-radiation/pyRTE-RRTMGP/issues/new/choose). +Please file a bug report on the [GitHub page](https://github.com/earth-system-radiation/pyRTE-RRTMGP/issues/new/choose). If possible, your issue should include a [minimal, complete, and verifiable example](https://stackoverflow.com/help/mcve) of the bug. ## How to Propose a New Feature? -Please file a feature request on the [Github page](https://github.com/earth-system-radiation/pyRTE-RRTMGP/issues/new/choose). +Please file a feature request on the [GitHub page](https://github.com/earth-system-radiation/pyRTE-RRTMGP/issues/new/choose). ## How to Set up a Local Development Environment? @@ -51,13 +51,13 @@ pytest tests ## How to Contribute a Patch That Fixes a Bug? Please fork this repository, branch from `main`, make your changes, and open a -Github [pull request](https://github.com/earth-system-radiation/pyRTE-RRTMTP/pulls) +GitHub [pull request](https://github.com/earth-system-radiation/pyRTE-RRTMTP/pulls) against the `main` branch. ## How to Contribute New Features? Please fork this repository, branch from `main`, make your changes, and open a -Github [pull request](https://github.com/earth-system-radiation/pyRTE-RRTMTP/pulls) +GitHub [pull request](https://github.com/earth-system-radiation/pyRTE-RRTMTP/pulls) against the `main` branch. Pull Requests for new features should include tests and documentation. diff --git a/README.md b/README.md index 0eb9112..1d1a736 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,7 @@ Currently, the following functions are available in the `pyrte_rrtmgp` package: ### Prerequisites -pyRTE-RRTMGP is currently only tested on x86_64 architecture with Linux (and, to some extend, macOS on Intel processors). +pyRTE-RRTMGP is currently only tested on x86_64 architecture with Linux (and, to some extent, macOS on Intel processors). The package source code is hosted [on GitHub](https://github.com/earth-system-radiation/pyRTE-RRTMGP) and uses git submodules to include the [RTE+RRTMGP Fortran software](https://earth-system-radiation.github.io/rte-rrtmgp/). The easiest way to install pyRTE-RRTMGP is to use `git`. You can install git from [here](https://git-scm.com/downloads). diff --git a/conda.recipe/meta.yaml b/conda.recipe/meta.yaml index 376d39b..ed50e38 100644 --- a/conda.recipe/meta.yaml +++ b/conda.recipe/meta.yaml @@ -38,6 +38,8 @@ requirements: run: - python - numpy >=1.21.0 + - xarray >=2023.5.0 + - netcdf4 >=1.5.7 test: imports: @@ -45,6 +47,8 @@ test: requires: - numpy>=1.21 - pytest>=7.4 + - xarray >=2023.5.0 + - netcdf4 >=1.5.7 source_files: - tests commands: diff --git a/docs/source/conf.py b/docs/source/conf.py index 31babeb..089cd01 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -4,6 +4,24 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html import datetime as dt +import os +import sys +import tomllib + +sys.path.insert(0, os.path.abspath("../../")) + + +def get_version_from_toml(): + if os.path.isfile("../../pyproject.toml"): + # read pyproject.toml + with open("../../pyproject.toml", "rb") as f: + pyproject = tomllib.load(f) + # get version from pyproject.toml + version = pyproject.get("project", {}).get("version", {}) + else: + version = None + return version if version else "dev" + # -- Project information ----------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information @@ -11,7 +29,8 @@ project = "pyRTE-RRTMGP" copyright = f"{dt.datetime.now().year}, Atmospheric and Environmental Research" author = "Atmospheric and Environmental Research" -release = "0.1" +version = get_version_from_toml() +release = version # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration @@ -19,14 +38,25 @@ extensions = [ "myst_parser", "sphinx_rtd_theme", + "sphinx.ext.autodoc", + "sphinx.ext.napoleon", ] templates_path = ["_templates"] exclude_patterns = [] +autodoc_mock_imports = [ + "pyrte_rrtmgp.pyrte_rrtmgp", + "numpy", + "xarray", +] # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output html_theme = "sphinx_rtd_theme" html_static_path = ["_static"] + +html_theme_options = { + "display_version": True, +} diff --git a/docs/source/contributor_guide/roadmap.md b/docs/source/contributor_guide/roadmap.md index d81cf0d..9c3236e 100644 --- a/docs/source/contributor_guide/roadmap.md +++ b/docs/source/contributor_guide/roadmap.md @@ -11,6 +11,6 @@ It fits into the broader scientific Python ecosystem and is designed to be intui - Documentation for all functions, multiple usage examples - Comprehensive test suite - Continuous integration and deployment -- Support for Windows, macOS ARM and other platforms +- Support for Windows, macOS ARM, and other platforms - Package installable via PyPI - Support for vertical (GPU) and horizontal (parallelization, e.g. with Dask) scaling diff --git a/docs/source/index.rst b/docs/source/index.rst index 9f6d20c..b90122e 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -9,9 +9,9 @@ Welcome to pyRTE-RRTMGP's documentation! pyRTE-RRTMGP provides a **Python interface** to the `RTE+RRTMGP `_ Fortran software package. -This package uses [pybind11](https://github.com/pybind/pybind11) to create a Python interface to a subset of the RTE+RRTMGP functions available in Fortran. +This package uses `pybind11 `_ to create a Python interface to a subset of the RTE+RRTMGP functions available in Fortran. -The RTE+RRTMGP package is a set of libraries for for computing radiative fluxes +The RTE+RRTMGP package is a set of libraries for computing radiative fluxes in planetary atmospheres. RTE+RRTMGP is described in a `paper `_ in `Journal of Advances in Modeling Earth Systems `_. @@ -26,6 +26,8 @@ in planetary atmospheres. RTE+RRTMGP is described in a user_guide/installation user_guide/usage + reference_guide/pyrte_rrtmgp_low_level + reference_guide/pyrte_rrtmgp_python_modules .. toctree:: :maxdepth: 2 diff --git a/docs/source/reference_guide/pyrte_rrtmgp_low_level.md b/docs/source/reference_guide/pyrte_rrtmgp_low_level.md new file mode 100644 index 0000000..4b50cc5 --- /dev/null +++ b/docs/source/reference_guide/pyrte_rrtmgp_low_level.md @@ -0,0 +1,20 @@ +(low_level_interface)= +# Low-Level Interface + +The `pyrte_rrtmgp.pyrte_rrtmgp` module provides a low-level interface to the RTE-RRTMGP and RRTMGP functions. It is generated by pybind11 from C headers of RTE-RRTMGP. + +See the [RTE-RRTMGP repository on GitHub](https://github.com/earth-system-radiation/pyRTE-RRTMGP) and the [RTE-RRTMGP documentation](https://earth-system-radiation.github.io/rte-rrtmgp/) for more information about these functions. + +The following RTE functions from [RTE-RRTMGP](https://github.com/earth-system-radiation/pyRTE-RRTMGP) are currently available in the `pyrte_rrtmgp.pyrte_rrtmgp` module: + +```{include} ../../../README.md +:start-after: +:end-before: +``` + +The following RRTMGP functions from [RTE-RRTMGP](https://github.com/earth-system-radiation/pyRTE-RRTMGP) are currently available in the `pyrte_rrtmgp.pyrte_rrtmgp` module: + +```{include} ../../../README.md +:start-after: +:end-before: +``` diff --git a/docs/source/reference_guide/pyrte_rrtmgp_python_modules.rst b/docs/source/reference_guide/pyrte_rrtmgp_python_modules.rst new file mode 100644 index 0000000..356c5eb --- /dev/null +++ b/docs/source/reference_guide/pyrte_rrtmgp_python_modules.rst @@ -0,0 +1,33 @@ +.. _module_ref: + +Python Modules +============== + +The documentation below provides details for several of the high-level functions wrapping the underlying Fortran code. + +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 +------------------------ + +.. automodule:: pyrte_rrtmgp.rte + :members: + :undoc-members: + :show-inheritance: + +pyrte\_rrtmgp.rrtmgp module +--------------------------- + +.. automodule:: pyrte_rrtmgp.rrtmgp + :members: + :undoc-members: + :show-inheritance: + +pyrte\_rrtmgp.utils module +-------------------------- + +.. automodule:: pyrte_rrtmgp.utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/user_guide/usage.md b/docs/source/user_guide/usage.md index 3855e27..c7c1714 100644 --- a/docs/source/user_guide/usage.md +++ b/docs/source/user_guide/usage.md @@ -2,17 +2,27 @@ This section provides a brief overview of how to use `pyrte_rrtmgp` with Python. -## Importing the Package +## Python Package Structure -To use any of the RTE-RRTMGP functions in Python, you must first import the `pyrte_rrtmgp` package: +The `pyrte_rrtmgp` package contains the following submodules: -```python -import pyrte_rrtmgp.pyrte_rrtmgp +- `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.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). ``` -This gives you access to [RTE-RRTMGP](https://github.com/earth-system-radiation/pyRTE-RRTMGP)'s Fortran functions directly in Python. +## Importing the Package + +To use any of the RTE-RRTMGP functions in Python, you first need to import `pyrte_rrtmgp` and the respective submodule you want to use. +For example: ``import pyrte_rrtmgp.pyrte_rrtmgp`` or ``import pyrte_rrtmgp.rrtmgp as rrtmgp``. -For example: +The example below uses the `pyrte_rrtmgp.pyrte_rrtmgp` submodule, which gives you low-level access to [RTE-RRTMGP](https://github.com/earth-system-radiation/pyRTE-RRTMGP)'s Fortran functions in Python: ```python import pyrte_rrtmgp.pyrte_rrtmgp as py @@ -21,25 +31,3 @@ args = list_of_arguments py.rte_lw_solver_noscat(*args) ``` - -## RTE Functions - -The following RTE functions from [RTE-RRTMGP](https://github.com/earth-system-radiation/pyRTE-RRTMGP) are currently available in the `pyrte_rrtmgp` package: - -```{include} ../../../README.md -:start-after: -:end-before: -``` - -## RRTMGP Functions - -The following RRTMGP functions from [RTE-RRTMGP](https://github.com/earth-system-radiation/pyRTE-RRTMGP) are currently available in the `pyrte_rrtmgp` package: - -```{include} ../../../README.md -:start-after: -:end-before: -``` - -```{seealso} -See the [RTE-RRTMGP repository on GitHub](https://github.com/earth-system-radiation/pyRTE-RRTMGP) and the [RTE-RRTMGP documentation](https://earth-system-radiation.github.io/rte-rrtmgp/) for more information about these functions. -``` diff --git a/examples/lw_noscat.ipynb b/examples/lw_noscat.ipynb new file mode 100644 index 0000000..3397db3 --- /dev/null +++ b/examples/lw_noscat.ipynb @@ -0,0 +1,80 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "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", + "\n", + "ERROR_TOLERANCE = 1e-4\n", + "\n", + "rte_rrtmgp_dir = \"../rrtmgp-data\"\n", + "clear_sky_example_files = f\"{rte_rrtmgp_dir}/examples/rfmip-clear-sky/inputs\"\n", + "\n", + "rfmip = xr.load_dataset(\n", + " f\"../tests/test_python_frontend/multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc\"\n", + ")\n", + "rfmip = rfmip.sel(expt=0) # only one experiment\n", + "\n", + "kdist = xr.load_dataset(f\"../tests/test_python_frontend/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", + "\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", + ")\n", + "\n", + "rlu = xr.load_dataset(\"../tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc\")\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", + "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", + "assert np.isclose(solver_flux_down, ref_flux_down, atol=ERROR_TOLERANCE).all()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/pyproject.toml b/pyproject.toml index 34ae06a..0f1828d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,9 +4,7 @@ version = "0.0.2" description = "A Python interface to the RTE+RRTMGP Fortran software package." readme = "README.md" requires-python = ">=3.7" -dependencies = [ - "numpy>=1.21.0" -] +dependencies = ["numpy>=1.21.0", "xarray>=2023.5.0", "netcdf4>=1.5.7"] classifiers = [ "Development Status :: 4 - Beta", "License :: OSI Approved :: MIT License", @@ -26,7 +24,7 @@ build-backend = "scikit_build_core.build" [project.optional-dependencies] -test = ["pytest", "numpy>=1.21.0"] +test = ["pytest", "numpy>=1.21.0", "xarray>=2023.5.0", "netcdf4>=1.5.7"] [tool.scikit-build] @@ -38,9 +36,7 @@ minversion = "6.0" addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"] xfail_strict = true log_cli_level = "INFO" -filterwarnings = [ - "error", -] +filterwarnings = ["error"] testpaths = ["tests"] diff --git a/pyrte_rrtmgp/__init__.py b/pyrte_rrtmgp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyrte_rrtmgp/gas_optics.py b/pyrte_rrtmgp/gas_optics.py new file mode 100644 index 0000000..f171e89 --- /dev/null +++ b/pyrte_rrtmgp/gas_optics.py @@ -0,0 +1,216 @@ +from pyrte_rrtmgp.rrtmgp import ( + compute_planck_source, + compute_tau_absorption, + interpolation, +) +from pyrte_rrtmgp.utils import ( + extract_gas_names, + flavors_from_kdist, + get_idx_minor, + gpoint_flavor_from_kdist, + rfmip_2_col_gas, +) + + +class GasOptics: + + def __init__(self, kdist, rfmip, gases_to_use=None): + self.kdist = kdist + self.rfmip = rfmip + + kdist_gas_names = extract_gas_names(kdist["gas_names"].values) + self.kdist_gas_names = kdist_gas_names + rfmip_vars = list(rfmip.keys()) + + gas_names = {n: n + "_GM" for n in kdist_gas_names if n + "_GM" in rfmip_vars} + + # Create a dict that maps the gas names in the kdist gas names to the gas names in the rfmip dataset + gas_names.update( + { + "co": "carbon_monoxide_GM", + "ch4": "methane_GM", + "o2": "oxygen_GM", + "n2o": "nitrous_oxide_GM", + "n2": "nitrogen_GM", + "co2": "carbon_dioxide_GM", + "ccl4": "carbon_tetrachloride_GM", + "cfc22": "hcfc22_GM", + "h2o": "water_vapor", + "o3": "ozone", + "no2": "no2", + } + ) + + # sort gas names based on kdist + gas_names = {g: gas_names[g] for g in kdist_gas_names if g in gas_names} + + if gases_to_use is not None: + gas_names = {g: gas_names[g] for g in gases_to_use} + + self.gas_names = gas_names + + self.tlay = rfmip["temp_layer"].values + self.play = rfmip["pres_layer"].values + self.col_gas = rfmip_2_col_gas(rfmip, list(gas_names.values()), dry_air=True) + + @property + def source_is_internal(self): + variables = self.kdist.data_vars + return "totplnk" in variables and "plank_fraction" in variables + + def gas_optics(self): + self.gas_optics_int() + self.compute_planck() + self.compute_gas_taus() + return self.tau, self.lay_src, self.lev_src, self.sfc_src, self.sfc_src_jac + + def gas_optics_int(self): + neta = len(self.kdist["mixing_fraction"]) + press_ref = self.kdist["press_ref"].values + temp_ref = self.kdist["temp_ref"].values + + press_ref_trop = self.kdist["press_ref_trop"].values.item() + + # dry air is zero + vmr_idx = [ + i for i, g in enumerate(self.kdist_gas_names, 1) if g in self.gas_names + ] + vmr_idx = [0] + vmr_idx + vmr_ref = self.kdist["vmr_ref"].sel(absorber_ext=vmr_idx).values.T + + # just the unique sets of gases + flavor = flavors_from_kdist(self.kdist) + + ( + self.jtemp, + self.fmajor, + self.fminor, + self.col_mix, + self.tropo, + self.jeta, + self.jpress, + ) = interpolation( + neta=neta, + flavor=flavor, + press_ref=press_ref, + temp_ref=temp_ref, + press_ref_trop=press_ref_trop, + vmr_ref=vmr_ref, + play=self.play, + tlay=self.tlay, + col_gas=self.col_gas, + ) + + def compute_planck(self): + + tlay = self.rfmip["temp_layer"].values + tlev = self.rfmip["temp_level"].values + tsfc = self.rfmip["surface_temperature"].values + pres_layers = self.rfmip["pres_layer"]["layer"] + top_at_1 = pres_layers[0] < pres_layers[-1] + band_lims_gpt = self.kdist["bnd_limits_gpt"].values.T + pfracin = self.kdist["plank_fraction"].values.transpose(0, 2, 1, 3) + temp_ref_min = self.kdist["temp_ref"].values.min() + temp_ref_max = self.kdist["temp_ref"].values.max() + totplnk = self.kdist["totplnk"].values.T + + gpoint_flavor = gpoint_flavor_from_kdist(self.kdist) + + self.sfc_src, self.lay_src, self.lev_src, self.sfc_src_jac = ( + compute_planck_source( + tlay, + tlev, + tsfc, + top_at_1, + self.fmajor, + self.jeta, + self.tropo, + self.jtemp, + self.jpress, + band_lims_gpt, + pfracin, + temp_ref_min, + temp_ref_max, + totplnk, + gpoint_flavor, + ) + ) + + def compute_gas_taus(self): + + idx_h2o = list(self.gas_names).index("h2o") + 1 + + gpoint_flavor = gpoint_flavor_from_kdist(self.kdist) + + kmajor = self.kdist["kmajor"].values + kminor_lower = self.kdist["kminor_lower"].values + kminor_upper = self.kdist["kminor_upper"].values + minor_limits_gpt_lower = self.kdist["minor_limits_gpt_lower"].values.T + minor_limits_gpt_upper = self.kdist["minor_limits_gpt_upper"].values.T + + minor_scales_with_density_lower = self.kdist[ + "minor_scales_with_density_lower" + ].values.astype(bool) + minor_scales_with_density_upper = self.kdist[ + "minor_scales_with_density_upper" + ].values.astype(bool) + scale_by_complement_lower = self.kdist[ + "scale_by_complement_lower" + ].values.astype(bool) + scale_by_complement_upper = self.kdist[ + "scale_by_complement_upper" + ].values.astype(bool) + + gas_name_list = list(self.gas_names.keys()) + + band_lims_gpt = self.kdist["bnd_limits_gpt"].values.T + + minor_gases_lower = extract_gas_names(self.kdist["minor_gases_lower"].values) + minor_gases_upper = extract_gas_names(self.kdist["minor_gases_upper"].values) + # check if the index is correct + idx_minor_lower = get_idx_minor(gas_name_list, minor_gases_lower) + idx_minor_upper = get_idx_minor(gas_name_list, minor_gases_upper) + + minor_scaling_gas_lower = extract_gas_names( + self.kdist["scaling_gas_lower"].values + ) + minor_scaling_gas_upper = extract_gas_names( + self.kdist["scaling_gas_upper"].values + ) + + idx_minor_scaling_lower = get_idx_minor(gas_name_list, minor_scaling_gas_lower) + idx_minor_scaling_upper = get_idx_minor(gas_name_list, minor_scaling_gas_upper) + + kminor_start_lower = self.kdist["kminor_start_lower"].values + kminor_start_upper = self.kdist["kminor_start_upper"].values + + self.tau = compute_tau_absorption( + idx_h2o, + gpoint_flavor, + band_lims_gpt, + kmajor, + kminor_lower, + kminor_upper, + minor_limits_gpt_lower, + minor_limits_gpt_upper, + minor_scales_with_density_lower, + minor_scales_with_density_upper, + scale_by_complement_lower, + scale_by_complement_upper, + idx_minor_lower, + idx_minor_upper, + idx_minor_scaling_lower, + idx_minor_scaling_upper, + kminor_start_lower, + kminor_start_upper, + self.tropo, + self.col_mix, + self.fmajor, + self.fminor, + self.play, + self.tlay, + self.col_gas, + self.jeta, + self.jtemp, + self.jpress, + ) diff --git a/pyrte_rrtmgp/rrtmgp.py b/pyrte_rrtmgp/rrtmgp.py new file mode 100644 index 0000000..117e641 --- /dev/null +++ b/pyrte_rrtmgp/rrtmgp.py @@ -0,0 +1,405 @@ +from typing import Tuple + +import numpy as np +import numpy.typing as npt + +from pyrte_rrtmgp.pyrte_rrtmgp import ( + rrtmgp_compute_Planck_source, + rrtmgp_compute_tau_absorption, + rrtmgp_compute_tau_rayleigh, + rrtmgp_interpolation, +) + + +def interpolation( + neta: int, + flavor: npt.NDArray, + press_ref: npt.NDArray, + temp_ref: npt.NDArray, + press_ref_trop: float, + vmr_ref: npt.NDArray, + play: npt.NDArray, + tlay: npt.NDArray, + col_gas: npt.NDArray, +) -> Tuple[ + npt.NDArray, + npt.NDArray, + npt.NDArray, + npt.NDArray, + npt.NDArray, + npt.NDArray, + npt.NDArray, +]: + """Interpolate the RRTMGP coefficients. + + Args: + neta (int): Number of mixing_fraction. + flavor (np.ndarray): Index into vmr_ref of major gases for each flavor. + press_ref (np.ndarray): Reference pressure grid. + temp_ref (np.ndarray): Reference temperature grid. + press_ref_trop (float): Reference pressure at the tropopause. + vmr_ref (np.ndarray): Reference volume mixing ratio. + play (np.ndarray): Pressure layers. + tlay (np.ndarray): Temperature layers. + col_gas (np.ndarray): Gas concentrations. + + Returns: + Tuple: A tuple containing the following arrays: + - jtemp (np.ndarray): Temperature interpolation index. + - fmajor (np.ndarray): Major gas interpolation fraction. + - fminor (np.ndarray): Minor gas interpolation fraction. + - col_mix (np.ndarray): Mixing fractions. + - tropo (np.ndarray): Use lower (or upper) atmosphere tables. + - jeta (np.ndarray): Index for binary species interpolation. + - jpress (np.ndarray): Pressure interpolation index. + """ + npres = press_ref.shape[0] + ntemp = temp_ref.shape[0] + ncol, nlay, ngas = col_gas.shape + ngas = ngas - 1 # Fortran uses index 0 here + nflav = flavor.shape[1] + + press_ref_log = np.log(press_ref) + press_ref_log_delta = (press_ref_log.min() - press_ref_log.max()) / ( + len(press_ref_log) - 1 + ) + press_ref_trop_log = np.log(press_ref_trop) + + temp_ref_min = temp_ref.min() + temp_ref_delta = (temp_ref.max() - temp_ref.min()) / (len(temp_ref) - 1) + + # outputs + jtemp = np.ndarray([nlay, ncol], dtype=np.int32) + fmajor = np.ndarray([nflav, nlay, ncol, 2, 2, 2], dtype=np.float64) + fminor = np.ndarray([nflav, nlay, ncol, 2, 2], dtype=np.float64) + col_mix = np.ndarray([nflav, nlay, ncol, 2], dtype=np.float64) + tropo = np.ndarray([nlay, ncol], dtype=np.int32) + jeta = np.ndarray([nflav, nlay, ncol, 2], dtype=np.int32) + jpress = np.ndarray([nlay, ncol], dtype=np.int32) + + args = [ + ncol, + nlay, + ngas, + nflav, + neta, + npres, + ntemp, + flavor.flatten("F"), + press_ref_log.flatten("F"), + temp_ref.flatten("F"), + press_ref_log_delta, + temp_ref_min, + temp_ref_delta, + press_ref_trop_log, + vmr_ref.flatten("F"), + play.flatten("F"), + tlay.flatten("F"), + col_gas.flatten("F"), + jtemp, + fmajor, + fminor, + col_mix, + tropo, + jeta, + jpress, + ] + + rrtmgp_interpolation(*args) + + tropo = tropo != 0 # Convert to boolean + return jtemp.T, fmajor.T, fminor.T, col_mix.T, tropo.T, jeta.T, jpress.T + + +def compute_planck_source( + tlay, + tlev, + tsfc, + top_at_1, + fmajor, + jeta, + tropo, + jtemp, + jpress, + band_lims_gpt, + pfracin, + temp_ref_min, + temp_ref_max, + totplnk, + gpoint_flavor, +): + """Compute the Planck source function for a radiative transfer calculation. + + Args: + tlay (numpy.ndarray): Temperature at layer centers (K), shape (ncol, nlay). + tlev (numpy.ndarray): Temperature at layer interfaces (K), shape (ncol, nlay+1). + tsfc (numpy.ndarray): Surface temperature, shape (ncol,). + top_at_1 (bool): Flag indicating if the top layer is at index 0. + sfc_lay (int): Index of the surface layer. + fmajor (numpy.ndarray): Interpolation weights for major gases, shape (2, 2, 2, ncol, nlay, nflav). + jeta (numpy.ndarray): Interpolation indexes in eta, shape (2, ncol, nlay, nflav). + tropo (numpy.ndarray): Use upper- or lower-atmospheric tables, shape (ncol, nlay). + jtemp (numpy.ndarray): Interpolation indexes in temperature, shape (ncol, nlay). + jpress (numpy.ndarray): Interpolation indexes in pressure, shape (ncol, nlay). + band_lims_gpt (numpy.ndarray): Start and end g-point for each band, shape (2, nbnd). + pfracin (numpy.ndarray): Fraction of the Planck function in each g-point, shape (ntemp, neta, npres+1, ngpt). + temp_ref_min (float): Minimum reference temperature for Planck function interpolation. + totplnk (numpy.ndarray): Total Planck function by band at each temperature, shape (nPlanckTemp, nbnd). + gpoint_flavor (numpy.ndarray): Major gas flavor (pair) by upper/lower, g-point, shape (2, ngpt). + + Returns: + sfc_src (numpy.ndarray): Planck emission from the surface, shape (ncol, ngpt). + lay_src (numpy.ndarray): Planck emission from layer centers, shape (ncol, nlay, ngpt). + lev_src (numpy.ndarray): Planck emission from layer boundaries, shape (ncol, nlay+1, ngpt). + sfc_source_Jac (numpy.ndarray): Jacobian (derivative) of the surface Planck source with respect to surface temperature, shape (ncol, ngpt). + """ + + _, ncol, nlay, nflav = jeta.shape + nPlanckTemp, nbnd = totplnk.shape + ntemp, neta, npres_e, ngpt = pfracin.shape + npres = npres_e - 1 + + sfc_lay = nlay if top_at_1 else 1 + + gpoint_bands = [] + + totplnk_delta = (temp_ref_max - temp_ref_min) / (nPlanckTemp - 1) + + # outputs + sfc_src = np.ndarray([ngpt, ncol], dtype=np.float64) + lay_src = np.ndarray([ngpt, nlay, ncol], dtype=np.float64) + lev_src = np.ndarray([ngpt, nlay + 1, ncol], dtype=np.float64) + sfc_src_jac = np.ndarray([ngpt, ncol], dtype=np.float64) + + args = [ + ncol, + nlay, + nbnd, + ngpt, + nflav, + neta, + npres, + ntemp, + nPlanckTemp, + tlay.flatten("F"), + tlev.flatten("F"), + tsfc.flatten("F"), + sfc_lay, + fmajor.flatten("F"), + jeta.flatten("F"), + tropo.flatten("F"), + jtemp.flatten("F"), + jpress.flatten("F"), + gpoint_bands, + band_lims_gpt.flatten("F"), + pfracin.flatten("F"), + temp_ref_min, + totplnk_delta, + totplnk.flatten("F"), + gpoint_flavor.flatten("F"), + sfc_src, + lay_src, + lev_src, + sfc_src_jac, + ] + + rrtmgp_compute_Planck_source(*args) + + return sfc_src.T, lay_src.T, lev_src.T, sfc_src_jac.T + + +def compute_tau_absorption( + idx_h2o, + gpoint_flavor, + band_lims_gpt, + kmajor, + kminor_lower, + kminor_upper, + minor_limits_gpt_lower, + minor_limits_gpt_upper, + minor_scales_with_density_lower, + minor_scales_with_density_upper, + scale_by_complement_lower, + scale_by_complement_upper, + idx_minor_lower, + idx_minor_upper, + idx_minor_scaling_lower, + idx_minor_scaling_upper, + kminor_start_lower, + kminor_start_upper, + tropo, + col_mix, + fmajor, + fminor, + play, + tlay, + col_gas, + jeta, + jtemp, + jpress, +): + """Compute the absorption optical depth for a set of atmospheric profiles. + + Args: + idx_h2o (int): Index of the water vapor gas species. + gpoint_flavor (np.ndarray): Spectral g-point flavor indices. + band_lims_gpt (np.ndarray): Spectral band limits in g-point space. + kmajor (np.ndarray): Major gas absorption coefficients. + kminor_lower (np.ndarray): Minor gas absorption coefficients for the lower atmosphere. + kminor_upper (np.ndarray): Minor gas absorption coefficients for the upper atmosphere. + minor_limits_gpt_lower (np.ndarray): Spectral g-point limits for minor contributors in the lower atmosphere. + minor_limits_gpt_upper (np.ndarray): Spectral g-point limits for minor contributors in the upper atmosphere. + minor_scales_with_density_lower (np.ndarray): Flags indicating if minor contributors in the lower atmosphere scale with density. + minor_scales_with_density_upper (np.ndarray): Flags indicating if minor contributors in the upper atmosphere scale with density. + scale_by_complement_lower (np.ndarray): Flags indicating if minor contributors in the lower atmosphere should be scaled by the complement. + scale_by_complement_upper (np.ndarray): Flags indicating if minor contributors in the upper atmosphere should be scaled by the complement. + idx_minor_lower (np.ndarray): Indices of minor contributors in the lower atmosphere. + idx_minor_upper (np.ndarray): Indices of minor contributors in the upper atmosphere. + idx_minor_scaling_lower (np.ndarray): Indices of minor contributors in the lower atmosphere that require scaling. + idx_minor_scaling_upper (np.ndarray): Indices of minor contributors in the upper atmosphere that require scaling. + kminor_start_lower (np.ndarray): Starting indices of minor absorption coefficients in the lower atmosphere. + kminor_start_upper (np.ndarray): Starting indices of minor absorption coefficients in the upper atmosphere. + tropo (np.ndarray): Flags indicating if a layer is in the troposphere. + col_mix (np.ndarray): Column-dependent gas mixing ratios. + fmajor (np.ndarray): Major gas absorption coefficient scaling factors. + fminor (np.ndarray): Minor gas absorption coefficient scaling factors. + play (np.ndarray): Pressure in each layer. + tlay (np.ndarray): Temperature in each layer. + col_gas (np.ndarray): Column-dependent gas concentrations. + jeta (np.ndarray): Indices of temperature/pressure levels. + jtemp (np.ndarray): Indices of temperature levels. + jpress (np.ndarray): Indices of pressure levels. + + Returns: + np.ndarray): tau Absorption optical depth. + """ + + ntemp, npres_e, neta, ngpt = kmajor.shape + npres = npres_e - 1 + nbnd = band_lims_gpt.shape[1] + _, ncol, nlay, nflav = jeta.shape + ngas = col_gas.shape[2] - 1 + nminorlower = minor_scales_with_density_lower.shape[0] + nminorupper = minor_scales_with_density_upper.shape[0] + nminorklower = kminor_lower.shape[2] + nminorkupper = kminor_upper.shape[2] + + # outputs + tau = np.zeros([ngpt, nlay, ncol], dtype=np.float64) + + args = [ + ncol, + nlay, + nbnd, + ngpt, + ngas, + nflav, + neta, + npres, + ntemp, + nminorlower, + nminorklower, + nminorupper, + nminorkupper, + idx_h2o, + gpoint_flavor.flatten("F"), # correct + band_lims_gpt.flatten("F"), + kmajor.transpose(0, 2, 1, 3).flatten("F"), + kminor_lower.flatten("F"), + kminor_upper.flatten("F"), + minor_limits_gpt_lower.flatten("F"), + minor_limits_gpt_upper.flatten("F"), + minor_scales_with_density_lower.flatten("F"), + minor_scales_with_density_upper.flatten("F"), + scale_by_complement_lower.flatten("F"), + scale_by_complement_upper.flatten("F"), + idx_minor_lower.flatten("F"), + idx_minor_upper.flatten("F"), + idx_minor_scaling_lower.flatten("F"), + idx_minor_scaling_upper.flatten("F"), + kminor_start_lower.flatten("F"), + kminor_start_upper.flatten("F"), + tropo.flatten("F"), + col_mix.flatten("F"), + fmajor.flatten("F"), + fminor.flatten("F"), + play.flatten("F"), + tlay.flatten("F"), + col_gas.flatten("F"), + jeta.flatten("F"), + jtemp.flatten("F"), + jpress.flatten("F"), + tau, + ] + + rrtmgp_compute_tau_absorption(*args) + + return tau.T + + +def compute_tau_rayleigh( + npres, + gpoint_flavor, + band_lims_gpt, + krayl, + idx_h2o, + col_dry, + col_gas, + fminor, + jeta, + tropo, + jtemp, +): + """Compute Rayleigh optical depth. + + Args: + npres (int): Number of pressure values. + gpoint_flavor (numpy.ndarray): Major gas flavor (pair) by upper/lower, g-point (shape: (2, ngpt)). + band_lims_gpt (numpy.ndarray): Start and end g-point for each band (shape: (2, nbnd)). + krayl (numpy.ndarray): Rayleigh scattering coefficients (shape: (ntemp, neta, ngpt, 2)). + idx_h2o (int): Index of water vapor in col_gas. + col_dry (numpy.ndarray): Column amount of dry air (shape: (ncol, nlay)). + col_gas (numpy.ndarray): Input column gas amount (molecules/cm^2) (shape: (ncol, nlay, 0:ngas)). + fminor (numpy.ndarray): Interpolation weights for major gases - computed in interpolation() (shape: (2, 2, ncol, nlay, nflav)). + jeta (numpy.ndarray): Interpolation indexes in eta - computed in interpolation() (shape: (2, ncol, nlay, nflav)). + tropo (numpy.ndarray): Use upper- or lower-atmospheric tables? (shape: (ncol, nlay)). + jtemp (numpy.ndarray): Interpolation indexes in temperature - computed in interpolation() (shape: (ncol, nlay)). + + Returns: + numpy.ndarray: Rayleigh optical depth (shape: (ncol, nlay, ngpt)). + """ + + ncol, nlay, ngas = col_gas.shape + ntemp, neta, ngpt, _ = krayl.shape + nflav = jeta.shape[3] + nbnd = band_lims_gpt.shape[1] + + # outputs + tau_rayleigh = np.ndarray((ncol, nlay, ngpt), dtype=np.float64) + + args = [ + ncol, + nlay, + nbnd, + ngpt, + ngas, + nflav, + neta, + npres, + ntemp, + gpoint_flavor, + band_lims_gpt, + krayl, + idx_h2o, + col_dry, + col_gas, + fminor, + jeta, + tropo, + jtemp, + tau_rayleigh, + ] + + rrtmgp_compute_tau_rayleigh(*args) + + return tau_rayleigh diff --git a/pyrte_rrtmgp/rte.py b/pyrte_rrtmgp/rte.py new file mode 100644 index 0000000..09c932f --- /dev/null +++ b/pyrte_rrtmgp/rte.py @@ -0,0 +1,164 @@ +from typing import Optional, Tuple + +import numpy as np +import numpy.typing as npt + +from pyrte_rrtmgp.pyrte_rrtmgp import rte_lw_solver_noscat, rte_sw_solver_noscat + +GAUSS_DS = np.array( + [ + [1.66, 0.0, 0.0, 0.0], # Diffusivity angle, not Gaussian angle + [1.18350343, 2.81649655, 0.0, 0.0], + [1.09719858, 1.69338507, 4.70941630, 0.0], + [1.06056257, 1.38282560, 2.40148179, 7.15513024], + ] +) + + +GAUSS_WTS = np.array( + [ + [0.5, 0.0, 0.0, 0.0], + [0.3180413817, 0.1819586183, 0.0, 0.0], + [0.2009319137, 0.2292411064, 0.0698269799, 0.0], + [0.1355069134, 0.2034645680, 0.1298475476, 0.0311809710], + ] +) + + +def lw_solver_noscat( + tau: npt.NDArray, + lay_source: npt.NDArray, + lev_source: npt.NDArray, + sfc_emis: npt.NDArray, + sfc_src: npt.NDArray, + top_at_1: bool = True, + nmus: int = 1, + inc_flux: Optional[npt.NDArray] = None, + ds: Optional[npt.NDArray] = None, + weights: Optional[npt.NDArray] = None, + do_broadband: Optional[bool] = True, + do_Jacobians: Optional[bool] = False, + sfc_src_jac: Optional[npt.NDArray] = [], + do_rescaling: Optional[bool] = False, + ssa: Optional[npt.NDArray] = None, + g: Optional[np.ndarray] = None, +) -> Tuple: + """ + Perform longwave radiation transfer calculations without scattering. + + Args: + top_at_1 (bool): Flag indicating whether the top of the atmosphere is at level 1. + nmus (int): Number of quadrature points. + tau (npt.NDArray): Array of optical depths. + lay_source (npt.NDArray): Array of layer sources. + lev_source (npt.NDArray): Array of level sources. + sfc_emis (npt.NDArray): Array of surface emissivities. + sfc_src (npt.NDArray): Array of surface sources. + inc_flux (npt.NDArray): Array of incoming fluxes. + ds (Optional[npt.NDArray], optional): Array of integration weights. Defaults to None. + weights (Optional[npt.NDArray], optional): Array of Gaussian quadrature weights. Defaults to None. + do_broadband (Optional[bool], optional): Flag indicating whether to compute broadband fluxes. Defaults to None. + do_Jacobians (Optional[bool], optional): Flag indicating whether to compute Jacobians. Defaults to None. + sfc_src_jac (Optional[npt.NDArray], optional): Array of surface source Jacobians. Defaults to None. + do_rescaling (Optional[bool], optional): Flag indicating whether to perform flux rescaling. Defaults to None. + ssa (Optional[npt.NDArray], optional): Array of single scattering albedos. Defaults to None. + g (Optional[np.ndarray], optional): Array of asymmetry parameters. Defaults to None. + + Returns: + Tuple: A tuple containing the following arrays: + - flux_up_jac (np.ndarray): Array of upward flux Jacobians. + - broadband_up (np.ndarray): Array of upward broadband fluxes. + - broadband_dn (np.ndarray): Array of downward broadband fluxes. + - flux_up (np.ndarray): Array of upward fluxes. + - flux_dn (np.ndarray): Array of downward fluxes. + """ + + ncol, nlay, ngpt = tau.shape + + # default values + n_quad_angs = nmus + + if inc_flux is None: + inc_flux = np.zeros(sfc_src.shape) + + if ds is None: + ds = np.empty((ncol, ngpt, n_quad_angs)) + for imu in range(n_quad_angs): + for igpt in range(ngpt): + for icol in range(ncol): + ds[icol, igpt, imu] = GAUSS_DS[imu, n_quad_angs - 1] + + if weights is None: + weights = GAUSS_WTS[0:n_quad_angs, n_quad_angs - 1] + + ssa = ssa or tau + g = g or tau + + # outputs + flux_up_jac = np.full([nlay + 1, ncol], np.nan, dtype=np.float64) + broadband_up = np.full([nlay + 1, ncol], np.nan, dtype=np.float64) + broadband_dn = np.full([nlay + 1, ncol], np.nan, dtype=np.float64) + flux_up = np.full([ngpt, nlay + 1, ncol], np.nan, dtype=np.float64) + flux_dn = np.full([ngpt, nlay + 1, ncol], np.nan, dtype=np.float64) + + args = [ + ncol, + nlay, + ngpt, + top_at_1, + nmus, + ds.flatten("F"), + weights.flatten("F"), + tau.flatten("F"), + lay_source.flatten("F"), + lev_source.flatten("F"), + sfc_emis.flatten("F"), + sfc_src.flatten("F"), + inc_flux.flatten("F"), + flux_up, + flux_dn, + do_broadband, + broadband_up, + broadband_dn, + do_Jacobians, + sfc_src_jac.flatten("F"), + flux_up_jac, + do_rescaling, + ssa.flatten("F"), + g.flatten("F"), + ] + + rte_lw_solver_noscat(*args) + + return flux_up_jac.T, broadband_up.T, broadband_dn.T, flux_up.T, flux_dn.T + + +def sw_solver_noscat( + top_at_1, + tau, + mu0, + inc_flux_dir, +): + """ + Computes the direct-beam flux for a shortwave radiative transfer problem without scattering. + + Args: + top_at_1 (bool): Logical flag indicating if the top layer is at index 1. + tau (numpy.ndarray): Absorption optical thickness of size (ncol, nlay, ngpt). + mu0 (numpy.ndarray): Cosine of solar zenith angle of size (ncol, nlay). + inc_flux_dir (numpy.ndarray): Direct beam incident flux of size (ncol, ngpt). + + Returns: + numpy.ndarray: Direct-beam flux of size (ncol, nlay+1, ngpt). + """ + + ncol, nlay, ngpt = tau.shape + + # outputs + flux_dir = np.ndarray((ncol, nlay + 1, ngpt), dtype=np.float64) + + args = [ncol, nlay, ngpt, top_at_1, tau, mu0, inc_flux_dir, flux_dir] + + rte_lw_solver_noscat(*args) + + return flux_dir diff --git a/pyrte_rrtmgp/utils.py b/pyrte_rrtmgp/utils.py new file mode 100644 index 0000000..7d054ba --- /dev/null +++ b/pyrte_rrtmgp/utils.py @@ -0,0 +1,195 @@ +from typing import List + +import numpy as np +import numpy.typing as npt +import xarray as xr + +# Constants +HELMERT1 = 9.80665 +HELMERT2 = 0.02586 +M_DRY = 0.028964 +M_H2O = 0.018016 +AVOGAD = 6.02214076e23 + + +def flavors_from_kdist(kdist: xr.Dataset) -> npt.NDArray: + """Get the unique flavors from the k-distribution file. + + Args: + kdist (xr.Dataset): K-distribution file. + + Returns: + np.ndarray: Unique flavors. + """ + key_species = kdist["key_species"].values + tot_flav = len(kdist["bnd"]) * len(kdist["atmos_layer"]) + npairs = len(kdist["pair"]) + all_flav = np.reshape(key_species, (tot_flav, npairs)) + # (0,0) becomes (2,2) because absorption coefficients for these g-points will be 0. + all_flav[np.all(all_flav == [0, 0], axis=1)] = [2, 2] + # we do that instead of unique to preserv the order + _, idx = np.unique(all_flav, axis=0, return_index=True) + return all_flav[np.sort(idx)].T + + +def get_col_dry(vmr_h2o, plev, latitude=None): + """Calculate the dry column of the atmosphere + + Args: + vmr_h2o (np.ndarray): Water vapor volume mixing ratio + plev (np.ndarray): Pressure levels + latitude (np.ndarray): Latitude of the location + + Returns: + np.ndarray: Dry column of the atmosphere + """ + ncol = plev.shape[0] + nlev = plev.shape[1] + col_dry = np.zeros((ncol, nlev - 1)) + + if latitude is not None: + g0 = HELMERT1 - HELMERT2 * np.cos(2.0 * np.pi * latitude / 180.0) + else: + g0 = np.full(ncol, HELMERT1) # Assuming grav is a constant value + + # TODO: use numpy instead of loops + for ilev in range(nlev - 1): + for icol in range(ncol): + delta_plev = abs(plev[icol, ilev] - plev[icol, ilev + 1]) + fact = 1.0 / (1.0 + vmr_h2o[icol, ilev]) + m_air = (M_DRY + M_H2O * vmr_h2o[icol, ilev]) * fact + col_dry[icol, ilev] = ( + 10.0 * delta_plev * AVOGAD * fact / (1000.0 * m_air * 100.0 * g0[icol]) + ) + return col_dry + + +def rfmip_2_col_gas(rfmip: xr.Dataset, gas_names: List[str], dry_air: bool = False): + """Convert RFMIP data to column gas concentrations. + + Args: + rfmip (xr.Dataset): RFMIP data. + gas_names (list): List of gas names. + dry_air (bool, optional): Include dry air. Defaults to False. + + Returns: + np.ndarray: Column gas concentrations. + """ + + ncol = len(rfmip["site"]) + nlay = len(rfmip["layer"]) + col_gas = [] + for gas_name in gas_names: + # if gas_name is not available, fill it with zeros + if gas_name not in rfmip.data_vars.keys(): + gas_values = np.zeros((ncol, nlay)) + else: + try: + scale = float(rfmip[gas_name].units) + except AttributeError: + scale = 1.0 + gas_values = rfmip[gas_name].values * scale + + if gas_values.ndim == 0: + gas_values = np.full((ncol, nlay), gas_values) + col_gas.append(gas_values) + + if dry_air: + if "h2o" not in gas_names and "water_vapor" not in gas_names: + raise ValueError( + "h2o gas must be included in gas_names to calculate dry air" + ) + if "h2o" in gas_names: + h2o_idx = gas_names.index("h2o") + else: + h2o_idx = gas_names.index("water_vapor") + vmr_h2o = col_gas[h2o_idx] + col_dry = get_col_dry(vmr_h2o, rfmip["pres_level"].values, latitude=None) + col_gas = [col_dry] + col_gas + + col_gas = np.stack(col_gas, axis=-1).astype(np.float64) + col_gas[:, :, 1:] = col_gas[:, :, 1:] * col_gas[:, :, :1] + + return col_gas + + +def gpoint_flavor_from_kdist(kdist: xr.Dataset) -> npt.NDArray: + """Get the g-point flavors from the k-distribution file. + + Each g-point is associated with a flavor, which is a pair of key species. + + Args: + kdist (xr.Dataset): K-distribution file. + + Returns: + np.ndarray: G-point flavors. + """ + key_species = kdist["key_species"].values + flavors = flavors_from_kdist(kdist) + + band_ranges = [ + [i] * (r.values[1] - r.values[0] + 1) + for i, r in enumerate(kdist["bnd_limits_gpt"], 1) + ] + gpoint_bands = np.concatenate(band_ranges) + + key_species_rep = key_species.copy() + key_species_rep[np.all(key_species_rep == [0, 0], axis=2)] = [2, 2] + + # unique flavors + flist = flavors.T.tolist() + + def key_species_pair2flavor(key_species_pair): + return flist.index(key_species_pair.tolist()) + 1 + + flavors_bands = np.apply_along_axis( + key_species_pair2flavor, 2, key_species_rep + ).tolist() + gpoint_flavor = np.array([flavors_bands[gp - 1] for gp in gpoint_bands]).T + + return gpoint_flavor + + +def extract_gas_names(gas_names): + """Extract gas names from the gas_names array, decoding and removing the suffix + + Args: + gas_names (np.ndarray): Gas names + + Returns: + list: List of gas names + """ + output = [] + for gas in gas_names: + output.append(gas.tobytes().decode().strip().split("_")[0]) + return output + + +def get_idx_minor(gas_names, minor_gases): + """Index of each minor gas in col_gas + + Args: + gas_names (list): Gas names + minor_gases (list): List of minor gases + + Returns: + list: Index of each minor gas in col_gas + """ + idx_minor_gas = [] + for gas in minor_gases: + try: + gas_idx = gas_names.index(gas) + 1 + except ValueError: + gas_idx = -1 + idx_minor_gas.append(gas_idx) + return np.array(idx_minor_gas, dtype=np.int32) + + +def calculate_mu0(solar_zenith_angle): + """Calculate the cosine of the solar zenith angle + + Args: + solar_zenith_angle (np.ndarray): Solar zenith angle in degrees + """ + usecol_values = solar_zenith_angle < 90.0 - 2.0 * np.spacing(90.0) + return np.where(usecol_values, np.cos(np.radians(solar_zenith_angle)), 1.0) diff --git a/tests/test_python_frontend/multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc b/tests/test_python_frontend/multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc new file mode 100644 index 0000000..5e7ac60 Binary files /dev/null and b/tests/test_python_frontend/multiple_input4MIPs_radiation_RFMIP_UColorado-RFMIP-1-2_none.nc differ diff --git a/tests/test_python_frontend/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc b/tests/test_python_frontend/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc new file mode 100644 index 0000000..5c2201b Binary files /dev/null and b/tests/test_python_frontend/rld_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc differ diff --git a/tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc b/tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc new file mode 100644 index 0000000..14ab6b0 Binary files /dev/null and b/tests/test_python_frontend/rlu_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc differ diff --git a/tests/test_python_frontend/rrtmgp-gas-lw-g256.nc b/tests/test_python_frontend/rrtmgp-gas-lw-g256.nc new file mode 100644 index 0000000..c66d3a5 Binary files /dev/null and b/tests/test_python_frontend/rrtmgp-gas-lw-g256.nc differ