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/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..b2bb204 100644 --- a/src/lib/libthermo.cpp +++ b/src/lib/libthermo.cpp @@ -2,6 +2,41 @@ 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 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)); + + 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 +176,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 +188,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[], const 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 a085711..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,14 +18,10 @@ 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 + cdef cppclass wind_components[T]: T u, v wind_components() noexcept @@ -48,9 +34,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 @@ -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/__init__.py b/src/nzthermo/__init__.py index b1af963..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,10 +40,14 @@ "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", # .utils @@ -101,7 +103,9 @@ downdraft_cape, el, lfc, + mixed_layer, mixing_ratio, + most_unstable_cape_cin, most_unstable_parcel, most_unstable_parcel_index, ) 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 1687c67..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( @@ -154,8 +232,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: 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: @@ -225,11 +300,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. @@ -289,7 +365,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. ]) """ @@ -298,6 +374,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: @@ -345,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]( @@ -375,25 +453,21 @@ cdef void parcel_profile_1d( T[:] pressure, # (Z,) T temperature, T dewpoint, - ) 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](p0, t0, dewpoint) + lcl = C.lcl[T](pressure[0], temperature, dewpoint) - # [dry ascent] + # [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) - + out[i] = C.dry_lapse(pressure[i], p0, temperature) # [ moist ascent ] if stop != Z: @@ -405,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'): @@ -423,53 +494,32 @@ 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 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, - double step = 1000.0, - double eps = 0.1, - size_t max_iters = 5, + np.ndarray where = None, ): cdef: size_t N, Z BroadcastMode mode np.ndarray out - (pressure, temperature, dewpoint), mode = pressure_mode(pressure, temperature, dewpoint) + (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 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, - ) + out = np.empty((N, Z), dtype=dtype) + if dtype == np.float64: + out[...] = parcel_profile_2d[double](pressure, temperature, dewpoint, mode) else: - raise ValueError("Invalid strategy.") + out[...] = parcel_profile_2d[float](pressure, temperature, dewpoint, mode) return out @@ -497,40 +547,47 @@ 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) + + 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 - # [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) - - ep[:stop] = pressure[:stop] - et[:stop] = temperature[:stop] - etd[:stop] = dewpoint[:stop] - - # [ 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 # [ 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) @@ -541,20 +598,19 @@ cdef T[:, :, :] parcel_profile_with_lcl_2d( BroadcastMode mode, ) noexcept: cdef: - size_t N, Z, i - T[:, :, :] out + size_t N, Z, i, idx + T[:, :] ep, et, etd, pt 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) - + 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,44 +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, :], - pressure[i, :], - temperature[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] @@ -613,11 +665,11 @@ 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] - 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) @@ -684,7 +736,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) @@ -697,7 +748,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]) @@ -722,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: @@ -750,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.pyi b/src/nzthermo/_ufunc.pyi index 10899de..8b77d82 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: ... @@ -45,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 5dfd3f5..70efd0e 100644 --- a/src/nzthermo/_ufunc.pyx +++ b/src/nzthermo/_ufunc.pyx @@ -13,16 +13,24 @@ that generates the stub file from the c++ header file. # cython: cdivision=True # pyright: reportGeneralTypeIssues=false +from typing import TypeVar + +import numpy as np cimport cython cimport numpy as np +from libcpp.cmath cimport fabs, isnan cimport nzthermo._C as C -from nzthermo._C cimport epsilon +from nzthermo._C cimport Md, Mw 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 @@ -160,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 @@ -172,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 @@ -198,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 bddd251..7577de1 100644 --- a/src/nzthermo/core.py +++ b/src/nzthermo/core.py @@ -12,11 +12,11 @@ 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 _core as core, functional as F +from . import functional as F +from ._core import moist_lapse, parcel_profile_with_lcl from ._ufunc import ( - between_or_close, dewpoint as _dewpoint, dry_lapse, equivalent_potential_temperature, @@ -24,6 +24,7 @@ lcl, lcl_pressure, mixing_ratio, + pressure_vector, saturation_mixing_ratio, saturation_vapor_pressure, vapor_pressure, @@ -32,27 +33,47 @@ ) from .const import Rd from .typing import Kelvin, N, Pascal, Z, shape -from .utils import Vector1d, broadcast_nz +from .utils import Axis, Vector1d, broadcast_nz -float_ = TypeVar("float_", bound=np.floating[Any], covariant=True) +_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 -z_axis: Final[tuple[slice, None]] = np.s_[:, newaxis] -N_AXIS: Final[tuple[None, slice]] = np.s_[newaxis, :] FASTPATH: dict[str, Any] = {"__fastpath": True} +def parcel_mixing_ratio( + 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[_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 + + # ------------------------------------------------------------------------------------------------- # downdraft_cape # ------------------------------------------------------------------------------------------------- @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[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). Calculate the downward convective available potential energy (DCAPE) of a given upper air @@ -74,29 +95,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) - 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=where, + out=np.full_like(temperature, np.inf), + ) + zx = theta_e.argmin(axis=1) - 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 = 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) + 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) @@ -109,10 +130,11 @@ 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, which: L["bottom", "top"] = "bottom", @@ -127,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: @@ -141,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) # ------------------------------------------------------------------------------------------------- @@ -149,18 +168,17 @@ 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[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_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[float_]] | None = None, -) -> tuple[Vector1d[float_], Vector1d[float_]] | Vector1d[float_]: - if parcel_temperature_profile is None: - pressure, temperature, dewpoint, parcel_temperature_profile = core.parcel_profile_with_lcl( + dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, +) -> tuple[Vector1d[_T], Vector1d[_T]] | Vector1d[_T]: + if parcel_profile is None: + pressure, temperature, dewpoint, parcel_profile = parcel_profile_with_lcl( pressure, temperature, dewpoint ) @@ -173,42 +191,44 @@ def _el_lfc( LCL = Vector1d.from_func(lcl, p0, t0, td0).unsqueeze() - pressure, parcel_temperature_profile, temperature = ( - pressure[:, 1:], - parcel_temperature_profile[:, 1:], - temperature[:, 1:], + pressure, parcel_profile, temperature = ( + pressure[aloft], + parcel_profile[aloft], + temperature[aloft], ) - 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.Z) - 1 + left_of_env = (parcel_profile[top_idx] <= temperature[top_idx])[:, newaxis] + EL = F.intersect_nz( + pressure, + parcel_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, - parcel_temperature_profile, + parcel_profile, temperature, "increasing", 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), + np.where(LCL.is_below(pressure, close=True), parcel_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,35 +244,34 @@ 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[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_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[float_]: - return _el_lfc( - "EL", pressure, temperature, dewpoint, parcel_temperature_profile, which_el=which - ) +) -> Vector1d[_T]: + return _el_lfc("EL", pressure, temperature, dewpoint, parcel_profile, which_el=which) @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[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_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[float_]] | None = None, -) -> Vector1d[float_]: + dewpoint_start: np.ndarray[shape[N], np.dtype[_T]] | None = None, +) -> Vector1d[_T]: return _el_lfc( "LFC", pressure, temperature, dewpoint, - parcel_temperature_profile, + parcel_profile, which_lfc=which, dewpoint_start=dewpoint_start, ) @@ -260,22 +279,22 @@ 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[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, - parcel_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[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, temperature, dewpoint, - parcel_temperature_profile, + parcel_profile, which_lfc=which_lfc, which_el=which_el, dewpoint_start=dewpoint_start, @@ -283,55 +302,125 @@ def el_lfc( # ------------------------------------------------------------------------------------------------- -# cape_cin +# nzthermo.core.most_unstable_parcel # ------------------------------------------------------------------------------------------------- +@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, bottom: float | None = None, -): +) -> np.ndarray[shape[N], np.dtype[np.intp]]: 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[surface] 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] = 30000.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, + ) + + +# ------------------------------------------------------------------------------------------------- +# 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[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[pressure_vector[shape[N, Z], np.dtype[_T]]], + temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], + dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[_T]]], /, parcel_profile: np.ndarray, + *, which_lfc: L["bottom", "top"] = "bottom", which_el: L["bottom", "top"] = "top", ) -> tuple[np.ndarray, np.ndarray]: - 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), + + parcel_profile = virtual_temperature( + parcel_profile, parcel_mixing_ratio(pressure, temperature, dewpoint) ) - # 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 +429,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[float_]]], - temperature: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], - dewpoint: Kelvin[np.ndarray[shape[N, Z], np.dtype[float_]]], +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, + *, + depth: Pascal[float] = 30000.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_]]], - 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 - - 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] +) -> tuple[np.ndarray, np.ndarray]: + idx = most_unstable_parcel_index( + pressure, + temperature, + dewpoint, + depth, + bottom, + **FASTPATH, + ) + mask = np.arange(pressure.shape[1]) >= idx[:, newaxis] - # theta_e = equivalent_potential_temperature(p, t, td) - # idx = np.argmax(theta_e, axis=1) + p, t, td, mu_profile = parcel_profile_with_lcl( + pressure, + temperature, + dewpoint, + where=mask, + ) - # 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 cab81dc..09c4f01 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 @@ -33,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. @@ -80,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: @@ -209,12 +230,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..18b9062 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, @@ -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, @@ -46,22 +52,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 +75,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 +92,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 +109,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 +120,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, *, @@ -148,14 +148,41 @@ 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) 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,29 +205,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: - 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_]]]: + 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": @@ -210,10 +233,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: @@ -223,17 +246,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: @@ -253,44 +273,39 @@ 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[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, ], - _T, + _S, ], ) -> Callable[ - Concatenate[ - Pascal[ArrayLike[np.floating[Any]]], - Kelvin[ArrayLike[np.floating[Any]]], - Kelvin[ArrayLike[np.floating[Any]]], - _P, - ], - _T, + Concatenate[Pascal[ArrayLike[_T]], Kelvin[ArrayLike[_T]], Kelvin[ArrayLike[_T]], _P], + _S, ]: @functools.wraps(f) def wrapper( - pressure: ArrayLike[np.floating[Any]], - temperature: ArrayLike[np.floating[Any]], - dewpoint: ArrayLike[np.floating[Any]], + 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 + # 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_vector(pressure), temperature, dewpoint, *args, **kwargs) return wrapper @@ -305,39 +320,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_]: @@ -346,9 +361,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) @@ -365,9 +378,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") @@ -386,20 +398,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") @@ -498,7 +508,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 @@ -515,15 +526,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, @@ -546,12 +557,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]", @@ -578,9 +589,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) diff --git a/tests/core_test.py b/tests/core_test.py index 1419f9b..38fc7b1 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -1,6 +1,7 @@ # noqa +from __future__ import annotations + import itertools -import warnings from typing import Any import metpy.calc as mpcalc @@ -17,11 +18,14 @@ downdraft_cape, el, lfc, + mixed_layer, + most_unstable_cape_cin, + most_unstable_parcel, most_unstable_parcel_index, ) np.set_printoptions( - precision=3, + precision=6, suppress=True, threshold=150, linewidth=150, @@ -56,6 +60,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): @@ -301,8 +312,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]): @@ -317,28 +337,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) # =============================================================================================== # @@ -347,16 +370,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)) @@ -681,7 +705,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, @@ -1026,10 +1050,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 @@ -1055,6 +1089,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.""" @@ -1166,28 +1209,13 @@ def test_lfc_profile_nan() -> None: assert_almost_equal(lfc_t, 9.6977, 0) # RETURNS: 9.58921545636997 -@pytest.mark.cape -@pytest.mark.mu_cape @pytest.mark.regression -@pytest.mark.parametrize("depth", [30000.0]) -def test_most_unstable_parcel_index(depth) -> None: - assert_array_equal( - most_unstable_parcel_index(P, T, Td, depth=depth), - [ - mpcalc.most_unstable_parcel(P * Pa, T[i] * K, Td[i] * K, depth=depth * Pa)[-1] - for i in range(T.shape[0]) - ], - ) - - @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}") + 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, @@ -1200,9 +1228,123 @@ def test_lfc_metpy_regression(which) -> None: assert_allclose(lfc_t[i], lfc_t_.m, atol=1.0) +# ............................................................................................... # +# nzthermo.core.most_unstable_parcel +# ............................................................................................... # +@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( + most_unstable_parcel_index(P, T, Td, depth=depth), + [ + mpcalc.most_unstable_parcel(P * Pa, T[i] * K, Td[i] * K, depth=depth * Pa)[-1] + for i in range(T.shape[0]) + ], + ) + + +# ------------------------------------------------------------------------------------------------- +# nzthermo.core.mixed_layer +# ------------------------------------------------------------------------------------------------- +@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(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.regression +@pytest.mark.most_unstable_parcel +@pytest.mark.parametrize("depth", [30000.0]) +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]): + 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_) + + +# ............................................................................................... # +# 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 # ............................................................................................... # +@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( @@ -1240,3 +1382,42 @@ 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.broadcasting +@pytest.mark.most_unstable_cape_cin +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, + ), + ) + + +@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, + 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=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, )