diff --git a/src/include/libthermo.hpp b/src/include/libthermo.hpp index f312c54..a238146 100644 --- a/src/include/libthermo.hpp +++ b/src/include/libthermo.hpp @@ -95,7 +95,7 @@ class lcl { const size_t max_iters = DEFAULT_ITERS ) noexcept; constexpr T wet_bulb_temperature(const T pressure, const T step = DEFAULT_STEP) noexcept; - constexpr size_t index(const T pressure[], const size_t size) noexcept; + constexpr size_t index_on(const T pressure[], const size_t size) noexcept; }; template diff --git a/src/lib/functional.cpp b/src/lib/functional.cpp index f26c7c4..9a0c1b4 100644 --- a/src/lib/functional.cpp +++ b/src/lib/functional.cpp @@ -158,7 +158,7 @@ constexpr T rk2(Fn fn, T x0, T x1, T y, T step /* 1000.0 (Pa) */) noexc T k1, delta, abs_delta; size_t N = 1; - abs_delta = fabs(delta = (x1 - x0)); + abs_delta = std::fabs(delta = (x1 - x0)); if (abs_delta > step) { N = (size_t)ceil(abs_delta / step); delta = delta / (T)N; @@ -206,9 +206,9 @@ constexpr T fixed_point( p2 = fn(p1, x0, args...); delta = p2 - 2.0 * p1 + p0; - p2 = delta ? p0 - pow(p1 - p0, 2) / delta : p2; /* delta squared */ + p2 = delta ? p0 - std::pow(p1 - p0, 2) / delta : p2; /* delta squared */ - if ((p0 ? fabs((p2 - p0) / p0) : p2 /* absolute relative error */) < eps) + if ((p0 ? std::fabs((p2 - p0) / p0) : p2 /* absolute relative error */) < eps) return p2; p0 = p2; @@ -289,8 +289,8 @@ constexpr point intersect_1d( x0 = x[i0]; x1 = x[i1]; if (log_x) { - x0 = log(x0); - x1 = log(x1); + x0 = std::log(x0); + x1 = std::log(x1); } a0 = a[i0]; @@ -305,7 +305,7 @@ constexpr point intersect_1d( y_intercept = ((x_intercept - x0) / LIMIT_ZERO(x1 - x0)) * (a1 - a0) + a0; if (log_x) - x_intercept = exp(x_intercept); + x_intercept = std::exp(x_intercept); return {x_intercept, y_intercept}; }; diff --git a/src/lib/libthermo.cpp b/src/lib/libthermo.cpp index 35e6f1b..41217dd 100644 --- a/src/lib/libthermo.cpp +++ b/src/lib/libthermo.cpp @@ -2,17 +2,15 @@ namespace libthermo { -// helper functions - /** * @brief given a pressure array of decreasing values, find the index of the pressure level that * corresponds to the given value. This function is optimized for evenly distributed pressure * and will typically find the index in O(1) time. * * @tparam T - * @param levels - * @param value - * @param size + * @param levels pressure levels array + * @param value value to find + * @param size size of the array * @return constexpr size_t */ template @@ -92,13 +90,17 @@ constexpr T potential_temperature(const T pressure, const T temperature) noexcep } /** - * @brief theta_e + * @brief Equivalent potential temperature, `theta-e (𝜃𝑒)`, is a quantity that is conserved during + * changes to an air parcel's pressure (that is, during vertical motions in the atmosphere), even + * if water vapor condenses during that pressure change. It is therefore more conserved than the + * ordinary potential temperature, which remains constant only for unsaturated vertical motions + * (pressure changes). * * @tparam T - * @param pressure - * @param temperature - * @param dewpoint - * @return constexpr T + * @param pressure `(Pa)` + * @param temperature `(K)` + * @param dewpoint `(K)` + * @return `theta-e` - `(K)` */ template constexpr T equivalent_potential_temperature( @@ -106,20 +108,20 @@ constexpr T equivalent_potential_temperature( ) noexcept { const T r = saturation_mixing_ratio(pressure, dewpoint); const T e = saturation_vapor_pressure(dewpoint); - const T t_l = 56 + 1.0 / (1.0 / (dewpoint - 56) + log(temperature / dewpoint) / 800.0); + const T t_l = 56. + 1. / (1. / (dewpoint - 56.) + log(temperature / dewpoint) / 800.); const T th_l = potential_temperature(pressure - e, temperature) * pow(temperature / t_l, 0.28 * r); - return th_l * exp(r * (1 + 0.448 * r) * (3036.0 / t_l - 1.78)); + return th_l * exp(r * (1. + 0.448 * r) * (3036. / t_l - 1.78)); } /** * @brief theta_w * * @tparam T - * @param pressure - * @param temperature - * @param dewpoint + * @param pressure `(Pa)` + * @param temperature `(K)` + * @param dewpoint `(K)` * @return constexpr T */ template @@ -188,7 +190,7 @@ constexpr T lcl::wet_bulb_temperature(const T pressure, const T step) noexcep } template -constexpr size_t lcl::index(const T pressure[], const size_t size) noexcept { +constexpr size_t lcl::index_on(const T pressure[], const size_t size) noexcept { return index_pressure(pressure, this->pressure, size); } diff --git a/src/nzthermo/_C.pxd b/src/nzthermo/_C.pxd index 4910ea7..4be7b36 100644 --- a/src/nzthermo/_C.pxd +++ b/src/nzthermo/_C.pxd @@ -54,7 +54,7 @@ cdef extern from "libthermo.cpp" namespace "libthermo" nogil: lcl() noexcept lcl(T pressure, T temperature) noexcept lcl(T pressure, T temperature, T dewpoint) noexcept - size_t index(T* levels, size_t size) noexcept + size_t index_on(T* levels, size_t size) noexcept # 1x1 T saturation_vapor_pressure[T](T temperature) noexcept diff --git a/src/nzthermo/_core.pyi b/src/nzthermo/_core.pyi index 8990eb6..b58304c 100644 --- a/src/nzthermo/_core.pyi +++ b/src/nzthermo/_core.pyi @@ -14,7 +14,6 @@ _Dtype_t = TypeVar("_Dtype_t", bound=np.generic) _DualArrayLike = ( SupportsArray[_Dtype_t] | NestedSequence[SupportsArray[_Dtype_t]] | _T | NestedSequence[_T] ) -_float = TypeVar("_float", np.float32, np.float64) OPENMP_ENABLED: bool = ... diff --git a/src/nzthermo/_core.pyx b/src/nzthermo/_core.pyx index d4cd3f5..870f43b 100644 --- a/src/nzthermo/_core.pyx +++ b/src/nzthermo/_core.pyx @@ -103,16 +103,6 @@ kappa = C.kappa # ............................................................................................... # # helpers # ............................................................................................... # -cdef cvarray nzarray((size_t, size_t) shape, size_t size): - return cvarray( - shape, - itemsize=size, - format='f' if size == sizeof(float) else 'd', - mode='c', - allocate_buffer=True, - ) - - def broadcast_and_nanmask( np.ndarray pressure not None, np.ndarray temperature not None, @@ -284,7 +274,7 @@ cdef T[:, :] moist_lapse_2d( T[:, :] out N, Z = temperature.shape[0], pressure.shape[1] - out = nzarray((N, Z), pressure.itemsize) + out = np.empty((N, Z), dtype=np.dtype(f"f{pressure.itemsize}")) with nogil, parallel(): if BROADCAST is mode: for i in prange(N, schedule='dynamic'): @@ -466,7 +456,7 @@ cdef void parcel_profile_1d( # [dry ascent] # stop the dry ascent at the LCL - stop = lcl.index(&pressure[0], Z) + stop = lcl.index_on(&pressure[0], Z) for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent out[i] = C.dry_lapse(pressure[i], p0, temperature) @@ -549,8 +539,7 @@ cdef void parcel_profile_with_lcl_1d( lcl = C.lcl[T](p0, t0, td0) # [dry ascent] .. parcel temperature from the surface up to the LCL .. - - stop = lcl.index(&pressure[0], Z) + stop = lcl.index_on(&pressure[0], Z) if stop: ep[:stop] = pressure[:stop] @@ -774,7 +763,7 @@ cdef intersect_2d( T[:, :] out N, Z = y0.shape[0], y1.shape[1] - out = nzarray((2, N), x.itemsize) + out = np.empty((2, N), f'f{x.itemsize}') with nogil, parallel(): if BROADCAST is mode: for i in prange(N, schedule='dynamic'): diff --git a/src/nzthermo/_ufunc.pyi b/src/nzthermo/_ufunc.pyi index 1ce4e75..68b194f 100644 --- a/src/nzthermo/_ufunc.pyi +++ b/src/nzthermo/_ufunc.pyi @@ -13,9 +13,9 @@ _DType_T_co = TypeVar("_DType_T_co", bound=np.dtype[Any]) class pressure_vector(np.ndarray[_S, _DType_T_co]): @overload - def __new__(cls, array: np.ndarray[_S, _DType_T_co]) -> pressure_vector[_S, _DType_T_co]: ... + def __new__(cls, __x: np.ndarray[_S, _DType_T_co], /) -> pressure_vector[_S, _DType_T_co]: ... @overload - def __new__(cls, array: ArrayLike) -> pressure_vector[Any, np.dtype[Any]]: ... + def __new__(cls, __x: ArrayLike, /) -> pressure_vector[Any, np.dtype[Any]]: ... def is_above( self, bottom: ArrayLike, close: bool = ... ) -> np.ndarray[_S, np.dtype[np.bool_]]: ... diff --git a/src/nzthermo/core.py b/src/nzthermo/core.py index 15c9913..325f575 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -9,7 +9,7 @@ from __future__ import annotations -from typing import Any, Final, Literal as L, TypeVar +from typing import Any, Literal as L, TypeVar import numpy as np from numpy.typing import ArrayLike, NDArray @@ -34,19 +34,16 @@ wet_bulb_temperature, ) from .typing import Kelvin, N, Pascal, Z, shape -from .utils import Axis, Vector1d, broadcast_nz +from .utils import Parcel, broadcast_nz _S = TypeVar("_S", bound=shape) _T = TypeVar("_T", bound=np.floating[Any], covariant=True) -newaxis: Final[None] = np.newaxis -surface: Final[tuple[slice, slice]] = np.s_[:, :1] -aloft: Final[tuple[slice, slice]] = np.s_[:, 1:] -NaN = np.nan FASTPATH: dict[str, Any] = {"__fastpath": True} +@broadcast_nz def parcel_mixing_ratio( pressure: Pascal[pressure_vector[_S, np.dtype[_T]]], temperature: Kelvin[np.ndarray[_S, np.dtype[_T]]], @@ -55,16 +52,114 @@ def parcel_mixing_ratio( *, where: np.ndarray[_S, np.dtype[np.bool_]] | None = None, ) -> np.ndarray[_S, np.dtype[_T]]: + r"""Calculate the parcel mixing ratio. + + Calculate the mixing ratio of a parcel given the pressure, temperature, and dewpoint. The + mixing ratio is the ratio of the mass of water vapor to the mass of dry air in a given + volume of air. This function calculates the mixing ratio of a parcel of air given the + pressure, temperature, and dewpoint of the parcel. + + + Parameters + ---------- + pressure : array_like[(N, Z), floating] + Atmospheric pressure profile (Pa). This array must be from high to low pressure. + temperature : array_like[(N, Z), floating] + Temperature (K) at the levels given by `pressure` + dewpoint : array_like[(N, Z), floating] + Dewpoint (K) at the levels given by `pressure` + where : array_like[(N, Z), bool], optional + Where ``True`` the saturation mixing ratio is calculated, using the dewpoint and where + ``False`` the mixing ratio is calculated using the temperature. If not provided the + mixing ratio is calculated using the dewpoint upto the LCL (lifting condensation level) + and the temperature above the LCL. + + Returns + ------- + parcel_mixing_ratio : array_like[(N, Z), floating] + Parcel mixing ratio (kg/kg) + + """ if where is None: where = pressure.is_below( - lcl_pressure(pressure[surface], temperature[surface], dewpoint[surface]) + lcl_pressure(pressure[:, :1], temperature[:, :1], dewpoint[:, :1]) ) - r = saturation_mixing_ratio(pressure, dewpoint, out=np.empty_like(temperature), where=where) - r = saturation_mixing_ratio(pressure, temperature, out=r, where=~where) + + r = saturation_mixing_ratio( + pressure, + temperature, + out=saturation_mixing_ratio(pressure, dewpoint, where=where), + where=~where, + ) return r +# ------------------------------------------------------------------------------------------------- +# convective condensation level +# ------------------------------------------------------------------------------------------------- +@broadcast_nz +def ccl( + pressure: Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + /, + *, + height: ArrayLike | None = None, + mixed_layer_depth: ArrayLike | None = None, + which: L["bottom", "top"] = "bottom", +) -> tuple[Pascal[NDArray[_T]], Kelvin[NDArray[_T]], Kelvin[NDArray[_T]]]: + r"""Calculate convective condensation level (CCL) and convective temperature. + + The Convective Condensation Level (CCL) is the level at which condensation will occur if + sufficient afternoon heating causes rising parcels of air to reach saturation. The CCL is + greater than or equal in height (lower or equal pressure level) than the LCL. The CCL and the + LCL are equal when the atmosphere is saturated. The CCL is found at the intersection of the + saturation mixing ratio line (through the surface dewpoint) and the environmental temperature. + + Parameters + ---------- + pressure : array_like[(N, Z), floating] + Atmospheric pressure profile (Pa). This array must be from high to low pressure. + temperature : array_like[(N, Z), floating] + Temperature (K) at the levels given by `pressure` + dewpoint : array_like[(N, Z), floating] + Dewpoint (K) at the levels given by `pressure` + height : array_like[(N,), floating], optional + + Returns + ------- + TODO : ... + + Examples + -------- + TODO : ... + + """ + if mixed_layer_depth is None: + r = mixing_ratio(saturation_vapor_pressure(dewpoint[:, :1]), pressure[:, :1]) + else: + r = mixed_layer( + pressure, + mixing_ratio(saturation_vapor_pressure(dewpoint), pressure), + height=height, + depth=mixed_layer_depth, + # TODO: + # the metpy implementation forces interpolation here which is not implemented + # currently with our mixed layer function + interpolate=True, + )[0][:, np.newaxis] + + if height is not None: + raise NotImplementedError + + rt_profile = _dewpoint(vapor_pressure(pressure, r)) + + p, t = F.intersect_nz(pressure, rt_profile, temperature, "increasing", log_x=True).pick(which) + + return p, t, dry_lapse(pressure[:, 0], t, p) + + # ------------------------------------------------------------------------------------------------- # downdraft_cape # ------------------------------------------------------------------------------------------------- @@ -88,14 +183,18 @@ def downdraft_cape( Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` - where : `array_like[[N, Z], bool]`, optional + where : array_like[(N, Z), bool], optional + Returns + ------- + DCAPE : array_like[(N,), floating] + Downward Convective Available Potential Energy (J/kg) Examples -------- @@ -112,20 +211,17 @@ def downdraft_cape( pressure, temperature, dewpoint, - # masking values with inf will alow us to call argmin without worrying about nan where=where, out=np.full_like(temperature, np.inf), ) zx = theta_e.argmin(axis=1) - p_top = pressure[nx, zx] if pressure.shape == temperature.shape else pressure[0, zx] + p_top = pressure[0 if pressure.shape[0] == 1 else nx, zx] t_top = temperature[nx, zx] # (N,) td_top = dewpoint[nx, zx] # (N,) wb_top = wet_bulb_temperature(p_top, t_top, td_top) # (N,) - # our moist_lapse rate function has nan ignoring capabilities - # pressure = pressure.where(pressure >= p_top[:, newaxis], NaN) - pressure = pressure.where(pressure.is_below(p_top[:, newaxis], close=True), NaN) + pressure = pressure.where(pressure.is_below(p_top[:, np.newaxis], close=True), np.nan) e_vt = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) # (N, Z) trace = moist_lapse(pressure, wb_top, p_top) # (N, Z) p_vt = virtual_temperature(trace, saturation_mixing_ratio(pressure, trace)) # (N, Z) @@ -135,61 +231,6 @@ def downdraft_cape( return DCAPE -# ------------------------------------------------------------------------------------------------- -# convective condensation level -# ------------------------------------------------------------------------------------------------- -@broadcast_nz -def ccl( - pressure: Pascal[NDArray[_T]], - temperature: Kelvin[NDArray[_T]], - dewpoint: Kelvin[NDArray[_T]], - /, - *, - height=None, - mixed_layer_depth=None, - which: L["bottom", "top"] = "bottom", -) -> tuple[Pascal[NDArray[_T]], Kelvin[NDArray[_T]], Kelvin[NDArray[_T]]]: - r"""Calculate convective condensation level (CCL). - - The Convective Condensation Level (CCL) is the level at which condensation will occur if - sufficient afternoon heating causes rising parcels of air to reach saturation. The CCL is - greater than or equal in height (lower or equal pressure level) than the LCL. The CCL and the - LCL are equal when the atmosphere is saturated. The CCL is found at the intersection of the - saturation mixing ratio line (through the surface dewpoint) and the environmental temperature. - - Parameters - ---------- - pressure : `array_like[[N, Z] | [Z], floating]` - Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` - Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` - Dewpoint (K) at the levels given by `pressure` - TODO : ... - - Returns - ------- - TODO : ... - - Examples - -------- - TODO : ... - - """ - if mixed_layer_depth is None: - r = mixing_ratio(saturation_vapor_pressure(dewpoint[surface]), pressure[surface]) - else: - raise NotImplementedError - if height is not None: - raise NotImplementedError - - rt_profile = _dewpoint(vapor_pressure(pressure, r)) - - p, t = F.intersect_nz(pressure, rt_profile, temperature, "increasing", log_x=True).pick(which) - - return p, t, dry_lapse(pressure[:, 0], t, p) - - # ------------------------------------------------------------------------------------------------- # el & lfc # ------------------------------------------------------------------------------------------------- @@ -203,7 +244,8 @@ def _el_lfc( which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, -) -> tuple[Vector1d[_T], Vector1d[_T]] | Vector1d[_T]: +) -> tuple[Parcel[_T], Parcel[_T]] | Parcel[_T]: + r"""short cut for el, lfc, and el_lfc functions.""" if parcel_profile is None: pressure, temperature, dewpoint, parcel_profile = parcel_profile_with_lcl( pressure, temperature, dewpoint @@ -216,17 +258,17 @@ def _el_lfc( else: td0 = dewpoint_start - LCL = Vector1d.from_func(lcl, p0, t0, td0).unsqueeze() + LCL = Parcel.from_func(lcl, p0, t0, td0).unsqueeze() pressure, parcel_profile, temperature = ( - pressure[aloft], - parcel_profile[aloft], - temperature[aloft], + pressure[:, 1:], + parcel_profile[:, 1:], + temperature[:, 1:], ) if pick != "LFC": # find the Equilibrium Level (EL) - top_idx = np.arange(N), np.argmin(~np.isnan(pressure), Axis.Z) - 1 - left_of_env = (parcel_profile[top_idx] <= temperature[top_idx])[:, newaxis] + top_idx = np.arange(N), np.argmin(~np.isnan(pressure), axis=1) - 1 + left_of_env = (parcel_profile[top_idx] <= temperature[top_idx])[:, np.newaxis] EL = F.intersect_nz( pressure, parcel_profile, @@ -253,7 +295,7 @@ def _el_lfc( is_lcl = no_lfc & greater_or_close( # the mask only needs to be applied to either the temperature or parcel_temperature_profile - np.where(LCL.is_below(pressure, close=True), parcel_profile, NaN), + np.where(LCL.is_below(pressure, close=True), parcel_profile, np.nan), temperature, ).any(axis=1, out=np.empty((N, 1), dtype=np.bool_), keepdims=True) @@ -278,7 +320,7 @@ def el( parcel_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, *, which: L["top", "bottom"] = "top", -) -> Vector1d[_T]: +) -> Parcel[_T]: r"""Calculate the equilibrium level (EL). This works by finding the last intersection of the ideal parcel path and the measured @@ -287,11 +329,11 @@ def el( Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` parcel_profile : array_like[[N, Z], floating], optional The parcel's temperature profile from which to calculate the EL. Defaults to the @@ -332,7 +374,7 @@ def lfc( *, which: L["top", "bottom"] = "top", dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, -) -> Vector1d[_T]: +) -> Parcel[_T]: return _el_lfc( "LFC", pressure, @@ -355,16 +397,16 @@ def el_lfc( which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, -) -> tuple[Vector1d[_T], Vector1d[_T]]: +) -> tuple[Parcel[_T], Parcel[_T]]: r"""TODO ... Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` TODO : ... @@ -389,7 +431,73 @@ def el_lfc( # ------------------------------------------------------------------------------------------------- -# nzthermo.core.most_unstable_parcel +# cape_cin +# ------------------------------------------------------------------------------------------------- +@broadcast_nz +def cape_cin( + pressure: Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + /, + parcel_profile: np.ndarray, + *, + which_lfc: L["bottom", "top"] = "bottom", + which_el: L["bottom", "top"] = "top", +) -> tuple[np.ndarray, np.ndarray]: + r"""Calculate CAPE and CIN. + + Parameters + ---------- + pressure : array_like[(N, Z), floating] + Atmospheric pressure profile (Pa). This array must be from high to low pressure. + temperature : array_like[(N, Z), floating] + Temperature (K) at the levels given by `pressure` + dewpoint : array_like[(N, Z), floating] + Dewpoint (K) at the levels given by `pressure` + parcel_profile : array_like[(N, Z), floating] + + Returns + ------- + TODO : ... + + Examples + -------- + TODO : ... + """ + # The mixing ratio of the parcel comes from the dewpoint below the LCL, is saturated + # based on the temperature above the LCL + + parcel_profile = virtual_temperature( + parcel_profile, parcel_mixing_ratio(pressure, temperature, dewpoint, **FASTPATH) + ) + temperature = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) + # Calculate the EL limit of integration + (EL, _), (LFC, _) = _el_lfc( + "BOTH", + pressure, + temperature, + dewpoint, + parcel_profile, + which_lfc, + which_el, + ) + EL, LFC = np.reshape((EL, LFC), (2, -1, 1)) # reshape for broadcasting + + tzx = F.zero_crossings(pressure, parcel_profile - temperature) # temperature zero crossings + + p, t = tzx.where_between(LFC, EL, close=True) + CAPE = Rd * F.nantrapz(t, np.log(p), axis=1) + CAPE[CAPE < 0.0] = 0.0 + + p, t = tzx.where_below(LFC, close=True) + CIN = Rd * F.nantrapz(t, np.log(p), axis=1) + CIN[CIN > 0.0] = 0.0 + + return CAPE, CIN + + +# ------------------------------------------------------------------------------------------------- +# most_unstable # ------------------------------------------------------------------------------------------------- @broadcast_nz def most_unstable_parcel_index( @@ -405,11 +513,11 @@ def most_unstable_parcel_index( Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` TODO : ... @@ -424,7 +532,7 @@ def most_unstable_parcel_index( if height is not None: raise NotImplementedError("height argument is not implemented") - pbot = (pressure[surface] if bottom is None else np.asarray(bottom)).reshape(-1, 1) + pbot = (pressure[:, :1] if bottom is None else np.asarray(bottom)).reshape(-1, 1) ptop = pbot - depth theta_e = equivalent_potential_temperature( @@ -457,11 +565,11 @@ def most_unstable_parcel( Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` TODO : ... @@ -485,75 +593,6 @@ def most_unstable_parcel( ) -# ------------------------------------------------------------------------------------------------- -# cape_cin -# ------------------------------------------------------------------------------------------------- -@broadcast_nz -def cape_cin( - pressure: Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], - /, - parcel_profile: np.ndarray, - *, - which_lfc: L["bottom", "top"] = "bottom", - which_el: L["bottom", "top"] = "top", -) -> tuple[np.ndarray, np.ndarray]: - r"""TODO ... - - Parameters - ---------- - pressure : `array_like[[N, Z] | [Z], floating]` - Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` - Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` - Dewpoint (K) at the levels given by `pressure` - TODO : ... - - Returns - ------- - TODO : ... - - Examples - -------- - TODO : ... - """ - # The mixing ratio of the parcel comes from the dewpoint below the LCL, is saturated - # based on the temperature above the LCL - - parcel_profile = virtual_temperature( - parcel_profile, parcel_mixing_ratio(pressure, temperature, dewpoint) - ) - temperature = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) - # Calculate the EL limit of integration - (EL, _), (LFC, _) = _el_lfc( - "BOTH", - pressure, - temperature, - dewpoint, - parcel_profile, - which_lfc, - which_el, - ) - EL, LFC = np.reshape((EL, LFC), (2, -1, 1)) # reshape for broadcasting - - tzx = F.zero_crossings(pressure, parcel_profile - temperature) # temperature zero crossings - - p, t = tzx.where_between(LFC, EL, close=True) - CAPE = Rd * F.nantrapz(t, np.log(p), axis=1) - CAPE[CAPE < 0.0] = 0.0 - - p, t = tzx.where_below(LFC, close=True) - CIN = Rd * F.nantrapz(t, np.log(p), axis=1) - CIN[CIN > 0.0] = 0.0 - - return CAPE, CIN - - -# ------------------------------------------------------------------------------------------------- -# nzthermo.core.most_unstable -# ------------------------------------------------------------------------------------------------- @broadcast_nz def most_unstable_cape_cin( pressure: Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], @@ -564,15 +603,15 @@ def most_unstable_cape_cin( depth: Pascal[float] = 30000.0, bottom: Pascal[float] | None = None, ) -> tuple[np.ndarray, np.ndarray]: - r"""TODO ... + r"""Calculate most unstable CAPE and CIN. Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` TODO : ... @@ -592,7 +631,7 @@ def most_unstable_cape_cin( bottom, **FASTPATH, ) - mask = np.arange(pressure.shape[1]) >= idx[:, newaxis] + mask = np.arange(pressure.shape[1]) >= idx[:, np.newaxis] p, t, td, mu_profile = parcel_profile_with_lcl( pressure, @@ -605,32 +644,27 @@ def most_unstable_cape_cin( # ------------------------------------------------------------------------------------------------- -# nzthermo.core.mixed_layer +# mixed_layer # ------------------------------------------------------------------------------------------------- -@broadcast_nz def mixed_layer( - pressure: Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - *, - depth: float | NDArray[np.floating[Any]] = 10000.0, + *args: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + depth: ArrayLike = 10000.0, height: ArrayLike | None = None, bottom: ArrayLike | None = None, - interpolate=False, -) -> tuple[ - np.ndarray[shape[N], np.dtype[_T]], - Kelvin[np.ndarray[shape[N], np.dtype[_T]]], -]: - r"""TODO ... + where: ArrayLike | None = None, + interpolate: bool = False, +) -> tuple[np.ndarray[shape[N], np.dtype[_T]], ...]: + r"""Calculate the mean temperature and dewpoint of a mixed layer. Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` TODO : ... @@ -642,15 +676,19 @@ def mixed_layer( -------- TODO : ... """ + pressure, *args = F.exactly_2d(pressure, *args) + if height is not None: raise NotImplementedError("height argument is not implemented") if interpolate: raise NotImplementedError("interpolate argument is not implemented") - bottom = (pressure[surface] if bottom is None else np.asarray(bottom)).reshape(-1, 1) - top = bottom - depth - - where = pressure.is_between(bottom, top) + if where is None: + bottom = (pressure[:, :1] if bottom is None else np.asarray(bottom)).reshape(-1, 1) + top = bottom - np.asarray(depth) + where = pressure.view(pressure_vector).is_between(bottom, top) + else: + where = np.asarray(where, dtype=np.bool_) depth = np.asarray( # use asarray otherwise the depth is cast to pressure_vector which doesn't @@ -659,9 +697,7 @@ def mixed_layer( - np.min(pressure, initial=np.inf, axis=1, where=where) ) - T, Td = F.nantrapz([temperature, dewpoint], pressure, axis=-1, where=where) / -depth - - return T, Td + return tuple(F.nantrapz(args, pressure, axis=-1, where=where) / -depth) @broadcast_nz @@ -681,15 +717,15 @@ def mixed_parcel( np.ndarray[shape[N], np.dtype[_T]], Kelvin[np.ndarray[shape[N], np.dtype[_T]]], ]: - r"""TODO ... + r"""Calculate the mean temperature and dewpoint of a mixed parcel. Parameters ---------- - pressure : `array_like[[N, Z] | [Z], floating]` + pressure : array_like[(N, Z), floating] Atmospheric pressure profile (Pa). This array must be from high to low pressure. - temperature : `array_like[[N, Z], floating]` + temperature : array_like[(N, Z), floating] Temperature (K) at the levels given by `pressure` - dewpoint : `array_like[[N, Z], floating]` + dewpoint : array_like[(N, Z), floating] Dewpoint (K) at the levels given by `pressure` TODO : ... @@ -718,14 +754,11 @@ def mixed_parcel( height=height, depth=depth, interpolate=interpolate, - **FASTPATH, ) mean_temperature = mean_theta * exner_function(parcel_start_pressure) - mean_vapor_pressure = vapor_pressure(parcel_start_pressure, mean_mixing_ratio) - mean_dewpoint = _dewpoint(mean_vapor_pressure) + mean_dewpoint = _dewpoint(vapor_pressure(parcel_start_pressure, mean_mixing_ratio)) return ( - # parcel_start_pressure, np.broadcast_to(parcel_start_pressure, mean_temperature.shape), mean_temperature, mean_dewpoint, diff --git a/src/nzthermo/functional.py b/src/nzthermo/functional.py index 22e98a3..f90710d 100644 --- a/src/nzthermo/functional.py +++ b/src/nzthermo/functional.py @@ -1,6 +1,7 @@ from __future__ import annotations import functools +from typing import NamedTuple, Generic from typing import ( Any, Callable, @@ -22,7 +23,7 @@ from numpy.typing import ArrayLike, NDArray from .typing import N, Z, shape -from .utils import Vector2d, exactly_2d +from .utils import Profile, exactly_2d _T = TypeVar("_T", bound=np.floating[Any]) _P = ParamSpec("_P") @@ -50,14 +51,19 @@ def sort_nz( *args: _P.args, **kwargs: _P.kwargs, ): - sort = ( - np.arange(pressure.shape[0])[:, np.newaxis], - np.argsort(pressure, axis=1, kind="quicksort"), - ) - - pressure = pressure[sort][:, ::-1] - temperature = temperature[sort][:, ::-1] - dewpoint = dewpoint[sort][:, ::-1] + p_sort = np.argsort(pressure, axis=1, kind="quicksort") + if pressure.shape != temperature.shape: + sort = np.s_[:, p_sort] + pressure = pressure[sort][:, ::-1] + sort = np.arange(pressure.shape[0])[:, np.newaxis], p_sort + temperature = temperature[sort][:, ::-1] + dewpoint = dewpoint[sort][:, ::-1] + else: + sort = np.arange(pressure.shape[0])[:, np.newaxis], p_sort + pressure = pressure[sort][:, ::-1] + temperature = temperature[sort][:, ::-1] + dewpoint = dewpoint[sort][:, ::-1] + if callable(where): where = where(pressure, *args, **kwargs) @@ -181,7 +187,7 @@ def intersect_nz( b: np.ndarray[shape[N, Z], np.dtype[_T]], direction: L["increasing", "decreasing"] = "increasing", log_x: bool = False, -) -> Vector2d[_T]: +) -> Profile[_T]: """ NOTE: this function will leave a trailing `nan` value at the end of the array, and clip the array to the length of the longest column. @@ -242,13 +248,13 @@ def intersect_nz( clip = max(np.argmax(np.isnan(x), axis=1)) + 1 - return Vector2d(x[:, :clip], y[:, :clip]) + return Profile(x[:, :clip], y[:, :clip]) def zero_crossings( X: np.ndarray[shape[N, Z], np.dtype[_T]] | np.ndarray[shape[N], np.dtype[_T]], Y: np.ndarray[shape[N, Z], np.dtype[_T]], -) -> Vector2d[_T]: +) -> Profile[_T]: """ This function targets the `metpy.thermo._find_append_zero_crossings` function but with support for 2D arrays. @@ -280,4 +286,25 @@ def zero_crossings( # clip axis 1 to the last non nan value. clip = max(np.argmax(np.isnan(x), axis=1)) + 1 - return Vector2d(x[:, :clip], y[:, :clip]) + return Profile(x[:, :clip], y[:, :clip]) + + +class TopK(NamedTuple, Generic[_T]): + values: NDArray[_T] + indices: NDArray[np.intp] + + +def topk( + values: NDArray[_T], k: int, axis: int = -1, absolute: bool = False, sort: bool = False +) -> TopK[_T]: + values = np.asarray(values) + arg = np.abs(values) if absolute else values + arg[np.isnan(values)] = -np.inf + indices = np.argpartition(arg, -k, axis=axis)[-k:] + values = np.take_along_axis(values, indices, axis=axis) + if sort: + idx = np.flip(np.argsort(values, axis=axis), axis=axis) + indices = np.take_along_axis(indices, idx, axis=axis) + values = np.take_along_axis(values, idx, axis=axis) + + return TopK(values, indices) diff --git a/src/nzthermo/utils.py b/src/nzthermo/utils.py index 257f0c4..dc57886 100644 --- a/src/nzthermo/utils.py +++ b/src/nzthermo/utils.py @@ -80,7 +80,7 @@ class Axis(enum.IntEnum): Z = 1 -class PVectorNd(NamedTuple, Generic[_S, _T]): +class ParcelProfile(NamedTuple, Generic[_S, _T]): pressure: Pascal[np.ndarray[_S, np.dtype[_T]]] temperature: Kelvin[np.ndarray[_S, np.dtype[_T]]] @@ -111,7 +111,7 @@ def where( def is_below( self, pressure: Pascal[NDArray[np.floating[Any]]] | Self, *, close: bool = False ) -> NDArray[np.bool_]: - if isinstance(pressure, PVectorNd): + if isinstance(pressure, ParcelProfile): pressure = pressure.pressure if not close: return self.pressure > pressure @@ -129,9 +129,9 @@ def where_below( return self.where(self.is_below(pressure, close=close), x_fill, y_fill) def is_above( - self, pressure: Pascal[NDArray[np.floating[Any]]] | PVectorNd, *, close: bool = False + self, pressure: Pascal[NDArray[np.floating[Any]]] | ParcelProfile, *, close: bool = False ) -> NDArray[np.bool_]: - if isinstance(pressure, PVectorNd): + if isinstance(pressure, ParcelProfile): pressure = pressure.pressure if not close: return self.pressure < pressure @@ -140,7 +140,7 @@ def is_above( def where_above( self, - pressure: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + pressure: Pascal[NDArray[np.floating[Any]]] | ParcelProfile, x_fill: ArrayLike[np.floating[Any]] = np.nan, y_fill: ArrayLike[np.floating[Any]] | None = None, *, @@ -150,14 +150,14 @@ def where_above( def is_between( self, - bottom: Pascal[NDArray[np.floating[Any]]] | PVectorNd, - top: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + bottom: Pascal[NDArray[np.floating[Any]]] | ParcelProfile, + top: Pascal[NDArray[np.floating[Any]]] | ParcelProfile, *, close: bool = False, ): - if isinstance(bottom, PVectorNd): + if isinstance(bottom, ParcelProfile): bottom = bottom.pressure - if isinstance(top, PVectorNd): + if isinstance(top, ParcelProfile): top = top.pressure if not close: return (self.pressure > bottom) & (self.pressure < top) @@ -166,8 +166,8 @@ def is_between( def where_between( self, - bottom: Pascal[NDArray[np.floating[Any]]] | PVectorNd, - top: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + bottom: Pascal[NDArray[np.floating[Any]]] | ParcelProfile, + top: Pascal[NDArray[np.floating[Any]]] | ParcelProfile, x_fill: ArrayLike[np.floating[Any]] = np.nan, y_fill: ArrayLike[np.floating[Any]] | None = None, *, @@ -216,41 +216,56 @@ def reshape(self, *shape: int) -> tuple[Pascal[NDArray[_T]], Kelvin[NDArray[_T]] return p, t -class Vector1d(PVectorNd[shape[N], _T]): - def unsqueeze(self) -> Vector2d[_T]: +class Parcel(ParcelProfile[shape[N], _T]): + r"""class for containing a (N,) parcel. + + The vertical coordinate for a parcel temperature is the pressure value. + """ + + def unsqueeze(self) -> Profile[_T]: s = np.s_[:, np.newaxis] - return Vector2d(self.pressure[s], self.temperature[s]) + return Profile(self.pressure[s], self.temperature[s]) + + +class Profile(ParcelProfile[shape[N, Z], _T]): + r"""class for containing a (N, Z) profile. + The vertical coordinate for a profile temperature is the pressure value. It is assumed that + pressure is monotonically decreasing with height, and that any nan values are at the top + of the profile ie the end of the array. + """ + + def pick(self, which: L["bottom", "top"]) -> Parcel[_T]: + if which not in {"bottom", "top"}: + raise ValueError(f"which must be either 'bottom' or 'top', got {which!r}") -class Vector2d(PVectorNd[shape[N, Z], _T]): - def pick(self, which: L["bottom", "top"]) -> Vector1d[_T]: p, t = self.pressure, self.temperature nx = np.arange(p.shape[0]) if which == "bottom": idx = np.argmin(~np.isnan(p), axis=1) - 1 # the last non-nan value - return Vector1d(p[nx, idx], t[nx, idx]) + return Parcel(p[nx, idx], t[nx, idx]) elif which == "top": - return Vector1d(p[nx, 0], t[nx, 0]) # the first value is the uppermost value + return Parcel(p[nx, 0], t[nx, 0]) # the first value is the uppermost value - def bottom(self) -> Vector1d[_T]: + def bottom(self) -> Parcel[_T]: return self.pick("bottom") - def top(self) -> Vector1d[_T]: + def top(self) -> Parcel[_T]: return self.pick("top") - def sort(self) -> Self: + def sort(self) -> Profile[_T]: N = self.pressure.shape[0] sort = np.arange(N)[:, np.newaxis], np.argsort(self.pressure, axis=1, kind="quicksort") - return self.__class__(self.pressure[sort], self.temperature[sort]) + return Profile(self.pressure[sort], self.temperature[sort]) -@overload -def exactly_2d(x: NDArray[_T], /) -> np.ndarray[shape[N, Z], np.dtype[_T]]: ... @overload def exactly_2d( *args: NDArray[_T], ) -> tuple[np.ndarray[shape[N, Z], np.dtype[_T]], ...]: ... +@overload +def exactly_2d(__x: NDArray[_T], /) -> np.ndarray[shape[N, Z], np.dtype[_T]]: ... def exactly_2d( *args: NDArray[_T], ) -> np.ndarray[shape[N, Z], np.dtype[_T]] | tuple[np.ndarray[shape[N, Z], np.dtype[_T]], ...]: diff --git a/tests/conftest.py b/tests/conftest.py index 032d26e..8b13789 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,42 +1 @@ -from __future__ import annotations -from typing import Any - -import numpy as np -import pytest -import xarray as xr - -import nzthermo as nzt - -_DATASET = None - - -def get_data(**sel: Any): - global _DATASET - if _DATASET is None: - _DATASET = xr.open_dataset( - "data/hrrr.t00z.wrfprsf00.grib2", - engine="cfgrib", - backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}}, - ) - - ds = _DATASET.sel(**sel) - T = ds["t"].to_numpy() # (K) (Z, Y, X) - Z, Y, X = T.shape - N = Y * X - T = T.reshape(Z, N).transpose() # (N, Z) - P = ds["isobaricInhPa"].to_numpy().astype(np.float32) * 100.0 # (Pa) - Q = ds["q"].to_numpy() # (kg/kg) (Z, Y, X) - Q = Q.reshape(Z, N).transpose() # (N, Z) - Td = nzt.dewpoint_from_specific_humidity(P, Q) - # lat = ds["latitude"].to_numpy() - # lon = (ds["longitude"].to_numpy() + 180) % 360 - 180 - # timestamp = datetime.datetime.fromisoformat(ds["time"].to_numpy().astype(str).item()) - # extent = [lon.min(), lon.max(), lat.min(), lat.max()] - - return (P, T, Td), (Z, Y, X) - - -@pytest.fixture -def isobaric(): - return get_data() diff --git a/tests/core_test.py b/tests/core_test.py index a8d42b7..50007e9 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -371,6 +371,49 @@ def test_parcel_profile_with_lcl_metpy_regression() -> None: # =============================================================================================== # # nzthermo.core # =============================================================================================== # +# ............................................................................................... # +# nzthermo.core.ccl +# ............................................................................................... # +@pytest.mark.ccl +@pytest.mark.parametrize( + "dtype,mixed_layer_depth,which", + itertools.product([np.float64, np.float32], [None, 10000], ["bottom", "top"]), # type: ignore +) +def test_ccl_metpy_regression(dtype, mixed_layer_depth, which) -> None: + if mixed_layer_depth is not None: + # this is dependent on the interpolation of the mixed_layer function which is not implemented + with pytest.raises(NotImplementedError): + CCL_P, CCL_T, CT = ccl( + P.astype(dtype), + T.astype(dtype), + Td.astype(dtype), + mixed_layer_depth=mixed_layer_depth, + which=which, + ) + else: + CCL_P, CCL_T, CT = ccl( + P.astype(dtype), + T.astype(dtype), + Td.astype(dtype), + mixed_layer_depth=mixed_layer_depth, + which=which, + ) + for i in range(T.shape[1]): + CCL_P_, CCL_T_, CT_ = mpcalc.ccl( + P.astype(dtype) * Pa, + T[i].astype(dtype) * K, + Td[i].astype(dtype) * K, + mixed_layer_depth=mixed_layer_depth + if mixed_layer_depth is None + else mixed_layer_depth * Pa, + which=which, + ) + + assert_allclose(CCL_P[i], CCL_P_.m, atol=2) + assert_allclose(CCL_T[i], CCL_T_.m, atol=2) + assert_allclose(CT[i], CT_.m, atol=2) + + # ............................................................................................... # # nzthermo.core.downdraft_cape # ............................................................................................... # @@ -397,42 +440,6 @@ def test_downdraft_cape_metpy_regression(dtype) -> None: assert_allclose(DCAPE[i], DCAPE_, rtol=1e-2) -# ............................................................................................... # -# nzthermo.core.ccl -# ............................................................................................... # -@pytest.mark.ccl -@pytest.mark.parametrize("dtype", [np.float64, np.float32]) -def test_ccl_metpy_regression(dtype) -> None: - P = ( - ([993.0, 957.0, 925.0, 886.0, 850.0, 813.0, 798.0, 732.0, 716.0, 700.0] * units.hPa) - .to("pascal") - .m.astype(dtype) - ) - T = ([34.6, 31.1, 27.8, 24.3, 21.4, 19.6, 18.7, 13, 13.5, 13] * C).to("K").m.astype(dtype) - Td = ([19.6, 18.7, 17.8, 16.3, 12.4, -0.4, -3.8, -6, -13.2, -11] * C).to("K").m.astype(dtype) - ccl_t, ccl_p, ct = ccl(P, T, Td, which="bottom") - - def get_metpy_ccl(p, t, td): - return [x.m for x in mpcalc.ccl(p * Pa, t * K, td * K, which="bottom")] - - assert_allclose( - np.ravel((ccl_t, ccl_p, ct)), - get_metpy_ccl(P, T, Td), - atol=1e-6, - ) - P = np.array([P, P, P]) - T = np.array([T, T - 1, T]) - Td = np.array([Td, Td, Td - 0.5]) - ccl_t, ccl_p, ct = ccl(P, T, Td, which="bottom") - - for i in range(len(P)): - assert_allclose( - (ccl_t[i], ccl_p[i], ct[i]), - get_metpy_ccl(P[i], T[i], Td[i]), - atol=1e-6, - ) - - # ............................................................................................... # # nzthermo.core.el # ............................................................................................... # @@ -1065,9 +1072,9 @@ def test_el_profile_nan_with_parcel_profile() -> None: dewpoints.to(K).m, parcel_temps, ) - print(el_pressure, el_temperature) - # assert_almost_equal(el_pressure, 673.0104 * hPa, 3) - # assert_almost_equal(el_temperature, 5.8853 * C, 3) + + assert_almost_equal(el_pressure, 673.0104 * hPa, 3) + assert_almost_equal(el_temperature, 5.8853 * C, 3) @pytest.mark.el @@ -1304,7 +1311,6 @@ def test_most_unstable_parcel_index_broadcasting(depth) -> None: assert_array_equal( most_unstable_parcel_index(P, T, Td, depth=depth), most_unstable_parcel_index(np.broadcast_to(P, T.shape), T, Td, depth=depth), - err_msg="most_unstable_parcel_index failed to on broadcasted pressure input.", ) @@ -1331,7 +1337,6 @@ def test_most_unstable_parcel_broadcasting(depth) -> None: assert_array_equal( most_unstable_parcel(P, T, Td, depth=depth), most_unstable_parcel(np.broadcast_to(P, T.shape), T, Td, depth=depth), - err_msg="most_unstable_parcel failed to on broadcasted pressure input.", ) @@ -1356,15 +1361,11 @@ def test_most_unstable_parcel_regression(depth) -> None: # ............................................................................................... # @pytest.mark.broadcasting @pytest.mark.most_unstable_cape_cin -def test_most_unstable_cape_cin_broadcasting(): +@pytest.mark.parametrize("depth", [30000.0]) +def test_most_unstable_cape_cin_broadcasting(depth) -> None: assert_array_equal( - most_unstable_cape_cin(P, T, Td, depth=30000.0), - most_unstable_cape_cin( - np.broadcast_to(P, T.shape), - T, - Td, - depth=30000.0, - ), + most_unstable_cape_cin(P, T, Td, depth=depth), + most_unstable_cape_cin(np.broadcast_to(P, T.shape), T, Td, depth=depth), ) @@ -1399,41 +1400,33 @@ def test_most_unstable_cape_cin_metpy_regression(depth) -> None: @pytest.mark.broadcasting @pytest.mark.mixed_layer def test_mixed_layer_broadcasting() -> None: - """ - NOTE: using assert_array_equal I'm not entirely sure wy broadcasting the pressure - is causing causing some 1e-5 differences in the results, but atol of 1e-5 is well within - and acceptable range for the test to pass. - - ```bash - E Mismatched elements: 233 / 1080 (21.6%) - E Max absolute difference among violations: 0.000031 - E Max relative difference among violations: 0. - ``` - """ - - assert_allclose( + assert_array_equal( mixed_layer(P, T, Td), mixed_layer(np.broadcast_to(P, T.shape), T, Td), - atol=TEMPERATURE_ABSOLUTE_TOLERANCE, ) @pytest.mark.regression @pytest.mark.mixed_layer -def test_mixed_layer_regression() -> None: - t, td = mixed_layer(P, T, Td) - for i in range(T.shape[0]): - t_, td_ = mpcalc.mixed_layer(P * Pa, T[i] * K, Td[i] * K, interpolate=False) - assert_allclose( - t[i], - t_.m, - atol=TEMPERATURE_ABSOLUTE_TOLERANCE, - ) - assert_allclose( - td[i], - td_.m, - atol=TEMPERATURE_ABSOLUTE_TOLERANCE, - ) +@pytest.mark.parametrize("interpolate", [True, False]) +def test_mixed_layer_regression(interpolate) -> None: + if interpolate: + with pytest.raises(NotImplementedError): + mixed_layer(P, T, Td, interpolate=interpolate) + else: + t, td = mixed_layer(P, T, Td, interpolate=interpolate) + for i in range(T.shape[0]): + t_, td_ = mpcalc.mixed_layer(P * Pa, T[i] * K, Td[i] * K, interpolate=interpolate) + assert_allclose( + t[i], + t_.m, + atol=TEMPERATURE_ABSOLUTE_TOLERANCE, + ) + assert_allclose( + td[i], + td_.m, + atol=TEMPERATURE_ABSOLUTE_TOLERANCE, + ) # ............................................................................................... # @@ -1442,22 +1435,9 @@ def test_mixed_layer_regression() -> None: @pytest.mark.broadcasting @pytest.mark.mixed_parcel def test_mixed_parcel_broadcasting() -> None: - """ - NOTE: using assert_array_equal I'm not entirely sure wy broadcasting the pressure - is causing causing some 1e-5 differences in the results, but atol of 1e-5 is well within - and acceptable range for the test to pass. - - ```bash - E Mismatched elements: 233 / 1080 (21.6%) - E Max absolute difference among violations: 0.000031 - E Max relative difference among violations: 0. - ``` - """ - - assert_allclose( + assert_array_equal( mixed_parcel(P, T, Td), mixed_parcel(np.broadcast_to(P, T.shape), T, Td), - atol=TEMPERATURE_ABSOLUTE_TOLERANCE, ) @@ -1489,41 +1469,21 @@ def test_mixed_parcel_regression() -> None: # ............................................................................................... # # nzthermo.core.mixed_layer_cape_cin # ............................................................................................... # -# @pytest.mark.skip @pytest.mark.broadcasting @pytest.mark.mixed_layer_cape_cin def test_mixed_layer_cape_cin_broadcasting() -> None: - """ - NOTE: using assert_array_equal I'm not entirely sure wy broadcasting the pressure - is causing causing some 1e-5 differences in the results, but atol of 1e-5 is well within - and acceptable range for the test to pass. - - ```bash - E Mismatched elements: 233 / 1080 (21.6%) - E Max absolute difference among violations: 0.000031 - E Max relative difference among violations: 0. - ``` - """ - - assert_allclose( + assert_array_equal( mixed_layer_cape_cin(P, T, Td), mixed_layer_cape_cin(np.broadcast_to(P, T.shape), T, Td), - atol=TEMPERATURE_ABSOLUTE_TOLERANCE, ) -# @pytest.mark.skip @pytest.mark.regression @pytest.mark.mixed_layer_cape_cin def test_mixed_layer_cape_cin_regression() -> None: CAPE, CIN = mixed_layer_cape_cin(P, T, Td) - CAPE_ = [] - CIN_ = [] for i in range(T.shape[0]): - cape_, cin_ = mpcalc.mixed_layer_cape_cin(P * Pa, T[i] * K, Td[i] * K, interpolate=False) - CAPE_.append(cape_.m) - CIN_.append(cin_.m) - - assert_allclose(CAPE, CAPE_, atol=20) - assert_allclose(CIN, CIN_, atol=200) + CAPE_, CIN__ = mpcalc.mixed_layer_cape_cin(P * Pa, T[i] * K, Td[i] * K, interpolate=False) + assert_allclose(CAPE[i], CAPE_.m, atol=20) + assert_allclose(CIN[i], CIN__.m, atol=200)