prophys 0.3.0

prophys: probabilistic-physics modelling toolkit

A differentiable, JAX-native DSL for spatial probability and risk models.

Overview

prophys lets you describe an uncertainty structure (geometry, physics, distributions) as a symbolic Python model, and compiles it into a single differentiable JAX graph with log_prob, sample, expectation, tail-risk metrics, gradient-based calibration, and design optimization.

It is not a distribution library like NumPyro or Distrax. It is the layer above one. The distinguishing capability is that geometry, interpolated fields, physics surrogates, and probability live in the same differentiable graph: a pipeline route vertex, an emission abatement fraction, or a copula correlation can all be trained or optimized with the same jax.grad pass that evaluates the model.

It combines:

• A symbolic expression core (Expr, Param, Input, RandomVariable) with NumPy-style operator overloading. Models are built with ordinary arithmetic and compile to pure jax.numpy.
• Differentiable geometry structures in shared coordinate frames: PointList, LineList, Polygon/PolygonList (signed distance, containment, area), Polyhedron (3D halfspace volumes), raster Fields with bilinear interpolation, and evaluation Grids.
• Transformations: linear/affine/polynomial, smoothable PiecewiseLinear, Decay, Logistic dose-response, and interpolated TableLookup1D/2D characteristic curves.
• Distributions: 12+ univariate families, circular (VonMises, WrappedGaussian), multivariate Gaussian (Cholesky-parameterized), mixtures, empirical/KDE, and Gaussian copulas with always-valid learnable correlation.
• Probabilistic inputs anywhere in the graph: a RandomVariable (e.g. wind speed/direction) can feed the physics, not just the output noise. The compiled model marginalizes it by reparameterized Monte Carlo with gradients intact.
• A training engine (Optax): calibrate one attribute, finetune the whole model jointly across attributes with shared parameters, fit_distribution for standalone/multivariate distribution learning, with trainable/frozen parameter control.
• Design optimization with penalty constraints, where bounded parameters are constrained structurally (sigmoid/softplus bijectors) rather than by projection.
• An observation model: left- and right-censored, interval-censored and missing records each contribute their correct likelihood term instead of a substituted number, and prophys.io reads them out of CSV, wind roses, GeoTIFF and NetCDF with units converted at the boundary.
• Uncertainty on every number: Monte-Carlo standard errors on each sampled statistic (batch means, Maritz-Jarrett quantiles, influence-function CVaR), and observed-information standard errors, confidence intervals, profile likelihoods and delta-method propagation on each fitted parameter.
• Sensitivity analysis: exact local derivatives and dimensionless elasticities by autodiff, and global Sobol variance decomposition separating direct effects from interactions.
• Model validation: AIC/AICc/BIC, WAIC and PSIS-LOO with Pareto-k diagnostics, model comparison with the standard error of the difference, PIT calibration, CRPS, and posterior predictive checks.
• Posterior inference through NumPyro: NUTS over the compiled model, with split-r-hat, effective sample sizes and divergence reporting.
• Units that are checked — dimensions, SI prefixes, compound expressions and affine temperature scales, enforced at frame and attribute boundaries at model-build time.
• A standardized export format (ModelPackage): quantized PMF, histogram, or JAX-free sampler representations, JSON round-trippable, consumer-agnostic, carrying a run manifest that records what produced the result.

Installation

pip install prophys

Ships as a prebuilt binary wheel, including the compiled native core (prophys._core), so no Rust toolchain or compilation step is needed.

Free tier / licensing: model calls touching up to 250 structural objects (points, polygon vertices, segments, raster cells, ...) run without any license. This does not count Monte-Carlo sample counts, observation counts, or training iterations, only declared model structure. Larger models require a signed license (set via the PROPHYS_LICENSE environment variable). Contact RhineQC GmbH to obtain one.

With optional extras:

pip install "prophys[viz]"        # matplotlib plotting helpers (prophys.plot)
pip install "prophys[geo]"        # shapely/pyproj geodata import
pip install "prophys[inference]"  # NumPyro bridge

Requirements: Python ≥ 3.11 (JAX and Optax install automatically as dependencies). Python 3.10 is not supported: jaxlib has not published a cp310 wheel since 0.9.0.

Quick Start

import jax
import jax.numpy as jnp
import optax
import prophys as prp

# Geometry in a shared frame: the pipeline route (optimizable) and a receptor
site = prp.Frame("site", units="m")
route = prp.Polyline(
    prp.Param("route", shape=(4, 2), init=jnp.array([[0.0, 0.0], [300.0, 60.0], [650.0, -30.0], [900.0, 0.0]])),
    site,
)
house = prp.Point(jnp.array([500.0, 250.0]), site)

# Physics: a leak can occur anywhere along the route, so its location is a
# continuous, reparameterized random variable pushed through the route's
# differentiable point_at(t), not a hand-picked set of candidate points
leak_frac = prp.RandomVariable("leak_frac", prp.Uniform(low=0.0, high=1.0))
leak_point = route.point_at(leak_frac)

# Gas dispersion from the leak to the house, then concentration -> a
# probability of health issues (dose-response curve)
distance = prp.norm(house.coords[0] - leak_point)
concentration = prp.Param("emission_rate", init=1.5) * prp.exp(-distance / 150.0)
health_risk = prp.PiecewiseLinear([0.0, 0.05, 0.2, 0.5], [0.0, 0.05, 0.3, 0.85], smooth=0.02)(concentration)

# Probabilistic attribute + standardized model object
incident = prp.UncertainAttribute("health_incident", prp.Bernoulli(prob=prp.clip(health_risk, 1e-4, 1 - 1e-4)))
model = prp.ProbabilityModel(incident)
compiled = model.compile(mode="opt", n_samples=200)

# Probability of a health incident, marginalized over every leak location
# along the route, with a gradient with respect to every vertex of the route
p_incident = compiled.expectation("health_incident")
grad_wrt_route = jax.grad(
    lambda params: compiled.expectation("health_incident", params=params)
)(compiled.default_params())

# Calibrate against observations ...
result = prp.calibrate(compiled, "health_incident", observations, n_steps=300)

# ... or optimize the route itself, directly against the marginalized
# probability, with plain JAX/optax (prp.optimize is only a convenience
# wrapper around the same pattern, not required)
opt = optax.adam(0.05)
raw_params = compiled.default_params()
opt_state = opt.init(raw_params)

for _ in range(200):
    loss, grads = jax.value_and_grad(
        lambda params: compiled.expectation("health_incident", params=params)
    )(raw_params)
    updates, opt_state = opt.update(grads, opt_state)
    raw_params = optax.apply_updates(raw_params, updates)

Uncertainty at the inputs, not just the output

wind_speed = prp.RandomVariable("v", prp.Weibull(scale=6.0, concentration=2.0))
wind_dir   = prp.RandomVariable("phi", prp.VonMises(loc=prp.deg2rad(35.0), kappa=3.0))

# Any expression downstream of these is a random quantity. The compiled
# model marginalizes them via reparameterized Monte Carlo, inside one
# JAX trace, so gradients flow back through the marginalization.

Training the model

# Fit one attribute
prp.calibrate(compiled, "exposure", observations)

# Jointly train the entire model. Shared parameters pool information
# across every attribute's observations
prp.finetune(compiled, {"exposure": obs_e, "complaints": obs_c}, frozen=["emission_rate"])

# Learn a standalone (also multivariate) distribution from data
mvn = prp.MultivariateGaussian(
    mean=prp.Param("mu", shape=(2,), init=jnp.zeros(2)),
    cholesky=prp.Param("L", shape=(2, 2), init=jnp.eye(2)),
)
prp.fit_distribution(mvn, data)

Records that are not exact numbers

A concentration below a detection limit, a specimen still unfailed at the end of a test, a reading binned at recording, a failed sensor: each carries a different likelihood term, and substituting a number for any of them changes the fitted parameters.

obs = prp.Observations.detection_limit(concentrations, limit=0.5)   # "< 0.5"
obs = prp.Observations.right_censored(times, censored=still_running)
obs = prp.Observations.interval(lower_edges, lower_edges + width)
obs = prp.Observations.from_arrays(values, missing=np.isnan(values))

prp.finetune(compiled, {"c": obs})    # accepted anywhere an array was

prophys.io reads them straight out of the formats they arrive in — a CSV with <0.5 and ND entries, a wind rose in degrees, a GeoTIFF, a NetCDF extract — with units declared per column and converted at the boundary.

The uncertainty of the numbers

Every reported statistic comes with its own error, and every fitted parameter with its own interval.

compiled.assess("loss", "cvar", level=0.99)   # estimate + Monte-Carlo error
prp.required_samples(estimate, 0.01)          # draws needed for 1% precision

u = prp.parameter_uncertainty(compiled, {"y": data}, result=result)
u.interval("sigma")                           # standard errors + intervals
u.collinear_pairs()                           # what the data cannot separate
prp.propagate(u, lambda p: p["mu"] + 1.645 * p["sigma"])   # into any quantity

The Monte-Carlo error uses batch means for a mean, the Maritz-Jarrett order-statistic estimator for a quantile, and the influence function for a CVaR — the last including the uncertainty of the VaR threshold itself. Parameter errors come from the observed information at the optimum, mapped into constrained space so a positive scale never gets an interval straddling zero; an unidentified direction is reported as such rather than as numbers.

What the answer depends on

prp.local_sensitivity(compiled, "power")      # exact derivatives, elasticities
prp.sobol_sensitivity(compiled, "power")      # variance shares + interactions
compiled.conditional_expectation("power", {"wind": np.array([4., 8., 12.])})

Whether the model is any good

prp.information_criteria(compiled, {"y": data}, result=result)   # AIC/AICc/BIC
prp.waic(compiled, "y", data, draws)                             # predictive
prp.loo(compiled, "y", data, draws)                              # + Pareto k
prp.compare({"simple": a, "richer": b})                          # with SE of Δ
prp.calibration(compiled, "y", data)      # under-dispersed? biased?
prp.crps(compiled, "y", data)             # the whole distribution, in its units
prp.posterior_predictive_check(compiled, "y", data, statistic="skew")

prp.report_card(compiled, "y", data, result=result)   # all of the above

The posterior, not just the mode

from prophys.inference import sample_posterior     # pip install "prophys[inference]"

posterior = sample_posterior(compiled, {"y": data}, init_params=result.raw_params)
print(posterior.summary())        # credible intervals, r_hat, ESS, divergences
prp.waic(compiled, "y", data, posterior.thin(400))

Units that are checked

site = prp.Frame("site", units="m")
plan = prp.Frame("plan", units="km")
prp.FrameTransform(plan, site).apply(pts)   # converts; it is not the identity

power = prp.UncertainAttribute("power", dist, unit="kW")
power.convert(1000.0, "MW")   # 1.0
power.convert(1.0, "kWh")     # UnitError: power is not energy

Export

pkg = prp.export.export_model(
    compiled,
    params=result.raw_params,
    observations={"y": data},
    diagnostics=prp.report_card(compiled, "y", data, result=result),
)
pkg.save("model.json")       # versioned, consumer-agnostic interchange format

Alongside the quantized PMF / histogram / JAX-free sampler, format 1.1 carries a run manifest: the prophys version, a structure hash of the model, the parameters in both constrained and unconstrained form, the seed and Monte-Carlo settings, the licensing regime, the observation counts by kind, and any attached diagnostics. A later run can verify against it:

loaded = prp.export.ModelPackage.load("model.json")
loaded.verify(model.compile())                     # raises if the model changed
compiled = prp.export.reproduce(loaded.manifest, model)

Downstream consumers implement their own import against the documented ModelPackage shape. prophys carries no consumer-specific dependencies.

Examples

• examples/gas_dispersion/: a complete community-health-risk model. Gaussian plume physics, wind as random inputs, terrain-corrected effective stack height, dose-response calibration, and constrained abatement optimization. Every primitive family appears once.

Documentation

Full docs are at docs.rhineqc.com/distribution/prophys, covering concepts (symbolic graphs, frames, eval/opt modes, marginalization), the full structure/distribution/transformation catalogs with plots, units and dimensional checking, reading measurement data, calibration and design optimization, Monte-Carlo and parameter uncertainty, sensitivity analysis, model validation, posterior inference, step-by-step worked examples, and the export format specification.

License

prophys is proprietary software by RhineQC GmbH. It is intended for internal or explicitly licensed distribution.