From 3f241cc773991fab676a2aff9adf4777cd23eabb Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Mon, 17 Jun 2024 20:11:23 -0500 Subject: [PATCH 1/6] =?UTF-8?q?=E2=9B=88=EF=B8=8F=20clean=20up?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/nzthermo/_C.pxd | 6 +++--- src/nzthermo/_core.pyx | 2 +- src/nzthermo/core.py | 30 +++++++++++++++--------------- src/nzthermo/functional.py | 9 --------- src/nzthermo/utils.py | 27 ++++++++++----------------- 5 files changed, 29 insertions(+), 45 deletions(-) diff --git a/src/nzthermo/_C.pxd b/src/nzthermo/_C.pxd index a085711..cbd3b6f 100644 --- a/src/nzthermo/_C.pxd +++ b/src/nzthermo/_C.pxd @@ -36,6 +36,9 @@ cdef inline size_t search_pressure(Float[:] pressure, Float value) noexcept nogi cdef extern from "wind.cpp" namespace "libthermo" nogil: + T wind_direction[T](T, T) noexcept + T wind_magnitude[T](T, T) noexcept + cdef cppclass wind_components[T]: T u, v wind_components() noexcept @@ -48,9 +51,6 @@ cdef extern from "wind.cpp" namespace "libthermo" nogil: wind_vector(T, T) noexcept wind_vector(wind_components[T]) noexcept - T wind_direction[T](T, T) noexcept - T wind_magnitude[T](T, T) noexcept - cdef extern from "libthermo.cpp" namespace "libthermo" nogil: const double T0 # `(J/kg*K)` - freezing point in kelvin diff --git a/src/nzthermo/_core.pyx b/src/nzthermo/_core.pyx index 1687c67..8c732fd 100644 --- a/src/nzthermo/_core.pyx +++ b/src/nzthermo/_core.pyx @@ -289,7 +289,7 @@ def moist_lapse( [1013.93, 1000, 975, 950, 925, 900, ...], [np.nan, np.nan, 975, 950, 925, 900, ...] ]) * 100.0 # (N, Z) :: pressure profile - >>> reference_pressure = pressure[np.arange(len(pressure)), np.argmin(np.isnan(pressure), axis=1)] + >>> reference_pressure = pressure[np.arange(pressure.shape[0]), np.argmin(np.isnan(pressure), axis=1)] >>> reference_pressure array([101312., 101393., 97500. ]) """ diff --git a/src/nzthermo/core.py b/src/nzthermo/core.py index bddd251..f0ecf0a 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -179,22 +179,22 @@ def _el_lfc( temperature[:, 1:], ) - top_idx = np.arange(N), np.argmin(~np.isnan(pressure), axis=1) - 1 - # find the Equilibrium Level (EL) - left_of_env = (parcel_temperature_profile[top_idx] <= temperature[top_idx])[:, newaxis] - EL = F.intersect_nz( - pressure, - parcel_temperature_profile, - temperature, - "decreasing", - log_x=True, - ).where( - # If the top of the sounding parcel is warmer than the environment, there is no EL - lambda el: el.is_above(LCL) & left_of_env - ) + if pick != "LFC": # find the Equilibrium Level (EL) + top_idx = np.arange(N), np.argmin(~np.isnan(pressure), axis=1) - 1 + left_of_env = (parcel_temperature_profile[top_idx] <= temperature[top_idx])[:, newaxis] + EL = F.intersect_nz( + pressure, + parcel_temperature_profile, + temperature, + "decreasing", + log_x=True, + ).where( + # If the top of the sounding parcel is warmer than the environment, there is no EL + lambda el: el.is_above(LCL) & left_of_env + ) - if pick == "EL": - return EL.pick(which_el) + if pick == "EL": + return EL.pick(which_el) LFC = F.intersect_nz( pressure, diff --git a/src/nzthermo/functional.py b/src/nzthermo/functional.py index cab81dc..06db101 100644 --- a/src/nzthermo/functional.py +++ b/src/nzthermo/functional.py @@ -209,12 +209,3 @@ def zero_crossings( clip = max(np.argmax(np.isnan(x), axis=1)) return Vector2d(x[:, :clip], y[:, :clip]) - - -def pressure_top( - n: int, pressure: np.ndarray[shape[N, Z], np.dtype[_T]] -) -> tuple[ - np.ndarray[shape[N], np.dtype[np.int_]], - np.ndarray[shape[Z], np.dtype[np.int_]], -]: - return np.arange(n)[:, np.newaxis], np.nanargmin(~np.isnan(pressure), axis=1) - 1 diff --git a/src/nzthermo/utils.py b/src/nzthermo/utils.py index 975b4e1..28ceaff 100644 --- a/src/nzthermo/utils.py +++ b/src/nzthermo/utils.py @@ -178,15 +178,11 @@ def __str__(self) -> str: @classmethod def from_func( cls, - func: Callable[ - _P, - tuple[Pascal[NDArray[float_]], Kelvin[NDArray[float_]]], - ], + func: Callable[_P, tuple[Pascal[NDArray[float_]], Kelvin[NDArray[float_]]]], *args: _P.args, **kwargs: _P.kwargs, ) -> Self: - x, y = func(*args, **kwargs) - return cls(x, y) + return cls(*func(*args, **kwargs)) def reshape(self, *shape: int) -> tuple[Pascal[NDArray[float_]], Kelvin[NDArray[float_]]]: p, t = np.reshape([self.pressure, self.temperature], (2, *shape)) @@ -262,35 +258,32 @@ def broadcast_nz( ], ) -> Callable[ Concatenate[ - Pascal[ArrayLike[np.floating[Any]]], - Kelvin[ArrayLike[np.floating[Any]]], - Kelvin[ArrayLike[np.floating[Any]]], - _P, + Pascal[ArrayLike[float_]], Kelvin[ArrayLike[float_]], Kelvin[ArrayLike[float_]], _P ], _T, ]: @functools.wraps(f) def wrapper( - pressure: ArrayLike[np.floating[Any]], - temperature: ArrayLike[np.floating[Any]], - dewpoint: ArrayLike[np.floating[Any]], + pressure: ArrayLike[float_], + temperature: ArrayLike[float_], + dewpoint: ArrayLike[float_], *args: _P.args, **kwargs: _P.kwargs, ) -> _T: + if kwargs.pop("__fastpath", False): + return f(pressure, temperature, dewpoint, *args, **kwargs) # type: ignore + # TODO # - add support for squeezing what would have been a 1d input # - add support for reshaping: # (T, Z, Y, X) -> (N, Z) # (Z, Y, X) -> (N, Z)' - if kwargs.pop("__fastpath", False): - return f(pressure, temperature, dewpoint, *args, **kwargs) # type: ignore - pressure, temperature, dewpoint = exactly_2d( magnitude(pressure, "pascal"), magnitude(temperature, "kelvin"), magnitude(dewpoint, "kelvin"), ) - return f(pressure, temperature, dewpoint, *args, **kwargs) # type: ignore + return f(pressure, temperature, dewpoint, *args, **kwargs) return wrapper From ead01c799244ffdd0276ac1183d8fed2d2aac129 Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Tue, 18 Jun 2024 09:26:20 -0500 Subject: [PATCH 2/6] =?UTF-8?q?=E2=9B=88=EF=B8=8F=20type=20annotations=20a?= =?UTF-8?q?nd=20pressure=20search?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/include/libthermo.hpp | 12 +-- src/lib/libthermo.cpp | 49 +++++++++-- src/nzthermo/_C.pxd | 18 +--- src/nzthermo/_core.pyx | 15 ++-- src/nzthermo/core.py | 84 +++++++++--------- src/nzthermo/utils.py | 175 +++++++++++++++++--------------------- 6 files changed, 176 insertions(+), 177 deletions(-) diff --git a/src/include/libthermo.hpp b/src/include/libthermo.hpp index 93e081a..792a1a6 100644 --- a/src/include/libthermo.hpp +++ b/src/include/libthermo.hpp @@ -28,13 +28,12 @@ static constexpr double Md = 28.96546; // `(g/mol)` - molecular weight of dry a static constexpr double epsilon = Mw / Md; // `Mw / Md` - molecular weight ratio static constexpr double kappa = Rd / Cp; // `Rd / Cp` - ratio of gas constants -/* ........................................{ struct }........................................... */ +/* ........................................{ helper }........................................... */ + template -struct Parcel { - T pressure; - T temperature; - T dewpoint; -}; +constexpr size_t index_pressure(const T x[], const T value, const size_t size) noexcept; + +/* ........................................{ struct }........................................... */ template constexpr T mixing_ratio(const T partial_press, const T total_press) noexcept; @@ -111,6 +110,7 @@ class lcl { ) 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; }; template diff --git a/src/lib/libthermo.cpp b/src/lib/libthermo.cpp index ca08125..1aea9b9 100644 --- a/src/lib/libthermo.cpp +++ b/src/lib/libthermo.cpp @@ -2,6 +2,39 @@ 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 + * @return constexpr size_t + */ +template +constexpr size_t index_pressure(const T levels[], const T value, const size_t size) noexcept { + const size_t N = size - 1; + const T p0 = levels[0]; + const T p1 = levels[N]; + const T step = ((p1 - p0) / N); + + size_t idx = (size_t)((value / step) - (p0 / step)); + + if (idx >= N) + return N; + + while ((idx < N) && (value < levels[idx])) + idx++; + + return idx; +} + +// thermodynamic functions + template constexpr T mixing_ratio(const T partial_press, const T total_press) noexcept { return epsilon * partial_press / (total_press - partial_press); @@ -141,14 +174,11 @@ constexpr T lcl_pressure( template constexpr lcl::lcl( - const T pressure, const T temperature, const T dewpoint, const T eps, const size_t max_iters + const T pressure_, const T temperature_, const T dewpoint_, const T eps, const size_t max_iters ) noexcept { - const T r = mixing_ratio(saturation_vapor_pressure(dewpoint), pressure); - const T lcl_p = fixed_point(find_lcl, max_iters, eps, pressure, temperature, r); - const T lcl_t = libthermo::dewpoint(lcl_p, r); - - this->pressure = lcl_p; - this->temperature = lcl_t; + const T r = mixing_ratio(saturation_vapor_pressure(dewpoint_), pressure_); + pressure = fixed_point(find_lcl, max_iters, eps, pressure_, temperature_, r); + temperature = dewpoint(pressure, r); } template @@ -156,6 +186,11 @@ constexpr T lcl::wet_bulb_temperature(const T pressure, const T step) noexcep return moist_lapse(this->pressure, pressure, this->temperature, step); } +template +constexpr size_t lcl::index(const T pressure[], size_t size) noexcept { + return index_pressure(pressure, this->pressure, size); +} + template constexpr T wet_bulb_temperature( const T pressure, diff --git a/src/nzthermo/_C.pxd b/src/nzthermo/_C.pxd index cbd3b6f..d0bf950 100644 --- a/src/nzthermo/_C.pxd +++ b/src/nzthermo/_C.pxd @@ -1,16 +1,6 @@ # cython: boundscheck=False # pyright: reportGeneralTypeIssues=false -ctypedef fused Float: - float - double - - -cdef extern from "" namespace "std" nogil: - cdef cppclass pair[T, U]: - T pressure "first" - U temperature "second" - cdef extern from "functional.cpp" namespace "libthermo" nogil: cdef cppclass point[T]: @@ -28,13 +18,6 @@ cdef extern from "functional.cpp" namespace "libthermo" nogil: size_t search_sorted[T](T* x, T value, size_t size, bint inverted) noexcept -cdef inline size_t search_pressure(Float[:] pressure, Float value) noexcept nogil: - cdef size_t Z = pressure.shape[0] - if pressure[Z - 1] > value: - return Z - return search_sorted(&pressure[0], value, Z, True) - - cdef extern from "wind.cpp" namespace "libthermo" nogil: T wind_direction[T](T, T) noexcept T wind_magnitude[T](T, T) noexcept @@ -71,6 +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 # 1x1 T saturation_vapor_pressure[T](T temperature) noexcept diff --git a/src/nzthermo/_core.pyx b/src/nzthermo/_core.pyx index 8c732fd..b7aa19d 100644 --- a/src/nzthermo/_core.pyx +++ b/src/nzthermo/_core.pyx @@ -375,7 +375,6 @@ cdef void parcel_profile_1d( T[:] pressure, # (Z,) T temperature, T dewpoint, - ) noexcept nogil: cdef: size_t Z, i, stop @@ -387,10 +386,9 @@ cdef void parcel_profile_1d( t0 = out[0] = temperature lcl = C.lcl[T](p0, t0, dewpoint) - # [dry ascent] # stop the dry ascent at the LCL - stop = C.search_pressure(pressure, lcl.pressure) + stop = lcl.index(&pressure[0], Z) for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent out[i] = C.dry_lapse(pressure[i], p0, t0) @@ -497,16 +495,13 @@ 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 = C.search_pressure(pressure, lcl.pressure) - for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent - pt[i] = C.dry_lapse(pressure[i], p0, t0) - + stop = lcl.index(&pressure[0], Z) ep[:stop] = pressure[:stop] et[:stop] = temperature[:stop] etd[:stop] = dewpoint[:stop] + for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent + pt[i] = C.dry_lapse(pressure[i], p0, t0) # [ lcl ] ep[stop] = lcl.pressure @@ -525,12 +520,12 @@ cdef void parcel_profile_with_lcl_1d( dewpoint[stop] ) pt[stop] = lcl.temperature + # [ moist ascent ] .. parcel temperature from the LCL to the top of the atmosphere .. if stop != Z: ep[stop + 1:] = pressure[stop:] et[stop + 1:] = temperature[stop:] etd[stop + 1:] = dewpoint[stop:] - moist_lapse_1d(pt[stop + 1:], pressure[stop:], lcl.pressure, lcl.temperature) diff --git a/src/nzthermo/core.py b/src/nzthermo/core.py index f0ecf0a..5cefcf4 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -32,13 +32,11 @@ ) from .const import Rd from .typing import Kelvin, N, Pascal, Z, shape -from .utils import Vector1d, broadcast_nz +from .utils import Vector1d, broadcast_nz, Axis -float_ = TypeVar("float_", bound=np.floating[Any], covariant=True) +_T = TypeVar("_T", bound=np.floating[Any], covariant=True) newaxis: Final[None] = np.newaxis NaN = np.nan -z_axis: Final[tuple[slice, None]] = np.s_[:, newaxis] -N_AXIS: Final[tuple[None, slice]] = np.s_[newaxis, :] FASTPATH: dict[str, Any] = {"__fastpath": True} @@ -49,10 +47,10 @@ # ------------------------------------------------------------------------------------------------- @broadcast_nz def downdraft_cape( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], -) -> np.ndarray[shape[N], np.dtype[float_]]: + pressure: Pascal[np.ndarray[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]]], +) -> np.ndarray[shape[N], np.dtype[_T]]: """Calculate downward CAPE (DCAPE). Calculate the downward convective available potential energy (DCAPE) of a given upper air @@ -109,9 +107,9 @@ def downdraft_cape( # ------------------------------------------------------------------------------------------------- @broadcast_nz def ccl( - pressure: Pascal[NDArray[float_]], - temperature: Kelvin[NDArray[float_]], - dewpoint: Kelvin[NDArray[float_]], + pressure: Pascal[NDArray[_T]], + temperature: Kelvin[NDArray[_T]], + dewpoint: Kelvin[NDArray[_T]], /, height=None, mixed_layer_depth=None, @@ -149,16 +147,16 @@ def ccl( # ------------------------------------------------------------------------------------------------- def _el_lfc( pick: L["EL", "LFC", "BOTH"], - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + pressure: Pascal[np.ndarray[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_temperature_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", - dewpoint_start: np.ndarray[shape[N], np.dtype[float_]] | None = None, -) -> tuple[Vector1d[float_], Vector1d[float_]] | Vector1d[float_]: + dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, +) -> tuple[Vector1d[_T], Vector1d[_T]] | Vector1d[_T]: if parcel_temperature_profile is None: pressure, temperature, dewpoint, parcel_temperature_profile = core.parcel_profile_with_lcl( pressure, temperature, dewpoint @@ -180,7 +178,7 @@ def _el_lfc( ) if pick != "LFC": # find the Equilibrium Level (EL) - top_idx = np.arange(N), np.argmin(~np.isnan(pressure), axis=1) - 1 + top_idx = np.arange(N), np.argmin(~np.isnan(pressure), Axis.Z) - 1 left_of_env = (parcel_temperature_profile[top_idx] <= temperature[top_idx])[:, newaxis] EL = F.intersect_nz( pressure, @@ -204,11 +202,13 @@ def _el_lfc( log_x=True, ).where_above(LCL) - is_lcl = (no_lfc := LFC.is_nan().all(axis=1, keepdims=True)) & greater_or_close( + no_lfc = LFC.is_nan().all(Axis.Z, out=np.empty((N, 1), dtype=np.bool_), keepdims=True) + + 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_temperature_profile, NaN), temperature, - ).any(axis=1, keepdims=True) + ).any(Axis.Z, out=np.empty((N, 1), dtype=np.bool_), keepdims=True) LFC = LFC.select( [~no_lfc, is_lcl], @@ -224,14 +224,14 @@ def _el_lfc( @broadcast_nz def el( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + pressure: Pascal[np.ndarray[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_temperature_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, which: L["top", "bottom"] = "top", -) -> Vector1d[float_]: +) -> Vector1d[_T]: return _el_lfc( "EL", pressure, temperature, dewpoint, parcel_temperature_profile, which_el=which ) @@ -239,14 +239,14 @@ def el( @broadcast_nz def lfc( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + pressure: Pascal[np.ndarray[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_temperature_profile: np.ndarray | None = None, which: L["top", "bottom"] = "top", - dewpoint_start: np.ndarray[shape[N], np.dtype[float_]] | None = None, -) -> Vector1d[float_]: + dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, +) -> Vector1d[_T]: return _el_lfc( "LFC", pressure, @@ -260,16 +260,16 @@ def lfc( @broadcast_nz def el_lfc( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + pressure: Pascal[np.ndarray[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_temperature_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", - dewpoint_start: np.ndarray[shape[N], np.dtype[float_]] | None = None, -) -> tuple[Vector1d[float_], Vector1d[float_]]: + dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, +) -> tuple[Vector1d[_T], Vector1d[_T]]: return _el_lfc( "BOTH", pressure, @@ -310,9 +310,9 @@ def most_unstable_parcel_index( @broadcast_nz def cape_cin( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + pressure: Pascal[np.ndarray[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", @@ -362,16 +362,16 @@ def cape_cin( @broadcast_nz def most_unstable_parcel( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + pressure: Pascal[np.ndarray[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]]], /, depth: Pascal[float] = 30_000.0, bottom: Pascal[float] | None = None, ) -> tuple[ - Pascal[np.ndarray[shape[N], np.dtype[float_]]], - Kelvin[np.ndarray[shape[N], np.dtype[float_]]], - Kelvin[np.ndarray[shape[N], np.dtype[float_]]], + Pascal[np.ndarray[shape[N], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N], np.dtype[_T]]], np.ndarray[shape[N, Z], np.dtype[np.intp]], ]: depth = 100.0 if depth is None else depth diff --git a/src/nzthermo/utils.py b/src/nzthermo/utils.py index 28ceaff..de6ec2f 100644 --- a/src/nzthermo/utils.py +++ b/src/nzthermo/utils.py @@ -1,6 +1,6 @@ from __future__ import annotations -import abc +import enum import functools from typing import ( TYPE_CHECKING, @@ -46,22 +46,22 @@ ) -_T = TypeVar("_T") +_S = TypeVar("_S") _P = ParamSpec("_P") -float_ = TypeVar("float_", bound=np.floating[Any], covariant=True) +_T = TypeVar("_T", bound=np.floating[Any], covariant=True) ArrayLike: TypeAlias = ( - "SupportsArray[float_] | NestedSequence[SupportsArray[float_]] | NestedSequence[float] | float" + "SupportsArray[_T] | NestedSequence[SupportsArray[_T]] | NestedSequence[float] | float" ) if TYPE_CHECKING: - def magnitude(x: ArrayLike[float_], unit: str) -> NDArray[float_]: ... + def magnitude(x: ArrayLike[_T], unit: str) -> NDArray[_T]: ... elif pint is not None: def magnitude(x, unit): if isinstance(x, pint.Quantity): - return x.to(unit).magnitude + x = x.to(unit).magnitude return np.asarray(x) else: @@ -69,20 +69,14 @@ def magnitude(x, unit): return np.asarray(x) -class pressure_vector(abc.ABC): - @abc.abstractmethod - def is_below( - self, pressure: Pascal[NDArray[np.floating[Any]]], *, close: bool = False - ) -> NDArray[np.bool_]: ... - @abc.abstractmethod - def is_above( - self, pressure: Pascal[NDArray[np.floating[Any]]], *, close: bool = False - ) -> NDArray[np.bool_]: ... +class Axis(enum.IntEnum): + N = 0 + Z = 1 -class PVectorNd(NamedTuple, Generic[_T, float_]): - pressure: Pascal[np.ndarray[_T, np.dtype[float_]]] - temperature: Kelvin[np.ndarray[_T, np.dtype[float_]]] +class PVectorNd(NamedTuple, Generic[_S, _T]): + pressure: Pascal[np.ndarray[_S, np.dtype[_T]]] + temperature: Kelvin[np.ndarray[_S, np.dtype[_T]]] def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): if method == "__call__": @@ -92,8 +86,8 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): def where( self, - condition: np.ndarray[_T, np.dtype[np.bool_]] - | Callable[[Self], np.ndarray[_T, np.dtype[np.bool_]]], + condition: np.ndarray[_S, np.dtype[np.bool_]] + | Callable[[Self], np.ndarray[_S, np.dtype[np.bool_]]], x_fill: ArrayLike[np.floating[Any]] = np.nan, y_fill: ArrayLike[np.floating[Any]] | None = None, ) -> Self: @@ -109,7 +103,7 @@ def where( ) def is_below( - self, pressure: Pascal[NDArray[np.floating[Any]]] | PVectorNd, *, close: bool = False + self, pressure: Pascal[NDArray[np.floating[Any]]] | Self, *, close: bool = False ) -> NDArray[np.bool_]: if isinstance(pressure, PVectorNd): pressure = pressure.pressure @@ -120,7 +114,7 @@ def is_below( def where_below( self, - pressure: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + pressure: Pascal[NDArray[np.floating[Any]]] | Self, x_fill: ArrayLike[np.floating[Any]] = np.nan, y_fill: ArrayLike[np.floating[Any]] | None = None, *, @@ -154,8 +148,8 @@ def is_nan(self) -> NDArray[np.bool_]: def select( self, condlist: Sequence[NDArray[np.bool_]], - x_choice: Sequence[NDArray[float_]], - y_choice: Sequence[NDArray[float_]], + x_choice: Sequence[NDArray[_T]], + y_choice: Sequence[NDArray[_T]], x_default: ArrayLike[np.floating[Any]] = np.nan, y_default: ArrayLike[np.floating[Any]] | None = None, ) -> Self: @@ -178,25 +172,25 @@ def __str__(self) -> str: @classmethod def from_func( cls, - func: Callable[_P, tuple[Pascal[NDArray[float_]], Kelvin[NDArray[float_]]]], + func: Callable[_P, tuple[Pascal[NDArray[_T]], Kelvin[NDArray[_T]]]], *args: _P.args, **kwargs: _P.kwargs, ) -> Self: return cls(*func(*args, **kwargs)) - def reshape(self, *shape: int) -> tuple[Pascal[NDArray[float_]], Kelvin[NDArray[float_]]]: + def reshape(self, *shape: int) -> tuple[Pascal[NDArray[_T]], Kelvin[NDArray[_T]]]: p, t = np.reshape([self.pressure, self.temperature], (2, *shape)) return p, t -class Vector1d(PVectorNd[shape[N], float_]): - def unsqueeze(self) -> Vector2d[float_]: +class Vector1d(PVectorNd[shape[N], _T]): + def unsqueeze(self) -> Vector2d[_T]: s = np.s_[:, np.newaxis] return Vector2d(self.pressure[s], self.temperature[s]) -class Vector2d(PVectorNd[shape[N, Z], float_]): - def pick(self, which: L["bottom", "top"]) -> Vector1d[float_]: +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": @@ -206,10 +200,10 @@ def pick(self, which: L["bottom", "top"]) -> Vector1d[float_]: elif which == "top": return Vector1d(p[nx, 0], t[nx, 0]) # the first value is the uppermost value - def bottom(self) -> Vector1d[float_]: + def bottom(self) -> Vector1d[_T]: return self.pick("bottom") - def top(self) -> Vector1d[float_]: + def top(self) -> Vector1d[_T]: return self.pick("top") def sort(self) -> Self: @@ -219,17 +213,14 @@ def sort(self) -> Self: @overload -def exactly_2d(x: NDArray[float_], /) -> np.ndarray[shape[N, Z], np.dtype[float_]]: ... +def exactly_2d(x: NDArray[_T], /) -> np.ndarray[shape[N, Z], np.dtype[_T]]: ... @overload def exactly_2d( - *args: NDArray[float_], -) -> tuple[np.ndarray[shape[N, Z], np.dtype[float_]], ...]: ... + *args: NDArray[_T], +) -> tuple[np.ndarray[shape[N, Z], np.dtype[_T]], ...]: ... def exactly_2d( - *args: NDArray[float_], -) -> ( - np.ndarray[shape[N, Z], np.dtype[float_]] - | tuple[np.ndarray[shape[N, Z], np.dtype[float_]], ...] -): + *args: NDArray[_T], +) -> np.ndarray[shape[N, Z], np.dtype[_T]] | tuple[np.ndarray[shape[N, Z], np.dtype[_T]], ...]: values = [] for x in args: if x.ndim == 0: @@ -249,27 +240,25 @@ def exactly_2d( def broadcast_nz( f: Callable[ Concatenate[ - Pascal[np.ndarray[shape[N, Z], np.dtype[float_]]], - Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], + Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], _P, ], - _T, + _S, ], ) -> Callable[ - Concatenate[ - Pascal[ArrayLike[float_]], Kelvin[ArrayLike[float_]], Kelvin[ArrayLike[float_]], _P - ], - _T, + Concatenate[Pascal[ArrayLike[_T]], Kelvin[ArrayLike[_T]], Kelvin[ArrayLike[_T]], _P], + _S, ]: @functools.wraps(f) def wrapper( - pressure: ArrayLike[float_], - temperature: ArrayLike[float_], - dewpoint: ArrayLike[float_], + pressure: ArrayLike[_T], + temperature: ArrayLike[_T], + dewpoint: ArrayLike[_T], *args: _P.args, **kwargs: _P.kwargs, - ) -> _T: + ) -> _S: if kwargs.pop("__fastpath", False): return f(pressure, temperature, dewpoint, *args, **kwargs) # type: ignore @@ -298,39 +287,39 @@ def wrapper( julian = np.float64 datetime64 = np.datetime64 +_T1 = TypeVar("_T1", datetime64, calendar, julian, unix) +_T2 = TypeVar("_T2", datetime64, calendar, julian, unix) -_DT1 = TypeVar("_DT1", datetime64, calendar, julian, unix) -_DT2 = TypeVar("_DT2", datetime64, calendar, julian, unix) -DType_DT1 = np.dtype[_DT1] | type[_DT1] | SupportsDType[_DT1] +DType_T = np.dtype[_T1] | type[_T1] | SupportsDType[_T1] @overload -def isdtype(x: NDArray[Any], /, dtype: DType_DT1[_DT1]) -> TypeGuard[NDArray[_DT1]]: ... +def isdtype(x: NDArray[Any], /, dtype: DType_T[_T1]) -> TypeGuard[NDArray[_T1]]: ... @overload -def isdtype(x: np.dtype[_DT1], /, dtype: DType_DT1[_DT1]) -> TypeGuard[np.dtype[_DT1]]: ... +def isdtype(x: np.dtype[_T1], /, dtype: DType_T[_T1]) -> TypeGuard[np.dtype[_T1]]: ... @overload -def isdtype(x: type[_DT1], /, dtype: DType_DT1[_DT1]) -> TypeGuard[np.dtype[_DT1]]: ... +def isdtype(x: type[_T1], /, dtype: DType_T[_T1]) -> TypeGuard[np.dtype[_T1]]: ... def isdtype( - x: NDArray[Any] | np.dtype[_DT1] | type[_DT1], /, dtype: DType_DT1[_DT1] -) -> TypeGuard[NDArray[_DT1]] | TypeGuard[np.dtype[_DT1]]: + x: NDArray[Any] | np.dtype[_T1] | type[_T1], /, dtype: DType_T[_T1] +) -> TypeGuard[NDArray[_T1]] | TypeGuard[np.dtype[_T1]]: if isinstance(x, np.dtype): - x_dtype = x + arg = x elif isinstance(x, type): - x_dtype = np.dtype(x) + arg = np.dtype(x) else: - x_dtype = x.dtype + arg = x.dtype if dtype is calendar: return ( - np.issubdtype(x_dtype, calendar) - and x_dtype.names is not None + np.issubdtype(arg, calendar) + and arg.names is not None and ( {"year", "month", "day", "hour", "minute", "second", "microsecond"}.issubset( - x_dtype.names + arg.names ) ) ) - return np.issubdtype(x_dtype, dtype) + return np.issubdtype(arg, dtype) def leap_year(x: NDArray[datetime64 | calendar | julian | unix]) -> NDArray[np.bool_]: @@ -339,9 +328,7 @@ def leap_year(x: NDArray[datetime64 | calendar | julian | unix]) -> NDArray[np.b return (year % 4 == 0) & ((year % 100 != 0) | (year % 400 == 0)) -def leap_day( - x: NDArray[datetime64 | calendar | julian | unix], -) -> NDArray[np.bool_]: +def leap_day(x: NDArray[datetime64 | calendar | julian | unix]) -> NDArray[np.bool_]: x = to_date(x) year, month, day = x["year"], x["month"], x["day"] return leap_year(year) & (month == 2) & (day == 29) @@ -358,9 +345,8 @@ def to_date( elif isdtype(x, unix) or isdtype(x, julian): x = to_datetime64(x) - out = np.recarray( - x.shape, dtype=np.dtype([("year", np.int32), ("month", np.int32), ("day", np.int32)]) - ) + dtype = np.dtype([("year", np.int32), ("month", np.int32), ("day", np.int32)]) + out = np.recarray(x.shape, dtype) Y, M, D = (x.astype(f"datetime64[{s}]") for s in "YMD") @@ -379,20 +365,18 @@ def to_calendar( elif isdtype(x, unix) or isdtype(x, julian): x = to_datetime64(x) - out = np.recarray( - x.shape, - dtype=np.dtype( - [ - ("year", np.int32), - ("month", np.int32), - ("day", np.int32), - ("hour", np.int32), - ("minute", np.int32), - ("second", np.int32), - ("microsecond", np.int32), - ] - ), + dtype = np.dtype( + [ + ("year", np.int32), + ("month", np.int32), + ("day", np.int32), + ("hour", np.int32), + ("minute", np.int32), + ("second", np.int32), + ("microsecond", np.int32), + ] ) + out = np.recarray(x.shape, dtype=dtype) Y, M, D, h, m, s = (x.astype(f"datetime64[{s}]") for s in "YMDhms") @@ -491,7 +475,8 @@ def to_julian_day(x: NDArray[datetime64 | calendar | julian | unix]) -> NDArray[ + D - 32075 ) - hms = (h * 3600 + m * 60 + (s + (ms / 1_000_000))) / 86400 + + hms = (h * 3600 + m * 60 + (s + (ms / 1e6))) / 86400 return ymd + hms - 0.5 @@ -508,15 +493,15 @@ def to_julian_day(x: NDArray[datetime64 | calendar | julian | unix]) -> NDArray[ def cast_to( - x: NDArray[datetime64 | calendar | julian | unix], dtype: type[_DT2] | str -) -> NDArray[_DT2]: + x: NDArray[datetime64 | calendar | julian | unix], dtype: type[_T2] | str +) -> NDArray[_T2]: if isinstance(dtype, str): return _function_map[_string_map[dtype]](x) return _function_map[dtype](x) -class timeseries(np.ndarray[Any, np.dtype[_DT1]]): +class timeseries(np.ndarray[Any, np.dtype[_T1]]): @overload def __new__( cls, @@ -539,12 +524,12 @@ def __new__( def __new__( cls, data: Any, - dtype: np.dtype[_DT1] = ..., - ) -> timeseries[_DT1]: ... + dtype: np.dtype[_T1] = ..., + ) -> timeseries[_T1]: ... def __new__( cls, data: Any, - dtype: np.dtype[_DT1] + dtype: np.dtype[_T1] | L[ "datetime64", "datetime64[Y]", @@ -571,9 +556,9 @@ def to(self, dtype: L["calendar"]) -> timeseries[calendar]: ... @overload def to(self, dtype: L["julian"]) -> timeseries[julian]: ... @overload - def to(self, dtype: type[_DT2]) -> timeseries[_DT2]: ... + def to(self, dtype: type[_T2]) -> timeseries[_T2]: ... def to( - self, dtype: type[_DT2] | L["unix", "datetime64", "calendar", "julian"] + self, dtype: type[_T2] | L["unix", "datetime64", "calendar", "julian"] ) -> timeseries[Any]: return cast_to(self, dtype).view(timeseries) From 64dec55cccd4fec1a9a8e6c22ff22ee08d6bd2cb Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Tue, 18 Jun 2024 17:59:37 -0500 Subject: [PATCH 3/6] =?UTF-8?q?=E2=9B=88=EF=B8=8F=20most=20unstable=20cape?= =?UTF-8?q?=20cin?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/lib/libthermo.cpp | 4 +- src/nzthermo/_core.pyi | 6 ++ src/nzthermo/_core.pyx | 89 +++++++++++------- src/nzthermo/_ufunc.pyi | 21 ++++- src/nzthermo/_ufunc.pyx | 55 ++++++++--- src/nzthermo/core.py | 181 +++++++++++++++++++++++-------------- src/nzthermo/functional.py | 11 +++ src/nzthermo/utils.py | 39 +++++++- tests/core_test.py | 147 +++++++++++++++++++++++++++--- 9 files changed, 423 insertions(+), 130 deletions(-) diff --git a/src/lib/libthermo.cpp b/src/lib/libthermo.cpp index 1aea9b9..5da2669 100644 --- a/src/lib/libthermo.cpp +++ b/src/lib/libthermo.cpp @@ -18,8 +18,10 @@ namespace libthermo { template constexpr size_t index_pressure(const T levels[], const T value, const size_t size) noexcept { const size_t N = size - 1; - const T p0 = levels[0]; const T p1 = levels[N]; + if (isnan(p1)) + return index_pressure(levels, value, N); + const T p0 = levels[0]; const T step = ((p1 - p0) / N); size_t idx = (size_t)((value / step) - (p0 / step)); diff --git a/src/nzthermo/_core.pyi b/src/nzthermo/_core.pyi index b4b1290..d4848a8 100644 --- a/src/nzthermo/_core.pyi +++ b/src/nzthermo/_core.pyi @@ -54,20 +54,26 @@ def moist_lapse[T: np.floating[Any]]( pressure: Pascal[np.ndarray[Any, np.dtype[T]]], temperature: Kelvin[np.ndarray[Any, np.dtype[np.floating[Any]]]], reference_pressure: Pascal[np.ndarray[Any, np.dtype[np.floating[Any]]]] | None = ..., + /, *, dtype: type[T | float] | L["float32", "float64"] | None = ..., + where: np.ndarray[shape[N], np.dtype[np.bool_]] | None = ..., ) -> Kelvin[np.ndarray[Any, np.dtype[T]]]: ... def parcel_profile[T: np.floating[Any]]( pressure: Pascal[np.ndarray[shape[Z], np.dtype[T]] | np.ndarray[shape[N, Z], np.dtype[T]]], temperature: Kelvin[np.ndarray[shape[N], np.dtype[np.floating[Any]]]], dewpoint: Kelvin[np.ndarray[shape[N], np.dtype[np.floating[Any]]]], /, + *, + where: np.ndarray[shape[N], np.dtype[np.bool_]] | None = ..., ) -> Kelvin[np.ndarray[shape[N, Z], np.dtype[T]]]: ... def parcel_profile_with_lcl[T: np.floating[Any]]( pressure: Pascal[np.ndarray[shape[Z], np.dtype[T]] | np.ndarray[shape[N, Z], np.dtype[T]]], temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]], /, + *, + where: np.ndarray[shape[N], np.dtype[np.bool_]] | None = ..., ) -> tuple[ Pascal[np.ndarray[shape[N, Z], np.dtype[T]]], Kelvin[np.ndarray[shape[N, Z], np.dtype[T]]], diff --git a/src/nzthermo/_core.pyx b/src/nzthermo/_core.pyx index b7aa19d..47ed4e8 100644 --- a/src/nzthermo/_core.pyx +++ b/src/nzthermo/_core.pyx @@ -225,11 +225,12 @@ cdef T[:, :] moist_lapse_2d( def moist_lapse( - np.ndarray pressure, - np.ndarray temperature, + np.ndarray pressure not None, + np.ndarray temperature not None, np.ndarray reference_pressure = None, *, object dtype = None, + np.ndarray where = None, ): """ Calculate the moist adiabatic lapse rate. @@ -298,6 +299,9 @@ def moist_lapse( BroadcastMode mode np.ndarray out + if where is not None: + raise NotImplementedError("where argument is not supported.") + if dtype is None: dtype = pressure.dtype else: @@ -426,11 +430,11 @@ cdef T[:, :] parcel_profile_2d( def parcel_profile( - np.ndarray pressure, - np.ndarray temperature, - np.ndarray dewpoint, + np.ndarray pressure not None, + np.ndarray temperature not None, + np.ndarray dewpoint not None, *, - ProfileStrategy strategy = SURFACE_BASED, + np.ndarray where = None, double step = 1000.0, double eps = 0.1, size_t max_iters = 5, @@ -441,33 +445,34 @@ def parcel_profile( np.ndarray out (pressure, temperature, dewpoint), mode = pressure_mode(pressure, temperature, dewpoint) + if where is not None: + raise NotImplementedError("where argument is not supported.") - N, Z = temperature.shape[0], pressure.shape[1] + N, Z = temperature.shape[0], pressure.shape[1] + out = np.empty((N, Z), dtype=pressure.dtype) - if strategy == SURFACE_BASED: - if pressure.dtype == np.float64: - out[...] = parcel_profile_2d[double]( - pressure.astype(np.float64), - temperature.astype(np.float64), - dewpoint.astype(np.float64), - mode, - step, - eps, - max_iters, - ) - else: - out[...] = parcel_profile_2d[float]( - pressure.astype(np.float32), - temperature.astype(np.float32), - dewpoint.astype(np.float32), - mode, - step, - eps, - max_iters, - ) + if pressure.dtype == np.float64: + out[...] = parcel_profile_2d[double]( + pressure.astype(np.float64), + temperature.astype(np.float64), + dewpoint.astype(np.float64), + mode, + step, + eps, + max_iters, + ) else: - raise ValueError("Invalid strategy.") + out[...] = parcel_profile_2d[float]( + pressure.astype(np.float32), + temperature.astype(np.float32), + dewpoint.astype(np.float32), + mode, + step, + eps, + max_iters, + ) + return out @@ -529,6 +534,9 @@ cdef void parcel_profile_with_lcl_1d( moist_lapse_1d(pt[stop + 1:], pressure[stop:], lcl.pressure, lcl.temperature) + + + cdef T[:, :, :] parcel_profile_with_lcl_2d( T[:, :] pressure, T[:, :] temperature, @@ -536,12 +544,17 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( BroadcastMode mode, ) noexcept: cdef: - size_t N, Z, i + size_t N, Z, i, idx T[:, :, :] out N, Z = temperature.shape[0], pressure.shape[1] + 1 - out = np.empty((4, N, Z), dtype=np.float64 if sizeof(double) == pressure.itemsize else np.float32) - + out = np.full( + (4, N, Z), + fill_value=NaN, + dtype=np.float64 if sizeof(double) == pressure.itemsize else np.float32 + ) + # cdef long[:] indcies = np.argmax(np.isnan(pressure), axis=1) + # print(indcies) with nogil, parallel(): if BROADCAST is mode: for i in prange(N, schedule='dynamic'): @@ -556,15 +569,23 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( ) else: # MATRIX for i in prange(N, schedule='dynamic'): + + + parcel_profile_with_lcl_1d( out[0, i, :], out[1, i, :], out[2, i, :], out[3, i, :], - pressure[i, :], - temperature[i, :], + pressure[i, :], + temperature[i, :], dewpoint[i, :], ) + + print( + 'nanvalues', + np.argmin(~np.isnan(pressure), axis=1).any() + ) return out diff --git a/src/nzthermo/_ufunc.pyi b/src/nzthermo/_ufunc.pyi index 10899de..8b2c25a 100644 --- a/src/nzthermo/_ufunc.pyi +++ b/src/nzthermo/_ufunc.pyi @@ -1,10 +1,29 @@ -from typing import Annotated, ParamSpec, TypeVar +from typing import Annotated, Any, ParamSpec, TypeVar + +import numpy as np +from numpy.typing import ArrayLike from ._typing import _ufunc1x1, _ufunc2x1, _ufunc2x2, _ufunc3x1, _ufunc3x2 from .typing import Dimensionless, Kelvin, Pascal _P = ParamSpec("_P") +_S = TypeVar("_S") _T = TypeVar("_T") +_DType_T_co = TypeVar("_DType_T_co", bound=np.dtype[Any]) + +class pressure_vector(np.ndarray[_S, _DType_T_co]): + def is_above( + self, bottom: ArrayLike, close: bool = ... + ) -> np.ndarray[_S, np.dtype[np.bool_]]: ... + def is_below( + self, top: ArrayLike, close: bool = ... + ) -> np.ndarray[_S, np.dtype[np.bool_]]: ... + def is_between( + self, bottom: ArrayLike, top: ArrayLike, close: bool = ... + ) -> np.ndarray[_S, np.dtype[np.bool_]]: ... + def where( + self, condition: np.ndarray[_S, np.dtype[np.bool_]], fill: ArrayLike = ... + ) -> pressure_vector[_S, _DType_T_co]: ... @_ufunc2x1 def less_or_close(x: float, y: float) -> bool: ... diff --git a/src/nzthermo/_ufunc.pyx b/src/nzthermo/_ufunc.pyx index 5dfd3f5..294b1b0 100644 --- a/src/nzthermo/_ufunc.pyx +++ b/src/nzthermo/_ufunc.pyx @@ -14,15 +14,23 @@ that generates the stub file from the c++ header file. # pyright: reportGeneralTypeIssues=false +# c imports cimport cython cimport numpy as np - +from libcpp.cmath cimport fabs, isnan cimport nzthermo._C as C from nzthermo._C cimport epsilon +import numpy as np +from typing import TypeVar + np.import_array() np.import_ufunc() +_S = TypeVar("_S") +_T = TypeVar("_T") + + ctypedef fused T: float double @@ -32,31 +40,56 @@ ctypedef fused integer: long -cdef T abs(T x) noexcept nogil: - return x if x > 0 else -x - @cython.ufunc cdef bint less_or_close(T x, T y) noexcept nogil: return ( - x == x and y == y # nan check - and (x < y or abs(x - y) <= (1.0e-05 * abs(y))) + not isnan(x) and not isnan(y) + and (x < y or fabs(x - y) <= (1.0e-05 * fabs(y))) ) @cython.ufunc cdef bint greater_or_close(T x, T y) noexcept nogil: return ( - x == x and y == y # nan check - and (x > y or abs(x - y) <= (1.0e-05 * abs(y))) + not isnan(x) and not isnan(y) + and (x > y or fabs(x - y) <= (1.0e-05 * fabs(y))) ) @cython.ufunc cdef bint between_or_close(T x, T y0, T y1) noexcept nogil: return ( - x == x and y0 == y0 and y1 == y1 # nan check - and (x > y0 or abs(x - y0) <= (1.0e-05 * abs(y0))) - and (x < y1 or abs(x - y1) <= (1.0e-05 * abs(y1))) + not isnan(x) and not isnan(y0) and not isnan(y1) + and (x > y0 or fabs(x - y0) <= (1.0e-05 * fabs(y0))) + and (x < y1 or fabs(x - y1) <= (1.0e-05 * fabs(y1))) ) +class pressure_vector(np.ndarray[_S, np.dtype[_T]]): + def __new__(cls, pressure): + return np.asarray(pressure).view(cls) + + def is_above(self, bottom, close=True): + bottom = np.asarray(bottom) + if not close: + return np.asarray(self > bottom, np.bool_) + + return np.asarray(less_or_close(self, bottom), np.bool_) + + def is_below(self, top, close=True): + top = np.asarray(top) + if not close: + return np.asarray(self < top, np.bool_) + + return np.asarray(greater_or_close(self, top), np.bool_) + + def is_between(self, bottom, top, close=True): + bottom, top = np.asarray(bottom), np.asarray(top) + if not close: + return np.asarray((self > bottom) & (self < top), np.bool_) + + return np.asarray(between_or_close(self, top, bottom), np.bool_) + + def where(self, condition, fill=np.nan): + return np.where(condition, self, fill).view(pressure_vector) + # ............................................................................................... # # - wind diff --git a/src/nzthermo/core.py b/src/nzthermo/core.py index 5cefcf4..a003831 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -16,7 +16,6 @@ from . import _core as core, functional as F from ._ufunc import ( - between_or_close, dewpoint as _dewpoint, dry_lapse, equivalent_potential_temperature, @@ -24,6 +23,7 @@ lcl, lcl_pressure, mixing_ratio, + pressure_vector, saturation_mixing_ratio, saturation_vapor_pressure, vapor_pressure, @@ -32,7 +32,7 @@ ) from .const import Rd from .typing import Kelvin, N, Pascal, Z, shape -from .utils import Vector1d, broadcast_nz, Axis +from .utils import Axis, Vector1d, broadcast_nz _T = TypeVar("_T", bound=np.floating[Any], covariant=True) newaxis: Final[None] = np.newaxis @@ -42,12 +42,29 @@ FASTPATH: dict[str, Any] = {"__fastpath": True} +def mask_layers( + mask: np.ndarray[shape[N, Z], np.dtype[np.bool_]], + 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]]], +): + if broadcasted := pressure.shape == temperature.shape: + p_layer, t_layer, td_layer = np.where( + mask[newaxis, :, :], [pressure, temperature, dewpoint], NaN + ) + else: + p_layer = np.where(mask, pressure, NaN) + t_layer, td_layer = np.where(mask[newaxis, :, :], [temperature, dewpoint], NaN) + + return p_layer, t_layer, td_layer, broadcasted + + # ------------------------------------------------------------------------------------------------- # downdraft_cape # ------------------------------------------------------------------------------------------------- @broadcast_nz def downdraft_cape( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + 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]]], ) -> np.ndarray[shape[N], np.dtype[_T]]: @@ -74,25 +91,25 @@ def downdraft_cape( # Tims suggestion was to allow for the parcel to potentially be conditionally based mask = (pressure <= 70000.0) & (pressure >= 50000.0) - if broadcasted := pressure.shape == temperature.shape: - p_layer, t_layer, td_layer = np.where( - mask[newaxis, :, :], [pressure, temperature, dewpoint], NaN - ) - else: - p_layer = np.where(mask, pressure, NaN) - t_layer, td_layer = np.where(mask[newaxis, :, :], [temperature, dewpoint], NaN) - - theta_e = equivalent_potential_temperature(p_layer, t_layer, td_layer) - - zx = np.nanargmin(theta_e, axis=1) + theta_e = equivalent_potential_temperature( + pressure, + temperature, + dewpoint, + # masking values with inf will alow us to call argmin without worrying about nan + where=mask, + out=np.full(temperature.shape, np.inf, dtype=temperature.dtype), + ) + zx = theta_e.argmin(axis=1) + p_layer, t_layer, td_layer, broadcasted = mask_layers(mask, pressure, temperature, dewpoint) p_top = p_layer[nx, zx] if broadcasted else p_layer[0, zx] t_top = t_layer[nx, zx] # (N,) td_top = td_layer[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 = np.where(pressure >= p_top[:, newaxis], pressure, NaN) + # pressure = pressure.where(pressure >= p_top[:, newaxis], NaN) + pressure = pressure.where(pressure.is_below(p_top[:, newaxis], close=True), NaN) e_vt = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) # (N, Z) trace = core.moist_lapse(pressure, wb_top, p_top) # (N, Z) p_vt = virtual_temperature(trace, saturation_mixing_ratio(pressure, trace)) # (N, Z) @@ -147,7 +164,7 @@ def ccl( # ------------------------------------------------------------------------------------------------- def _el_lfc( pick: L["EL", "LFC", "BOTH"], - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + 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]]], /, @@ -224,7 +241,7 @@ def _el_lfc( @broadcast_nz def el( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + 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]]], /, @@ -239,7 +256,7 @@ def el( @broadcast_nz def lfc( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + 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]]], /, @@ -260,7 +277,7 @@ def lfc( @broadcast_nz def el_lfc( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + 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]]], /, @@ -285,10 +302,11 @@ def el_lfc( # ------------------------------------------------------------------------------------------------- # cape_cin # ------------------------------------------------------------------------------------------------- +@broadcast_nz def most_unstable_parcel_index( - pressure, - temperature, - dewpoint, + 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]]], /, depth: float = 30000.0, height: float | None = None, @@ -297,20 +315,49 @@ def most_unstable_parcel_index( if height is not None: raise NotImplementedError("height argument is not implemented") - pressure = np.atleast_2d(pressure) - p0 = pressure[:, 0] if bottom is None else np.asarray(bottom) # .reshape(-1, 1) - top = p0 - depth - (mask,) = np.nonzero(between_or_close(pressure, top, p0).astype(np.bool_).squeeze()) - t_layer = temperature[:, mask] - td_layer = dewpoint[:, mask] - p_layer = pressure[:, mask] - theta_e = equivalent_potential_temperature(p_layer, t_layer, td_layer) + pbot = (pressure[:, 0] if bottom is None else np.asarray(bottom)).reshape(-1, 1) + ptop = pbot - depth + + theta_e = equivalent_potential_temperature( + pressure, + temperature, + dewpoint, + where=pressure.is_between(pbot, ptop), + out=np.full(temperature.shape, -np.inf, dtype=temperature.dtype), + ) + return np.argmax(theta_e, axis=1) +@broadcast_nz +def most_unstable_parcel( + 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]]], + /, + depth: Pascal[float] = 30_000.0, + bottom: Pascal[float] | None = None, +) -> tuple[ + Pascal[np.ndarray[shape[N], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N], np.dtype[_T]]], + np.ndarray[shape[N, Z], np.dtype[np.intp]], +]: + idx = most_unstable_parcel_index( + pressure, temperature, dewpoint, depth=depth, bottom=bottom, **FASTPATH + ) + + return ( + pressure[np.arange(pressure.shape[0]), idx], + temperature[np.arange(temperature.shape[0]), idx], + dewpoint[np.arange(dewpoint.shape[0]), idx], + idx, + ) + + @broadcast_nz def cape_cin( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + 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]]], /, @@ -318,20 +365,23 @@ def cape_cin( which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", ) -> tuple[np.ndarray, np.ndarray]: - lcl_p = lcl_pressure(pressure[:, 0], temperature[:, 0], dewpoint[:, 0]) # ✔️ - # The mixing ratio of the parcel comes from the dewpoint below the LCL, is saturated # based on the temperature above the LCL - parcel_mixing_ratio = np.where( - pressure > lcl_p[:, newaxis], # below_lcl - saturation_mixing_ratio(pressure, dewpoint), - saturation_mixing_ratio(pressure, temperature), + below_lcl = pressure.is_below( + lcl_pressure(pressure[:, 0], temperature[:, 0], dewpoint[:, 0])[:, newaxis] + ) + parcel_mixing_ratio = saturation_mixing_ratio( + pressure, dewpoint, out=np.empty_like(temperature), where=below_lcl + ) + parcel_mixing_ratio = saturation_mixing_ratio( + pressure, temperature, out=parcel_mixing_ratio, where=~below_lcl ) + # Convert the temperature/parcel profile to virtual temperature temperature = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) parcel_profile = virtual_temperature(parcel_profile, parcel_mixing_ratio) # Calculate the EL limit of integration - (el_p, _), (lfc_p, _) = _el_lfc( + (EL, _), (LFC, _) = _el_lfc( "BOTH", pressure, temperature, @@ -340,51 +390,46 @@ def cape_cin( which_lfc, which_el, ) + EL, LFC = np.reshape((EL, LFC), (2, -1, 1)) # reshape for broadcasting - el_p[np.isnan(el_p)] = np.nanmin(pressure, axis=1) - - lfc_p, el_p = np.reshape((lfc_p, el_p), (2, -1, 1)) # reshape for broadcasting + tzx = F.zero_crossings(pressure, parcel_profile - temperature) # temperature zero crossings - X, Y = F.zero_crossings(pressure, parcel_profile - temperature) # ((N, Z), ...) - - mask = between_or_close(X, el_p, lfc_p).astype(np.bool_) - x, y = np.where(mask[newaxis, ...], [X, Y], NaN) - CAPE = Rd * F.nantrapz(y, np.log(x), axis=1) + 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 - mask = greater_or_close(X, lfc_p).astype(np.bool_) - x, y = np.where(mask[newaxis, ...], [X, Y], NaN) - CIN = Rd * F.nantrapz(y, np.log(x), axis=1) + 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 @broadcast_nz -def most_unstable_parcel( - pressure: Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], +def most_unstable_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]]], /, depth: Pascal[float] = 30_000.0, bottom: Pascal[float] | None = None, -) -> tuple[ - Pascal[np.ndarray[shape[N], np.dtype[_T]]], - Kelvin[np.ndarray[shape[N], np.dtype[_T]]], - Kelvin[np.ndarray[shape[N], np.dtype[_T]]], - np.ndarray[shape[N, Z], np.dtype[np.intp]], -]: - depth = 100.0 if depth is None else depth - p0 = (pressure[:, 0] if bottom is None else np.asarray(bottom)).reshape(-1, 1) - top = p0 - depth +) -> tuple[np.ndarray, np.ndarray]: + mask = ( + np.arange(pressure.shape[1]) + >= most_unstable_parcel_index(pressure, temperature, dewpoint, depth, bottom, **FASTPATH)[ + :, newaxis + ] + ) + pressure = pressure.where(mask, -np.inf) - n = np.arange(temperature.shape[0]) - mask = between_or_close(pressure, p0, top).astype(np.bool_) - p = pressure[n, mask] - t = temperature[n, mask] - td = dewpoint[n, mask] + sort = np.arange(mask.shape[0])[:, newaxis], np.argsort(pressure, axis=1, kind="quicksort") + pressure = pressure[sort][:, ::-1] + pressure[np.isneginf(pressure)] = np.nan - # theta_e = equivalent_potential_temperature(p, t, td) - # idx = np.argmax(theta_e, axis=1) + p, t, td, mu_profile = core.parcel_profile_with_lcl( + pressure, + temperature[sort][:, ::-1], + dewpoint[sort][:, ::-1], + ) - # return p[n, idx], t[n, idx], td[n, idx], np.array([n, idx]) + return cape_cin(p.view(pressure_vector), t, td, parcel_profile=mu_profile, **FASTPATH) diff --git a/src/nzthermo/functional.py b/src/nzthermo/functional.py index 06db101..f153327 100644 --- a/src/nzthermo/functional.py +++ b/src/nzthermo/functional.py @@ -16,6 +16,17 @@ _T = TypeVar("_T", bound=np.floating[Any]) +def nanwhere( + mask: np.ndarray[shape[N, Z], np.dtype[np.bool_]], + x: np.ndarray[shape[N, Z], np.dtype[_T]], + *args: np.ndarray[shape[N, Z], np.dtype[_T]], +) -> tuple[np.ndarray[shape[N, Z], np.dtype[_T]], ...]: + if x.shape == args[0].shape: + return tuple(np.where(mask[np.newaxis, :, :], np.nan, [x, *args])) + + return (np.where(mask, np.nan, x),) + tuple(np.where(mask[np.newaxis, :, :], np.nan, args)) + + @overload def nanroll_2d(__x: NDArray[_T]) -> np.ndarray[shape[N, Z], np.dtype[_T]]: ... @overload diff --git a/src/nzthermo/utils.py b/src/nzthermo/utils.py index de6ec2f..18b9062 100644 --- a/src/nzthermo/utils.py +++ b/src/nzthermo/utils.py @@ -23,7 +23,13 @@ import numpy as np from numpy.typing import NDArray -from ._ufunc import delta_t, greater_or_close, less_or_close +from ._ufunc import ( + between_or_close, + delta_t, + greater_or_close, + less_or_close, + pressure_vector, +) from .typing import ( Kelvin, N, @@ -142,6 +148,33 @@ def where_above( ) -> Self: return self.where(self.is_above(pressure, close=close), x_fill, y_fill) + def is_between( + self, + bottom: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + top: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + *, + close: bool = False, + ): + if isinstance(bottom, PVectorNd): + bottom = bottom.pressure + if isinstance(top, PVectorNd): + top = top.pressure + if not close: + return (self.pressure > bottom) & (self.pressure < top) + + return between_or_close(self.pressure, top, bottom).astype(np.bool_) + + def where_between( + self, + bottom: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + top: Pascal[NDArray[np.floating[Any]]] | PVectorNd, + x_fill: ArrayLike[np.floating[Any]] = np.nan, + y_fill: ArrayLike[np.floating[Any]] | None = None, + *, + close: bool = False, + ) -> Self: + return self.where(self.is_between(bottom, top, close=close), x_fill, y_fill) + def is_nan(self) -> NDArray[np.bool_]: return np.isnan(self.pressure) @@ -240,7 +273,7 @@ def exactly_2d( def broadcast_nz( f: Callable[ Concatenate[ - Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], _P, @@ -272,7 +305,7 @@ def wrapper( magnitude(temperature, "kelvin"), magnitude(dewpoint, "kelvin"), ) - return f(pressure, temperature, dewpoint, *args, **kwargs) + return f(pressure_vector(pressure), temperature, dewpoint, *args, **kwargs) return wrapper diff --git a/tests/core_test.py b/tests/core_test.py index 1419f9b..acf7c01 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -1,4 +1,6 @@ # noqa +from __future__ import annotations + import itertools import warnings from typing import Any @@ -11,12 +13,15 @@ import nzthermo._core as _C from nzthermo._core import moist_lapse +from nzthermo._ufunc import pressure_vector from nzthermo.core import ( cape_cin, ccl, downdraft_cape, el, lfc, + most_unstable_cape_cin, + most_unstable_parcel, most_unstable_parcel_index, ) @@ -1166,6 +1171,29 @@ def test_lfc_profile_nan() -> None: assert_almost_equal(lfc_t, 9.6977, 0) # RETURNS: 9.58921545636997 +@pytest.mark.lfc +@pytest.mark.regression +@pytest.mark.parametrize("which", ["top", "bottom"]) +def test_lfc_metpy_regression(which) -> None: + prof = _C.parcel_profile(P, T[:, 0], Td[:, 0]) + + lfc_p, lfc_t = lfc(P, T, Td, prof, which) + print(f"\nlfc {which}") + for i in range(T.shape[0]): + lfc_p_, lfc_t_ = mpcalc.lfc( + P * Pa, + T[i] * K, + Td[i] * K, + prof[i] * K, + which=which, + ) + assert_allclose(lfc_p[i], lfc_p_.m, atol=500.0) # type: ignore + assert_allclose(lfc_t[i], lfc_t_.m, atol=1.0) + + +# ............................................................................................... # +# nzthermo.core.most_unstable_parcel +# ............................................................................................... # @pytest.mark.cape @pytest.mark.mu_cape @pytest.mark.regression @@ -1178,26 +1206,85 @@ def test_most_unstable_parcel_index(depth) -> None: for i in range(T.shape[0]) ], ) + assert_array_equal( + most_unstable_parcel_index(P, T, Td, depth=depth), + most_unstable_parcel_index(np.broadcast_to(P, (T.shape[0], P.size)), T, Td, depth=depth), + err_msg="most_unstable_parcel_index failed to on broadcasted pressure input.", + ) -@pytest.mark.lfc +@pytest.mark.cape +@pytest.mark.mu_cape @pytest.mark.regression -@pytest.mark.parametrize("which", ["top", "bottom"]) -def test_lfc_metpy_regression(which) -> None: - prof = _C.parcel_profile(P, T[:, 0], Td[:, 0]) +@pytest.mark.parametrize("depth", [30000.0]) +def test_most_unstable_parcel(depth) -> None: + p, t, td, idx = most_unstable_parcel( + np.broadcast_to(P, (T.shape[0], P.size)), T, Td, depth=depth + ) - lfc_p, lfc_t = lfc(P, T, Td, prof, which) - print(f"\nlfc {which}") for i in range(T.shape[0]): - lfc_p_, lfc_t_ = mpcalc.lfc( - P * Pa, + p_, t_, td_, idx_ = mpcalc.most_unstable_parcel( + P * Pa, T[i] * K, Td[i] * K, depth=depth * Pa + ) + assert_array_equal(p[i], p_.m) + assert_array_equal(t[i], t_.m) + assert_array_equal(td[i], td_.m) + assert_array_equal(idx[i], idx_) + + assert_array_equal( + (p, t, td, idx), + most_unstable_parcel(np.broadcast_to(P, (T.shape[0], P.size)), T, Td, depth=depth), + err_msg="most_unstable_parcel failed to on broadcasted pressure input.", + ) + + pressure = P.reshape(1, -1).view(pressure_vector) + mask = np.arange(pressure.shape[1]) >= most_unstable_parcel_index(P, T, Td)[:, np.newaxis] + + p_masked = pressure.where(mask, -np.inf) # np.where(mask, pressure, -np.inf) + + sort = ( + np.arange(mask.shape[0])[:, np.newaxis], + np.argsort(p_masked, axis=1, kind="quicksort"), + ) + + p_masked = p_masked[sort][:, ::-1] + nan_mask = np.isneginf(p_masked) + p_masked[nan_mask] = np.nan + + p, t, td, mu_profile = _C.parcel_profile_with_lcl( + p_masked, + T[sort][:, ::-1], + Td[sort][:, ::-1], + ) + + for i in range(T.shape[0]): + _, _, _, idx = mpcalc.most_unstable_parcel( + pressure.squeeze() * Pa, T[i] * K, Td[i] * K, - prof[i] * K, - which=which, + depth=depth * Pa, ) - assert_allclose(lfc_p[i], lfc_p_.m, atol=500.0) # type: ignore - assert_allclose(lfc_t[i], lfc_t_.m, atol=1.0) + + p_, t_, td_, mu_profile_ = mpcalc.parcel_profile_with_lcl( + pressure.squeeze()[idx:] * Pa, + T[i, idx:] * K, + Td[i, idx:] * K, + ) + + m = np.s_[:] + p_ = p_.m[m] + t_ = t_.m[m] + td_ = td_.m[m] + mu_profile_ = mu_profile_.m[m] + if np.isnan(t_[0]): + continue + + assert_allclose(p[i, : len(p_)], p_, atol=100.0) + assert_allclose(t[i, : len(t_)], t_, atol=0.5) + assert_allclose(td[i, : len(td_)], td_, atol=0.5) + assert_allclose(mu_profile[i, : len(mu_profile_)], mu_profile_, atol=0.5) + + most_unstable_cape_cin(P, T, Td, depth=30000.0) # ............................................................................................... # @@ -1240,3 +1327,39 @@ def test_cape_cin_metpy_regression(which_lfc, which_el) -> None: assert_allclose(CAPE[i], CAPE_.m, atol=10) assert_allclose(CIN[i], CIN_.m, atol=10) + + +# ............................................................................................... # +# nzthermo.core.most_unstable_cape_cin +# ............................................................................................... # +@pytest.mark.most_unstable_cape_cin +@pytest.mark.regression +@pytest.mark.parametrize( + "depth", + [30000.0], +) +def test_most_unstable_cape_cin_metpy_regression(depth) -> None: + """ + TODO currently this test is passing on 95% of the cases, need to investigate the. + there error appears to be something in the logic block of the el_lfc function. + + The current test cases run 500 samples and we are failing on 17 of them specifically when + `which_el=bottom` parameter is used. realistically using the lower EL is not a typical use + case but it should still be tested. + """ + + CAPE, CIN = most_unstable_cape_cin( + P, + T, + Td, + depth=depth, + ) + + for i in range(T.shape[0]): + CAPE_, CIN_ = mpcalc.most_unstable_cape_cin( + P * Pa, + T[i] * K, + Td[i] * K, + depth=depth * Pa, + ) + assert_allclose(CAPE[i], CAPE_.m, atol=1000) From 2e1e3d08889a1ab3fb887fa9fbfd0bd720a3a83e Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Tue, 18 Jun 2024 21:17:52 -0500 Subject: [PATCH 4/6] =?UTF-8?q?=E2=9B=88=EF=B8=8F=20itemsize=20;?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/nzthermo/_core.pyx | 33 +++++++-------------------------- tests/core_test.py | 12 ++++++++---- 2 files changed, 15 insertions(+), 30 deletions(-) diff --git a/src/nzthermo/_core.pyx b/src/nzthermo/_core.pyx index 47ed4e8..9f993bc 100644 --- a/src/nzthermo/_core.pyx +++ b/src/nzthermo/_core.pyx @@ -154,7 +154,7 @@ cdef T[:] dispatch( T[:] out N, Z = temperature.shape[0], pressure.shape[1] - out = np.empty((N,), dtype=np.float64 if sizeof(double) == pressure.itemsize else np.float32) + out = np.empty((N,), dtype=np.dtype(f"f{pressure.itemsize}")) with nogil: if BROADCAST is mode: @@ -534,9 +534,6 @@ cdef void parcel_profile_with_lcl_1d( moist_lapse_1d(pt[stop + 1:], pressure[stop:], lcl.pressure, lcl.temperature) - - - cdef T[:, :, :] parcel_profile_with_lcl_2d( T[:, :] pressure, T[:, :] temperature, @@ -548,13 +545,8 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( T[:, :, :] out N, Z = temperature.shape[0], pressure.shape[1] + 1 - out = np.full( - (4, N, Z), - fill_value=NaN, - dtype=np.float64 if sizeof(double) == pressure.itemsize else np.float32 - ) - # cdef long[:] indcies = np.argmax(np.isnan(pressure), axis=1) - # print(indcies) + out = np.full((4, N, Z), fill_value=NaN, dtype=np.dtype(f"f{pressure.itemsize}")) + with nogil, parallel(): if BROADCAST is mode: for i in prange(N, schedule='dynamic'): @@ -569,9 +561,6 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( ) else: # MATRIX for i in prange(N, schedule='dynamic'): - - - parcel_profile_with_lcl_1d( out[0, i, :], out[1, i, :], @@ -581,12 +570,7 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( temperature[i, :], dewpoint[i, :], ) - - print( - 'nanvalues', - np.argmin(~np.isnan(pressure), axis=1).any() - ) - + return out @@ -606,7 +590,6 @@ def parcel_profile_with_lcl(np.ndarray pressure, np.ndarray temperature, np.ndar temperature.astype(np.float64), dewpoint.astype(np.float64), mode, - ) else: out[...] = parcel_profile_with_lcl_2d[float]( @@ -614,7 +597,6 @@ def parcel_profile_with_lcl(np.ndarray pressure, np.ndarray temperature, np.ndar temperature.astype(np.float32), dewpoint.astype(np.float32), mode, - ) return out[0], out[1], out[2], out[3] @@ -633,7 +615,7 @@ cdef T[:] _interpolate_nz( T[:] out N, Z = x.shape[0], xp.shape[0] - out = np.empty(N, dtype=np.float32 if sizeof(float) == x.itemsize else np.float64) + out = np.empty(N, dtype=np.dtype(f"f{x.itemsize}")) with nogil, parallel(): for n in prange(N, schedule='runtime'): out[n] = C.interpolate_1d(x[n], &xp[0], &fp[n, 0], Z) @@ -700,7 +682,6 @@ def interpolate_nz( array([296.63569648, 296.79664494, 296.74736566, 297.07070398, 297.54936596]), array([295.07855875, 294.79437914, 295.27081714, 295.4858194, 296.31665617]) ) - """ dtype = __x.dtype cdef np.ndarray xp = np.asarray(__xp, dtype=dtype) @@ -713,7 +694,7 @@ def interpolate_nz( else: out[i] = _interpolate_nz[float](__x, xp, fp[i], log_x) - if interp_nan: + if interp_nan: mask = np.isnan(out[i]) out[i, mask] = np.interp(np.flatnonzero(mask), np.flatnonzero(~mask), __x[~mask]) @@ -770,7 +751,7 @@ def intersect( if increasing is False and direction == 'increasing': increasing = True - + out = np.empty((2, a.shape[0]), x.dtype) if x.dtype == np.float64: out[...] = intersect_2d[double](x, a, b, mode, log_x, increasing, bottom) diff --git a/tests/core_test.py b/tests/core_test.py index acf7c01..bbaefb0 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -1194,8 +1194,8 @@ def test_lfc_metpy_regression(which) -> None: # ............................................................................................... # # nzthermo.core.most_unstable_parcel # ............................................................................................... # +@pytest.mark.mu @pytest.mark.cape -@pytest.mark.mu_cape @pytest.mark.regression @pytest.mark.parametrize("depth", [30000.0]) def test_most_unstable_parcel_index(depth) -> None: @@ -1213,9 +1213,10 @@ def test_most_unstable_parcel_index(depth) -> None: ) +@pytest.mark.mu @pytest.mark.cape -@pytest.mark.mu_cape @pytest.mark.regression +@pytest.mark.most_unstable_cape @pytest.mark.parametrize("depth", [30000.0]) def test_most_unstable_parcel(depth) -> None: p, t, td, idx = most_unstable_parcel( @@ -1276,6 +1277,7 @@ def test_most_unstable_parcel(depth) -> None: t_ = t_.m[m] td_ = td_.m[m] mu_profile_ = mu_profile_.m[m] + if np.isnan(t_[0]): continue @@ -1332,8 +1334,10 @@ def test_cape_cin_metpy_regression(which_lfc, which_el) -> None: # ............................................................................................... # # nzthermo.core.most_unstable_cape_cin # ............................................................................................... # -@pytest.mark.most_unstable_cape_cin +@pytest.mark.mu +@pytest.mark.cape_cin @pytest.mark.regression +@pytest.mark.most_unstable_cape_cin @pytest.mark.parametrize( "depth", [30000.0], @@ -1362,4 +1366,4 @@ def test_most_unstable_cape_cin_metpy_regression(depth) -> None: Td[i] * K, depth=depth * Pa, ) - assert_allclose(CAPE[i], CAPE_.m, atol=1000) + assert_allclose(CAPE[i], CAPE_.m, atol=900) From a9fba5cee21cd8909d532969e3e75a49c3fd69aa Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Wed, 19 Jun 2024 11:02:51 -0500 Subject: [PATCH 5/6] =?UTF-8?q?=E2=9B=88=EF=B8=8F=20most=20unstable=20cape?= =?UTF-8?q?=20cin?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 1 + src/lib/libthermo.cpp | 2 +- src/nzthermo/__init__.py | 2 + src/nzthermo/_core.pyx | 269 +++++++++++++++++++++++---------------- src/nzthermo/_ufunc.pyx | 83 ++++++++++++ src/nzthermo/core.py | 132 +++++++++++-------- tests/core_test.py | 233 +++++++++++++++++---------------- tests/functional_test.py | 4 +- 8 files changed, 446 insertions(+), 280 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ce8acfd..8fb4b02 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ markers = [ "mu_cape", "parcel_profile", "regression", + "broadcasting", ] pythonpath = ["src"] diff --git a/src/lib/libthermo.cpp b/src/lib/libthermo.cpp index 5da2669..b2bb204 100644 --- a/src/lib/libthermo.cpp +++ b/src/lib/libthermo.cpp @@ -189,7 +189,7 @@ constexpr T lcl::wet_bulb_temperature(const T pressure, const T step) noexcep } template -constexpr size_t lcl::index(const T pressure[], size_t size) noexcept { +constexpr size_t lcl::index(const T pressure[], const size_t size) noexcept { return index_pressure(pressure, this->pressure, size); } diff --git a/src/nzthermo/__init__.py b/src/nzthermo/__init__.py index b1af963..42b936c 100644 --- a/src/nzthermo/__init__.py +++ b/src/nzthermo/__init__.py @@ -48,6 +48,7 @@ "downdraft_cape", "most_unstable_parcel", "most_unstable_parcel_index", + "most_unstable_cape_cin", # .utils "timeseries", ] @@ -104,5 +105,6 @@ mixing_ratio, most_unstable_parcel, most_unstable_parcel_index, + most_unstable_cape_cin, ) from .utils import timeseries diff --git a/src/nzthermo/_core.pyx b/src/nzthermo/_core.pyx index 9f993bc..8f2e56e 100644 --- a/src/nzthermo/_core.pyx +++ b/src/nzthermo/_core.pyx @@ -18,6 +18,7 @@ template T fn(T ...){...} ``` """ +import warnings from cython.parallel cimport parallel, prange from cython.view cimport array as cvarray @@ -110,21 +111,98 @@ cdef cvarray nzarray((size_t, size_t) shape, size_t size): ) -cdef pressure_mode( - np.ndarray pressure, - np.ndarray temperature, - np.ndarray dewpoint, +def broadcast_and_nanmask( + np.ndarray pressure not None, + np.ndarray temperature not None, + np.ndarray dewpoint, + long ndim = 2, + object where = None, + object dtype = None, ): - if pressure.ndim == 1: + """ + TODO: ideally this function needs to be exposed to the `core.py` or implmented as part of the + decorator, so that it is only ever called once. As part of the masking operation it would be + useful to properly mask and sort the data once per function call. + + + >>> where &= np.isfinite(temperature) & np.isfinite(dewpoint) & np.isfinite(pressure) + + Validates the input arrays and determines the broadcast mode. by brodcasting the ``pressure``, + ``temperature``, and ``dewpoint`` arrays to a suitable shape for the calling function. + + In some cases temperature and dewpoint arrays can be 1D, if this is the case the ``ndim`` + argument must be set to 1. + + masks the data with nan values and sorts in descending pressure order (bottom -> top) with nan + values at the end, this allows us to short circuit some computations as some of the functions + will stop at the first nan value. + + Insures that all arays are in ``C_CONTIGUOUS`` memory layout. This is important for the Cython + memoryview to work correctly. + """ + cdef: + size_t N, Z + BroadcastMode mode + + if temperature.ndim != ndim: + raise ValueError(f"temperature must be a {ndim}D array.") + elif dewpoint.ndim != ndim: + raise ValueError(f"dewpoint must be a {ndim}D array.") + elif ndim == 2: + if temperature.shape[0] != dewpoint.shape[0]: + raise ValueError("temperature and dewpoint must have the same number of rows.") + elif temperature.shape[1] != dewpoint.shape[1]: + raise ValueError("temperature and dewpoint must have the same number of columns.") + elif ndim == 1: + if temperature.shape[0] != dewpoint.shape[0]: + raise ValueError("temperature and dewpoint must have the same number of elements.") + + if pressure.ndim == 1 or pressure.shape[0] == 1: pressure = pressure.reshape(1, -1) mode = BROADCAST else: mode = MATRIX + + if where is not None: + if not isinstance(where, np.ndarray): + raise ValueError("where must be a numpy array.") + + + mode = MATRIX + N = temperature.shape[0] + Z = temperature.shape[1] + + where = np.atleast_2d(where) + # where &= np.isfinite(temperature) & np.isfinite(dewpoint) & np.isfinite(pressure) + + pressure = np.broadcast_to(pressure.squeeze(), (N, Z)) + pressure = np.where(where, pressure, -np.inf) + + sort = np.arange(N)[:, np.newaxis], np.argsort(pressure, axis=1) + pressure = pressure[sort][:, ::-1] + temperature = temperature[sort][:, ::-1] + dewpoint = dewpoint[sort][:, ::-1] + m = np.isneginf(pressure) + pressure[m] = np.nan + temperature[m] = np.nan + dewpoint[m] = np.nan + + if dtype is None: + dtype = temperature.dtype + + if dtype != np.float32 and dtype != np.float64: + + warnings.warn(f"the dtype {dtype} is not supported, defaulting to np.float64.") + dtype = np.float64 - return (pressure, temperature, dewpoint), mode + pressure = np.ascontiguousarray(pressure, dtype=dtype) + temperature = np.ascontiguousarray(temperature, dtype=dtype) + dewpoint = np.ascontiguousarray(dewpoint, dtype=dtype) + return (pressure, temperature, dewpoint), mode, dtype -# need to figuoure out a way to possibly pass in **kwargs maybe via a options struct + +# need to figure out a way to possibly pass in **kwargs maybe via a options struct ctypedef T (*Dispatch)(const T*, const T*, const T*, size_t) noexcept nogil cdef T[:] dispatch( @@ -155,7 +233,6 @@ cdef T[:] dispatch( N, Z = temperature.shape[0], pressure.shape[1] out = np.empty((N,), dtype=np.dtype(f"f{pressure.itemsize}")) - with nogil: if BROADCAST is mode: for i in prange(N, schedule='dynamic'): @@ -200,14 +277,12 @@ cdef T[:, :] moist_lapse_2d( T[:] reference_pressure, T[:] temperature, BroadcastMode mode, - ) noexcept: cdef: size_t N, Z, i T[:, :] out N, Z = temperature.shape[0], pressure.shape[1] - out = nzarray((N, Z), pressure.itemsize) with nogil, parallel(): if BROADCAST is mode: @@ -349,7 +424,6 @@ def moist_lapse( raise ValueError("Unable to determine the broadcast mode.") Z = pressure.shape[1] - out = np.empty((N, Z), dtype=dtype) if np.float32 == dtype: out[...] = moist_lapse_2d[float]( @@ -382,20 +456,18 @@ cdef void parcel_profile_1d( ) noexcept nogil: cdef: size_t Z, i, stop - T p0, t0, reference_pressure, next_pressure + T p0, reference_pressure, next_pressure C.lcl[T] lcl - Z = pressure.shape[0] + Z = out.shape[0] p0 = pressure[0] - t0 = out[0] = temperature + lcl = C.lcl[T](pressure[0], temperature, dewpoint) - lcl = C.lcl[T](p0, t0, dewpoint) - # [dry ascent] + # [dry ascent] # stop the dry ascent at the LCL stop = lcl.index(&pressure[0], Z) for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent - out[i] = C.dry_lapse(pressure[i], p0, t0) - + out[i] = C.dry_lapse(pressure[i], p0, temperature) # [ moist ascent ] if stop != Z: @@ -407,16 +479,13 @@ cdef T[:, :] parcel_profile_2d( T[:] temperature, T[:] dewpoint, BroadcastMode mode, - T step, - T eps, - size_t max_iters, ) noexcept: cdef: size_t N, Z, i 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'): @@ -425,7 +494,7 @@ cdef T[:, :] parcel_profile_2d( else: # MATRIX for i in prange(N, schedule='dynamic'): parcel_profile_1d(out[i], pressure[i, :], temperature[i], dewpoint[i]) - + return out @@ -435,44 +504,22 @@ def parcel_profile( np.ndarray dewpoint not None, *, np.ndarray where = None, - double step = 1000.0, - double eps = 0.1, - size_t max_iters = 5, ): cdef: size_t N, Z BroadcastMode mode np.ndarray out - (pressure, temperature, dewpoint), mode = pressure_mode(pressure, temperature, dewpoint) - if where is not None: - raise NotImplementedError("where argument is not supported.") - + (pressure, temperature, dewpoint), mode, dtype = broadcast_and_nanmask( + pressure, temperature, dewpoint, ndim=1, where=where + ) N, Z = temperature.shape[0], pressure.shape[1] - - out = np.empty((N, Z), dtype=pressure.dtype) - if pressure.dtype == np.float64: - out[...] = parcel_profile_2d[double]( - pressure.astype(np.float64), - temperature.astype(np.float64), - dewpoint.astype(np.float64), - mode, - step, - eps, - max_iters, - ) + out = np.empty((N, Z), dtype=dtype) + if dtype == np.float64: + out[...] = parcel_profile_2d[double](pressure, temperature, dewpoint, mode) else: - out[...] = parcel_profile_2d[float]( - pressure.astype(np.float32), - temperature.astype(np.float32), - dewpoint.astype(np.float32), - mode, - step, - eps, - max_iters, - ) - + out[...] = parcel_profile_2d[float](pressure, temperature, dewpoint, mode) return out @@ -501,30 +548,40 @@ 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) - ep[:stop] = pressure[:stop] - et[:stop] = temperature[:stop] - etd[:stop] = dewpoint[:stop] - for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent - pt[i] = C.dry_lapse(pressure[i], p0, t0) - - # [ lcl ] - ep[stop] = lcl.pressure - et[stop] = C.linear_interpolate( - lcl.pressure, - pressure[stop - 1], - pressure[stop], - temperature[stop - 1], - temperature[stop] - ) - etd[stop] = C.linear_interpolate( - lcl.pressure, - pressure[stop - 1], - pressure[stop], - dewpoint[stop - 1], - dewpoint[stop] - ) - pt[stop] = lcl.temperature + + if stop: + ep[:stop] = pressure[:stop] + et[:stop] = temperature[:stop] + etd[:stop] = dewpoint[:stop] + for i in prange(0, stop, schedule='dynamic'): # parallelize the dry ascent + pt[i] = C.dry_lapse(pressure[i], p0, t0) + + # [ lcl ] + ep[stop] = lcl.pressure + et[stop] = C.linear_interpolate( + lcl.pressure, + pressure[stop - 1], + pressure[stop], + temperature[stop - 1], + temperature[stop] + ) + etd[stop] = C.linear_interpolate( + lcl.pressure, + pressure[stop - 1], + pressure[stop], + dewpoint[stop - 1], + dewpoint[stop] + ) + pt[stop] = lcl.temperature + else: + # the lcl was found at the surface which is a little odd. In the metpy implementation + # this causes interpolation warnings, but we can just set the values to the surface values + ep[0] = lcl.pressure + et[0] = t0 + etd[0] = td0 + pt[0] = lcl.temperature # [ moist ascent ] .. parcel temperature from the LCL to the top of the atmosphere .. if stop != Z: @@ -542,19 +599,18 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( ) noexcept: cdef: size_t N, Z, i, idx - T[:, :, :] out + T[:, :] ep, et, etd, pt N, Z = temperature.shape[0], pressure.shape[1] + 1 - out = np.full((4, N, Z), fill_value=NaN, dtype=np.dtype(f"f{pressure.itemsize}")) - + ep, et, etd, pt = np.full((4, N, Z), fill_value=NaN, dtype=np.dtype(f"f{pressure.itemsize}")) with nogil, parallel(): if BROADCAST is mode: for i in prange(N, schedule='dynamic'): parcel_profile_with_lcl_1d( - out[0, i, :], - out[1, i, :], - out[2, i, :], - out[3, i, :], + ep[i, :], + et[i, :], + etd[i, :], + pt[i, :], pressure[0, :], # broadcast 1d pressure array temperature[i, :], dewpoint[i, :], @@ -562,42 +618,40 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( else: # MATRIX for i in prange(N, schedule='dynamic'): parcel_profile_with_lcl_1d( - out[0, i, :], - out[1, i, :], - out[2, i, :], - out[3, i, :], + ep[i, :], + et[i, :], + etd[i, :], + pt[i, :], pressure[i, :], temperature[i, :], dewpoint[i, :], ) - return out + return np.asarray([ep, et, etd, pt]) -def parcel_profile_with_lcl(np.ndarray pressure, np.ndarray temperature, np.ndarray dewpoint): +def parcel_profile_with_lcl( + np.ndarray pressure not None, + np.ndarray temperature not None, + np.ndarray dewpoint not None, + *, + np.ndarray where = None, +): cdef: size_t N, Z BroadcastMode mode np.ndarray out - (pressure, temperature, dewpoint), mode = pressure_mode(pressure, temperature, dewpoint) - N, Z = temperature.shape[0], pressure.shape[1] + (pressure, temperature, dewpoint), mode, dtype = broadcast_and_nanmask( + pressure, temperature, dewpoint, ndim=2, where=where + ) - out = np.empty((4, N, Z + 1), dtype=pressure.dtype) - if pressure.dtype == np.float64: - out[...] = parcel_profile_with_lcl_2d[double]( - pressure.astype(np.float64), - temperature.astype(np.float64), - dewpoint.astype(np.float64), - mode, - ) + N, Z = temperature.shape[0], pressure.shape[1] + out = np.empty((4, N, Z + 1), dtype=dtype) + if dtype == np.float64: + out[...] = parcel_profile_with_lcl_2d[double](pressure, temperature, dewpoint, mode) else: - out[...] = parcel_profile_with_lcl_2d[float]( - pressure.astype(np.float32), - temperature.astype(np.float32), - dewpoint.astype(np.float32), - mode, - ) + out[...] = parcel_profile_with_lcl_2d[float](pressure, temperature, dewpoint, mode) return out[0], out[1], out[2], out[3] @@ -611,7 +665,7 @@ cdef T[:] _interpolate_nz( bint log_x = 0, ) noexcept: cdef: - size_t N, Z, n + size_t N, Z, n T[:] out N, Z = x.shape[0], xp.shape[0] @@ -719,7 +773,6 @@ cdef intersect_2d( T[:, :] out N, Z = y0.shape[0], y1.shape[1] - out = nzarray((2, N), x.itemsize) with nogil, parallel(): if BROADCAST is mode: @@ -747,12 +800,12 @@ def intersect( BroadcastMode mode np.ndarray out - (x, a, b), mode = pressure_mode(x, a, b) + (x, a, b), mode, dtype = broadcast_and_nanmask(x, a, b) if increasing is False and direction == 'increasing': increasing = True - out = np.empty((2, a.shape[0]), x.dtype) + out = np.empty((2, a.shape[0]), dtype) if x.dtype == np.float64: out[...] = intersect_2d[double](x, a, b, mode, log_x, increasing, bottom) else: diff --git a/src/nzthermo/_ufunc.pyx b/src/nzthermo/_ufunc.pyx index 294b1b0..3a78789 100644 --- a/src/nzthermo/_ufunc.pyx +++ b/src/nzthermo/_ufunc.pyx @@ -231,6 +231,89 @@ cdef T dry_lapse(T pressure, T temperature, T reference_pressure) noexcept nogil @cython.ufunc # theta_e cdef T equivalent_potential_temperature(T pressure, T temperature, T dewpoint) noexcept nogil: + """ + Parameters + ---------- + x : array_like + pressure (Pa) values. + x1 : array_like + temperature (K) values. + x2 : array_like + dewpoint (K) values. + out : ndarray, None, or tuple of ndarray and None, optional + A location into which the result is stored. If provided, it must have + a shape that the inputs broadcast to. If not provided or None, + a freshly-allocated array is returned. A tuple (possible only as a + keyword argument) must have length equal to the number of outputs. + where : array_like, optional + This condition is broadcast over the input. At locations where the + condition is True, the `out` array will be set to the ufunc result. + Elsewhere, the `out` array will retain its original value. + Note that if an uninitialized `out` array is created via the default + ``out=None``, locations within it where the condition is False will + remain uninitialized. + **kwargs + For other keyword-only arguments, see the + :ref:`ufunc docs `. + + Returns + ------- + theta_e : ndarray + Equivalent potential temperature (K). + + Examples + ------- + >>> import numpy as np + >>> import nzthermo as nzt + >>> data = np.load("tests/data.npz", allow_pickle=False) + >>> pressure = data['P'] + >>> temperature = data['T'] + >>> dewpoint = data['Td'] + >>> assert pressure.ndim == 1 and pressure.shape != temperature.shape + >>> mask = (pressure <= 70000.0) & (pressure >= 50000.0) + >>> theta_e = nzt.equivalent_potential_temperature( + ... pressure, + ... temperature, + ... dewpoint, + ... where=mask, # masking values with inf will alow us to call argmin without worrying about nan + ... out=np.full_like(temperature, np.inf), + ... ) + >>> theta_e.shape + (540, 40) + >>> theta_e.argmin(axis=1) + array([13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 15, 21, 21, 21, 21, 21, 21, + 20, 13, 18, 14, 14, 14, 14, 20, 20, 14, 14, 16, 18, 13, 13, 13, 13, + 13, 13, 13, 13, 13, 13, 21, 21, 17, 15, 21, 18, 21, 13, 13, 13, 18, + 13, 14, 13, 16, 13, 19, 18, 18, 20, 13, 13, 15, 14, 13, 13, 13, 13, + 13, 14, 21, 18, 21, 21, 13, 21, 20, 21, 14, 13, 19, 20, 13, 16, 13, + 18, 16, 18, 21, 20, 13, 13, 14, 16, 16, 14, 13, 13, 13, 19, 21, 21, + 21, 21, 20, 17, 20, 21, 21, 13, 13, 20, 13, 14, 18, 13, 13, 13, 13, + 14, 13, 13, 13, 13, 13, 13, 13, 13, 14, 15, 19, 18, 18, 20, 19, 19, + 13, 20, 21, 13, 14, 20, 18, 18, 16, 13, 13, 13, 16, 13, 13, 13, 13, + 13, 13, 14, 13, 13, 15, 15, 15, 15, 13, 13, 16, 16, 20, 18, 15, 21, + 21, 13, 16, 16, 14, 13, 13, 13, 13, 13, 13, 13, 14, 13, 14, 15, 13, + 13, 13, 14, 21, 21, 21, 16, 14, 15, 13, 17, 18, 13, 20, 18, 18, 20, + 14, 18, 14, 13, 13, 13, 19, 18, 14, 14, 13, 15, 15, 18, 21, 20, 19, + 21, 20, 21, 21, 14, 14, 18, 20, 15, 18, 13, 16, 14, 16, 14, 16, 18, + 13, 13, 20, 13, 18, 18, 18, 16, 17, 19, 19, 18, 20, 21, 20, 18, 21, + 17, 17, 19, 18, 16, 18, 13, 13, 14, 13, 16, 16, 16, 16, 18, 16, 14, + 14, 16, 18, 18, 19, 18, 17, 18, 20, 21, 21, 20, 20, 21, 15, 19, 17, + 18, 18, 13, 15, 16, 13, 13, 16, 15, 13, 13, 14, 13, 13, 18, 18, 16, + 19, 19, 16, 16, 19, 19, 18, 20, 19, 21, 20, 18, 20, 18, 18, 13, 15, + 15, 17, 18, 16, 13, 13, 13, 13, 14, 13, 13, 16, 16, 18, 18, 16, 16, + 17, 18, 20, 19, 16, 19, 13, 14, 14, 18, 17, 16, 15, 18, 18, 13, 13, + 13, 14, 13, 13, 13, 14, 13, 16, 16, 19, 17, 14, 14, 15, 16, 17, 18, + 15, 13, 14, 13, 15, 13, 13, 18, 13, 13, 14, 14, 15, 14, 13, 13, 13, + 13, 13, 13, 14, 16, 19, 15, 18, 15, 13, 15, 15, 16, 16, 13, 13, 19, + 17, 13, 13, 13, 13, 13, 13, 15, 19, 13, 13, 13, 13, 13, 13, 13, 13, + 13, 15, 16, 16, 13, 13, 18, 16, 16, 15, 14, 13, 14, 13, 15, 17, 16, + 13, 13, 13, 16, 15, 18, 13, 13, 13, 13, 13, 13, 13, 13, 14, 16, 17, + 13, 17, 17, 16, 16, 14, 13, 13, 15, 16, 16, 15, 15, 17, 18, 13, 15, + 15, 14, 14, 13, 13, 13, 13, 13, 13, 15, 14, 15, 16, 14, 17, 17, 16, + 16, 13, 13, 14, 20, 17, 17, 14, 16, 16, 13, 13, 17, 16, 15, 14, 13, + 13, 13, 13, 13, 13, 13, 14, 20, 18, 18, 15, 17, 13, 14, 13, 13, 13, + 14, 13, 13, 13, 15, 20, 18, 13, 14, 19, 13, 16, 13]) + """ return C.equivalent_potential_temperature(pressure, temperature, dewpoint) diff --git a/src/nzthermo/core.py b/src/nzthermo/core.py index a003831..a0737c4 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -14,7 +14,8 @@ import numpy as np from numpy.typing import NDArray -from . import _core as core, functional as F +from . import functional as F +from ._core import parcel_profile_with_lcl, moist_lapse from ._ufunc import ( dewpoint as _dewpoint, dry_lapse, @@ -43,22 +44,43 @@ def mask_layers( - mask: np.ndarray[shape[N, Z], np.dtype[np.bool_]], 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]]], -): + /, + *, + where: np.ndarray[shape[N, Z], np.dtype[np.bool_]], + fill_value: float = NaN, +) -> tuple[ + Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + bool, +]: if broadcasted := pressure.shape == temperature.shape: p_layer, t_layer, td_layer = np.where( - mask[newaxis, :, :], [pressure, temperature, dewpoint], NaN + where[newaxis, :, :], [pressure, temperature, dewpoint], fill_value ) else: - p_layer = np.where(mask, pressure, NaN) - t_layer, td_layer = np.where(mask[newaxis, :, :], [temperature, dewpoint], NaN) + p_layer = np.where(where, pressure, fill_value) + t_layer, td_layer = np.where(where[newaxis, :, :], [temperature, dewpoint], fill_value) return p_layer, t_layer, td_layer, broadcasted +def parcel_mixing_ratio( + 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]]], + /, + *, + where: np.ndarray[shape[N, Z], np.dtype[np.bool_]], +): + r = saturation_mixing_ratio(pressure, dewpoint, out=np.empty_like(temperature), where=where) + r = saturation_mixing_ratio(pressure, temperature, out=r, where=~where) + return r + + # ------------------------------------------------------------------------------------------------- # downdraft_cape # ------------------------------------------------------------------------------------------------- @@ -67,6 +89,8 @@ def downdraft_cape( 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]]], + /, + where: np.ndarray[shape[N, Z], np.dtype[np.bool_]] | None = None, ) -> np.ndarray[shape[N], np.dtype[_T]]: """Calculate downward CAPE (DCAPE). @@ -89,29 +113,29 @@ def downdraft_cape( N, _ = temperature.shape nx = np.arange(N) # Tims suggestion was to allow for the parcel to potentially be conditionally based - mask = (pressure <= 70000.0) & (pressure >= 50000.0) + if where is None: + where = (pressure <= 70000.0) & (pressure >= 50000.0) theta_e = equivalent_potential_temperature( pressure, temperature, dewpoint, # masking values with inf will alow us to call argmin without worrying about nan - where=mask, - out=np.full(temperature.shape, np.inf, dtype=temperature.dtype), + where=where, + out=np.full_like(temperature, np.inf), ) zx = theta_e.argmin(axis=1) - p_layer, t_layer, td_layer, broadcasted = mask_layers(mask, pressure, temperature, dewpoint) - p_top = p_layer[nx, zx] if broadcasted else p_layer[0, zx] - t_top = t_layer[nx, zx] # (N,) - td_top = td_layer[nx, zx] # (N,) + p_top = pressure[nx, zx] if pressure.shape == temperature.shape else pressure[0, 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) e_vt = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) # (N, Z) - trace = core.moist_lapse(pressure, wb_top, p_top) # (N, Z) + trace = moist_lapse(pressure, wb_top, p_top) # (N, Z) p_vt = virtual_temperature(trace, saturation_mixing_ratio(pressure, trace)) # (N, Z) DCAPE = Rd * F.nantrapz(p_vt - e_vt, np.log(pressure), axis=1) @@ -128,6 +152,7 @@ def ccl( temperature: Kelvin[NDArray[_T]], dewpoint: Kelvin[NDArray[_T]], /, + *, height=None, mixed_layer_depth=None, which: L["bottom", "top"] = "bottom", @@ -168,14 +193,13 @@ def _el_lfc( temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_temperature_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] - | None = None, + parcel_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, 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]: - if parcel_temperature_profile is None: - pressure, temperature, dewpoint, parcel_temperature_profile = core.parcel_profile_with_lcl( + if parcel_profile is None: + pressure, temperature, dewpoint, parcel_profile = parcel_profile_with_lcl( pressure, temperature, dewpoint ) @@ -188,18 +212,18 @@ def _el_lfc( LCL = Vector1d.from_func(lcl, p0, t0, td0).unsqueeze() - pressure, parcel_temperature_profile, temperature = ( + pressure, parcel_profile, temperature = ( pressure[:, 1:], - parcel_temperature_profile[:, 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_temperature_profile[top_idx] <= temperature[top_idx])[:, newaxis] + left_of_env = (parcel_profile[top_idx] <= temperature[top_idx])[:, newaxis] EL = F.intersect_nz( pressure, - parcel_temperature_profile, + parcel_profile, temperature, "decreasing", log_x=True, @@ -213,7 +237,7 @@ def _el_lfc( LFC = F.intersect_nz( pressure, - parcel_temperature_profile, + parcel_profile, temperature, "increasing", log_x=True, @@ -223,7 +247,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_temperature_profile, NaN), + np.where(LCL.is_below(pressure, close=True), parcel_profile, NaN), temperature, ).any(Axis.Z, out=np.empty((N, 1), dtype=np.bool_), keepdims=True) @@ -245,13 +269,11 @@ def el( temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_temperature_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] - | None = None, + parcel_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, + *, which: L["top", "bottom"] = "top", ) -> Vector1d[_T]: - return _el_lfc( - "EL", pressure, temperature, dewpoint, parcel_temperature_profile, which_el=which - ) + return _el_lfc("EL", pressure, temperature, dewpoint, parcel_profile, which_el=which) @broadcast_nz @@ -260,7 +282,8 @@ def lfc( temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_temperature_profile: np.ndarray | None = None, + parcel_profile: np.ndarray | None = None, + *, which: L["top", "bottom"] = "top", dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, ) -> Vector1d[_T]: @@ -269,7 +292,7 @@ def lfc( pressure, temperature, dewpoint, - parcel_temperature_profile, + parcel_profile, which_lfc=which, dewpoint_start=dewpoint_start, ) @@ -281,8 +304,8 @@ def el_lfc( temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_temperature_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] - | None = None, + parcel_profile: Kelvin[np.ndarray[shape[N, Z], np.dtype[np.floating[Any]]]] | None = None, + *, which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, @@ -292,7 +315,7 @@ def el_lfc( pressure, temperature, dewpoint, - parcel_temperature_profile, + parcel_profile, which_lfc=which_lfc, which_el=which_el, dewpoint_start=dewpoint_start, @@ -335,6 +358,7 @@ def most_unstable_parcel( temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, + *, depth: Pascal[float] = 30_000.0, bottom: Pascal[float] | None = None, ) -> tuple[ @@ -362,6 +386,7 @@ def cape_cin( 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]: @@ -370,16 +395,11 @@ def cape_cin( below_lcl = pressure.is_below( lcl_pressure(pressure[:, 0], temperature[:, 0], dewpoint[:, 0])[:, newaxis] ) - parcel_mixing_ratio = saturation_mixing_ratio( - pressure, dewpoint, out=np.empty_like(temperature), where=below_lcl - ) - parcel_mixing_ratio = saturation_mixing_ratio( - pressure, temperature, out=parcel_mixing_ratio, where=~below_lcl - ) - # Convert the temperature/parcel profile to virtual temperature + parcel_profile = virtual_temperature( + parcel_profile, parcel_mixing_ratio(pressure, temperature, dewpoint, where=below_lcl) + ) temperature = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) - parcel_profile = virtual_temperature(parcel_profile, parcel_mixing_ratio) # Calculate the EL limit of integration (EL, _), (LFC, _) = _el_lfc( "BOTH", @@ -411,25 +431,25 @@ def most_unstable_cape_cin( temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - depth: Pascal[float] = 30_000.0, + *, + depth: Pascal[float] = 30000.0, bottom: Pascal[float] | None = None, ) -> tuple[np.ndarray, np.ndarray]: - mask = ( - np.arange(pressure.shape[1]) - >= most_unstable_parcel_index(pressure, temperature, dewpoint, depth, bottom, **FASTPATH)[ - :, newaxis - ] + idx = most_unstable_parcel_index( + pressure, + temperature, + dewpoint, + depth, + bottom, + **FASTPATH, ) - pressure = pressure.where(mask, -np.inf) - - sort = np.arange(mask.shape[0])[:, newaxis], np.argsort(pressure, axis=1, kind="quicksort") - pressure = pressure[sort][:, ::-1] - pressure[np.isneginf(pressure)] = np.nan + mask = np.arange(pressure.shape[1]) >= idx[:, newaxis] - p, t, td, mu_profile = core.parcel_profile_with_lcl( + p, t, td, mu_profile = parcel_profile_with_lcl( pressure, - temperature[sort][:, ::-1], - dewpoint[sort][:, ::-1], + temperature, + dewpoint, + where=mask, ) return cape_cin(p.view(pressure_vector), t, td, parcel_profile=mu_profile, **FASTPATH) diff --git a/tests/core_test.py b/tests/core_test.py index bbaefb0..c3531e0 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -2,7 +2,6 @@ from __future__ import annotations import itertools -import warnings from typing import Any import metpy.calc as mpcalc @@ -13,7 +12,6 @@ import nzthermo._core as _C from nzthermo._core import moist_lapse -from nzthermo._ufunc import pressure_vector from nzthermo.core import ( cape_cin, ccl, @@ -26,7 +24,7 @@ ) np.set_printoptions( - precision=3, + precision=6, suppress=True, threshold=150, linewidth=150, @@ -61,6 +59,13 @@ def assert_nan(value: np.ndarray, value_units=None): P: np.ndarray = data["P"] T: np.ndarray = data["T"][step] Td: np.ndarray = data["Td"][step] +# In very rare cases the data accessed from the HRRR model had dewpoint temperatures greater than +# the actual temperature. This is not physically possible and is likely due to rounding errors. +# This also makes testing quite difficult because in many cases metpy will report a nan values +# and throw interpolation warnings. To avoid this we will set the dewpoint temperature to be less +# than the actual temperature. +_super_saturation = Td > T +Td[_super_saturation] = T[_super_saturation] def pressure_levels(sfc=1013.25, dtype: Any = np.float64): @@ -306,8 +311,17 @@ def test_moist_lapse(dtype): # ............................................................................................... # # nzthermo._core.parcel_profile # ............................................................................................... # +@pytest.mark.broadcasting @pytest.mark.parcel_profile +def test_parcel_profile_broadcasting() -> None: + assert_array_equal( + _C.parcel_profile(P, T[:, 0], Td[:, 0]), + _C.parcel_profile(np.broadcast_to(P, T.shape), T[:, 0], Td[:, 0]), + ) + + @pytest.mark.regression +@pytest.mark.parcel_profile def test_parcel_profile_metpy_regression() -> None: prof = _C.parcel_profile(P, T[:, 0], Td[:, 0]) for i in range(T.shape[0]): @@ -322,28 +336,31 @@ def test_parcel_profile_metpy_regression() -> None: # ............................................................................................... # # nzthermo._core.parcel_profile_with_lcl # ............................................................................................... # +@pytest.mark.broadcasting @pytest.mark.parcel_profile +def test_parcel_profile_with_lcl_broadcasting() -> None: + p, t, td, tp = _C.parcel_profile_with_lcl(P, T, Td) + p_, t_, td_, tp = _C.parcel_profile_with_lcl(np.broadcast_to(P, T.shape), T, Td) + assert_array_equal(p, p_) + assert_array_equal(t, t_) + assert_array_equal(td, td_) + assert_array_equal(tp, tp) + + @pytest.mark.regression +@pytest.mark.parcel_profile def test_parcel_profile_with_lcl_metpy_regression() -> None: ep, et, etd, ptp = _C.parcel_profile_with_lcl(P, T[:, :], Td[:, :]) for i in range(ep.shape[0]): - with warnings.catch_warnings(): # UserWarning - warnings.simplefilter("ignore") # Interpolation point out of data bounds encountered - ep_, et_, etd_, pt_ = mpcalc.parcel_profile_with_lcl( - P * Pa, - T[i, :] * K, - Td[i, :] * K, - ) + ep_, et_, etd_, pt_ = mpcalc.parcel_profile_with_lcl( + P * Pa, + T[i, :] * K, + Td[i, :] * K, + ) assert_allclose(ep[i], ep_.m, rtol=1e-3) - if np.isnan(et_.m[0]): # warning throw by metpy sometimes caused the first value to be nan - # and the rest of the values to be shifted by one - assert_allclose(et[i, 1:], et_.m[1:], rtol=1e-3) - assert_allclose(etd[i, 1:], etd_.m[1:], rtol=1e-3) - assert_allclose(ptp[i, 1:], pt_.m[1:], rtol=1e-2) - else: - assert_allclose(et[i], et_.m, rtol=1e-3) - assert_allclose(etd[i], etd_.m, rtol=1e-3) - assert_allclose(ptp[i], pt_.m, rtol=1e-3) + assert_allclose(et[i], et_.m, rtol=1e-3) + assert_allclose(etd[i], etd_.m, rtol=1e-3) + assert_allclose(ptp[i], pt_.m, rtol=1e-3) # =============================================================================================== # @@ -352,16 +369,17 @@ def test_parcel_profile_with_lcl_metpy_regression() -> None: # ............................................................................................... # # nzthermo.core.downdraft_cape # ............................................................................................... # +@pytest.mark.broadcasting @pytest.mark.downdraft_cape def test_downdraft_cape_with_broadcasted_pressure() -> None: - assert_allclose( - downdraft_cape(np.broadcast_to(P, (T.shape[0], P.size)), T, Td), + assert_array_equal( + downdraft_cape(np.broadcast_to(P, T.shape), T, Td), downdraft_cape(P, T, Td), ) -@pytest.mark.downdraft_cape @pytest.mark.regression +@pytest.mark.downdraft_cape @pytest.mark.parametrize("dtype", [np.float64, np.float32]) def test_downdraft_cape_metpy_regression(dtype) -> None: DCAPE = downdraft_cape(P.astype(dtype), T.astype(dtype), Td.astype(dtype)) @@ -686,7 +704,7 @@ def test_el_lfc_equals_lcl() -> None: @pytest.mark.el @pytest.mark.skip -def test_el_small_surface_instability(): +def test_el_small_surface_instability() -> None: """Test that no EL is found when there is a small pocket of instability at the sfc.""" levels = [ 959.0, @@ -1031,10 +1049,20 @@ def test_el_profile_nan_with_parcel_profile() -> None: levels = np.array([959.0, 779.2, 751.3, 724.3, 700.0, 269.0]) * hPa temperatures = np.array([22.2, 14.6, np.nan, 9.4, 7.0, -38.0]) * C dewpoints = np.array([19.0, -11.2, -10.8, -10.4, np.nan, -53.2]) * C - parcel_temps = _C.parcel_profile(levels, temperatures[0], dewpoints[0]).to("degC") - el_pressure, el_temperature = el(levels, temperatures, dewpoints, parcel_temps) - assert_almost_equal(el_pressure, 673.0104 * hPa, 3) - assert_almost_equal(el_temperature, 5.8853 * C, 3) + parcel_temps = _C.parcel_profile( + levels.to(Pa).m, + temperatures[:1].to(K).m, + dewpoints[:1].to(K).m, + ) # .to("degC") + el_pressure, el_temperature = el( + levels.to(Pa).m, + temperatures.to(K).m, + 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) @pytest.mark.el @@ -1060,6 +1088,15 @@ def test_el_metpy_regression(which) -> None: # ............................................................................................... # # nzthermo.core.lfc # ............................................................................................... # +@pytest.mark.broadcasting +@pytest.mark.lfc +def test_lfc_broadcasting() -> None: + assert_array_equal( + lfc(P, T, Td), + lfc(np.broadcast_to(P, T.shape), T, Td), + ) + + @pytest.mark.lfc def test_lfc_basic() -> None: """Test LFC calculation.""" @@ -1171,14 +1208,13 @@ def test_lfc_profile_nan() -> None: assert_almost_equal(lfc_t, 9.6977, 0) # RETURNS: 9.58921545636997 -@pytest.mark.lfc @pytest.mark.regression +@pytest.mark.lfc @pytest.mark.parametrize("which", ["top", "bottom"]) def test_lfc_metpy_regression(which) -> None: prof = _C.parcel_profile(P, T[:, 0], Td[:, 0]) - lfc_p, lfc_t = lfc(P, T, Td, prof, which) - print(f"\nlfc {which}") + lfc_p, lfc_t = lfc(P, T, Td, prof, which=which) for i in range(T.shape[0]): lfc_p_, lfc_t_ = mpcalc.lfc( P * Pa, @@ -1194,9 +1230,19 @@ def test_lfc_metpy_regression(which) -> None: # ............................................................................................... # # nzthermo.core.most_unstable_parcel # ............................................................................................... # -@pytest.mark.mu -@pytest.mark.cape +@pytest.mark.broadcasting +@pytest.mark.most_unstable_parcel +@pytest.mark.parametrize("depth", [30000.0]) +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.", + ) + + @pytest.mark.regression +@pytest.mark.most_unstable_parcel @pytest.mark.parametrize("depth", [30000.0]) def test_most_unstable_parcel_index(depth) -> None: assert_array_equal( @@ -1206,22 +1252,24 @@ def test_most_unstable_parcel_index(depth) -> None: for i in range(T.shape[0]) ], ) + + +@pytest.mark.broadcasting +@pytest.mark.most_unstable_parcel +@pytest.mark.parametrize("depth", [30000.0]) +def test_most_unstable_parcel_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[0], P.size)), T, Td, depth=depth), - err_msg="most_unstable_parcel_index failed to on broadcasted pressure input.", + 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.", ) -@pytest.mark.mu -@pytest.mark.cape @pytest.mark.regression -@pytest.mark.most_unstable_cape +@pytest.mark.most_unstable_parcel @pytest.mark.parametrize("depth", [30000.0]) def test_most_unstable_parcel(depth) -> None: - p, t, td, idx = most_unstable_parcel( - np.broadcast_to(P, (T.shape[0], P.size)), T, Td, depth=depth - ) + p, t, td, idx = most_unstable_parcel(P, T, Td, depth=depth) for i in range(T.shape[0]): p_, t_, td_, idx_ = mpcalc.most_unstable_parcel( @@ -1232,66 +1280,24 @@ def test_most_unstable_parcel(depth) -> None: assert_array_equal(td[i], td_.m) assert_array_equal(idx[i], idx_) - assert_array_equal( - (p, t, td, idx), - most_unstable_parcel(np.broadcast_to(P, (T.shape[0], P.size)), T, Td, depth=depth), - err_msg="most_unstable_parcel failed to on broadcasted pressure input.", - ) - - pressure = P.reshape(1, -1).view(pressure_vector) - mask = np.arange(pressure.shape[1]) >= most_unstable_parcel_index(P, T, Td)[:, np.newaxis] - - p_masked = pressure.where(mask, -np.inf) # np.where(mask, pressure, -np.inf) - - sort = ( - np.arange(mask.shape[0])[:, np.newaxis], - np.argsort(p_masked, axis=1, kind="quicksort"), - ) - - p_masked = p_masked[sort][:, ::-1] - nan_mask = np.isneginf(p_masked) - p_masked[nan_mask] = np.nan - - p, t, td, mu_profile = _C.parcel_profile_with_lcl( - p_masked, - T[sort][:, ::-1], - Td[sort][:, ::-1], - ) - - for i in range(T.shape[0]): - _, _, _, idx = mpcalc.most_unstable_parcel( - pressure.squeeze() * Pa, - T[i] * K, - Td[i] * K, - depth=depth * Pa, - ) - - p_, t_, td_, mu_profile_ = mpcalc.parcel_profile_with_lcl( - pressure.squeeze()[idx:] * Pa, - T[i, idx:] * K, - Td[i, idx:] * K, - ) - - m = np.s_[:] - p_ = p_.m[m] - t_ = t_.m[m] - td_ = td_.m[m] - mu_profile_ = mu_profile_.m[m] - - if np.isnan(t_[0]): - continue - - assert_allclose(p[i, : len(p_)], p_, atol=100.0) - assert_allclose(t[i, : len(t_)], t_, atol=0.5) - assert_allclose(td[i, : len(td_)], td_, atol=0.5) - assert_allclose(mu_profile[i, : len(mu_profile_)], mu_profile_, atol=0.5) - - most_unstable_cape_cin(P, T, Td, depth=30000.0) - # ............................................................................................... # # nzthermo.core.cape_cin # ............................................................................................... # +@pytest.mark.cape_cin +@pytest.mark.broadcasting +def test_cape_cin_broadcasting(): + assert_array_equal( + cape_cin(P, T, Td, _C.parcel_profile(P, T[:, 0], Td[:, 0])), + cape_cin( + np.broadcast_to(P, T.shape), + T, + Td, + _C.parcel_profile(np.broadcast_to(P, T.shape), T[:, 0], Td[:, 0]), + ), + ) + + @pytest.mark.cape_cin @pytest.mark.regression @pytest.mark.parametrize( @@ -1334,24 +1340,24 @@ def test_cape_cin_metpy_regression(which_lfc, which_el) -> None: # ............................................................................................... # # nzthermo.core.most_unstable_cape_cin # ............................................................................................... # -@pytest.mark.mu -@pytest.mark.cape_cin -@pytest.mark.regression +@pytest.mark.broadcasting @pytest.mark.most_unstable_cape_cin -@pytest.mark.parametrize( - "depth", - [30000.0], -) -def test_most_unstable_cape_cin_metpy_regression(depth) -> None: - """ - TODO currently this test is passing on 95% of the cases, need to investigate the. - there error appears to be something in the logic block of the el_lfc function. +def test_most_unstable_cape_cin_broadcasting(): + 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, + ), + ) - The current test cases run 500 samples and we are failing on 17 of them specifically when - `which_el=bottom` parameter is used. realistically using the lower EL is not a typical use - case but it should still be tested. - """ +@pytest.mark.regression +@pytest.mark.most_unstable_cape_cin +@pytest.mark.parametrize("depth", [30000.0]) +def test_most_unstable_cape_cin_metpy_regression(depth) -> None: CAPE, CIN = most_unstable_cape_cin( P, T, @@ -1366,4 +1372,5 @@ def test_most_unstable_cape_cin_metpy_regression(depth) -> None: Td[i] * K, depth=depth * Pa, ) - assert_allclose(CAPE[i], CAPE_.m, atol=900) + assert_allclose(CAPE[i], CAPE_.m, atol=10) + assert_allclose(CIN[i], CIN_.m, atol=20) diff --git a/tests/functional_test.py b/tests/functional_test.py index d227818..aea078d 100644 --- a/tests/functional_test.py +++ b/tests/functional_test.py @@ -400,8 +400,8 @@ def test_insert_zero_crossings_specifically_for_cape_cin() -> None: temperature, dewpoint, parcel_profile, - "bottom", - "top", + which_lfc="bottom", + which_el="top", **FASTPATH, ) From fd96a3a576f9245f4502e54e0bd9ee1b8f10e079 Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Wed, 19 Jun 2024 13:13:31 -0500 Subject: [PATCH 6/6] =?UTF-8?q?=E2=9B=88=EF=B8=8Fmixed=20layer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/nzthermo/__init__.py | 12 ++-- src/nzthermo/_ufunc.pyi | 4 +- src/nzthermo/_ufunc.pyx | 14 ++--- src/nzthermo/core.py | 115 +++++++++++++++++++++---------------- src/nzthermo/functional.py | 12 +++- tests/core_test.py | 49 +++++++++++++++- 6 files changed, 143 insertions(+), 63 deletions(-) diff --git a/src/nzthermo/__init__.py b/src/nzthermo/__init__.py index 42b936c..3b132b4 100644 --- a/src/nzthermo/__init__.py +++ b/src/nzthermo/__init__.py @@ -32,9 +32,7 @@ "wet_bulb_temperature", "wet_bulb_potential_temperature", "dewpoint_from_specific_humidity", - "lfc", "parcel_profile", - "mixing_ratio", "vapor_pressure", "virtual_temperature", "wind_components", @@ -42,13 +40,16 @@ "wind_magnitude", "wind_vector", # .core - "el", "cape_cin", "ccl", "downdraft_cape", + "el", + "lfc", + "mixed_layer", + "mixing_ratio", + "most_unstable_cape_cin", "most_unstable_parcel", "most_unstable_parcel_index", - "most_unstable_cape_cin", # .utils "timeseries", ] @@ -102,9 +103,10 @@ downdraft_cape, el, lfc, + mixed_layer, mixing_ratio, + most_unstable_cape_cin, most_unstable_parcel, most_unstable_parcel_index, - most_unstable_cape_cin, ) from .utils import timeseries diff --git a/src/nzthermo/_ufunc.pyi b/src/nzthermo/_ufunc.pyi index 8b2c25a..8b77d82 100644 --- a/src/nzthermo/_ufunc.pyi +++ b/src/nzthermo/_ufunc.pyi @@ -64,7 +64,9 @@ def potential_temperature( pressure: Pascal[float], temperature: Kelvin[float] ) -> Theta[Kelvin[float]]: ... @_ufunc2x1 -def saturation_mixing_ratio(pressure: Pascal[float], temperature: Kelvin[float]) -> float: ... +def saturation_mixing_ratio( + pressure: Pascal[float], temperature: Kelvin[float] +) -> Dimensionless[float]: ... @_ufunc2x1 def virtual_temperature( temperature: Kelvin[float], vapor_pressure: Pascal[float] diff --git a/src/nzthermo/_ufunc.pyx b/src/nzthermo/_ufunc.pyx index 3a78789..70efd0e 100644 --- a/src/nzthermo/_ufunc.pyx +++ b/src/nzthermo/_ufunc.pyx @@ -13,16 +13,16 @@ that generates the stub file from the c++ header file. # cython: cdivision=True # pyright: reportGeneralTypeIssues=false +from typing import TypeVar + +import numpy as np -# c imports cimport cython cimport numpy as np from libcpp.cmath cimport fabs, isnan -cimport nzthermo._C as C -from nzthermo._C cimport epsilon -import numpy as np -from typing import TypeVar +cimport nzthermo._C as C +from nzthermo._C cimport Md, Mw np.import_array() np.import_ufunc() @@ -193,7 +193,7 @@ cdef T wobus(T temperature) noexcept nogil: # 2x1 @cython.ufunc cdef T mixing_ratio(T partial_pressure, T total_pressure) noexcept nogil: - return epsilon * partial_pressure / (total_pressure - partial_pressure) + return Mw / Md * partial_pressure / (total_pressure - partial_pressure) @cython.ufunc # theta @@ -205,7 +205,7 @@ cdef T potential_temperature(T pressure, T temperature) noexcept nogil: cdef T dewpoint_from_specific_humidity(T pressure, T specific_humidity) noexcept nogil: cdef T Q = specific_humidity or 1e-9 cdef T r = Q / (1 - Q) - return C.dewpoint(pressure * r / (epsilon + r)) + return C.dewpoint(pressure * r / (Mw / Md + r)) @cython.ufunc diff --git a/src/nzthermo/core.py b/src/nzthermo/core.py index a0737c4..7577de1 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -12,10 +12,10 @@ from typing import Any, Final, Literal as L, TypeVar import numpy as np -from numpy.typing import NDArray +from numpy.typing import ArrayLike, NDArray from . import functional as F -from ._core import parcel_profile_with_lcl, moist_lapse +from ._core import moist_lapse, parcel_profile_with_lcl from ._ufunc import ( dewpoint as _dewpoint, dry_lapse, @@ -35,47 +35,29 @@ from .typing import Kelvin, N, Pascal, Z, shape from .utils import Axis, Vector1d, 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} -def mask_layers( - 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]]], - /, - *, - where: np.ndarray[shape[N, Z], np.dtype[np.bool_]], - fill_value: float = NaN, -) -> tuple[ - Pascal[np.ndarray[shape[N, Z], np.dtype[_T]]], - Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], - Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], - bool, -]: - if broadcasted := pressure.shape == temperature.shape: - p_layer, t_layer, td_layer = np.where( - where[newaxis, :, :], [pressure, temperature, dewpoint], fill_value - ) - else: - p_layer = np.where(where, pressure, fill_value) - t_layer, td_layer = np.where(where[newaxis, :, :], [temperature, dewpoint], fill_value) - - return p_layer, t_layer, td_layer, broadcasted - - def parcel_mixing_ratio( - 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[pressure_vector[_S, np.dtype[_T]]], + temperature: Kelvin[np.ndarray[_S, np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[_S, np.dtype[_T]]], /, *, - where: np.ndarray[shape[N, Z], np.dtype[np.bool_]], -): + where: np.ndarray[_S, np.dtype[np.bool_]] | None = None, +) -> np.ndarray[_S, np.dtype[_T]]: + if where is None: + where = pressure.is_below( + lcl_pressure(pressure[surface], temperature[surface], dewpoint[surface]) + ) r = saturation_mixing_ratio(pressure, dewpoint, out=np.empty_like(temperature), where=where) r = saturation_mixing_ratio(pressure, temperature, out=r, where=~where) return r @@ -167,11 +149,8 @@ def ccl( saturation mixing ratio line (through the surface dewpoint) and the environmental temperature. """ - p0 = pressure[:, 0] # (N,) - td0 = dewpoint[:, 0] # (N,) - if mixed_layer_depth is None: - r = mixing_ratio(saturation_vapor_pressure(td0[:, newaxis]), p0[:, newaxis]) + r = mixing_ratio(saturation_vapor_pressure(dewpoint[surface]), pressure[surface]) else: raise NotImplementedError if height is not None: @@ -181,7 +160,7 @@ def ccl( p, t = F.intersect_nz(pressure, rt_profile, temperature, "increasing", log_x=True).pick(which) - return p, t, dry_lapse(p0, t, p) + return p, t, dry_lapse(pressure[:, 0], t, p) # ------------------------------------------------------------------------------------------------- @@ -213,9 +192,9 @@ def _el_lfc( LCL = Vector1d.from_func(lcl, p0, t0, td0).unsqueeze() pressure, parcel_profile, temperature = ( - pressure[:, 1:], - parcel_profile[:, 1:], - temperature[:, 1:], + pressure[aloft], + parcel_profile[aloft], + temperature[aloft], ) if pick != "LFC": # find the Equilibrium Level (EL) @@ -323,7 +302,7 @@ def el_lfc( # ------------------------------------------------------------------------------------------------- -# cape_cin +# nzthermo.core.most_unstable_parcel # ------------------------------------------------------------------------------------------------- @broadcast_nz def most_unstable_parcel_index( @@ -334,11 +313,11 @@ def most_unstable_parcel_index( depth: float = 30000.0, height: float | None = None, bottom: float | None = None, -): +) -> np.ndarray[shape[N], np.dtype[np.intp]]: if height is not None: raise NotImplementedError("height argument is not implemented") - pbot = (pressure[:, 0] if bottom is None else np.asarray(bottom)).reshape(-1, 1) + pbot = (pressure[surface] if bottom is None else np.asarray(bottom)).reshape(-1, 1) ptop = pbot - depth theta_e = equivalent_potential_temperature( @@ -359,7 +338,7 @@ def most_unstable_parcel( dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, *, - depth: Pascal[float] = 30_000.0, + depth: Pascal[float] = 30000.0, bottom: Pascal[float] | None = None, ) -> tuple[ Pascal[np.ndarray[shape[N], np.dtype[_T]]], @@ -379,6 +358,49 @@ def most_unstable_parcel( ) +# ------------------------------------------------------------------------------------------------- +# nzthermo.core.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]]], + /, + *, + depth: float | NDArray[np.floating[Any]] = 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]]], +]: + 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) + + depth = np.asarray( + # use asarray otherwise the depth is cast to pressure_vector which doesn't + # make sense for the temperature and dewpoint outputs + np.max(pressure, initial=-np.inf, axis=1, where=where) + - 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 + + +# ------------------------------------------------------------------------------------------------- +# cape_cin +# ------------------------------------------------------------------------------------------------- @broadcast_nz def cape_cin( pressure: Pascal[pressure_vector[shape[N, Z], np.dtype[_T]]], @@ -392,12 +414,9 @@ def cape_cin( ) -> tuple[np.ndarray, np.ndarray]: # The mixing ratio of the parcel comes from the dewpoint below the LCL, is saturated # based on the temperature above the LCL - below_lcl = pressure.is_below( - lcl_pressure(pressure[:, 0], temperature[:, 0], dewpoint[:, 0])[:, newaxis] - ) parcel_profile = virtual_temperature( - parcel_profile, parcel_mixing_ratio(pressure, temperature, dewpoint, where=below_lcl) + parcel_profile, parcel_mixing_ratio(pressure, temperature, dewpoint) ) temperature = virtual_temperature(temperature, saturation_mixing_ratio(pressure, dewpoint)) # Calculate the EL limit of integration diff --git a/src/nzthermo/functional.py b/src/nzthermo/functional.py index f153327..09c4f01 100644 --- a/src/nzthermo/functional.py +++ b/src/nzthermo/functional.py @@ -44,11 +44,16 @@ def nanroll_2d( return args +from numpy.typing import ArrayLike + + def nantrapz( y: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, x: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co | None = None, dx: float = 1.0, axis: SupportsIndex = -1, + *, + where: ArrayLike | None = None, ) -> NDArray[_T]: r""" This is a clone of the `numpy.trapz` function but with support for `nan` values. @@ -91,8 +96,13 @@ def nantrapz( The try-except block was removed because it was not necessary, for the use case of this of this library. """ + if where is not None: + y = np.where(where, y, np.nan) + if x is not None: + x = np.where(where, x, np.nan) - y = np.asanyarray(y) + else: + y = np.asanyarray(y) if x is None: d = dx else: diff --git a/tests/core_test.py b/tests/core_test.py index c3531e0..38fc7b1 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -18,6 +18,7 @@ downdraft_cape, el, lfc, + mixed_layer, most_unstable_cape_cin, most_unstable_parcel, most_unstable_parcel_index, @@ -1254,6 +1255,9 @@ def test_most_unstable_parcel_index(depth) -> None: ) +# ------------------------------------------------------------------------------------------------- +# nzthermo.core.mixed_layer +# ------------------------------------------------------------------------------------------------- @pytest.mark.broadcasting @pytest.mark.most_unstable_parcel @pytest.mark.parametrize("depth", [30000.0]) @@ -1268,7 +1272,7 @@ def test_most_unstable_parcel_broadcasting(depth) -> None: @pytest.mark.regression @pytest.mark.most_unstable_parcel @pytest.mark.parametrize("depth", [30000.0]) -def test_most_unstable_parcel(depth) -> None: +def test_most_unstable_parcel_regression(depth) -> None: p, t, td, idx = most_unstable_parcel(P, T, Td, depth=depth) for i in range(T.shape[0]): @@ -1281,6 +1285,49 @@ def test_most_unstable_parcel(depth) -> None: assert_array_equal(idx[i], idx_) +# ............................................................................................... # +# nzthermo.core.mixed_layer +# ............................................................................................... # +@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( + 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, + ) + + # ............................................................................................... # # nzthermo.core.cape_cin # ............................................................................................... #