⛈️mixed layer

src/nzthermo/__init__.py

@@ -32,23 +32,24 @@
     "wet_bulb_temperature",
     "wet_bulb_potential_temperature",
     "dewpoint_from_specific_humidity",
-    "lfc",
     "parcel_profile",
-    "mixing_ratio",
     "vapor_pressure",
     "virtual_temperature",
     "wind_components",
     "wind_direction",
     "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

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]

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

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

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: