diff --git a/docs/changelog.md b/docs/changelog.md index 267ae1d..3d1ed98 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -17,6 +17,7 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). * Deploy docs on [Read The Docs](https://joseki.readthedocs.io/). * Allow Numpy 2. * Support Python 3.13. +* Internalized ussa1976 dependency, which is no longer maintained. ### Fixed diff --git a/docs/requirements.txt b/docs/requirements.txt index 85ca808..a396d3f 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -56,12 +56,10 @@ charset-normalizer==3.4.2 click==8.1.8 ; python_full_version < '3.10' # via # joseki - # ussa1976 # uvicorn click==8.2.1 ; python_full_version >= '3.10' # via # joseki - # ussa1976 # uvicorn colorama==0.4.6 # via @@ -255,9 +253,7 @@ nbstripout==0.8.1 nest-asyncio==1.6.0 # via ipykernel netcdf4==1.7.2 - # via - # joseki - # ussa1976 + # via joseki notebook-shim==0.2.4 # via jupyterlab numpy==2.0.2 ; python_full_version < '3.10' @@ -270,7 +266,6 @@ numpy==2.0.2 ; python_full_version < '3.10' # pandas # scipy # seaborn - # ussa1976 # xarray numpy==2.2.6 ; python_full_version == '3.10.*' # via @@ -282,7 +277,6 @@ numpy==2.2.6 ; python_full_version == '3.10.*' # pandas # scipy # seaborn - # ussa1976 # xarray numpy==2.3.2 ; python_full_version >= '3.11' # via @@ -294,7 +288,6 @@ numpy==2.3.2 ; python_full_version >= '3.11' # pandas # scipy # seaborn - # ussa1976 # xarray overrides==7.7.0 # via jupyter-server @@ -415,17 +408,11 @@ rpds-py==0.26.0 # jsonschema # referencing scipy==1.13.1 ; python_full_version < '3.10' - # via - # joseki - # ussa1976 + # via joseki scipy==1.15.3 ; python_full_version == '3.10.*' - # via - # joseki - # ussa1976 + # via joseki scipy==1.16.0 ; python_full_version >= '3.11' - # via - # joseki - # ussa1976 + # via joseki seaborn==0.13.2 send2trash==1.8.3 # via jupyter-server @@ -546,8 +533,6 @@ uri-template==1.3.0 # via jsonschema urllib3==2.5.0 # via requests -ussa1976==0.3.4 - # via joseki uvicorn==0.35.0 # via sphinx-autobuild watchfiles==1.1.0 @@ -565,17 +550,11 @@ websocket-client==1.8.0 websockets==15.0.1 # via sphinx-autobuild xarray==2024.7.0 ; python_full_version < '3.10' - # via - # joseki - # ussa1976 + # via joseki xarray==2025.6.1 ; python_full_version == '3.10.*' - # via - # joseki - # ussa1976 + # via joseki xarray==2025.7.1 ; python_full_version >= '3.11' - # via - # joseki - # ussa1976 + # via joseki zipp==3.23.0 ; python_full_version < '3.10' # via # importlib-metadata diff --git a/pyproject.toml b/pyproject.toml index b421f18..37ba70c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,6 @@ dependencies = [ "pandas>=1.5.2", "scipy>=1.9.3", "xarray>=2022.12.0", - "ussa1976>=0.3.4", "attrs>=22.2.0", "importlib-resources>=5.10.2", ] diff --git a/src/joseki/constants.py b/src/joseki/constants.py index db58f27..bd33cc9 100644 --- a/src/joseki/constants.py +++ b/src/joseki/constants.py @@ -1,10 +1,22 @@ -from ussa1976.constants import F - from .units import ureg # Boltzmann constant (according to SciPy v1.10.1 Manual) K = 1.380649e-23 * ureg.joule / ureg.kelvin +# Sea level volume fractions below 86 km [dimensionless]. +F = { + "N2": 0.78084, + "O2": 0.209476, + "Ar": 0.00934, + "CO2": 0.000314, + "Ne": 0.00001818, + "He": 0.00000524, + "Kr": 0.00000114, + "Xe": 0.000000087, + "CH4": 0.000002, + "H2": 0.0000005, +} + # air main constituents molar fractions (values according to US Standard # Atmosphere 1976) AIR_MAIN_CONSTITUENTS_MOLAR_FRACTION = {m: F[m] for m in ["N2", "O2", "Ar"]} diff --git a/src/joseki/profiles/ussa_1976.py b/src/joseki/profiles/ussa_1976/__init__.py similarity index 90% rename from src/joseki/profiles/ussa_1976.py rename to src/joseki/profiles/ussa_1976/__init__.py index 3ff5ccb..5b06918 100644 --- a/src/joseki/profiles/ussa_1976.py +++ b/src/joseki/profiles/ussa_1976/__init__.py @@ -8,14 +8,14 @@ import typing as t import pint -import ussa1976 import xarray as xr from attrs import define -from .core import Profile -from .factory import factory -from .schema import history, schema -from ..units import to_quantity +from .core import compute as compute +from ..core import Profile +from ..factory import factory +from ..schema import history, schema +from ...units import to_quantity logger = logging.getLogger(__name__) @@ -56,11 +56,11 @@ def to_dataset( # compute profile if z is None: logging.debug("Computing profile with ussa1976 package") - ds = ussa1976.compute(variables=variables) + ds = compute(variables=variables) else: logging.debug("Computing profile with ussa1976 package") logging.debug("z=%s", z) - ds = ussa1976.compute(z=z.m_as("m"), variables=variables) + ds = compute(z=z.m_as("m"), variables=variables) # extract data coords = {"z": to_quantity(ds["z"]).to("km")} diff --git a/src/joseki/profiles/ussa_1976/constants.py b/src/joseki/profiles/ussa_1976/constants.py new file mode 100644 index 0000000..2ba3cf9 --- /dev/null +++ b/src/joseki/profiles/ussa_1976/constants.py @@ -0,0 +1,240 @@ +"""Constants module. + +As much as possible, constants' names are chosen to be as close as possible to +the notations used in :cite:`NASA1976USStandardAtmosphere`. + +Notes +----- +Constants' values are evaluated in the following set of units: +* length: meter +* time: second +* mass: kilogram +* temperature: kelvin +* quantity of matter: mole + +Note the following derived units: +* 1 Pa = 1 kg * m^-1 * s^-2 +* 1 Joule = 1 kg * m^2 * s^-2 +""" + +import numpy as np +from numpy.typing import NDArray + +K = 1.380622e-23 +"""Boltzmann constant [J * K^-1].""" + +M = { + "N2": 0.0280134, + "O2": 0.0319988, + "Ar": 0.039948, + "CO2": 0.04400995, + "Ne": 0.020183, + "He": 0.0040026, + "Kr": 0.08380, + "Xe": 0.13130, + "CH4": 0.01604303, + "H2": 0.00201594, + "O": 0.01599939, + "H": 0.00100797, +} +"""Molar masses of the individual species [kg * mole^-1].""" + +M0 = 0.028964425278793997 +"""Sea level mean air molar mass [kg * mole^-1].""" + +NA = 6.022169e23 +"""Avogadro number [mole^-1].""" + +R = 8.31432 +"""Universal gas constant [J * K^-1 * mole^-1].""" + +G0 = 9.80665 +"""Sea level gravity [m / s^-2].""" + +H: NDArray[np.float64] = np.array( + [ + 0.0, + 11e3, + 20e3, + 32e3, + 47e3, + 51e3, + 71e3, + 84852.05, + ] +) +"""Geopotential altitudes of the layers' boundaries (below 86 km) [m].""" + +LK: NDArray[np.float64] = np.array( + [ + -0.0065, + 0.0, + 0.0010, + 0.0028, + 0.0, + -0.0028, + -0.0020, + ] +) +"""Temperature gradients in the seven layers (below 86 km) [K * m^-1].""" + +P0 = 101325.0 +"""Pressure at sea level [Pa].""" + +R0 = 6.356766e6 +"""Effective Earth radius [m].""" + +T0 = 288.15 +"""Temperature at sea level [K].""" + +S = 110.4 +"""Sutherland constant in eq. 51 of :cite:`NASA1976USStandardAtmosphere` [K].""" + +BETA = 1.458e6 +""":math:`\\beta` constant in eq. 51 of :cite:`NASA1976USStandardAtmosphere` +[kg * m^-1 * s^-1 * K^-0.5].""" + +GAMMA = 1.40 +"""Ratio of specific heat of air at constant pressure to the specific heat of +air at constant volume [dimensionless].""" + +SIGMA = 3.65e-10 +"""Mean effective collision diameter [m].""" + +ALPHA = { + "N2": 0.0, + "O": 0.0, + "O2": 0.0, + "Ar": 0.0, + "He": -0.4, + "H": -0.25, +} +"""Thermal diffusion constants above 86 km [dimensionless].""" + +A = { + "O": 6.986e20, + "O2": 4.863e20, + "Ar": 4.487e20, + "He": 1.7e21, + "H": 3.305e21, +} +"""Thermal diffusion coefficients [m * s^-1].""" +B = { + "O": 0.75, + "O2": 0.75, + "Ar": 0.87, + "He": 0.691, + "H": 0.5, +} +"""Thermal diffusion constants [dimensionless].""" + +K_7 = 1.2e2 +"""Eddy diffusion coefficients [m^2 * s^-1].""" + +Q1 = { + "O": -5.809644e-13, + "O2": 1.366212e-13, + "Ar": 9.434079e-14, + "He": -2.457369e-13, +} +"""Vertical transport constants above 86 km [m^-3].""" + +Q2 = { + "O": -3.416248e-12, # warning: above 97 km, Q2 = 0 m^-3. + "O2": 0.0, + "Ar": 0.0, + "He": 0.0, +} +"""Vertical transport constants above 86 km [m^-3].""" + +U1 = { + "O": 56.90311e3, + "O2": 86e3, + "Ar": 86e3, + "He": 86e3, +} +"""Vertical transport constants above 86 km [m].""" + +U2 = { + "O": 97e3, +} +"""Vertical transport constants above 86 km [m].""" + +W1 = { + "O": 2.706240e-14, + "O2": 8.333333e-14, + "Ar": 8.333333e-14, + "He": 6.666667e-13, +} +"""Vertical transport constants above 86 km [m^-3].""" + +W2 = { + "O": 5.008765e-13, +} +"""Vertical transport constants above 86 km [m^-3].""" + +Z7 = 86e3 +"""Top altitude of the 7th layer [m].""" + +Z8 = 91e3 +"""Top altitude of the 8th layer [m].""" + +Z9 = 110e3 +"""Top altitude of the 9th layer [m].""" + +Z10 = 120e3 +"""Top altitude of the 10th layer [m].""" + +Z12 = 1000e3 +"""Top altitude of the 12nd layer [m].""" + +T7 = 186.8673 +"""Temperature at altitude :const:`.Z7` [K].""" + +T9 = 240.0 +"""Temperature at altitude :const:`.Z9` [K].""" + +T10 = 360.0 +"""Temperature at altitude :const:`.Z10` [K].""" + +T11 = 999.2356 +"""Temperature at altitude :const:`.Z11` [K].""" + +TINF = 1000.0 +"""Exospheric temperature [K].""" + +LAMBDA = 0.01875e-3 +""":math:`\\lambda` constant in eq. 32 of :cite:`NASA1976USStandardAtmosphere` +[m^-1].""" + +LK7 = 0.0 +"""Temperature gradient in the 8th layer [K * m^-1].""" + +LK9 = 12.0e-3 +"""Temperature gradient in the 10th layer [K * m^-1].""" + +N2_7 = 1.129794e20 +"""Molecular nitrogen number density at altitude :const:`.Z7` [m^-3].""" + +O_7 = 8.6e16 +"""Atomic oxygen number density at altitude :const:`.Z7` [m^-3].""" + +O2_7 = 3.030898e19 +"""Molecular oxygen number density at altitude :const:`.Z7` [m^-3].""" + +AR_7 = 1.351400e18 +"""Argon number density at altitude :const:`.Z7` [m^-3].""" + +HE_7 = 7.5817e14 +"""Helium number density at altitude :const:`.Z7` [m^-3]. + +Notes +----- +Assumes typo at page 13. +""" + +H_11 = 8.0e10 +"""Hydrogen number density at altitude :const:`.Z7` [m^-3].""" + +PHI = 7.2e11 +"""Vertical air particles flux [m^2 * s^-1].""" diff --git a/src/joseki/profiles/ussa_1976/core.py b/src/joseki/profiles/ussa_1976/core.py new file mode 100644 index 0000000..00f0d0f --- /dev/null +++ b/src/joseki/profiles/ussa_1976/core.py @@ -0,0 +1,1584 @@ +""" +U.S. Standard Atmosphere 1976 thermophysical model. + +The U.S. Standard Atmosphere 1976 model :cite:`NASA1976USStandardAtmosphere` +divides the atmosphere into two altitude regions: + +1. the low-altitude region, from 0 to 86 kilometers +2. the high-altitude region, from 86 to 1000 kilometers. + +A number of computational functions hereafter are specialized for one or +the other altitude region and is valid only in that altitude region, not in +the other. +Their name include a ``low_altitude`` or a ``high_altitude`` part to reflect +that they are valid only in the low altitude region and high altitude region, +respectively. +""" + +# IMPORTANT: This is a quick port of the external ussa1976 package, which is +# no longer maintained. The code was pasted here with minimal changes and has +# redundancies and inconsistencies with respect to the rest of the codebase. +# TODO: Refactor this module + +import datetime +from typing import Callable, List, Optional, Tuple, Union + +import numpy as np +import numpy.ma as ma +import xarray as xr +from numpy.typing import NDArray +from scipy.integrate import cumulative_trapezoid +from scipy.interpolate import interp1d + +from .constants import ( + ALPHA, + AR_7, + BETA, + G0, + GAMMA, + H_11, + HE_7, + K_7, + LAMBDA, + LK, + LK7, + LK9, + M0, + N2_7, + NA, + O2_7, + O_7, + P0, + PHI, + Q1, + Q2, + R0, + SIGMA, + T0, + T7, + T9, + T10, + T11, + TINF, + U1, + U2, + W1, + W2, + Z7, + Z8, + Z9, + Z10, + Z12, + A, + B, + H, + K, + M, + R, + S, +) +from ..._version import __version__ +from ...constants import F + +# List of all gas species +SPECIES = [ + "N2", + "O2", + "Ar", + "CO2", + "Ne", + "He", + "Kr", + "Xe", + "CH4", + "H2", + "O", + "H", +] + +# List of variables computed by the model +VARIABLES = [ + "t", + "p", + "n", + "n_tot", + "rho", + "mv", + "hp", + "v", + "mfp", + "f", + "cs", + "mu", + "nu", + "kt", +] + +# Variables standard names with respect to the Climate and Forecast (CF) +# convention +STANDARD_NAME = { + "t": "air_temperature", + "p": "air_pressure", + "n": "number_density", + "n_tot": "air_number_density", + "rho": "air_density", + "mv": "air_molar_volume", + "hp": "air_pressure_scale_height", + "v": "air_particles_mean_speed", + "mfp": "air_particles_mean_free_path", + "f": "air_particles_mean_collision_frequency", + "cs": "speed_of_sound_in_air", + "mu": "air_dynamic_viscosity", + "nu": "air_kinematic_viscosity", + "kt": "air_thermal_conductivity_coefficient", + "z": "altitude", + "h": "geopotential_height", +} + +# Variables long names +LONG_NAME = { + "t": "air temperature", + "p": "air pressure", + "n": "number_density", + "n_tot": "air number density", + "rho": "air density", + "mv": "air molar volume", + "hp": "air pressure scale height", + "v": "air particles mean speed", + "mfp": "air particles mean free path", + "f": "air particles mean collision frequency", + "cs": "speed of sound in air", + "mu": "air dynamic viscosity", + "nu": "air kinematic viscosity", + "kt": "air thermal conductivity coefficient", + "z": "altitude", + "h": "geopotential height", +} + +# Units of relevant quantities +UNITS = { + "t": "K", + "p": "Pa", + "n": "m^-3", + "n_tot": "m^-3", + "rho": "kg/m^3", + "mv": "m^3/mole", + "hp": "m", + "v": "m/s", + "mfp": "m", + "f": "s^-1", + "cs": "m/s", + "mu": "kg/(m*s)", + "nu": "m^2/s", + "kt": "W/(m*K)", + "z": "m", + "h": "m", + "s": "", +} + +# Variables dimensions +DIMS = { + "t": "z", + "p": "z", + "n": ("s", "z"), + "n_tot": "z", + "rho": "z", + "mv": "z", + "hp": "z", + "v": "z", + "mfp": "z", + "f": "z", + "cs": "z", + "mu": "z", + "nu": "z", + "kt": "z", +} + + +DEFAULT_Z = np.concatenate( + [ + np.arange(0.0, 11000.0, 50.0), + np.arange(11000.0, 32000.0, 100.0), + np.arange(32000.0, 50000.0, 200.0), + np.arange(50000.0, 100000.0, 500.0), + np.arange(100000.0, 300000.0, 1000.0), + np.arange(300000.0, 500000.0, 2000.0), + np.arange(500000.0, 1000001.0, 5000.0), + ] +) + + +def compute( + z: NDArray[np.float64] = DEFAULT_Z, + variables: Optional[List[str]] = None, +) -> xr.Dataset: + """Compute U.S. Standard Atmosphere 1976 data set on specified altitude grid. + + Parameters + ---------- + z: array + Altitude [m]. + + variables: list, optional + Names of the variables to compute. + + Returns + ------- + Dataset + Data set holding the values of the different atmospheric variables. + + Raises + ------ + ValueError + When altitude is out of bounds, or when variables are invalid. + """ + if np.any(z < 0.0): + raise ValueError("altitude values must be greater than or equal to zero") + + if np.any(z > 1000e3): + raise ValueError("altitude values must be less then or equal to 1000 km") + + if variables is None: + variables = VARIABLES + else: + for var in variables: + if var not in VARIABLES: + raise ValueError(var, " is not a valid variable name") + + # initialise data set + ds = init_data_set(z=z) + + z_delim = 86e3 + + # compute the model in the low-altitude region + low_altitude_region = ds.z.values <= z_delim + compute_low_altitude(data_set=ds, mask=low_altitude_region, inplace=True) + + # compute the model in the high-altitude region + high_altitude_region = ds.z.values > z_delim + compute_high_altitude(data_set=ds, mask=high_altitude_region, inplace=True) + + # replace all np.nan with 0.0 in number densities values + n = ds.n.values + n[np.isnan(n)] = 0.0 + ds.n.values = n + + # list names of variables to drop from the data set + names = [] + for var in ds.data_vars: # type: ignore + if var not in variables: + names.append(var) + + return ds.drop_vars(names) # type: ignore + + +def compute_low_altitude( + data_set: xr.Dataset, + mask: Optional[NDArray[bool]] = None, # type: ignore + inplace: bool = False, +) -> Optional[xr.Dataset]: + """Compute U.S. Standard Atmosphere 1976 in low-altitude region. + + Parameters + ---------- + data_set: Dataset + Data set to compute. + + mask: DataArray, optional + Mask to select the region of the data set to compute. + By default, the mask selects the entire data set. + + inplace: bool, default=False + If ``True``, modifies ``data_set`` in place, else returns a copy of + ``data_set``. + + Returns + ------- + Dataset + If ``inplace`` is ``True``, returns nothing, else returns a copy of + ``data_set``. + """ + if mask is None: + mask = np.full_like(data_set.coords["z"].values, True, dtype=bool) + + if inplace: + ds = data_set + else: + ds = data_set.copy(deep=True) + + z = ds.z[mask].values + + # compute levels temperature and pressure values + tb, pb = compute_levels_temperature_and_pressure_low_altitude() + + # compute geopotential height, temperature and pressure + h = to_geopotential_height(z) + t = compute_temperature_low_altitude(h=h, tb=tb) + p = compute_pressure_low_altitude(h=h, pb=pb, tb=tb) + + # compute the auxiliary atmospheric variables + n_tot = NA * p / (R * t) + rho = p * M0 / (R * t) + g = compute_gravity(z) + mu = BETA * np.power(t, 1.5) / (t + S) + + # assign data set with computed values + ds["t"].loc[dict(z=z)] = t + ds["p"].loc[dict(z=z)] = p + ds["n_tot"].loc[dict(z=z)] = n_tot + + species = ["N2", "O2", "Ar", "CO2", "Ne", "He", "Kr", "Xe", "CH4", "H2"] + for i, s in enumerate(SPECIES): + if s in species: + ds["n"][i].loc[dict(z=z)] = F[s] * n_tot + + ds["rho"].loc[dict(z=z)] = rho + ds["mv"].loc[dict(z=z)] = NA / n_tot + ds["hp"].loc[dict(z=z)] = R * t / (g * M0) + ds["v"].loc[dict(z=z)] = np.sqrt(8.0 * R * t / (np.pi * M0)) + ds["mfp"].loc[dict(z=z)] = np.sqrt(2.0) / ( + 2.0 * np.pi * np.power(SIGMA, 2.0) * n_tot + ) + ds["f"].loc[dict(z=z)] = ( + 4.0 + * NA + * np.power(SIGMA, 2.0) + * np.sqrt(np.pi * np.power(p, 2.0) / (R * M0 * t)) + ) + ds["cs"].loc[dict(z=z)] = np.sqrt(GAMMA * R * t / M0) + ds["mu"].loc[dict(z=z)] = mu + ds["nu"].loc[dict(z=z)] = mu / rho + ds["kt"].loc[dict(z=z)] = ( + 2.64638e-3 * np.power(t, 1.5) / (t + 245.4 * np.power(10.0, -12.0 / t)) + ) + + if not inplace: + return ds + else: + return None + + +def compute_high_altitude( + data_set: xr.Dataset, + mask: Optional[NDArray[bool]] = None, + inplace: bool = False, +) -> Optional[xr.Dataset]: + """Compute U.S. Standard Atmosphere 1976 in high-altitude region. + + Parameters + ---------- + data_set: Dataset + Data set to compute. + + mask: DataArray, optional + Mask to select the region of the data set to compute. + By default, the mask selects the entire data set. + + inplace: bool, default False + If ``True``, modifies ``data_set`` in place, else returns a copy of + ``data_set``. + + Returns + ------- + Dataset + If ``inplace`` is True, returns nothing, else returns a copy of + ``data_set``. + """ + if mask is None: + mask = np.full_like(data_set.coords["z"].values, True, dtype=bool) + + if inplace: + ds = data_set + else: + ds = data_set.copy(deep=True) + + z = ds.coords["z"][mask].values + if len(z) == 0: + return ds + + n = compute_number_densities_high_altitude(z) + species = ["N2", "O", "O2", "Ar", "He", "H"] + ni: NDArray[np.float64] = np.array([n.sel(s=s).values for s in species]) + n_tot = np.sum(ni, axis=0) + fi = ni / n_tot[np.newaxis, :] + mi: NDArray[np.float64] = np.array([M[s] for s in species]) + m = np.sum(fi * mi[:, np.newaxis], axis=0) + t = compute_temperature_high_altitude(z) + p = K * n_tot * t + rho = np.sum(ni * mi[:, np.newaxis], axis=0) / NA + g = compute_gravity(z) + + # assign data set with computed values + ds["t"].loc[dict(z=z)] = t + ds["p"].loc[dict(z=z)] = p + ds["n_tot"].loc[dict(z=z)] = n_tot + + for i, s in enumerate(SPECIES): + if s in species: + ds["n"][i].loc[dict(z=z)] = n.sel(s=s).values + + ds["rho"].loc[dict(z=z)] = rho + ds["mv"].loc[dict(z=z)] = NA / n_tot + ds["hp"].loc[dict(z=z)] = R * t / (g * m) + ds["v"].loc[dict(z=z)] = np.sqrt(8.0 * R * t / (np.pi * m)) + ds["mfp"].loc[dict(z=z)] = np.sqrt(2.0) / ( + 2.0 * np.pi * np.power(SIGMA, 2.0) * n_tot + ) + ds["f"].loc[dict(z=z)] = ( + 4.0 + * NA + * np.power(SIGMA, 2.0) + * np.sqrt(np.pi * np.power(p, 2.0) / (R * m * t)) + ) + + if not inplace: + return ds + else: + return None + + +def init_data_set(z: NDArray[np.float64]) -> xr.Dataset: # type: ignore + """Initialise data set. + + Parameters + ---------- + z: array + Altitudes [m]. + + Returns + ------- + Dataset + Initialised data set. + """ + data_vars = {} + for var in VARIABLES: + if var != "n": + data_vars[var] = ( + DIMS[var], + np.full(z.shape, np.nan), + { + "units": UNITS[var], + "long_name": LONG_NAME[var], + "standard_name": STANDARD_NAME[var], + }, + ) + else: + data_vars[var] = ( + DIMS[var], + np.full((len(SPECIES), len(z)), np.nan), + { + "units": UNITS[var], + "long_name": LONG_NAME[var], + "standard_name": STANDARD_NAME["n"], + }, + ) + + coords = { + "z": ( + "z", + z, + { + "standard_name": "altitude", + "long_name": "altitude", + "units": "m", + }, + ), + "s": ("s", SPECIES, {"long_name": "species", "standard_name": "species"}), + } + + attrs = { + "convention": "CF-1.9", + "title": "U.S. Standard Atmosphere 1976", + "history": ( + f"{datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')}" + f" - data set creation - ussa1976, version {__version__}" + ), + "source": f"ussa1976, version {__version__}", + "references": ( + "U.S. Standard Atmosphere, 1976, NASA-TM-X-74335NOAA-S/T-76-1562" + ), + } + + return xr.Dataset(data_vars, coords, attrs) # type: ignore + + +def compute_levels_temperature_and_pressure_low_altitude() -> Tuple[ + NDArray[np.float64], NDArray[np.float64] +]: + """Compute temperature and pressure at low-altitude region' levels. + + Returns + ------- + tuple of arrays: + Levels temperatures [K] and pressures [Pa]. + """ + tb = [T0] + pb = [P0] + for i in range(1, len(H)): + t_next = tb[i - 1] + LK[i - 1] * (H[i] - H[i - 1]) + tb.append(t_next) + if LK[i - 1] == 0: + p_next = compute_pressure_low_altitude_zero_gradient( + h=H[i], + hb=H[i - 1], + pb=pb[i - 1], + tb=tb[i - 1], + ) + else: + p_next = compute_pressure_low_altitude_non_zero_gradient( + h=H[i], + hb=H[i - 1], + pb=pb[i - 1], + tb=tb[i - 1], + lkb=LK[i - 1], + ) + pb.append(float(p_next)) + return np.array(tb, dtype=np.float64), np.array(pb, dtype=np.float64) + + +def compute_number_densities_high_altitude( + altitudes: NDArray[np.float64], +) -> xr.DataArray: + """Compute number density of individual species in high-altitude region. + + Parameters + ---------- + altitudes: array + Altitudes [m]. + + Returns + ------- + DataArray + Number densities of the individual species and total number density at + the given altitudes. + + Notes + ----- + A uniform altitude grid is generated and used for the computation of the + integral as well as for the computation of the number densities of the + individual species. This gridded data is then interpolated at the query + ``altitudes`` using a linear interpolation scheme in logarithmic space. + """ + # altitude grid + grid: NDArray[np.float64] = np.concatenate( + ( + np.linspace(start=Z7, stop=150e3, num=640, endpoint=False), + np.geomspace(start=150e3, stop=Z12, num=100, endpoint=True), + ) + ) + + # pre-computed variables + m = compute_mean_molar_mass_high_altitude(z=grid) + g = compute_gravity(z=grid) + t = compute_temperature_high_altitude(z=grid) + dt_dz = compute_temperature_gradient_high_altitude(z=grid) + below_115 = grid < 115e3 + k = eddy_diffusion_coefficient(grid[below_115]) + + n_grid = {} + + # ************************************************************************* + # Molecular nitrogen + # ************************************************************************* + + y = m * g / (R * t) # m^-1 + n_grid["N2"] = N2_7 * (T7 / t) * np.exp(-cumulative_trapezoid(y, grid, initial=0.0)) + + # ************************************************************************* + # Atomic oxygen + # ************************************************************************* + + d = thermal_diffusion_coefficient( + nb=n_grid["N2"][below_115], + t=t[below_115], + a=A["O"], + b=B["O"], + ) + y = thermal_diffusion_term_atomic_oxygen( + grid, + g, + t, + dt_dz, + d, + k, + ) + velocity_term_atomic_oxygen(grid) + n_grid["O"] = O_7 * (T7 / t) * np.exp(-cumulative_trapezoid(y, grid, initial=0.0)) + + # ************************************************************************* + # Molecular oxygen + # ************************************************************************* + + d = thermal_diffusion_coefficient( + nb=n_grid["N2"][below_115], + t=t[below_115], + a=A["O2"], + b=B["O2"], + ) + y = thermal_diffusion_term( + s="O2", + z_grid=grid, + g=g, + t=t, + dt_dz=dt_dz, + m=m, + d=d, + k=k, + ) + velocity_term("O2", grid) + n_grid["O2"] = O2_7 * (T7 / t) * np.exp(-cumulative_trapezoid(y, grid, initial=0.0)) + + # ************************************************************************* + # Argon + # ************************************************************************* + + background = ( + n_grid["N2"][below_115] + n_grid["O"][below_115] + n_grid["O2"][below_115] + ) + d = thermal_diffusion_coefficient( + nb=background, + t=t[below_115], + a=A["Ar"], + b=B["Ar"], + ) + y = thermal_diffusion_term( + s="Ar", + z_grid=grid, + g=g, + t=t, + dt_dz=dt_dz, + m=m, + d=d, + k=k, + ) + velocity_term("Ar", grid) + n_grid["Ar"] = AR_7 * (T7 / t) * np.exp(-cumulative_trapezoid(y, grid, initial=0.0)) + + # ************************************************************************* + # Helium + # ************************************************************************* + + background = ( + n_grid["N2"][below_115] + n_grid["O"][below_115] + n_grid["O2"][below_115] + ) + d = thermal_diffusion_coefficient( + nb=background, + t=t[below_115], + a=A["He"], + b=B["He"], + ) + y = thermal_diffusion_term( + s="He", + z_grid=grid, + g=g, + t=t, + dt_dz=dt_dz, + m=m, + d=d, + k=k, + ) + velocity_term("He", grid) + n_grid["He"] = HE_7 * (T7 / t) * np.exp(-cumulative_trapezoid(y, grid, initial=0.0)) + + # ************************************************************************* + # Hydrogen + # ************************************************************************* + + # below 500 km + mask = (grid >= 150e3) & (grid <= 500e3) + background = ( + n_grid["N2"][mask] + + n_grid["O"][mask] + + n_grid["O2"][mask] + + n_grid["Ar"][mask] + + n_grid["He"][mask] + ) + d = thermal_diffusion_coefficient( + background, + t[mask], + A["H"], + B["H"], + ) + alpha = ALPHA["H"] + _tau = tau_function(z_grid=grid[mask], below_500=True) + y = (PHI / d) * np.power(t[mask] / T11, 1 + alpha) * np.exp(_tau) # m^-4 + integral_values = cumulative_trapezoid( + y[::-1], grid[mask][::-1], initial=0.0 + ) # m^-3 + integral_values = integral_values[::-1] + n_below_500 = ( + (H_11 - integral_values) * np.power(T11 / t[mask], 1 + alpha) * np.exp(-_tau) + ) + + # above 500 km + _tau = tau_function( + z_grid=grid[grid > 500e3], + below_500=False, + ) + n_above_500 = H_11 * np.power(T11 / t[grid > 500e3], 1 + alpha) * np.exp(-_tau) + + n_grid["H"] = np.concatenate((n_below_500, n_above_500)) + + n = { + s: log_interp1d(grid, n_grid[s])(altitudes) + for s in ["N2", "O", "O2", "Ar", "He"] + } + + # Below 150 km, the number density of atomic hydrogen is zero. + n_h_below_150 = np.zeros(len(altitudes[altitudes < 150e3])) + n_h_above_150 = log_interp1d(grid[grid >= 150e3], n_grid["H"])( + altitudes[altitudes >= 150e3] + ) + n["H"] = np.concatenate((n_h_below_150, n_h_above_150)) + + n_concat: NDArray[np.float64] = np.array([n[species] for species in n]) + + return xr.DataArray( + n_concat, + dims=["s", "z"], + coords={ + "s": ("s", [s for s in n]), + "z": ("z", altitudes, dict(units="m")), + }, + attrs=dict(units="m^-3"), + ) + + +def compute_mean_molar_mass_high_altitude( + z: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute mean molar mass in high-altitude region. + + Parameters + ---------- + z: array + Altitude [m]. + + Returns + ------- + array + Mean molar mass [kg/mole]. + """ + return np.where(z <= 100e3, M0, M["N2"]) + + +def compute_temperature_high_altitude( + z: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute temperature in high-altitude region. + + Parameters + ---------- + z: array + Altitude [m]. + + Returns + ------- + array + Temperature [K]. + """ + a = -76.3232 # K + b = -19942.9 # m + tc = 263.1905 # K + + def t(z: float) -> float: + """Compute temperature at given altitude. + + Parameters + ---------- + z: float + Altitude [m]. + + Returns + ------- + float + Temperature [K]. + + Raises + ------ + ValueError + If the altitude is out of range. + """ + if Z7 <= z <= Z8: + return T7 + elif Z8 < z <= Z9: + return tc + a * float(np.sqrt(1.0 - np.power((z - Z8) / b, 2.0))) + elif Z9 < z <= Z10: + return T9 + LK9 * (z - Z9) + elif Z10 < z <= Z12: + return TINF - (TINF - T10) * float( + np.exp(-LAMBDA * (z - Z10) * (R0 + Z10) / (R0 + z)) + ) + else: + raise ValueError(f"altitude value '{z}' is out of range") + + temperature = np.vectorize(t)(z) + return np.array(temperature, dtype=np.float64) + + +def compute_temperature_gradient_high_altitude( + z: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute temperature gradient in high-altitude region. + + Parameters + ---------- + z: array + Altitude [m]. + + Returns + ------- + array + Temperature gradient [K/m]. + """ + a = -76.3232 # [dimensionless] + b = -19942.9 # m + + def gradient(z_value: float) -> float: + """Compute temperature gradient at given altitude. + + Parameters + ---------- + z_value: float + Altitude [m]. + + Raises + ------ + ValueError + When altitude is out of bounds. + + Returns + ------- + float + Temperature gradient [K / m]. + """ + if Z7 <= z_value <= Z8: + return LK7 + elif Z8 < z_value <= Z9: + return ( + -a + / b + * ((z_value - Z8) / b) + / float(np.sqrt(1 - np.square((z_value - Z8) / b))) + ) + elif Z9 < z_value <= Z10: + return LK9 + elif Z10 < z_value <= Z12: + zeta = (z_value - Z10) * (R0 + Z10) / (R0 + z_value) + return ( + LAMBDA + * (TINF - T10) + * float(np.square((R0 + Z10) / (R0 + z_value))) + * float(np.exp(-LAMBDA * zeta)) + ) + + else: + raise ValueError( + f"altitude z ({z_value}) out of range, should be in [{Z7}, {Z12}] m" + ) + + z_values: NDArray[np.float64] = np.array(z, dtype=float) + dt_dz = np.vectorize(gradient)(z_values) + return np.array(dt_dz, dtype=np.float64) + + +def thermal_diffusion_coefficient( + nb: NDArray[np.float64], + t: NDArray[np.float64], + a: float, + b: float, +) -> NDArray[np.float64]: + r"""Compute thermal diffusion coefficient values in high-altitude region. + + Parameters + ---------- + nb: array + Background number density [m^-3]. + + t: array + Temperature [K]. + + a: float + Thermal diffusion constant :math:`a` [m^-1 * s^-1]. + + b: float + Thermal diffusion constant :math:`b` [dimensionless]. + + Returns + ------- + array + Thermal diffusion coefficient [m^2 * s^-1]. + """ + k = (a / nb) * np.power(t / 273.15, b) + return np.array(k, dtype=np.float64) + + +def eddy_diffusion_coefficient(z: NDArray[np.float64]) -> NDArray[np.float64]: + r"""Compute Eddy diffusion coefficient in high-altitude region. + + Parameters + ---------- + z: array + Altitude [m]. + + Returns + ------- + array + Eddy diffusion coefficient [m^2 * s^-1]. + + Notes + ----- + Valid in the altitude region :math:`86 \leq z \leq 150` km. + """ + return np.where( + z < 95e3, + K_7, + K_7 * np.exp(1.0 - (4e8 / (4e8 - np.square(z - 95e3)))), + ) + + +def f_below_115_km( + g: NDArray[np.float64], + t: NDArray[np.float64], + dt_dz: NDArray[np.float64], + m: Union[float, NDArray[np.float64]], + mi: float, + alpha: float, + d: NDArray[np.float64], + k: NDArray[np.float64], +) -> NDArray[np.float64]: + r"""Evaluate function :math:`f` below 115 km altitude. + + Evaluates the function :math:`f` defined by equation (36) in + :cite:`NASA1976USStandardAtmosphere` in the altitude region :math:`86` km + :math:`\leq z \leq 115` km. + + Parameters + ---------- + g: array + Gravity values at the different altitudes [m * s^-2]. + + t: array + Temperature values at the different altitudes [K]. + + dt_dz: array + Temperature gradient values at the different altitudes [K * m^-1]. + + m: array + Molar mass [kg * mole^-1]. + + mi: float + Species molar masses [kg * mole^-1]. + + alpha: float + Alpha thermal diffusion constant [dimensionless]. + + d: array + Thermal diffusion coefficient values at the different altitudes + [m^2 * s^-1]. + + k: array + Eddy diffusion coefficient values at the different altitudes + [m^2 * s^-1]. + + Returns + ------- + array + Function :math:`f` at the different altitudes. + """ + term_1 = g * d / ((d + k) * (R * t)) + term_2 = mi + (m * k) / d + (alpha * R * dt_dz) / g + return term_1 * term_2 + + +def f_above_115_km( + g: NDArray[np.float64], + t: NDArray[np.float64], + dt_dz: NDArray[np.float64], + mi: float, + alpha: float, +) -> NDArray[np.float64]: + r"""Evaluate function :math:`f` above 115 km altitude. + + Evaluate the function :math:`f` defined by equation (36) in + :cite:`NASA1976USStandardAtmosphere` in the altitude region :math:`115` + :math:`\lt z \leq 1000` km. + + Parameters + ---------- + g: array + Gravity at the different altitudes [m * s^-2]. + + t: array + Temperature at the different altitudes [K]. + + dt_dz: array + Temperature gradient at the different altitudes [K * m^-1]. + + mi: float + Species molar masses [kg * mole^-1]. + + alpha: float + Alpha thermal diffusion constant [dimensionless]. + + Returns + ------- + array + Function :math:`f` at the different altitudes. + """ + return (g / (R * t)) * (mi + ((alpha * R) / g) * dt_dz) # type: ignore + + +def thermal_diffusion_term( + s: str, + z_grid: NDArray[np.float64], + g: NDArray[np.float64], + t: NDArray[np.float64], + dt_dz: NDArray[np.float64], + m: NDArray[np.float64], + d: NDArray[np.float64], + k: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute thermal diffusion term of given species in high-altitude region. + + Parameters + ---------- + s: str + Species. + + z_grid: array + Altitude grid [m]. + + g: array + Gravity values on the altitude grid [m * s^-2]. + + t: array + Temperature values on the altitude grid [K]. + + dt_dz: array + Temperature gradient values on the altitude grid [K * m^-1]. + + m: array + Values of the mean molar mass on the altitude grid [kg * mole^-1]. + + d: array + Molecular diffusion coefficient values on the altitude grid, + for altitudes strictly less than 115 km [m^2 * s^-1]. + + k: array + Eddy diffusion coefficient values on the altitude grid, for + altitudes strictly less than 115 km [m^2 * s^-1]. + + Returns + ------- + array + Thermal diffusion term [m^-1]. + """ + below_115_km = z_grid < 115e3 + fo1 = f_below_115_km( + g[below_115_km], + t[below_115_km], + dt_dz[below_115_km], + m[below_115_km], + M[s], + ALPHA[s], + d, + k, + ) + above_115_km = z_grid >= 115e3 + fo2 = f_above_115_km( + g[above_115_km], + t[above_115_km], + dt_dz[above_115_km], + M[s], + ALPHA[s], + ) + return np.concatenate((fo1, fo2)) + + +def thermal_diffusion_term_atomic_oxygen( + z_grid: NDArray[np.float64], + g: NDArray[np.float64], + t: NDArray[np.float64], + dt_dz: NDArray[np.float64], + d: NDArray[np.float64], + k: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute oxygen thermal diffusion term in high-altitude region. + + Parameters + ---------- + z_grid: array + Altitude grid [m]. + + g: array + Gravity values on the altitude grid [m * s^-2]. + + t: array + Temperature values on the altitude grid [K]. + + dt_dz: array + Temperature values gradient on the altitude grid [K * m^-1]. + + d: array + Thermal diffusion coefficient on the altitude grid [m^2 * s^-1]. + + k: array + Eddy diffusion coefficient values on the altitude grid [m^2 * s^-1]. + + Returns + ------- + array + Thermal diffusion term [-1]. + """ + mask1, mask2 = z_grid < 115e3, z_grid >= 115e3 + x1 = f_below_115_km( + g=g[mask1], + t=t[mask1], + dt_dz=dt_dz[mask1], + m=M["N2"], + mi=M["O"], + alpha=ALPHA["O"], + d=d, + k=k, + ) + x2 = f_above_115_km( + g=g[mask2], + t=t[mask2], + dt_dz=dt_dz[mask2], + mi=M["O"], + alpha=ALPHA["O"], + ) + return np.concatenate((x1, x2)) + + +def velocity_term_hump( + z: NDArray[np.float64], + q1: float, + q2: float, + u1: float, + u2: float, + w1: float, + w2: float, +) -> NDArray[np.float64]: + r"""Compute transport term. + + Compute the transport term given by equation (37) in + :cite:`NASA1976USStandardAtmosphere`. + + Parameters + ---------- + z: array + Altitude [m]. + + q1: float + Q constant [m^-3]. + + q2: float + q constant [m^-3]. + + u1: float + U constant [m]. + + u2: float + u constant [m]. + + w1: float + W constant [m^-3]. + + w2: float + w constant [m^-3]. + + Returns + ------- + array: + Transport term [m^-1]. + + Notes + ----- + Valid in the altitude region: 86 km :math:`\leq z \leq` 150 km. + """ + t = q1 * np.square(z - u1) * np.exp(-w1 * np.power(z - u1, 3.0)) + q2 * np.square( + u2 - z + ) * np.exp(-w2 * np.power(u2 - z, 3.0)) + return np.array(t, dtype=np.float64) + + +def velocity_term_no_hump( + z: NDArray[np.float64], q1: float, u1: float, w1: float +) -> NDArray[np.float64]: + r"""Compute transport term. + + Compute the transport term given by equation (37) in + :cite:`NASA1976USStandardAtmosphere` where the second term is zero. + + Parameters + ---------- + z: array + Altitude. + + q1: float + Q constant [m^-3]. + + u1: float + U constant [m]. + + w1: float + W constant [m^-3]. + + Returns + ------- + array + Transport term [m^-1]. + + Notes + ----- + Valid in the altitude region :math:`86` km :math:`\leq z \leq 150` km. + """ + t = q1 * np.square(z - u1) * np.exp(-w1 * np.power(z - u1, 3.0)) # m^-1 + return np.array(t, np.float64) + + +def velocity_term(s: str, z_grid: NDArray[np.float64]) -> NDArray[np.float64]: + """Compute velocity term of a given species in high-altitude region. + + Parameters + ---------- + s: str + Species. + + z_grid: array + Altitude grid [m]. + + Returns + ------- + array + Velocity term [m^-1]. + + Notes + ----- + Not valid for atomic oxygen. See :func:`velocity_term_atomic_oxygen`. + """ + x1 = velocity_term_no_hump( + z=z_grid[z_grid <= 150e3], + q1=Q1[s], + u1=U1[s], + w1=W1[s], + ) + + # Above 150 km, the velocity term is neglected, as indicated at p. 14 in + # :cite:`NASA1976USStandardAtmosphere` + x2 = np.zeros(len(z_grid[z_grid > 150e3])) + return np.concatenate((x1, x2)) + + +def velocity_term_atomic_oxygen( + grid: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute velocity term of atomic oxygen in high-altitude region. + + Parameters + ---------- + grid: array + Altitude grid [m]. + + Returns + ------- + array + Velocity term [m^-1]. + """ + mask1, mask2 = grid <= 150e3, grid > 150e3 + x1 = np.where( + grid[mask1] <= 97e3, + velocity_term_hump( + z=grid[mask1], + q1=Q1["O"], + q2=Q2["O"], + u1=U1["O"], + u2=U2["O"], + w1=W1["O"], + w2=W2["O"], + ), # m^-1 + velocity_term_no_hump( + z=grid[mask1], + q1=Q1["O"], + u1=U1["O"], + w1=W1["O"], + ), # m^-1 + ) + + x2 = np.zeros(len(grid[mask2])) + return np.concatenate((x1, x2)) + + +def tau_function( + z_grid: NDArray[np.float64], below_500: bool = True +) -> NDArray[np.float64]: + r"""Compute :math:`\tau` function. + + Compute integral given by equation (40) in + :cite:`NASA1976USStandardAtmosphere` at each point of an altitude grid. + + Parameters + ---------- + z_grid: array + Altitude grid (values sorted by ascending order) to use for integration + [m]. + + below_500: bool, default True + ``True`` if altitudes in ``z_grid`` are lower than 500 km, False + otherwise. + + Returns + ------- + array + Integral evaluations [dimensionless]. + + Notes + ----- + Valid for 150 km :math:`\leq z \leq` 500 km. + """ + if below_500: + z_grid = z_grid[::-1] + + y = ( + M["H"] + * compute_gravity(z=z_grid) + / (R * compute_temperature_high_altitude(z=z_grid)) + ) # m^-1 + integral_values: NDArray[np.float64] = np.array( + cumulative_trapezoid(y, z_grid, initial=0.0), dtype=np.float64 + ) + + if below_500: + values: NDArray[np.float64] = integral_values[::-1] + return values + else: + return integral_values + + +def log_interp1d( + x: NDArray[np.float64], y: NDArray[np.float64] +) -> Callable[[NDArray[np.float64]], NDArray[np.float64]]: + """Compute linear interpolation of :math:`y(x)` in logarithmic space. + + Parameters + ---------- + x: array + 1-D array of real values. + + y: array + N-D array of real values. The length of y along the interpolation axis + must be equal to the length of x. + + Returns + ------- + callable + Interpolating function. + """ + logx = np.log10(x) + logy = np.log10(y) + lin_interp = interp1d(logx, logy, kind="linear") + + def log_interp(z: NDArray[np.float64]) -> NDArray[np.float64]: + value = np.power(10.0, lin_interp(np.log10(z))) + return np.array(value, dtype=np.float64) + + return log_interp + + +def compute_pressure_low_altitude( + h: NDArray[np.float64], + pb: NDArray[np.float64], + tb: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute pressure in low-altitude region. + + Parameters + ---------- + h: array + Geopotential height [m]. + + pb: array + Levels pressure [Pa]. + + tb: array + Levels temperature [K]. + + Returns + ------- + array + Pressure [Pa]. + """ + # we create a mask for each layer + masks = [ + ma.masked_inside(h, H[i - 1], H[i]).mask # type: ignore + for i in range(1, len(H)) + ] + + # for each layer, we evaluate the pressure based on whether the + # temperature gradient is zero or non-zero + p = np.empty(len(h)) + for i, mask in enumerate(masks): + if LK[i] == 0: + p[mask] = compute_pressure_low_altitude_zero_gradient( + h=h[mask], + hb=H[i], + pb=pb[i], + tb=tb[i], + ) + else: + p[mask] = compute_pressure_low_altitude_non_zero_gradient( + h=h[mask], + hb=H[i], + pb=pb[i], + tb=tb[i], + lkb=LK[i], + ) + return p + + +def compute_pressure_low_altitude_zero_gradient( + h: Union[float, NDArray[np.float64]], + hb: float, + pb: float, + tb: float, +) -> NDArray[np.float64]: + """Compute pressure in low-altitude zero temperature gradient region. + + Parameters + ---------- + h: array + Geopotential height [m]. + + hb: float + Geopotential height at the bottom of the layer [m]. + + pb: float + Pressure at the bottom of the layer [Pa]. + + tb: float + Temperature at the bottom of the layer [K]. + + Returns + ------- + array + Pressure [Pa]. + """ + p = pb * np.exp(-G0 * M0 * (h - hb) / (R * tb)) + return np.array(p, dtype=np.float64) + + +def compute_pressure_low_altitude_non_zero_gradient( + h: Union[float, NDArray[np.float64]], + hb: float, + pb: float, + tb: float, + lkb: float, +) -> NDArray[np.float64]: + """Compute pressure in low-altitude non-zero temperature gradient region. + + Parameters + ---------- + h: array + Geopotential height [m]. + + hb: float + Geopotential height at the bottom of the layer [m]. + + pb: float + Pressure at the bottom of the layer [Pa]. + + tb: float + Temperature at the bottom of the layer [K]. + + lkb: float + Temperature gradient in the layer [K * m^-1]. + + Returns + ------- + array + Pressure [Pa]. + """ + p = pb * np.power(tb / (tb + lkb * (h - hb)), G0 * M0 / (R * lkb)) + return np.array(p, dtype=np.float64) + + +def compute_temperature_low_altitude( + h: NDArray[np.float64], + tb: NDArray[np.float64], +) -> NDArray[np.float64]: + """Compute temperature in low-altitude region. + + Parameters + ---------- + h: array + Geopotential height [m]. + + tb: array + Levels temperature [K]. + + Returns + ------- + array + Temperature [K]. + """ + # we create a mask for each layer + masks = [ + ma.masked_inside(h, H[i - 1], H[i]).mask # type: ignore + for i in range(1, len(H)) + ] + + # for each layer, we evaluate the pressure based on whether the + # temperature gradient is zero or not + t = np.empty(len(h)) + for i, mask in enumerate(masks): + if LK[i] == 0: + t[mask] = tb[i] + else: + t[mask] = tb[i] + LK[i] * (h[mask] - H[i]) + return t + + +def to_altitude(h: NDArray[np.float64]) -> NDArray[np.float64]: + """Convert geopotential height to (geometric) altitude. + + Parameters + ---------- + h: array + Geopotential altitude [m]. + + Returns + ------- + array + Altitude [m]. + """ + return R0 * h / (R0 - h) + + +def to_geopotential_height(z: NDArray[np.float64]) -> NDArray[np.float64]: + """Convert altitude to geopotential height. + + Parameters + ---------- + z: array + Altitude [m]. + + Returns + ------- + array + Geopotential height [m]. + """ + return R0 * z / (R0 + z) + + +def compute_gravity(z: NDArray[np.float64]) -> NDArray[np.float64]: + """Compute gravity. + + Parameters + ---------- + z : array + Altitude [m]. + + Returns + ------- + array + Gravity [m * s^-2]. + """ + return np.array(G0 * np.power((R0 / (R0 + z)), 2.0), dtype=np.float64) diff --git a/src/joseki/profiles/util.py b/src/joseki/profiles/util.py index d51ca8f..87a7c09 100644 --- a/src/joseki/profiles/util.py +++ b/src/joseki/profiles/util.py @@ -9,7 +9,6 @@ import xarray as xr from ..constants import MM, K -from ..units import ureg if sys.version_info[1] < 11: diff --git a/tests/profiles/test_ussa_1976_core.py b/tests/profiles/test_ussa_1976_core.py new file mode 100644 index 0000000..29ed4e1 --- /dev/null +++ b/tests/profiles/test_ussa_1976_core.py @@ -0,0 +1,491 @@ +"""Test cases for the core module.""" + +import typing as t + +import numpy as np +import pytest +import xarray as xr +from numpy.typing import NDArray + +from joseki.profiles.ussa_1976.constants import AR_7, M0, O2_7, O_7, H, M +from joseki.profiles.ussa_1976.core import ( + SPECIES, + VARIABLES, + compute, + compute_high_altitude, + compute_levels_temperature_and_pressure_low_altitude, + compute_low_altitude, + compute_mean_molar_mass_high_altitude, + compute_number_densities_high_altitude, + compute_temperature_gradient_high_altitude, + compute_temperature_high_altitude, + init_data_set, + to_altitude, +) + + +@pytest.fixture +def test_altitudes(): + """Test altitudes fixture.""" + return np.linspace(0.0, 100e3, 101) + + +def test_create(test_altitudes: NDArray[np.float64]): + """Creates a data set with expected data.""" + ds = compute(z=test_altitudes) + assert all([v in ds.data_vars for v in VARIABLES]) + + variables = ["p", "t", "n", "n_tot"] + ds = compute(z=test_altitudes, variables=variables) + + assert len(ds.dims) == 2 + assert "z" in ds.dims + assert "s" in ds.dims + assert len(ds.coords) == 2 + assert np.all(ds.z.values == test_altitudes) + assert list(ds.s.values) == SPECIES + assert all([var in ds for var in variables]) + assert all( + [ + x in ds.attrs + for x in ["convention", "title", "history", "source", "references"] + ] + ) + + +def test_create_invalid_variables( + test_altitudes: NDArray[np.float64], +): + """Raises when invalid variables are given.""" + invalid_variables = ["p", "t", "invalid", "n"] + with pytest.raises(ValueError): + compute(z=test_altitudes, variables=invalid_variables) + + +def test_create_invalid_z(): + """Raises when invalid altitudes values are given.""" + with pytest.raises(ValueError): + compute(z=np.array([-5.0])) # value is negative + + with pytest.raises(ValueError): + compute(z=np.array([1001e3])) # value is larger than 1000 km + + +def test_create_below_86_km_layers_boundary_altitudes(): + """ + Produces correct results. + + We test the computation of the atmospheric variables (pressure, + temperature and mass density) at the level altitudes, i.e. at the model + layer boundaries. We assert correctness by comparing their values with the + values from the table 1 of the U.S. Standard Atmosphere 1976 document. + """ + z = to_altitude(H) + ds = compute(z=z, variables=["p", "t", "rho"]) + + level_temperature: NDArray[np.float64] = np.array( + [288.15, 216.65, 216.65, 228.65, 270.65, 270.65, 214.65, 186.87] + ) + level_pressure: NDArray[np.float64] = np.array( + [101325.0, 22632.0, 5474.8, 868.01, 110.90, 66.938, 3.9564, 0.37338] + ) + level_mass_density: NDArray[np.float64] = np.array( + [ + 1.225, + 0.36392, + 0.088035, + 0.013225, + 0.0014275, + 0.00086160, + 0.000064261, + 0.000006958, + ] + ) + + assert np.allclose(ds.t.values, level_temperature, rtol=1e-4) + assert np.allclose(ds.p.values, level_pressure, rtol=1e-4) + assert np.allclose(ds.rho.values, level_mass_density, rtol=1e-3) + + +def test_create_below_86_km_arbitrary_altitudes(): + """ + Produces correct results. + + We test the computation of the atmospheric variables (pressure, + temperature and mass density) at arbitrary altitudes. We assert correctness + by comparing their values to the values from table 1 of the U.S. Standard + Atmosphere 1976 document. + """ + # The values below were selected arbitrarily from Table 1 of the document + # such that there is at least one value in each of the 7 temperature + # regions. + h: NDArray[np.float64] = np.array( + [ + 200.0, + 1450.0, + 5250.0, + 6500.0, + 9800.0, + 17900.0, + 24800.0, + 27100.0, + 37200.0, + 40000.0, + 49400.0, + 61500.0, + 79500.0, + 84000.0, + ] + ) + temperatures: NDArray[np.float64] = np.array( + [ + 286.850, + 278.725, + 254.025, + 245.900, + 224.450, + 216.650, + 221.450, + 223.750, + 243.210, + 251.050, + 270.650, + 241.250, + 197.650, + 188.650, + ] + ) + pressures: NDArray[np.float64] = np.array( + [ + 98945.0, + 85076.0, + 52239.0, + 44034.0, + 27255.0, + 7624.1, + 2589.6, + 1819.4, + 408.7, + 277.52, + 81.919, + 16.456, + 0.96649, + 0.43598, + ] + ) + mass_densities: NDArray[np.float64] = np.array( + [ + 1.2017, + 1.0633, + 0.71641, + 0.62384, + 0.42304, + 0.12259, + 0.040739, + 0.028328, + 0.0058542, + 0.0038510, + 0.0010544, + 0.00023764, + 0.000017035, + 0.0000080510, + ] + ) + + z = to_altitude(h) + ds = compute(z=z, variables=["t", "p", "rho"]) + + assert np.allclose(ds.t.values, temperatures, rtol=1e-4) + assert np.allclose(ds.p.values, pressures, rtol=1e-4) + assert np.allclose(ds.rho.values, mass_densities, rtol=1e-4) + + +@pytest.mark.parametrize( + "z_bounds", + [ + (0.0, 50e3), + (120e3, 650e3), + (70e3, 100e3), + ], +) +def test_init_data_set(z_bounds: t.Tuple[float, float]): + """Data set is initialized. + + Expected data variables are created and filled with nan values. + Expected dimensions and coordinates are present. + + Parameters + ---------- + z_bounds: tuple[float, float] + Altitude bound values. + """ + + def check_data_set(ds: xr.Dataset): + """Check a data set.""" + for var in VARIABLES: + assert var in ds + assert np.isnan(ds[var].values).all() + + assert ds.n.values.ndim == 2 + assert list(ds.s.values) == SPECIES + + z_start, z_stop = z_bounds + z = np.linspace(z_start, z_stop) + ds = init_data_set(z=z) + check_data_set(ds) + + +def test_compute_levels_temperature_and_pressure_low_altitude(): + """Computes correct level temperature and pressure values. + + The correct values are taken from :cite:`NASA1976USStandardAtmosphere`. + """ + tb, pb = compute_levels_temperature_and_pressure_low_altitude() + + level_temperature: NDArray[np.float64] = np.array( + [288.15, 216.65, 216.65, 228.65, 270.65, 270.65, 214.65, 186.87] + ) + level_pressure: NDArray[np.float64] = np.array( + [101325.0, 22632.0, 5474.8, 868.01, 110.90, 66.938, 3.9564, 0.37338] + ) + + assert np.allclose(tb, level_temperature, rtol=1e-3) + assert np.allclose(pb, level_pressure, rtol=1e-3) + + +def test_compute_number_density(): + """Computes correct number density values at arbitrary level altitudes. + + The correct values are taken from :cite:`NASA1976USStandardAtmosphere` + (table VIII, p. 210-215). + """ + # the following altitudes values are chosen arbitrarily + altitudes: NDArray[np.float64] = np.array( + [ + 86e3, + 90e3, + 95e3, + 100e3, + 110e3, + 120e3, + 150e3, + 200e3, + 300e3, + 400e3, + 500e3, + 600e3, + 700e3, + 800e3, + 900e3, + 1000e3, + ], + dtype=np.float64, + ) + + mask = altitudes > 150e3 + + # the corresponding number density values are from NASA (1976) - U.S. + # Standard Atmosphere, table VIII (p. 210-215) + values: t.Dict[str, NDArray[np.float64]] = { + "N2": np.array( + [ + 1.13e20, + 5.547e19, + 2.268e19, + 9.210e18, + 1.641e18, + 3.726e17, + 3.124e16, + 2.925e15, + 9.593e13, + 4.669e12, + 2.592e11, + 1.575e10, + 1.038e9, + 7.377e7, + 5.641e6, + 4.626e5, + ] + ), + "O": np.array( + [ + O_7, + 2.443e17, + 4.365e17, + 4.298e17, + 2.303e17, + 9.275e16, + 1.780e16, + 4.050e15, + 5.443e14, + 9.584e13, + 1.836e13, + 3.707e12, + 7.840e11, + 1.732e11, + 3.989e10, + 9.562e9, + ] + ), + "O2": np.array( + [ + O2_7, + 1.479e19, + 5.83e18, + 2.151e18, + 2.621e17, + 4.395e16, + 2.750e15, + 1.918e14, + 3.942e12, + 1.252e11, + 4.607e9, + 1.880e8, + 8.410e6, + 4.105e5, + 2.177e4, + 1.251e3, + ] + ), + "Ar": np.array( + [ + AR_7, + 6.574e17, + 2.583e17, + 9.501e16, + 1.046e16, + 1.366e15, + 5.0e13, + 1.938e12, + 1.568e10, + 2.124e8, + 3.445e6, + 6.351e4, + 1.313e3, + 3.027e1, + 7.741e-1, + 2.188e-2, + ] + ), + "He": np.array( + [ + 7.582e14, + 3.976e14, + 1.973e14, + 1.133e14, + 5.821e13, + 3.888e13, + 2.106e13, + 1.310e13, + 7.566e12, + 4.868e12, + 3.215e12, + 2.154e12, + 1.461e12, + 1.001e12, + 6.933e11, + 4.850e11, + ] + ), + "H": np.array( + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 3.767e11, + 1.630e11, + 1.049e11, + 8.960e10, + 8.0e10, + 7.231e10, + 6.556e10, + 5.961e10, + 5.434e10, + 4.967e10, + ] + ), + } + + n = compute_number_densities_high_altitude(altitudes=altitudes) + + assert np.allclose(n.sel(s="N2"), values["N2"], rtol=0.01) + # TODO: investigate the poor relative tolerance that is achieved here + assert np.allclose(n.sel(s="O"), values["O"], rtol=0.1) + assert np.allclose(n.sel(s="O2"), values["O2"], rtol=0.01) + assert np.allclose(n.sel(s="Ar"), values["Ar"], rtol=0.01) + assert np.allclose(n.sel(s="He"), values["He"], rtol=0.01) + assert np.allclose(n.sel(s="H")[mask], values["H"][mask], rtol=0.01) + + +def test_compute_mean_molar_mass(): + """Computes correct mean molar mass values. + + The correct values are taken from :cite:`NASA1976USStandardAtmosphere`. + """ + z: NDArray[np.float64] = np.linspace(86e3, 1000e3, 915) + assert np.allclose( + compute_mean_molar_mass_high_altitude(z=z), + np.where(z <= 100.0e3, M0, M["N2"]), + ) + + +def test_compute_temperature_above_86_km(): + """Compute correct temperature values. + + The correct values are taken from :cite:`NASA1976USStandardAtmosphere`. + """ + z: NDArray[np.float64] = np.array([90e3, 100e3, 110e3, 120e3, 130e3, 200e3, 500e3]) + assert np.allclose( + compute_temperature_high_altitude(z=z), + np.array([186.87, 195.08, 240.00, 360.0, 469.27, 854.56, 999.24]), + rtol=1e-3, + ) + + +def test_compute_temperature_above_86_km_invalid_altitudes(): + """Raises when altitude is out of range.""" + with pytest.raises(ValueError): + compute_temperature_high_altitude(z=np.array([10e3])) # 10 km < 86 km + + +def test_compute_high_altitude_no_mask(): + """Returns a Dataset.""" + z: NDArray[np.float64] = np.linspace(86e3, 1000e3) + ds = init_data_set(z=z) + compute_high_altitude(data_set=ds, mask=None, inplace=True) + assert isinstance(ds, xr.Dataset) + + +def test_compute_high_altitude_not_inplace(): + """Returns a Dataset.""" + z: NDArray[np.float64] = np.linspace(86e3, 1000e3) + ds1 = init_data_set(z=z) + ds2 = compute_high_altitude(data_set=ds1, mask=None, inplace=False) + assert ds1 != ds2 + assert isinstance(ds2, xr.Dataset) + + +def test_compute_low_altitude(): + """Returns a Dataset.""" + z: NDArray[np.float64] = np.linspace(0.0, 86e3) + ds = init_data_set(z=z) + compute_low_altitude(data_set=ds, mask=None, inplace=True) + assert isinstance(ds, xr.Dataset) + + +def test_compute_low_altitude_not_inplace(): + """Returns a Dataset.""" + z: NDArray[np.float64] = np.linspace(0.0, 86e3) + ds1 = init_data_set(z=z) + ds2 = compute_low_altitude(data_set=ds1, mask=None, inplace=False) + assert ds1 != ds2 + assert isinstance(ds2, xr.Dataset) + + +def test_compute_temperature_gradient_high_altitude(): + """Raises ValueError when altitude is out of bounds.""" + with pytest.raises(ValueError): + compute_temperature_gradient_high_altitude(z=np.array([1300e3])) diff --git a/uv.lock b/uv.lock index 177a75c..d3e30d0 100644 --- a/uv.lock +++ b/uv.lock @@ -1104,7 +1104,6 @@ dependencies = [ { name = "scipy", version = "1.13.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, { name = "scipy", version = "1.16.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "ussa1976" }, { name = "xarray", version = "2024.7.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, { name = "xarray", version = "2025.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -1149,7 +1148,6 @@ requires-dist = [ { name = "pandas", specifier = ">=1.5.2" }, { name = "pint", specifier = ">=0.20.1" }, { name = "scipy", specifier = ">=1.9.3" }, - { name = "ussa1976", specifier = ">=0.3.4" }, { name = "xarray", specifier = ">=2022.12.0" }, ] @@ -3612,29 +3610,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" }, ] -[[package]] -name = "ussa1976" -version = "0.3.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "click", version = "8.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "netcdf4" }, - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, - { name = "numpy", version = "2.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "scipy", version = "1.13.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, - { name = "scipy", version = "1.16.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "xarray", version = "2024.7.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, - { name = "xarray", version = "2025.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2c/83/6c3883fe969a7966cad1af0285ae6e4d292512e5ae894cc068b19e13c5ba/ussa1976-0.3.4.tar.gz", hash = "sha256:28eec80a65db5ab84ec79ea1732352bbb137408260f940374fa313c7440faa37", size = 15475, upload-time = "2023-01-31T11:25:43.222Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/77/3ac8a9bbef6c3be222c71bbe9587ff85adfa3fb85d83edba89c10823ae03/ussa1976-0.3.4-py3-none-any.whl", hash = "sha256:6722932c48234f02b4e1b0745f2dd3832954b14e7a2c6ece1211b50cb9e10bf6", size = 14793, upload-time = "2023-01-31T11:25:41.731Z" }, -] - [[package]] name = "uvicorn" version = "0.35.0"