Source code for openMWR.atm

"""
Functions for thermodynamic calculations such as saturation vapor pressure,
latent heat, absolute humidity, and dewpoint temperature.
"""

import numpy as np
import xarray as xr

from openMWR import consts

[docs] def saturation_vapor_pressure(T): """ Calculate saturation vapor pressure using the Magnus equation. Parameters ---------- T : float or ndarray Air temperature in kelvin. Returns ------- e_s : float or ndarray Saturation vapor pressure in pascals. """ # Convert temperature from kelvin to degrees Celsius T_C = T - consts.T_0 # Apply Magnus equation to compute saturation vapor pressure e_s = 611.2 * np.exp(17.62 * T_C / (T_C + 243.12)) return e_s
[docs] def latent_heat(T: np.ndarray) -> np.ndarray: """ Calculate the latent heat of vaporization as a function of temperature. Parameters ---------- T : ndarray Temperature in kelvin. Returns ------- L : ndarray Latent heat of vaporization in joules per kilogram (J/kg). Notes ----- Formula from Rogers and Yau, 1989: A Short Course in Cloud Physics (III.Ed.), Pergamon Press, 293p. Page 15 """ L = consts.L_0 - (consts.c_w - consts.c_pv) * (T - consts.T_0) return L
[docs] def calc_absolute_humidity(T, rh): """ Calculate absolute humidity from temperature and relative humidity. Parameters ---------- T : float or ndarray Air temperature in kelvin. rh : float or ndarray Relative humidity in percent (%). Returns ------- ah : float or ndarray Absolute humidity in kg/m³. """ ah = rh / 100 * saturation_vapor_pressure(T) / (consts.R_v * T) return ah
[docs] def dewpoint(e): """ Calculate dewpoint temperature from vapor pressure using the reversed Magnus equation. Parameters ---------- e : float, ndarray, or xr.DataArray Water vapor pressure in pascals (Pa). Returns ------- T_D : float, ndarray, or xr.DataArray Dewpoint temperature in kelvin (K). """ # Ensure non-positive vapor pressures become NaN if isinstance(e, xr.DataArray): e = e.where(e > 0) else: e = np.where(e > 0, e, np.nan) # Compute dewpoint using inverse Magnus formula log = np.log(e / 611.2) T_D = 243.12 * log / (17.62 - log) + consts.T_0 return T_D
[docs] def dewpoint_from_relative_humidity(T, rh): """ Calculate dewpoint temperature from temperature and relative humidity. Parameters ---------- T : float or ndarray Air temperature in kelvin. rh : float or ndarray Relative humidity in percent (%). Returns ------- T_D : float or ndarray Dewpoint temperature in kelvin (K). """ # Convert RH to fraction and compute vapor pressure e = rh / 100 * saturation_vapor_pressure(T) # Convert vapor pressure to dewpoint temperature T_D = dewpoint(e) return T_D
[docs] def dewpoint_from_absolute_humidity(T, ah): """ Calculate dewpoint temperature from temperature and absolute humidity. Parameters ---------- T : float or ndarray Air temperature in kelvin. ah : float or ndarray Absolute humidity in kg/m³. Returns ------- T_D : float or ndarray Dewpoint temperature in kelvin (K). """ # Convert absolute humidity to vapor pressure e = ah * consts.R_v * T # Convert vapor pressure to dewpoint temperature T_D = dewpoint(e) return T_D
[docs] def calc_pressure(surface_p, T, z): r""" Compute the hydrostatic pressure profile from surface pressure, temperature, and geometric height. Supports both NumPy arrays (1D calculation) and xarray.DataArray inputs (vectorized), using a self-dispatching mechanism. Parameters ---------- surface_p : float or xr.DataArray Surface pressure at the lowest level in pascals (Pa). T : ndarray or xr.DataArray Temperature profile in kelvin (K). Must contain a vertical dimension named ``'height'`` when used with xarray. z : ndarray or xr.DataArray Geometric height levels in meters (m). Must share the ``'height'`` dimension with `T` when using xarray. Returns ------- p : ndarray or xr.DataArray Pressure profile in pascals (Pa), evaluated at each height level. Notes ----- - For xarray inputs, the function automatically switches into a vectorized mode using :func:`xarray.apply_ufunc`. Inside this process, it calls itself again, but with NumPy arrays extracted from the DataArrays. - The underlying 1D computation follows the hydrostatic equation: .. math:: p(z_i) = p_0\, \exp\!\left( - \sum_{k=0}^{i-1} \frac{g\, \Delta z_k}{R_d\, T_{\text{mean},k}} \right) where :math:`T_{mean}` is the layer-mean temperature. """ # -------- CASE 1: xarray input → vectorized dispatch via apply_ufunc -------- if isinstance(surface_p, xr.DataArray) or isinstance(T, xr.DataArray) or isinstance(z, xr.DataArray): return xr.apply_ufunc( calc_pressure, # self-dispatching call surface_p, T, z, input_core_dims=[[], ['height'], ['height']], output_core_dims=[['height']], vectorize=True, dask="parallelized", output_dtypes=[float], ) # -------- CASE 2: NumPy inputs → perform actual 1D hydrostatic computation -------- delta_z = np.diff(z) # Layer-mean temperature T_mean = (T[:-1] + T[1:]) / 2 # Hydrostatic exponent for each layer exponent = -consts.g * delta_z / (consts.R_d * T_mean) # Initialize output array and integrate upward p = np.empty(z.size) p[0] = surface_p p[1:] = surface_p * np.exp(np.cumsum(exponent)) return p