Training & Design Optimization
Both parameter training (fit distribution/model parameters to data) and design optimization (choose a design variable, e.g. a position or an abatement fraction, to minimize/maximize an objective) reduce to the same operation: gradient descent through the compiled graph via Optax. The only difference is what plays the role of “the parameter being optimized” and what plays “the objective.”
Three training entry points, from narrowest to widest scope:
Function |
Scope |
|---|---|
One standalone distribution (univariate or multivariate), no model needed |
|
A compiled model, against observations of one attribute |
|
The entire model jointly, against observations of several attributes at once |
Calibration (MLE / MAP)
import prophys as prp
compiled = model.compile()
result = prp.calibrate(
compiled,
"acceptance_drop", # attribute name
observations, # array of observed values
n_steps=300,
learning_rate=0.05,
)
print(result.params) # fitted parameters, constrained space
print(result.final_loss)
Internally, calibrate() builds the negative
log-likelihood (optionally plus a prior term for MAP) as a function of the
model’s unconstrained parameters — the same space
Param.forward maps back from — so
Adam never has to worry about box constraints or positivity; those are
enforced structurally by the parameterization itself.
Freezing and selecting parameters
Both calibrate() and
finetune() accept trainable=[...] (whitelist) or
frozen=[...] (blacklist) of parameter names. Frozen parameters keep
their current values but still participate in the forward computation —
the standard workflow for staged fitting (fit the physics first, then the
noise model, then release everything for a final joint pass).
prp.calibrate(compiled, "complaints", records, trainable=["dose_k", "dose_x0"])
prp.finetune(compiled, all_observations, frozen=["emission_rate"])
Finetuning the entire model
finetune() maximizes the joint likelihood over
several attributes at once. Because attributes typically share upstream
parameters (a common emission rate, a shared monitoring sigma, a copula
correlation), joint fitting pools every observation’s information into
every shared parameter — with per-attribute weights if some data
sources are more trustworthy than others:
result = prp.finetune(
compiled,
{"exposure_res0": exposure_obs, "complaint_res0": complaint_records},
weights={"complaint_res0": 2.0},
)
Learning distributions directly
fit_distribution() trains the Param leaves of a
standalone distribution — no model required, and multivariate data works
the same way as univariate (the leading axis indexes observations):
# Univariate: recover Weibull parameters from samples
w = prp.Weibull(scale=prp.Param("scale", init=5.0, transform="softplus"),
concentration=prp.Param("k", init=1.0, transform="softplus"))
prp.fit_distribution(w, wind_measurements)
# Multivariate: learn a full covariance via its Cholesky factor
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, paired_measurements)
# Mixtures: learn component weights (softmax-normalized logits)
mix = prp.Mixture([prp.Gaussian(0.0, 1.0), prp.Gaussian(10.0, 1.0)],
logits=prp.Param("logits", shape=(2,), init=jnp.zeros(2)))
prp.fit_distribution(mix, samples)
The Cholesky and copula parameterizations are chosen so any unconstrained gradient step yields a valid covariance/correlation — no projections, no clipping, no failed updates.
Design optimization
result = prp.optimize(
objective, # a scalar Expr to minimize
wrt=[turbine_positions],
constraints=[min_spacing_violation], # each must evaluate <= 0
n_steps=200,
learning_rate=1.0,
)
print(result.params["pos"])
Constraints are enforced via a quadratic penalty; per-parameter bounds
don’t need a constraint at all, since a bounded Param
already can’t leave its valid range under any gradient step.
If the objective references a RandomVariable
(e.g. optimizing for a specific wind condition rather than marginalizing
over it), bind it via extra_env:
result = prp.optimize(
objective,
wrt=[abatement_param],
extra_env={"wind_speed": 6.0, "wind_direction": jnp.deg2rad(35.0)},
)
Using a compiled model in external JAX solvers
A CompiledModel is a bundle of pure JAX functions —
nothing about it is tied to prophys’ own Optax loops. Any quantity it
produces (an expectation, a log-likelihood, a CVaR) can be closed over in an
ordinary Python function of a parameter pytree and handed to any
JAX-gradient-based solver; gradients flow through the native kernels’
custom_vjp rules exactly as they do internally:
import jax
def objective(raw_params):
# any scalar the model can produce; raw_params is unconstrained space
return -compiled.log_prob("exposure", observations, params=raw_params).sum()
value_and_grad = jax.jit(jax.value_and_grad(objective))
# jaxopt
import jaxopt
result = jaxopt.LBFGS(fun=objective).run(compiled.default_params())
# optimistix
import optimistix as optx
sol = optx.minimise(lambda p, _: objective(p), optx.BFGS(rtol=1e-6, atol=1e-6),
compiled.default_params())
# scipy, via JAX-computed gradients
from scipy.optimize import minimize
import numpy as np
flat0 = np.asarray(compiled.default_params()["sigma"])
# ... flatten/unflatten with jax.flatten_util.ravel_pytree as usual
The same applies to symbolic expressions directly: expr.evaluate(env) is
a pure function of env, so jax.grad(lambda v: expr.evaluate({"x": v}))
works with any downstream solver, inside or outside jit.
Diagnostics
finite_difference_check() compares
jax.grad against central finite differences — useful whenever a
Surrogate wraps hand-written physics, to catch a
silent mismatch between the analytic and autodiff gradient.
check_identifiability() is a cheap heuristic flag for
“the loss was still improving meaningfully when optimization stopped.”
- prophys.engine.calibrate(compiled, attribute_name, observations, *, prior_neg_log_prob=None, init_params=None, trainable=None, frozen=None, n_steps=500, learning_rate=0.01, optimizer=None)[source]
Fit compiled’s parameters by maximizing the total log-likelihood of observations under the attribute attribute_name (MLE), optionally adding a prior_neg_log_prob(constrained_params) -> scalar term (MAP).
trainable/frozen restrict which Param leaves are updated; the rest keep their current values. For jointly fitting several attributes at once, use
finetune().
- prophys.engine.finetune(compiled, observations, *, weights=None, prior_neg_log_prob=None, init_params=None, trainable=None, frozen=None, n_steps=500, learning_rate=0.01, optimizer=None)[source]
Jointly fit the model against observations of several attributes.
observations maps attribute names to observation arrays; the loss is the (optionally weights-weighted) sum of every attribute’s negative log-likelihood. Because attributes typically share upstream parameters (a common emission rate, a shared noise sigma, a copula correlation), joint fitting propagates every observation’s information into every shared parameter — this is the “train the entire model” entry point.
- prophys.engine.fit_distribution(distribution, data, *, prior_neg_log_prob=None, init_params=None, n_steps=500, learning_rate=0.01, optimizer=None)[source]
Learn the Param leaves of a standalone Distribution directly from data — no ProbabilityModel required.
Works for univariate data (1-D array of scalar observations) and multivariate data alike: data’s leading axis indexes observations and each row is passed to log_prob as-is, so a
MultivariateGaussianwith a learnable Cholesky factor, aMixturewith learnable logits, or any distribution whose parameters are Param expressions can be fitted the same way.
- prophys.engine.optimize(objective, *, wrt, init_params=None, constraints=(), penalty_weight=1000.0, n_steps=500, learning_rate=0.01, optimizer=None, extra_env=None)[source]
Minimize objective (a scalar Expr) with respect to the named Param`s in `wrt, subject to constraints (each Expr must evaluate to
<= 0), enforced via a quadratic penalty.This performs plain (unconstrained-space) gradient descent; because Param.forward already maps bounded parameters through a sigmoid/ softplus, box constraints on individual parameters don’t need a penalty term at all — only cross-parameter constraints (e.g. minimum spacing) need constraints.
extra_env binds any non-Param leaves the objective/constraints reference (e.g. a RandomVariable fixed at a design point such as the median wind condition) — design optimization conditions on a scenario rather than marginalizing over it, unlike CompiledModel.expectation.
- prophys.engine.finite_difference_check(fn, params, eps=0.0001, rtol=0.01)[source]
Compare jax.grad(fn) against central finite differences for every scalar/array leaf in params. Returns, per parameter name, the max absolute and relative error between the two — use this to catch a hand-written custom gradient (e.g. in a Surrogate) that silently disagrees with autodiff.
- prophys.engine.check_identifiability(loss_history, tail_fraction=0.1)[source]
Cheap heuristic: if the loss over the final tail_fraction of steps hasn’t decreased meaningfully, warn that the fit may not have converged or that parameters may not be identifiable from the data.
- class prophys.engine.CompiledModel(model, mode, tau, n_samples, seed)[source]
A ProbabilityModel frozen into a concrete, jit/grad-friendly evaluation interface.
Parameters carrying the model’s Param leaves are passed as a flat dict of unconstrained values (matching Param.unconstrained_init() / Param.forward); this is exactly the representation Optax should optimize over, since it is unbounded.
- Parameters:
model (ProbabilityModel)
mode (str)
tau (float)
n_samples (int)
seed (int)
- constrained_params(params=None)[source]
Map unconstrained optimizer-space values to the constrained values the graph actually uses (identity for params without bounds/transform).
- log_prob(attribute_name, observation, params=None, key=None, n_samples=None)[source]
Log-density of observation under attribute.
If the attribute depends on upstream RandomVariable leaves (e.g. a noise level that depends on wind speed/direction), this marginalizes them out via reparameterized Monte Carlo:
log p(y) = logsumexp_i(log p(y | v_i)) - log(n_samples).
- sample(attribute_name, key, shape=(), params=None)[source]
Draw samples of attribute, first sampling any upstream RandomVariable leaves it depends on.
- expectation(attribute_name, params=None, key=None, n_samples=None)[source]
E[attribute], marginalizing any upstream RandomVariable leaves.
- prob(attribute_name, event, params=None, key=None, n_samples=None)[source]
P(attribute <= event)via Monte-Carlo sampling (works for any distribution, whether or not it has a closed-form CDF).
- quantile(attribute_name, q, params=None, key=None, n_samples=None)[source]
Empirical quantile from Monte-Carlo samples.