prophys 0.1.0

Prophys, a differentiable, JAX-native DSL for spatial probability and risk models

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.
• A standardized export format (ModelPackage): quantized PMF, histogram, or JAX-free sampler representations, JSON round-trippable, consumer-agnostic.

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)

Export

pkg = compiled.export()      # quantized PMF / histogram / JAX-free sampler
pkg.save("model.json")       # versioned, consumer-agnostic interchange format

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, calibration and design optimization, a step-by-step worked example, and the export format specification.

License

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