Linfeng - First PR of c++ thermo using pybind11

pyproject.toml

@@ -1,9 +1,9 @@
 [build-system]
-requires = ["setuptools>=61", "wheel", "setuptools_scm[toml]>=3.4"]
-build-backend = "setuptools.build_meta"
+requires = ["setuptools>=61", "setuptools-scm[toml]>=3.4", "scikit-build-core", "pybind11"]
+build-backend = "scikit_build_core.build"
 
 [project]
-name = "MetPy"
+name = "metpy"
 description = "Collection of tools for reading, visualizing and performing calculations with weather data."
 readme = "README.md"
 dynamic = ["version"]
@@ -174,5 +174,9 @@ max-complexity = 61
 [tool.ruff.lint.pydocstyle]
 convention = "numpy"
 
+[tool.scikit-build]
+metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
+cmake.source-dir = "src"
+
 [tool.setuptools_scm]
 version_scheme = "post-release"

src/CMakeLists.txt

@@ -0,0 +1,24 @@
+cmake_minimum_required(VERSION 3.15...3.26)
+project(
+  "${SKBUILD_PROJECT_NAME}"
+  LANGUAGES CXX
+  VERSION "${SKBUILD_PROJECT_VERSION}")
+
+# Enforce C++14 or higher
+#set(CMAKE_CXX_STANDARD 14)
+#set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
+find_package(pybind11 CONFIG REQUIRED)
+
+pybind11_add_module(_calc_mod MODULE
+  calcmod.cpp
+  constants.cpp
+  math.cpp
+  thermo.cpp
+)
+
+target_compile_definitions(_calc_mod
+                           PRIVATE VERSION_INFO=${PROJECT_VERSION})
+
+install(TARGETS _calc_mod DESTINATION metpy)

src/calcmod.cpp

@@ -0,0 +1,111 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+#include "constants.hpp"
+#include "thermo.hpp"
+#include <utility> // For std::pair
+#include <tuple>   // For std::make_tuple
+#include <iostream>
+
+namespace py = pybind11;
+
+int add(int i, int j) {
+    return i + j;
+}
+
+PYBIND11_MODULE(_calc_mod, m) {
+    m.doc() = "accelerator module docstring";
+
+    metpy_constants::load_constants_from_python();
+
+    m.def("add", &add, "Add two numbers");
+
+    m.def("moist_air_gas_constant", py::vectorize(MoistAirGasConstant),
+            "Calculate R_m, the gas constant for moist air.",
+            py::arg("specific_humidity"));
+
+    m.def("moist_air_specific_heat_pressure", py::vectorize(MoistAirSpecificHeatPressure),
+            "Calculate C_pm, the specific heat of moist air at constant pressure.",
+            py::arg("specific_humidity"));
+
+    m.def("water_latent_heat_vaporization", py::vectorize(WaterLatentHeatVaporization),
+            "Calculate water latent heat vaporization from temperature.",
+            py::arg("temperature"));
+
+    m.def("water_latent_heat_sublimation", py::vectorize(WaterLatentHeatSublimation),
+            "Calculate water latent heat sublimation from temperature.",
+            py::arg("temperature"));
+
+    m.def("relative_humidity_from_dewpoint", py::vectorize(RelativeHumidityFromDewPoint),
+            "Calculate relative humidity from temperature and dewpoint.",
+            py::arg("temperature"), py::arg("dewpoint"), py::arg("phase"));
+
+    m.def("dry_lapse", &DryLapseVectorized,
+            "Calculate the temperature along a pressure profile assuming dry adiabatic process.",
+            py::arg("pressure"), py::arg("ref_temperature"), py::arg("ref_pressure"));
+
+    m.def("dry_lapse_3d", &DryLapseVectorized_3D,
+            "Calculate the temperature along multiple pressure profiles assuming dry adiabatic process.",
+            py::arg("pressure"), py::arg("ref_temperature"), py::arg("ref_pressure"));
+
+    m.def("moist_lapse", &MoistLapseVectorized,
+          "Calculate the temperature along a pressure profile assuming saturated adiabatic process.",
+          py::arg("pressure"), py::arg("ref_temperature"), py::arg("ref_pressure"), py::arg("rk_nstep"));
+
+
+    m.def("lcl", &LCLVectorized,
+            "Calculate the lifting condensation level (LCL) from pressure, temperature and dewpoint.",
+            py::arg("pressure"), py::arg("temperature"), py::arg("dewpoint"));
+
+    m.def("parcel_profile",
+            [](py::array_t<double> pressure, double temperature, double dewpoint) {
+                // pressure.data() gives the beginning pointer of the data buffer
+                std::vector<double> pressure_vec(pressure.data(), pressure.data() + pressure.size());
+                std::vector<double> temp_prof = ParcelProfile(pressure_vec, temperature, dewpoint);
+                return py::array_t<double>(temp_prof.size(), temp_prof.data());
+            },
+            "Compute the parcel temperature profile as it rises from a given pressure and temperature.",
+            py::arg("pressure"), py::arg("temperature"), py::arg("dewpoint"));
+
+
+    m.def("saturation_vapor_pressure", py::vectorize(SaturationVaporPressure),
+            "Calculate saturation vapor pressure from temperature.",
+            py::arg("temperature"), py::arg("phase") = "liquid");
+
+    m.def("_saturation_vapor_pressure_liquid", py::vectorize(_SaturationVaporPressureLiquid),
+            "Calculate saturation vapor pressure from temperature.",
+            py::arg("temperature"));
+
+    m.def("_saturation_vapor_pressure_solid", py::vectorize(_SaturationVaporPressureSolid),
+            "Calculate saturation vapor pressure from temperature.",
+            py::arg("temperature"));
+
+    m.def("dewpoint", py::vectorize(DewPoint),
+            "Calculate dewpoint from water vapor partial pressure.",
+            py::arg("vapor_pressure"));
+
+    m.def("mixing_ratio", py::vectorize(MixingRatio),
+            "Calculate the mixing ratio of a gas.",
+            py::arg("partial_press"), py::arg("total_press"), py::arg("epsilon"));
+
+    m.def("saturation_mixing_ratio", py::vectorize(SaturationMixingRatio),
+            "Calculate the saturation mixing ratio of water vapor given total atmospheric pressure and temperature.",
+            py::arg("total_press"), py::arg("temperature"), py::arg("phase"));
+
+    m.def("specific_humidity_from_mixing_ratio", py::vectorize(SpecificHumidityFromMixingRatio),
+            "Calculate the specific humidity from the mixing ratio.",
+            py::arg("mixing_ratio"));
+
+    m.def("specific_humidity_from_dewpoint", py::vectorize(SpecificHumidityFromDewPoint),
+            "Calculate the specific humidity from the dewpoint temperature and pressure.",
+            py::arg("pressure"), py::arg("dewpoint"), py::arg("phase"));
+
+    m.def("virtual_temperature", py::vectorize(VirtualTemperature),
+            "Calculate virtual temperature from temperature and mixing ratio.",
+            py::arg("temperature"), py::arg("mixing_ratio"), py::arg("epsilon"));
+
+    m.def("virtual_temperature_from_dewpoint", py::vectorize(VirtualTemperatureFromDewpoint),
+            "Calculate virtual temperature from dewpoint.",
+            py::arg("pressure"), py::arg("temperature"), py::arg("dewpoint"), py::arg("epsilon"), py::arg("phase"));
+
+}

src/constants.cpp

@@ -0,0 +1,50 @@
+// constants.cpp
+#include "constants.hpp"
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+namespace metpy_constants {
+    double Mw;
+    double Rd;
+    double Rv;
+    double Cp_d;
+    double Cp_v;
+    double Cp_l;
+    double Lv;
+    double sat_pressure_0c;
+    double T0;
+    double Ls;
+    double Cp_i;
+    double zero_degc;
+    double epsilon;
+    double kappa;
+
+    void load_constants_from_python() {
+        py::object mod = py::module_::import("metpy.constants.nounit");
+
+//        Mw = mod.attr("Mw").attr("magnitude").cast<double>();
+//        Rv = mod.attr("Rv").attr("magnitude").cast<double>();
+//        Cp_v = mod.attr("Cp_v").attr("magnitude").cast<double>();
+//        Cp_l = mod.attr("Cp_l").attr("magnitude").cast<double>();
+//        Lv = mod.attr("Lv").attr("magnitude").cast<double>(); 
+//        sat_pressure_0c = mod.attr("sat_pressure_0c").cast<double>();
+//        T0 = mod.attr("T0").attr("magnitude").cast<double>();
+
+
+        Mw = mod.attr("Mw").cast<double>();
+        Rd = mod.attr("Rd").cast<double>();
+        Rv = mod.attr("Rv").cast<double>();
+        Cp_d = mod.attr("Cp_d").cast<double>();
+        Cp_v = mod.attr("Cp_v").cast<double>();
+        Cp_l = mod.attr("Cp_l").cast<double>();
+        Lv = mod.attr("Lv").cast<double>(); 
+        sat_pressure_0c = mod.attr("sat_pressure_0c").cast<double>();
+        T0 = mod.attr("T0").cast<double>();
+        Ls = mod.attr("Ls").cast<double>();
+        Cp_i = mod.attr("Cp_i").cast<double>();
+        zero_degc = mod.attr("zero_degc").cast<double>();
+        epsilon = mod.attr("epsilon").cast<double>();
+        kappa = mod.attr("kappa").cast<double>();
+    }
+}

src/constants.hpp

@@ -0,0 +1,49 @@
+#ifndef CONSTANTS_HPP
+#define CONSTANTS_HPP
+
+namespace metpy_constants {
+
+    // Gas constants (J / kg / K)
+//    extern double R;
+
+    // Dry air
+//    extern double Md;
+//    extern double Rd;
+//    extern double dry_air_spec_heat_ratio;
+//    extern double Cp_d = Rd * dry_air_spec_heat_ratio / (dry_air_spec_heat_ratio - 1.0); // (J / kg / K)
+//    extern double Cv_d = Cp_d / dry_air_spec_heat_ratio; // (J / kg / K)
+
+    // Water
+    extern double Mw;
+    extern double Rd;
+    extern double Rv;
+//    extern double wv_spec_heat_ratio = 1.33;
+    extern double Cp_d;
+    extern double Cp_v;
+//    extern double Cv_v = Cp_v / wv_spec_heat_ratio; // (J / kg / K)
+    extern double Cp_l;
+    extern double Lv;
+    extern double sat_pressure_0c;
+    extern double T0;
+    extern double Ls;
+    extern double Cp_i;
+    extern double zero_degc;
+    extern double epsilon;
+    extern double kappa;
+
+    // General Meteorological constants
+//    extern double epsilon = Mw / Md; // ≈ 0.622
+
+
+    // Standard gravity
+//    extern double g = 9.80665;       // m / s^2
+
+    // Reference pressure
+//    extern double P0 = 100000.0;     // Pa (hPa = 1000)
+
+
+    void load_constants_from_python();  // call once in your PYBIND11_MODULE
+}
+
+#endif
+

src/math.cpp

@@ -0,0 +1,32 @@
+#include <cmath>
+#include <iostream>   // for std::cerr
+#include "math.hpp"
+
+double lambert_wm1(double x, double tol, int max_iter) {
+    // Corless, R.M., Gonnet, G.H., Hare, D.E.G. et al. On the LambertW function.
+    // Adv Comput Math 5, 329–359 (1996). https://doi.org/10.1007/BF02124750 
+
+    if (x >= 0 || x < -1.0 / std::exp(1.0)) {
+        std::cerr << "Warning in function '" << __func__
+            << "': lambert_wm1 is only defined for -1/e < x < 0\n";
+        return std::numeric_limits<double>::quiet_NaN();
+    }
+
+    double L1 = std::log(-x);
+    double L2 = std::log(-L1);
+    double w = L1 - L2 + (L2 / L1);
+
+    for (int i = 0; i < max_iter; ++i) {
+        double ew = std::exp(w);
+        double w_ew = w * ew;
+        double diff = (w_ew - x) / (ew * (w + 1) - ((w + 2) * (w_ew - x)) / (2 * w + 2));
+        w -= diff;
+        if (std::abs(diff) < tol) {
+            return w;
+        }
+    }
+
+    std::cerr << "Warning in function '" << __func__
+        << "': lambert_wm1 did not converge.\n";
+    return std::numeric_limits<double>::quiet_NaN();
+}

src/math.hpp

@@ -0,0 +1,7 @@
+#ifndef MATH_HPP
+#define MATH_HPP
+
+double lambert_wm1(double x, double tol = 1e-10, int max_iter = 100);
+
+#endif // MATH_HPP
+