prophys.nn

Learned surrogates inside the probabilistic graph.

Surrogates are small to medium networks, tree ensembles and field transformers that are pretrained on their own data and then fine tuned jointly with every other parameter of a prophys model. All numerics run in optimized native kernels, with the same numerics on CPU and GPU.

Classes

DPOTNet	The whole DPOT network as one surrogate layer.
DPOTSurrogate	A field surrogate running a DPOT network.
Activation	Elementwise activation function applied as its own layer.
Bounded	Output head squashing into the open interval (low, high) with a sigmoid.
Choice	Softmax output head giving the choice probabilities of a discrete choice model.
CrossAttention	Multi head cross attention from the rows to a context set.
Dense	Fully connected layer: an affine map followed by an activation.
Derivatives	Outputs and input derivatives of a surrogate at collocation points.
FieldSurrogate	A surrogate whose input is a set or field of tokens rather than named scalars.
Fourier	Fourier feature embedding sin(2 pi B h) and cos(2 pi B h) (Tancik et al. 2020).
GaussianHead	Mean and scale output head for a heteroscedastic Gaussian likelihood.
HistogramHead	Histogram (bar distribution) output head for a full predictive distribution.
BarDistribution	Predictive distribution of a histogram head: a bar (Riemann) distribution with tails.
InContextRegressor	Prior fitted network that predicts the posterior predictive distribution from a context set in one pass.
InputConvex	Input convex neural network (Amos et al. 2017), convex in every input.
LayerNorm	Layer normalisation over the columns of each row, with learned gain and bias.
LinearPosterior	Gaussian posterior of a surrogate's linear head y = [phi, 1] w.
Lipschitz	Dense layer with a bounded Lipschitz constant in the maximum norm.
MixtureHead	Gaussian mixture output head for multimodal predictive distributions.
Monotone	Constrained monotone layer (Runje and Shankaranarayana 2023).
Physics	A layer whose body is a prophys symbolic expression.
PhysicsResidual	A physics informed residual loss term for pretrain().
PoolAttention	Attention pooling of the rows with a learned seed query (Lee et al. 2019).
PoolMean	Mean over the rows, a permutation invariant summary of a set.
PrefixAttention	Self attention in which every row attends only to the first keys rows.
PretrainResult	Training card returned by pretrain().
PriorTask	One synthetic regression task drawn from the mechanism prior.
QuantileDistribution	Predictive distribution of a quantile head: a piecewise linear quantile function with tails.
QuantileHead	Noncrossing quantile output head.
RFF	Random Fourier features of an RBF kernel (Rahimi and Recht 2007).
Reshape	Reinterpret each sample's tensor as rows by cols values.
Residual	Residual block h + f(h) with depth dense layers inside.
SelfAttention	Multi head self attention over the rows of a set or field.
SetBlock	Pre norm set transformer block: self attention and a GELU MLP, each with a residual connection.
SliceAttention	Transolver physics attention (Wu et al. 2024), linear in the number of rows.
SupportRegion	Region of input space covered by a surrogate's training data.
Surrogate	A trainable neural network surrogate with named, unit checked inputs, usable in a model.
SurrogateCall	Expression node evaluating a surrogate on input expressions.
TransolverBlock	Pre norm Transolver block: slice attention and a GELU MLP, each with a residual connection.
TreeEnsemble	A differentiable tree ensemble surrogate with named, unit checked inputs.
SampleSet	Expression node holding a fresh Monte Carlo sample [n, k] of an expression.

Functions

DPOT	Build a DPOT field surrogate with the hyperparameters of the reference model.
dpot_config	Hyperparameters of a DPOT checkpoint, read from its tensor shapes.
dpot_state	Leaf values of surrogate from a DPOT state dictionary.
dpot_tokens	Flatten fields [..., X, Y, T, C] to DPOT tokens [..., X*Y*T, C].
from_dpot	Build a DPOT surrogate carrying the weights of a released checkpoint.
Positive	Softplus output activation for a positive quantity.
field_tokens	Token rows of a raster field: cell centre coordinates followed by the cell values.
find_surrogates	Find every learned surrogate used by a probabilistic model.
fit_linear_posterior	Fit a Bayesian linear regression by maximising the model evidence.
geometry_tokens	Token rows of a geometry, one row [x0, y0, x1, y1] per segment or edge.
load_state_dict	Copy pretrained tensors from a PyTorch style state dictionary into a surrogate.
physics_attention_mapping	Map a SliceAttention layer's leaves to the keys of a Transolver physics attention module.
point_tokens	Token rows of a point set: coordinates followed by optional per point features.
sample_prior_task	Draw one synthetic regression task from the mechanism prior.
sample_set	Turn the Monte Carlo sample of an upstream quantity into a token set.
train_pfn	Train an InContextRegressor on fresh synthetic tasks from the mechanism prior.
transolver_mapping	Map a TransolverBlock layer's leaves to the keys of a Transolver block module.
fit_support	Fit the support region of a set of training rows.
pretrain	Train a surrogate on its own data before it enters a model.

Constants and objects

losses

Probabilistic training losses and distribution targets for surrogates.