Concepts

Symbolic graphs, not eager computation

Every quantity you build in prophys — a distance, a transformed value, a distribution parameter — is an Expr node, not a concrete number. Operators (+ - * / ** @, comparisons, indexing) are overloaded so ordinary Python arithmetic builds a graph:

import prophys as prp

a = prp.Param("a", init=2.0)
b = prp.Const(3.0)
expr = a * b + prp.exp(a)          # still symbolic
value = expr.evaluate({"a": 5.0}) # now a concrete jax.Array

evaluate(env) substitutes concrete values for every Param/Input leaf and recursively computes with jax.numpy — so any expression is automatically jit/grad/vmap-able, with no special handling required.

Leaves: Const, Input, Param, RandomVariable

  • Const — a fixed value (baked into the graph; folds away under jit).

  • Input — a named placeholder bound at evaluation time (e.g. observations passed into log_prob).

  • Param — a named, learnable/optimizable leaf. bounds=(lo, hi) or transform="softplus" constrain it through a smooth bijector, so unconstrained gradient descent (Adam, L-BFGS) never needs to know the parameter is bounded.

  • RandomVariable — a named stochastic leaf backed by a Distribution. Unlike the other leaves, uncertainty can enter a model at any point in the pipeline — not just at the very end. See Marginalizing over random inputs below.

Frames: a shared coordinate system

Every spatial structure (PointList, Polygon, Field, …) is annotated with a Frame. Structures declared in the same frame can be combined freely; combining structures from different, unregistered frames raises a FrameError immediately at model-build time (a plain Python check, not a runtime cost inside the JAX trace).

site = prp.Frame("site")
houses = prp.PointList(house_coords, site)
turbines = prp.PointList(turbine_coords, site)   # same frame -> fine

If two frames are related by a known affine transform (e.g. a local survey grid vs. a site-wide frame), register it once and structures in either frame interoperate transparently:

from prophys import FrameRegistry, FrameTransform

registry = FrameRegistry()
registry.register(FrameTransform(local_frame, site, offset=jnp.array([1000.0, 500.0])))

Eval mode vs. opt mode

Several primitives — the minimum over a set of distances, containment inside a polygon, a piecewise-linear kink — are not smooth at their decision boundary. Every such primitive exposes a mode argument:

  • mode="eval" (default): the exact, hard computation — a true minimum, a hard 0/1 containment indicator. Use this to read off a model’s current numbers.

  • mode="opt", with a temperature tau: a smooth approximation (softmin/softmax, a sigmoid-smoothed indicator, a soft segment blend for PiecewiseLinear) so gradients remain well-behaved when that quantity feeds a calibration or design optimization loop.

ProbabilityModel.compile(mode=...) sets this consistently for an entire model.

Differentiable discrete states

The same eval/opt split extends to discrete design decisions, not just geometry: BinaryState (a learnable on/off, e.g. a valve or a reinforcement decision) and CategoricalState (a learnable choice among k discrete options, e.g. a machine operating mode) evaluate to the exact hard state (0/1, one-hot) in eval mode, and to a temperature-relaxed soft assignment — the deterministic limit of a Gumbel-Softmax relaxation — in opt mode:

reinforce = prp.BinaryState("reinforce_line_3", init=0.3)
flow_capacity = base_capacity * (0.2 + 0.8 * reinforce)  # differentiable gate

pump_mode = prp.CategoricalState("pump_mode", k=3)
setting = prp.dot(pump_mode, jnp.array([0.0, 0.5, 1.0]))

Multiply or dot a state into any expression to gate that term; the same gradient pass used for continuous Params (via prp.optimize) then also searches over the relaxed discrete decision, with the mode/tau schedule following the model’s compile(mode=, tau=) setting exactly like the geometry primitives above.

Marginalizing over random inputs

Because a RandomVariable can sit anywhere in the graph (e.g. wind speed and direction feeding a dispersion model, not just the final noise term), evaluating an attribute’s log_prob or expectation in general requires integrating out that upstream randomness:

\[p(y) = \int p(y \mid v)\, p(v)\, dv, \qquad \mathbb{E}[Y] = \int \mathbb{E}[Y \mid v]\, p(v)\, dv\]

CompiledModel does this via reparameterized Monte Carlo, entirely inside one JAX trace: draw n_samples samples of every upstream RandomVariable, evaluate the attribute at each sample, and combine with logsumexp (for log_prob) or a plain mean (for expectation). Because sampling uses each distribution’s reparameterized form where available (Gaussian, Weibull, Exponential, Gumbel, LogNormal), gradients flow through the Monte-Carlo estimate back to every upstream Param — including parameters of the wind distribution itself.

A practical shape note: when an attribute’s distribution parameters have their own batch shape (e.g. one concentration value per residence), that shape must broadcast against the (n_samples,) shape of the marginalized samples. A per-receptor batch of size 1 broadcasts against any n_samples; a batch of, say, 4 residences generally will not unless n_samples happens to equal 4. See the gas-dispersion example for the standard pattern (build one scalar receptor per attribute that needs marginalizing).

The standardized model object

ProbabilityModel bundles a model’s UncertainAttributes (plus any AttributeInteractions) into one object with a single entry point: .compile(mode, tau, n_samples, seed) returns a CompiledModel exposing log_prob, sample, expectation, prob, quantile, cvar, and export — a stable interface regardless of how the underlying graph was built.