From fd96a3a576f9245f4502e54e0bd9ee1b8f10e079 Mon Sep 17 00:00:00 2001 From: Jason Leaver Date: Wed, 19 Jun 2024 13:13:31 -0500 Subject: [PATCH] =?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 # ............................................................................................... #