Development

pyproject.toml

@@ -49,6 +49,7 @@ markers = [
     "mu_cape",
     "parcel_profile",
     "regression",
+    "broadcasting",
 ]
 
 pythonpath = ["src"]

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 <floating T>
-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 <floating T>
 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 <floating T>

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 <floating T>
+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 <floating T>
 constexpr T mixing_ratio(const T partial_press, const T total_press) noexcept {
     return epsilon * partial_press / (total_press - partial_press);
@@ -141,21 +176,23 @@ constexpr T lcl_pressure(
 
 template <floating T>
 constexpr lcl<T>::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<T>, 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<T>, max_iters, eps, pressure_, temperature_, r);
+    temperature = dewpoint(pressure, r);
 }
 
 template <floating T>
 constexpr T lcl<T>::wet_bulb_temperature(const T pressure, const T step) noexcept {
     return moist_lapse(this->pressure, pressure, this->temperature, step);
 }
 
+template <floating T>
+constexpr size_t lcl<T>::index(const T pressure[], const size_t size) noexcept {
+    return index_pressure(pressure, this->pressure, size);
+}
+
 template <floating T>
 constexpr T wet_bulb_temperature(
   const T pressure,

src/nzthermo/_C.pxd

@@ -1,16 +1,6 @@
 # cython: boundscheck=False
 # pyright: reportGeneralTypeIssues=false
 
-ctypedef fused Float:
-    float
-    double
-
-
-cdef extern from "<utility>" 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

src/nzthermo/__init__.py

@@ -32,20 +32,22 @@
     "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",
     # .utils
@@ -101,7 +103,9 @@
     downdraft_cape,
     el,
     lfc,
+    mixed_layer,
     mixing_ratio,
+    most_unstable_cape_cin,
     most_unstable_parcel,
     most_unstable_parcel_index,
 )

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]]],