diff --git a/pyproject.toml b/pyproject.toml index 33bf3be94c4..04ef962d245 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt new file mode 100644 index 00000000000..d9be3d2855b --- /dev/null +++ b/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) diff --git a/src/calcmod.cpp b/src/calcmod.cpp new file mode 100644 index 00000000000..8c2cdbe64df --- /dev/null +++ b/src/calcmod.cpp @@ -0,0 +1,111 @@ +#include +#include +#include +#include "constants.hpp" +#include "thermo.hpp" +#include // For std::pair +#include // For std::make_tuple +#include + +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 pressure, double temperature, double dewpoint) { + // pressure.data() gives the beginning pointer of the data buffer + std::vector pressure_vec(pressure.data(), pressure.data() + pressure.size()); + std::vector temp_prof = ParcelProfile(pressure_vec, temperature, dewpoint); + return py::array_t(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")); + +} diff --git a/src/constants.cpp b/src/constants.cpp new file mode 100644 index 00000000000..6cc7bed6265 --- /dev/null +++ b/src/constants.cpp @@ -0,0 +1,50 @@ +// constants.cpp +#include "constants.hpp" +#include + +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(); +// Rv = mod.attr("Rv").attr("magnitude").cast(); +// Cp_v = mod.attr("Cp_v").attr("magnitude").cast(); +// Cp_l = mod.attr("Cp_l").attr("magnitude").cast(); +// Lv = mod.attr("Lv").attr("magnitude").cast(); +// sat_pressure_0c = mod.attr("sat_pressure_0c").cast(); +// T0 = mod.attr("T0").attr("magnitude").cast(); + + + Mw = mod.attr("Mw").cast(); + Rd = mod.attr("Rd").cast(); + Rv = mod.attr("Rv").cast(); + Cp_d = mod.attr("Cp_d").cast(); + Cp_v = mod.attr("Cp_v").cast(); + Cp_l = mod.attr("Cp_l").cast(); + Lv = mod.attr("Lv").cast(); + sat_pressure_0c = mod.attr("sat_pressure_0c").cast(); + T0 = mod.attr("T0").cast(); + Ls = mod.attr("Ls").cast(); + Cp_i = mod.attr("Cp_i").cast(); + zero_degc = mod.attr("zero_degc").cast(); + epsilon = mod.attr("epsilon").cast(); + kappa = mod.attr("kappa").cast(); + } +} diff --git a/src/constants.hpp b/src/constants.hpp new file mode 100644 index 00000000000..02271e28f61 --- /dev/null +++ b/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 + diff --git a/src/math.cpp b/src/math.cpp new file mode 100644 index 00000000000..12c89469a0f --- /dev/null +++ b/src/math.cpp @@ -0,0 +1,32 @@ +#include +#include // 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::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::quiet_NaN(); +} diff --git a/src/math.hpp b/src/math.hpp new file mode 100644 index 00000000000..c4c99ea2558 --- /dev/null +++ b/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 + diff --git a/src/metpy/calc/thermo.py b/src/metpy/calc/thermo.py index 40274192b28..e387bf0a439 100644 --- a/src/metpy/calc/thermo.py +++ b/src/metpy/calc/thermo.py @@ -16,6 +16,7 @@ import scipy.optimize as so from scipy.special import lambertw import xarray as xr +import metpy._calc_mod as _calc_mod from .exceptions import InvalidSoundingError from .tools import (_greater_or_close, _less_or_close, _remove_nans, find_bounding_indices, @@ -185,10 +186,8 @@ def water_latent_heat_vaporization(temperature): Eq 15, [Ambaum2020]_, using MetPy-defined constants in place of cited values. """ - return (mpconsts.nounit.Lv - - (mpconsts.nounit.Cp_l - mpconsts.nounit.Cp_v) - * (temperature - mpconsts.nounit.T0)) - + # Calling c++ calculation module + return _calc_mod.water_latent_heat_vaporization(temperature) @exporter.export @preprocess_and_wrap(wrap_like='temperature') @@ -227,10 +226,8 @@ def water_latent_heat_sublimation(temperature): Eq 18, [Ambaum2020]_, using MetPy-defined constants in place of cited values. """ - return (mpconsts.nounit.Ls - - (mpconsts.nounit.Cp_i - mpconsts.nounit.Cp_v) - * (temperature - mpconsts.nounit.T0)) - + # Calling c++ calculation module + return _calc_mod.water_latent_heat_sublimation(temperature) @exporter.export @preprocess_and_wrap(wrap_like='temperature') @@ -516,6 +513,30 @@ def dry_lapse(pressure, temperature, reference_pressure=None, vertical_dim=0): return temperature * (pressure / reference_pressure)**mpconsts.kappa + +@exporter.export +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'reference_pressure') +) +@process_units( + { + 'pressure': '[pressure]', + 'temperature': '[temperature]', + 'reference_pressure': '[pressure]' + }, + '[temperature]' +) +def dry_lapse_linfel(pressure, temperature, reference_pressure=None, vertical_dim=0): + """ + Linfeng's version of 'dry_lapse'. Added on Jun18 2025 + """ + if reference_pressure is None: + reference_pressure = pressure[0] + return _calc_mod.dry_lapse(pressure, temperature, reference_pressure) + + + @exporter.export @preprocess_and_wrap( wrap_like='temperature', @@ -654,6 +675,33 @@ def dt(p, t): return ret.squeeze() +@exporter.export +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'reference_pressure') +) +@process_units( + { + 'pressure': '[pressure]', + 'temperature': '[temperature]', + 'reference_pressure': '[pressure]' + }, + '[temperature]' +) +def moist_lapse_linfel(pressure, temperature, reference_pressure=None): + """ + Linfeng's version of 'moist_lapse'. Added on Jun 25 2025 + This function calculates the moist adiabatic profile for multiple starting + temperatures (2D surface) and a single communal starting pressure, along a + 1D pressure profile. + """ + if reference_pressure is None: + reference_pressure = pressure[0] + # nstep for RK4 is set to 30 + return _calc_mod.moist_lapse(pressure, temperature, reference_pressure, 30) + + + @exporter.export @preprocess_and_wrap() @process_units( @@ -739,6 +787,22 @@ def lcl(pressure, temperature, dewpoint, max_iters=None, eps=None): return p_lcl, t_lcl +@exporter.export +@preprocess_and_wrap() +@process_units( + {'pressure': '[pressure]', 'temperature': '[temperature]', 'dewpoint': '[temperature]'}, + ('[pressure]', '[temperature]') +) +def lcl_linfel(pressure, temperature, dewpoint, max_iters=None, eps=None): + """ + Linfeng's version of 'lcl'. Added on Jun23 2025 + """ + pressure, temperature, dewpoint = np.atleast_1d(pressure, temperature, dewpoint) + p_lcl, t_lcl = _calc_mod.lcl(pressure, temperature, dewpoint) + return p_lcl, t_lcl + + + @exporter.export @preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') @@ -1252,6 +1316,23 @@ def parcel_profile(pressure, temperature, dewpoint): return concatenate((t_l, t_u)) +@exporter.export +@preprocess_and_wrap(wrap_like='pressure') +@process_units( + { + 'pressure': '[pressure]', + 'temperature': '[temperature]', + 'dewpoint': '[temperature]' + }, + '[temperature]' +) # process units because no unit should be passed to c++ function +def parcel_profile_linfel(pressure, temperature, dewpoint): + """ + Linfeng's version of 'parcel_profile'. Added on Jun 24 2025 + """ + return _calc_mod.parcel_profile(pressure, temperature, dewpoint) + + @exporter.export @preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') @@ -1623,17 +1704,8 @@ def _saturation_vapor_pressure_liquid(temperature): .. math:: e = e_{s0} \frac{T_0}{T}^{(c_{pl} - c_{pv}) / R_v} \exp{ \frac{L_0}{R_v T_0} - \frac{L}{R_v T}} """ - latent_heat = water_latent_heat_vaporization._nounit(temperature) - heat_power = (mpconsts.nounit.Cp_l - mpconsts.nounit.Cp_v) / mpconsts.nounit.Rv - exp_term = ((mpconsts.nounit.Lv / mpconsts.nounit.T0 - latent_heat / temperature) - / mpconsts.nounit.Rv) - - return ( - mpconsts.nounit.sat_pressure_0c - * (mpconsts.nounit.T0 / temperature) ** heat_power - * np.exp(exp_term) - ) - + # Calling c++ calculation module + return _calc_mod._saturation_vapor_pressure_liquid(temperature) @preprocess_and_wrap(wrap_like='temperature') @process_units({'temperature': '[temperature]'}, '[pressure]') @@ -1660,17 +1732,8 @@ def _saturation_vapor_pressure_solid(temperature): .. math:: e_i = e_{i0} \frac{T_0}{T}^{(c_{pi} - c_{pv}) / R_v} \exp{ \frac{L_{s0}}{R_v T_0} - \frac{L_s}{R_v T}} """ - latent_heat = water_latent_heat_sublimation._nounit(temperature) - heat_power = (mpconsts.nounit.Cp_i - mpconsts.nounit.Cp_v) / mpconsts.nounit.Rv - exp_term = ((mpconsts.nounit.Ls / mpconsts.nounit.T0 - latent_heat / temperature) - / mpconsts.nounit.Rv) - - return ( - mpconsts.nounit.sat_pressure_0c - * (mpconsts.nounit.T0 / temperature) ** heat_power - * np.exp(exp_term) - ) - + # Calling c++ calculation module + return _calc_mod._saturation_vapor_pressure_solid(temperature) @exporter.export @preprocess_and_wrap(wrap_like='temperature', broadcast=('temperature', 'relative_humidity')) @@ -1750,9 +1813,8 @@ def dewpoint(vapor_pressure): Renamed ``e`` parameter to ``vapor_pressure`` """ - val = np.log(vapor_pressure / mpconsts.nounit.sat_pressure_0c) - return mpconsts.nounit.zero_degc + 243.5 * val / (17.67 - val) - + # Calling c++ calculation module + return _calc_mod.dewpoint(vapor_pressure) @exporter.export @preprocess_and_wrap(wrap_like='partial_press', broadcast=('partial_press', 'total_press')) @@ -1812,7 +1874,8 @@ def mixing_ratio(partial_press, total_press, molecular_weight_ratio=mpconsts.nou Renamed ``part_press``, ``tot_press`` parameters to ``partial_press``, ``total_press`` """ - return molecular_weight_ratio * partial_press / (total_press - partial_press) + # Calling c++ calculation module + return _calc_mod.mixing_ratio(partial_press, total_press, molecular_weight_ratio) @exporter.export @@ -2126,6 +2189,25 @@ def virtual_temperature( return temperature * ((mixing_ratio + molecular_weight_ratio) / (molecular_weight_ratio * (1 + mixing_ratio))) +@exporter.export +@preprocess_and_wrap(wrap_like='temperature', broadcast=('temperature', 'mixing_ratio')) +@process_units( + { + 'temperature': '[temperature]', + 'mixing_ratio': '[dimensionless]', + 'molecular_weight_ratio': '[dimensionless]' + }, + '[temperature]', + ignore_inputs_for_output=('molecular_weight_ratio',) +) +def virtual_temperature_linfel( + temperature, mixing_ratio, molecular_weight_ratio=mpconsts.nounit.epsilon): + """ + Linfeng's version of 'virtual_temperature'. Added on Jun 30 2025 + """ + return _calc_mod.virtual_temperature( + temperature, mixing_ratio, molecular_weight_ratio) + @exporter.export @preprocess_and_wrap(wrap_like='temperature', broadcast=('pressure', @@ -2195,6 +2277,38 @@ def virtual_temperature_from_dewpoint( return virtual_temperature(temperature, mixing_ratio, molecular_weight_ratio) +@exporter.export +@preprocess_and_wrap(wrap_like='temperature', broadcast=('pressure', + 'temperature', + 'dewpoint')) +@process_units( + { + 'pressure': '[pressure]', + 'temperature': '[temperature]', + 'dewpoint': '[temperature]', + 'molecular_weight_ratio': '[dimensionless]' + }, + '[temperature]', + ignore_inputs_for_output=('molecular_weight_ratio',) +) +def virtual_temperature_from_dewpoint_linfel( + pressure, + temperature, + dewpoint, + molecular_weight_ratio=mpconsts.nounit.epsilon, + *, + phase='liquid' +): + """ + Linfeng's version of 'virtual_temperature_from_dewpoint'. Added on Jun 30 2025 + """ + return _calc_mod.virtual_temperature_from_dewpoint(pressure, + temperature, + dewpoint, + molecular_weight_ratio, + phase) + + @exporter.export @preprocess_and_wrap( wrap_like='temperature', diff --git a/src/metpy/constants/nounit.py b/src/metpy/constants/nounit.py index 2fd6e8a78fa..d1b5e837303 100644 --- a/src/metpy/constants/nounit.py +++ b/src/metpy/constants/nounit.py @@ -6,6 +6,7 @@ from . import default from ..units import units +Mw = default.Mw.m_as('kg / mol') Rd = default.Rd.m_as('m**2 / K / s**2') Rv = default.Rv.m_as('m**2 / K / s**2') Lv = default.Lv.m_as('m**2 / s**2') diff --git a/src/thermo.cpp b/src/thermo.cpp new file mode 100644 index 00000000000..9f04ab82c81 --- /dev/null +++ b/src/thermo.cpp @@ -0,0 +1,464 @@ +#include +#include +#include +#include +#include +#include // For std::make_tuple +#include // For std::pair +#include +#include // for std::cerr +#include // for std::numeric_limits +#include "math.hpp" +#include "constants.hpp" +#include "thermo.hpp" + +namespace py = pybind11; +namespace mc = metpy_constants; + +double MoistAirGasConstant(double specific_humidity) { + return (1.0 - specific_humidity) * mc::Rd + specific_humidity * mc::Rv; +} + +double MoistAirSpecificHeatPressure(double specific_humidity) { + return (1.0 - specific_humidity) * mc::Cp_d + specific_humidity * mc::Cp_v; +} + +double WaterLatentHeatVaporization(double temperature) { + return mc::Lv - (mc::Cp_l - mc::Cp_v) * (temperature - mc::T0); +} + +double WaterLatentHeatSublimation(double temperature) { + return mc::Ls - (mc::Cp_i - mc::Cp_v) * (temperature - mc::T0); +} + +double RelativeHumidityFromDewPoint(double temperature, double dewpoint, std::string phase) { + double e_s = SaturationVaporPressure(temperature, phase); + double e = SaturationVaporPressure(dewpoint, phase); + return e / e_s; +} + +double DryLapse(double pressure, double ref_temperature, double ref_pressure) { + // calculate temperature at pressure given reference temperature and pressure + // assuming dry adiabatic process + return ref_temperature * pow(pressure / ref_pressure, mc::kappa); +} + +std::vector DryLapseProfile(const std::vector& pressure_profile, + double ref_temperature, + double ref_pressure) { + // Vectorized version of DryLapse for a pressure profile. C++ internally use. + + // calculate temperature profile of an air parcel lifting dry adiabatically + // through the given pressure profile. + std::vector temperature_profile; + temperature_profile.reserve(pressure_profile.size()); + + for (double p : pressure_profile) { + temperature_profile.push_back(DryLapse(p, ref_temperature, ref_pressure)); + } + return temperature_profile; +} + + +py::array_t DryLapseVectorized(py::array_t pressure, + py::array_t ref_temperature, + double ref_pressure) { + // This function calculates the dry adiabatic profile for multiple starting + // temperatures (2D surface) and a single communal starting pressure, along a + // 1D pressure profile. + // --- Step 1: Prepare the C++ vector for pressure levels --- + if (pressure.ndim() > 1) { + throw std::runtime_error("Input 'pressure' must be 1D array or a single value."); + } + std::vector pressure_vec(pressure.data(), pressure.data() + pressure.size()); + + // --- Step 2: Ensure the reference temperature array is contiguous --- + auto ref_temp_contig = py::array::ensure(ref_temperature, py::array::c_style); + + // --- Step 3: Define the shape of the output array: (N+1) dimension--- + std::vector out_shape; + for(int i = 0; i < ref_temp_contig.ndim(); ++i) { + out_shape.push_back(ref_temp_contig.shape(i)); + } + size_t profile_len = pressure_vec.size(); + out_shape.push_back(profile_len); + + auto out_array = py::array_t(out_shape); + + // --- Step 4: Get direct pointers to data buffers for fast access --- + const double* ref_temp_ptr = static_cast(ref_temp_contig.request().ptr); + double* out_array_ptr = out_array.mutable_data(); + size_t num_profiles = ref_temp_contig.size(); + + // --- Step 5: Loop through each reference temperature --- + for (size_t i = 0; i < num_profiles; ++i) { + for (size_t j = 0; j < profile_len; ++j) { + out_array_ptr[i * profile_len + j] = DryLapse(pressure_vec[j], ref_temp_ptr[i], ref_pressure); + } + } + + return out_array; +} + +py::array_t DryLapseVectorized_3D(py::array_t pressure, + py::array_t ref_temperature, + py::array_t ref_pressure) { + // --- Step 1: Ensure all input arrays are using a contiguous memory layout --- + auto p_contig = py::array::ensure(pressure, py::array::c_style); + auto ref_temp_contig = py::array::ensure(ref_temperature, py::array::c_style); + auto ref_press_contig = py::array::ensure(ref_pressure, py::array::c_style); + + // --- Step 2: Perform comprehensive shape validation --- + if (ref_temp_contig.ndim() != ref_press_contig.ndim()) { + throw std::runtime_error("Input 'ref_temperature' and 'ref_pressure' must have the same number of dimensions."); + } + if (p_contig.ndim() != ref_temp_contig.ndim() + 1) { + throw std::runtime_error("Input 'pressure' must have one more dimension than 'ref_temperature'."); + } + for (int i = 0; i < ref_temp_contig.ndim(); ++i) { + if (ref_temp_contig.shape(i) != ref_press_contig.shape(i) || + p_contig.shape(i+1) != ref_temp_contig.shape(i)) { + throw std::runtime_error("The horizontal dimensions of all input arrays must match."); + } + } + + // --- Step 3: Define the shape of the output array --- + // The output shape will be identical to the input pressure array's shape. + auto out_array = py::array_t(p_contig.request().shape); + + // --- Step 4: Get direct pointers to data buffers for fast access --- + const double* pressure_ptr = static_cast(p_contig.request().ptr); + const double* ref_temp_ptr = static_cast(ref_temp_contig.request().ptr); + const double* ref_press_ptr = static_cast(ref_press_contig.request().ptr); + double* out_array_ptr = out_array.mutable_data(); + + // --- Step 5: Define loop boundaries --- + size_t num_profiles = ref_temp_contig.size(); // Total number of horizontal points + size_t profile_len = p_contig.shape(0); // Length of the vertical dimension + + // --- Step 6: Loop through each horizontal point and its vertical profile --- + for (int j = 0; j < profile_len; ++j) { + for (int i = 0; i < num_profiles; ++i) { + // Calculate the index for the current point in the flattened arrays. + out_array_ptr[i+j*num_profiles] = DryLapse(pressure_ptr[i+j*num_profiles], + ref_temp_ptr[i], + ref_press_ptr[i]); + } + } + + return out_array; +} + +double CaldlnTdlnP(double temperature, double pressure) { + // Calculate dlnT/dlnP for a moist (saturated) adiabatic process. + double rs = SaturationMixingRatio(pressure, temperature, "liquid"); + + //double dlnT_dlnP_linfel = (mc::Rd + rs * mc::Rv) / (mc::Cp_d + rs * mc::Cp_v + + // (mc::Lv * mc::Lv * rs * mc::epsilon) / (mc::Rd * temperature * temperature)); + double dlnT_dlnP_Bakhshaii2013 = (mc::Rd + mc::Lv * rs / temperature) / (mc::Cp_d + + (mc::Lv * mc::Lv * rs * mc::epsilon) / (mc::Rd * temperature * temperature)); + + return dlnT_dlnP_Bakhshaii2013; +} + +double MoistLapse(double pressure, double ref_temperature, double ref_pressure, int rk_nstep) { + // calculate temperature at pressure given reference temperature and pressure + // assuming moist adiabatic expansion (vapor condenses and removed from the air + // parcel) + + double dlnP = log(pressure / ref_pressure) / (double)rk_nstep; + double T1 = ref_temperature; + double P1 = ref_pressure; + double k[4]; + + for (int i = 0; i < rk_nstep; ++i) { + k[0] = CaldlnTdlnP(T1, P1); + k[1] = CaldlnTdlnP(T1 * exp(k[0] * dlnP/2.), P1 * exp(dlnP/2.)); + k[2] = CaldlnTdlnP(T1 * exp(k[1] * dlnP/2.), P1 * exp(dlnP/2.)); + k[3] = CaldlnTdlnP(T1 * exp(k[2] * dlnP), P1 * exp(dlnP)); + + T1 = T1 * exp((k[0] + 2.0 * k[1] + 2.0 * k[2] + k[3]) * dlnP / 6.0); + P1 = P1 * exp(dlnP); + } + + return T1; // check final T1 P1 +} + +std::vector MoistLapseProfile(const std::vector& press_profile, + double ref_temperature, + double ref_pressure, + int rk_nstep) { + // MoistLapse for one full pressure profile given one ref_temperature. C++ internally use. + + // calculate temperature profile of an air parcel lifting saturated adiabatically + // through the given pressure profile. + std::vector temp_profile; + temp_profile.reserve(press_profile.size()); + + for (double p : press_profile) { + temp_profile.push_back(MoistLapse(p, ref_temperature, ref_pressure, rk_nstep)); + } + return temp_profile; +} + +py::array_t MoistLapseVectorized(py::array_t pressure, + py::array_t ref_temperature, + double ref_pressure, + int rk_nstep) { + // This function calculates the moist adiabatic profile for multiple starting + // temperatures (2D surface) and a single communal starting pressure, along a + // 1D pressure profile. + // --- Step 1: Prepare the C++ vector for pressure levels --- + if (pressure.ndim() > 1) { + throw std::runtime_error("Input 'pressure' must be 1D array or a single value."); + } + std::vector pressure_vec(pressure.data(), pressure.data() + pressure.size()); + + // --- Step 2: Ensure the reference temperature array is contiguous --- + auto ref_temp_contig = py::array::ensure(ref_temperature, py::array::c_style); + + // --- Step 3: Define the shape of the output array: (N+1) dimension--- + std::vector out_shape; + for(int i = 0; i < ref_temp_contig.ndim(); ++i) { + out_shape.push_back(ref_temp_contig.shape(i)); + } + size_t profile_len = pressure_vec.size(); + out_shape.push_back(profile_len); + + auto out_array = py::array_t(out_shape); + + // --- Step 4: Get direct pointers to data buffers for fast access --- + const double* ref_temp_ptr = static_cast(ref_temp_contig.request().ptr); + double* out_array_ptr = out_array.mutable_data(); + size_t num_profiles = ref_temp_contig.size(); + + // --- Step 5: Loop through each reference temperature --- + for (size_t i = 0; i < num_profiles; ++i) { + for (size_t j = 0; j < profile_len; ++j) { + out_array_ptr[i * profile_len + j] = MoistLapse(pressure_vec[j], ref_temp_ptr[i], ref_pressure, rk_nstep); + } + } + + return out_array; +} + + +std::pair LCL(double pressure, double temperature, double dewpoint) { + if (temperature < dewpoint) { + std::cerr << "Warning in function '" << __func__ + << "': Temperature must be greater than dew point for LCL calculation.\n"; + return {std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN()}; + } + + double q = SpecificHumidityFromDewPoint(pressure, dewpoint, "liquid"); + double moist_heat_ratio = MoistAirSpecificHeatPressure(q) / MoistAirGasConstant(q); + double spec_heat_diff = mc::Cp_l - mc::Cp_v; + + double a = moist_heat_ratio + spec_heat_diff / mc::Rv; + double b = -(mc::Lv + spec_heat_diff * mc::T0) / (mc::Rv * temperature); + double c = b / a; + + double rh = RelativeHumidityFromDewPoint(temperature, dewpoint, "liquid"); + double w_minus1 = lambert_wm1(pow(rh, 1.0 / a) * c * exp(c)); + double t_lcl = c / w_minus1 * temperature; + double p_lcl = pressure * pow(t_lcl / temperature, moist_heat_ratio); + + return {p_lcl, t_lcl}; +} + + +std::tuple, py::array_t> LCLVectorized(py::array_t pressure, + py::array_t temperature, + py::array_t dewpoint) { + + // This helper ensures the arrays are in C-style contiguous memory. + // If an input array is already contiguous, it's a zero-cost operation. + // If it's a slice or has a different memory layout, it creates a copy. + // This makes the subsequent looping simple and safe. + auto p_contig = py::array::ensure(pressure, py::array::c_style); + auto t_contig = py::array::ensure(temperature, py::array::c_style); + auto d_contig = py::array::ensure(dewpoint, py::array::c_style); + + // --- Step 1: Check that all input arrays have the same shape --- + if (p_contig.ndim() != t_contig.ndim() || p_contig.ndim() != d_contig.ndim()) { + throw std::runtime_error("Input arrays must have the same number of dimensions."); + } + for (int i = 0; i < p_contig.ndim(); ++i) { + if (p_contig.shape(i) != t_contig.shape(i) || p_contig.shape(i) != d_contig.shape(i)) { + throw std::runtime_error("Input arrays must have the same shape."); + } + } + + // --- Step 2: Create output arrays with the exact same N-D shape as the inputs --- + auto p_lcl = py::array_t(p_contig.request().shape); + auto t_lcl = py::array_t(p_contig.request().shape); + + // --- Step 3: Get the total number of elements to loop over --- + size_t size = p_contig.size(); + + // --- Step 4: Get direct pointers to the (now contiguous) data buffers --- + const double* p_ptr = static_cast(p_contig.request().ptr); + const double* t_ptr = static_cast(t_contig.request().ptr); + const double* d_ptr = static_cast(d_contig.request().ptr); + double* p_lcl_ptr = p_lcl.mutable_data(); + double* t_lcl_ptr = t_lcl.mutable_data(); + + // --- Step 5: Loop through the data as if it were a single flat 1D array --- + for (size_t i = 0; i < size; i++) { + // Call the scalar c++ function for each element + std::pair result = LCL(p_ptr[i], t_ptr[i], d_ptr[i]); + p_lcl_ptr[i] = result.first; + t_lcl_ptr[i] = result.second; + } + + // --- Step 6: Return a tuple of the two new, N-dimensional arrays --- + return std::make_tuple(p_lcl, t_lcl); +} + +bool _CheckPressure(const std::vector& pressure) { + for (size_t i = 0; i + 1 < pressure.size(); ++i) { + if (pressure[i] < pressure[i + 1]) { + return false; // pressure increases (invalid) + } + } + return true; // strictly non-increasing +} + + +std::vector ParcelProfile(const std::vector& pressure, + double temperature, + double dewpoint) { + // Returns a vector of temperatures corresponding to the input pressure levels. + + ParProStruct profile = _ParcelProfileHelper(pressure, temperature, dewpoint); + + // Combine lower and upper temperature profiles + std::vector combined_temp; + combined_temp.reserve(pressure.size()); + + combined_temp.insert(combined_temp.end(), profile.temp_lower.begin(), profile.temp_lower.end()); + combined_temp.insert(combined_temp.end(), profile.temp_upper.begin(), profile.temp_upper.end()); + + return combined_temp; +} + +ParProStruct _ParcelProfileHelper(const std::vector& pressure, double temperature, double dewpoint) { + // Check that pressure does not increase. + if (!_CheckPressure(pressure)) { + throw std::runtime_error( + "Pressure increases between at least two points in your sounding. " + "Using a smoothing filter (e.g., scipy.signal.medfilt) may fix this."); + } + + // Find the LCL + std::pair lcl_result = LCL(pressure[0], temperature, dewpoint); + double press_lcl = lcl_result.first; + double temp_lcl = lcl_result.second; + + // Establish profile below LCL + std::vector press_lower; + for (double p : pressure) { + if (p >= press_lcl) { + press_lower.push_back(p); + } + } + press_lower.push_back(press_lcl); + std::vector temp_lower = DryLapseProfile(press_lower, temperature, press_lower[0]); + + // Early return if profile ends before reaching above LCL + if (pressure.back() >= press_lcl) { + press_lower.pop_back(); + temp_lower.pop_back(); + return {press_lower, {}, press_lcl, temp_lower, {}, temp_lcl}; + } + + // Establish profile above LCL + std::vector press_upper; + press_upper.push_back(press_lcl); + for (double p : pressure) { + if (p < press_lcl) { + press_upper.push_back(p); + } + } + std::vector temp_upper = MoistLapseProfile(press_upper, temp_lower.back(), press_lcl, 30); + + press_lower.pop_back(); + temp_lower.pop_back(); + press_upper.erase(press_upper.begin()); + temp_upper.erase(temp_upper.begin()); + + return {press_lower, press_upper, press_lcl, temp_lower, temp_upper, temp_lcl}; +} + +double _SaturationVaporPressureLiquid(double temperature) { + double latent_heat = WaterLatentHeatVaporization(temperature); + double heat_power = (mc::Cp_l - mc::Cp_v) / mc::Rv; + double exp_term = (mc::Lv / mc::T0 - latent_heat / temperature) / mc::Rv; + + return mc::sat_pressure_0c * exp(exp_term) * pow(mc::T0 / temperature, heat_power); +} + +double _SaturationVaporPressureSolid(double temperature) { + double latent_heat = WaterLatentHeatSublimation(temperature); + double heat_power = (mc::Cp_i - mc::Cp_v) / mc::Rv; + double exp_term = (mc::Ls / mc::T0 - latent_heat / temperature) / mc::Rv; + + return mc::sat_pressure_0c * exp(exp_term) * pow(mc::T0 / temperature, heat_power); +} + +double SaturationVaporPressure(double temperature, std::string phase) { + if (phase == "liquid") { + return _SaturationVaporPressureLiquid(temperature); + } else if (phase == "solid") { + return _SaturationVaporPressureSolid(temperature); + } else if (phase == "auto") { + if (temperature > mc::T0) { + return _SaturationVaporPressureLiquid(temperature); + } else { + return _SaturationVaporPressureSolid(temperature); + } + } else { + throw std::invalid_argument("'" + phase + "' is not a valid option for phase. " + "Valid options are 'liquid', 'solid', or 'auto'."); + } +} + +double DewPoint(double vapor_pressure) { + double val = log(vapor_pressure / mc::sat_pressure_0c); + return mc::zero_degc + 243.5 * val / (17.67 - val); +} + +double MixingRatio(double partial_press, double total_press, double epsilon) { + return epsilon * partial_press / (total_press - partial_press); +} + +double SaturationMixingRatio(double total_press, double temperature, std::string phase) { + double e_s = SaturationVaporPressure(temperature, phase); + if (e_s >= total_press) { + std::cerr << "Warning in function '" << __func__ + << "': Total pressure must be greater than the saturation vapor pressure " + << "for liquid water to be in equilibrium.\n"; + return std::numeric_limits::quiet_NaN(); + } + return MixingRatio(e_s, total_press); +} + +double SpecificHumidityFromMixingRatio(double mixing_ratio) { + return mixing_ratio / (mixing_ratio + 1.0); +} + +double SpecificHumidityFromDewPoint(double pressure, double dewpoint, std::string phase) { + double mixing_ratio = SaturationMixingRatio(pressure, dewpoint, phase); + return SpecificHumidityFromMixingRatio(mixing_ratio); +} + +double VirtualTemperature(double temperature, double mixing_ratio, double epsilon) { + return temperature * (mixing_ratio + epsilon) / (epsilon * (1. + mixing_ratio)); +} + +double VirtualTemperatureFromDewpoint(double pressure, double temperature, + double dewpoint, double epsilon, + std::string phase) { + double mixing_ratio = SaturationMixingRatio(pressure, dewpoint, phase); + return VirtualTemperature(temperature, mixing_ratio, epsilon); +} diff --git a/src/thermo.hpp b/src/thermo.hpp new file mode 100644 index 00000000000..555bbc437ae --- /dev/null +++ b/src/thermo.hpp @@ -0,0 +1,80 @@ +#ifndef THERMO_HPP // if not defined +#define THERMO_HPP // define the header file + +#include +#include +#include +#include +#include // For std::pair +#include // For std::tuple +#include "constants.hpp" + +namespace py = pybind11; +namespace mc = metpy_constants; + +double MoistAirGasConstant(double specific_humidity); +double MoistAirSpecificHeatPressure(double specific_humidity); + +double WaterLatentHeatVaporization(double temperature); +double WaterLatentHeatSublimation(double temperature); + +double RelativeHumidityFromDewPoint(double temperature, double dewpoint, std::string phase="liquid"); + +double DryLapse(double pressure, double ref_temperature, double ref_pressure); +std::vector DryLapseProfile(const std::vector& pressure_profile, + double ref_temperature, + double ref_pressure); +py::array_t DryLapseVectorized(py::array_t pressure, + py::array_t ref_temperature, + double ref_pressure); +py::array_t DryLapseVectorized_3D(py::array_t pressure, + py::array_t ref_temperature, + py::array_t ref_pressure); + +double CaldlnTdlnP(double temperature, double pressure); +double MoistLapse(double pressure, double ref_temperature, double ref_pressure, int rk_nstep); +std::vector MoistLapseProfile(const std::vector& press_profile, + double ref_temperature, + double ref_pressure, + int rk_nstep); +py::array_t MoistLapseVectorized(py::array_t pressure, + py::array_t ref_temperature, + double ref_pressure, + int rk_nstep); + +std::pair LCL(double pressure, double temperature, double dewpoint); +std::tuple, py::array_t> LCLVectorized(py::array_t pressure, + py::array_t temperature, + py::array_t dewpoint); + +bool _CheckPressure(const std::vector& pressure); + +// Return struct for _ParcelProfileHelper +struct ParProStruct { + std::vector press_lower, press_upper; + double press_lcl; + std::vector temp_lower, temp_upper; + double temp_lcl; +}; + +std::vector ParcelProfile(const std::vector& pressure, + double temperature, + double dewpoint); + +ParProStruct _ParcelProfileHelper(const std::vector& pressure, double temperature, double dewpoint); + +double _SaturationVaporPressureLiquid(double temperature); +double _SaturationVaporPressureSolid(double temperature); +double SaturationVaporPressure(double temperature, std::string phase); + +double DewPoint(double vapor_pressure); +double MixingRatio(double partial_press, double total_press, double epsilon=mc::epsilon); +double SaturationMixingRatio(double total_press, double temperature, std::string phase); +double SpecificHumidityFromMixingRatio(double mixing_ratio); +double SpecificHumidityFromDewPoint(double pressure, double dew_point, std::string phase); + +double VirtualTemperature(double temperature, double mixing_ratio, double epsilon); +double VirtualTemperatureFromDewpoint(double pressure, double temperature, + double dewpoint, double epsilon, + std::string phase); +#endif // THERMO_HPP