metrust 0.4.3

Rust-powered MetPy-compatible calculation layer with optional GPU acceleration

metrust

MetPy-compatible calculation layer, powered by Rust.

150/150 metpy.calc functions implemented natively, plus 36 extras. Often faster on real-world workflows. Verified against MetPy on SounderPy, MetPy Cookbook examples, and synthetic global grids.

# The only change: swap the import
from metrust.calc import cape_cin, potential_temperature, vorticity
from metrust.units import units

Installation

pip install metrust

Optional GPU acceleration for supported calculations:

pip install "metrust[gpu]"

For plotting, xarray accessor, or Level2File (forwarded to MetPy):

pip install metrust metpy

What It Does

metrust implements every function in metpy.calc with a Rust backend compiled via PyO3. The Python API now matches MetPy's public metpy.calc signatures and is designed for MetPy-compatible units, return types, and runtime behavior on the shared calculation surface:

import numpy as np
from metrust.calc import (
    cape_cin, parcel_profile, bunkers_storm_motion,
    storm_relative_helicity, significant_tornado_parameter,
    vorticity, divergence, advection,
)
from metrust.units import units

# Sounding analysis (same API as MetPy)
p = np.array([1000, 925, 850, 700, 500, 300]) * units.hPa
T = np.array([25, 20, 15, 5, -15, -40]) * units.degC
Td = np.array([20, 15, 10, -5, -25, -50]) * units.degC

prof = parcel_profile(p, T[0], Td[0])
cape, cin = cape_cin(p, T, Td, prof)  # MetPy parcel_profile form works

# Grid kinematics with xarray (dx/dy inferred from lat/lon coords)
vort = vorticity(u_xarray, v_xarray)  # spherical metric corrections included
div = divergence(u_xarray, v_xarray)

Optional GPU Backend

Current users do not need to change anything. metrust stays on the Rust CPU backend by default.

When met-cu is installed, you can opt in explicitly:

import metrust.calc as mcalc

mcalc.set_backend("gpu")
theta = mcalc.potential_temperature(pressure, temperature)

with mcalc.use_backend("cpu"):
    theta_cpu = mcalc.potential_temperature(pressure, temperature)

The GPU backend currently targets the overlap where met-cu is already strong and verified:

• potential_temperature
• equivalent_potential_temperature
• dewpoint
• vorticity
• frontogenesis
• q_vector
• compute_cape_cin
• compute_srh
• compute_shear
• compute_pw
• composite_reflectivity_from_hydrometeors

Eligible dispatch currently focuses on scalar thermo plus uniform 2-D Cartesian grid workloads. Latitude/longitude-derived spacing, map-scale corrections, and other projection-aware cases fall back to the Rust CPU backend.

metrust still returns the same Pint/NumPy-facing API surface. Unsupported cases automatically stay on the Rust CPU path.

Speed

Benchmarked on real-world workflows (v0.3.3, validated by Codex against MetPy):

Workflow	Speedup	Notes
MetPy Cookbook sounding analysis	6.0x	Full severe weather stack
MetPy Cookbook 500 hPa grid	6.1x	Vorticity, smoothing, advection
MetPy Cookbook Q-vectors	6.1x	Q-vector divergence
SounderPy compute-heavy subset	29.7x	Thermo + wind + severe params
MetPy isentropic example	2.3x	Isentropic interpolation + Montgomery
Vorticity/divergence (global grid)	2.3x	Spherical corrections on 721x1440

Three-Way Benchmark: MetPy vs Rust vs CUDA

Real HRRR model output (40 isobaric levels, 1059 × 1799 grid, ~1.9 M points). RTX 5090, python tests/benchmark_gpu.py.

Scalar Thermodynamics (2D: 1059×1799)

Function	MetPy	Rust	CUDA	Rust/MetPy	CUDA/Rust
potential_temperature ★	11.2 ms	13.0 ms	7.9 ms	0.9x	1.6x
equiv_potential_temperature ★	303.1 ms	16.5 ms	9.1 ms	18x	1.8x
dewpoint ★	33.6 ms	10.8 ms	8.8 ms	3.1x	1.2x
saturation_vapor_pressure	66.7 ms	8.2 ms	—	8.1x	—
saturation_mixing_ratio	78.6 ms	13.7 ms	—	5.7x	—
dewpoint_from_rh	110.3 ms	7.6 ms	—	14x	—
rh_from_dewpoint	138.7 ms	10.9 ms	—	13x	—
virtual_temperature	31.1 ms	19.7 ms	—	1.6x	—
mixing_ratio	11.6 ms	8.5 ms	—	1.4x	—
wet_bulb_temperature	>10 min	26.9 ms	—	—	—

Grid Kinematics (2D: 1059×1799)

Function	MetPy	Rust	CUDA	Rust/MetPy	CUDA/Rust
vorticity ★	98.3 ms	92.8 ms	9.3 ms	1.1x	10x
divergence	96.1 ms	91.2 ms	—	1.1x	—
frontogenesis ★	733.0 ms	339.4 ms	12.2 ms	2.2x	28x
q_vector ★	390.3 ms	310.1 ms	10.8 ms	1.3x	29x
advection	161.8 ms	87.7 ms	—	1.8x	—

1D Sounding (40 levels, single column)

Function	MetPy	Rust	Rust/MetPy
parcel_profile	5.5 ms	0.074 ms	74x
cape_cin	1.4 ms	0.254 ms	5.4x
lcl	0.118 ms	0.065 ms	1.8x
lfc	6.7 ms	0.107 ms	62x
el	6.6 ms	0.112 ms	59x
precipitable_water	2.1 ms	0.057 ms	36x

Grid Composites (3D: 40×1059×1799 → 2D) — MetPy has no grid equivalents

Function	Rust	CUDA	CUDA/Rust
compute_cape_cin ★	2.96 s	674.5 ms	4.4x
compute_srh ★	223.5 ms	135.8 ms	1.6x
compute_shear ★	190.4 ms	166.5 ms	1.1x
compute_pw ★	191.6 ms	107.9 ms	1.8x
composite_refl_hydrometeors ★	154.2 ms	232.2 ms	0.7x

Summary

	MetPy	Rust	CUDA
Scalar thermo (×10)	785 ms	136 ms	26 ms
Grid kinematics (×5)	1.48 s	921 ms	32 ms
1D sounding (×6)	22 ms	0.67 ms	—
Grid composites (×5)	—	3.72 s	1.32 s
★ GPU-eligible total	—	4.50 s	1.37 s (3.3x)

★ = dispatches to CUDA when set_backend("gpu"). All other functions stay on Rust CPU regardless of backend setting.

Array Throughput

1M elements, 32-core Ryzen, rayon parallel:

Function	Time	Throughput
potential_temperature	1.8 ms	550 M/s
wet_bulb_temperature	7.3 ms	137 M/s
wind_speed	1.5 ms	660 M/s

Numerical Parity

Verified on the MetPy OUN 2011-05-22 12Z test sounding:

Metric	Difference from MetPy
CAPE	+4.0 J/kg
MUCAPE	+7.6 J/kg
SRH (0-1 km)	+0.3 m^2/s^2
Critical angle	+0.2 deg
Bunkers RM	+0.02 m/s
STP	+0.01
Montgomery streamfunction	corr 1.0000
Isentropic pressure	7e-13 hPa diff
Vorticity (global lat/lon)	corr 1.0000

Uses MetPy-exact physical constants (Rd, Cp, Lv, epsilon), MetPy's CAPE integration formula (g * dTv/Tv * dz), pressure-weighted Bunkers algorithm, Newton solver for isentropic interpolation, and spherical metric tensor corrections for lat/lon grid kinematics.

Coverage

• 150/150 metpy.calc functions (100% coverage)
• 36 extras not in MetPy (grid composites, fire weather indices, etc.)
• 28 Rust array bindings with rayon parallelism and GIL release
• Pint application registry shared with MetPy (no cross-registry errors)
• xarray support with coordinate inference and shape preservation

What's Not Native

These forward to MetPy when installed:

• metrust.plots (matplotlib-based plotting)
• metrust.xarray (xarray accessor)
• metrust.io.Level2File (NEXRAD Level II)

Core metrust.calc stays native Rust by default with no required MetPy dependency. A small parity-sensitive subset may delegate to MetPy when it is installed; otherwise those paths still fall back to local metrust implementations. The optional met-cu backend is an explicit accelerator, not a requirement.

Examples

See examples/ for complete drop-in demos:

• cookbook_sounding.py — MetPy Cookbook sounding analysis
• cookbook_500hpa_grid.py — MetPy Cookbook 500 hPa vorticity advection
• sounderpy_dropin.py — SounderPy-style full sounding pipeline

Testing

cargo test --workspace          # 1,186 Rust tests
python -m pytest tests/ -q      # 30 Python tests (including MetPy compatibility regression)

Documentation

Full docs at fahrenheitresearch.github.io/metrust-py, including:

• API reference for all 186 functions
• Beginner tutorials (Weather 101, soundings, grids, recipes)
• Migration guide from MetPy
• Performance benchmarks

License

MIT