mixle.models.pinn module

PINNRegression – a physics-informed neural network as a Mixle conditional-density model.

A NeuralGaussian fits p(y | x) = N(y; module(x), noise^2 I) from labeled (x, y) pairs alone. PINNRegression is the same model plus a residual penalty: at every M-step it also draws unlabeled collocation points from a box domain, evaluates a caller-supplied PDE/ODE residual on the module’s output via autograd, and adds residual_weight * mean(residual**2) to the training loss – the standard physics-informed-neural-network (PINN) loss, L = L_data + w * L_physics.

This makes the labeled-data term double as boundary/initial conditions (or scattered measurements) and the residual term enforce the governing equation between them, so the fitted module honors the physics even where it never saw labeled data – the whole point of a PINN over plain regression. With zero labeled data (suff_stat empty) the model still trains: pure PDE-residual fitting, a boundary-value/collocation solver.

The reported density (log_density()/seq_log_density(), inherited unchanged from NeuralGaussian) is the data-fit Gaussian NLL only – the model never claims the residual penalty as part of its probability model. mixle.inference.planning.certify() already caps a bare gradient-fit model like this at STATIONARY (no global-optimum claim), so penalized= adds nothing for a standalone fit; pass certify(structure, penalized="PINN residual") when this model is composed as one block of a larger structure that otherwise contains closed-form EM blocks, so the composite certificate records the residual-penalized training step as a gradient-based block (mirroring how mixle.ppl.core.ode_residual()’s soft-constraint fits are certified).

Requires torch. residual_fn(module, collocation_points) -> tensor computes the residual using torch.autograd.grad on the module’s output w.r.t. collocation_points (which arrive with requires_grad_(True) already set) – ordinary PINN practice, e.g. for a 1-D heat equation u_t = alpha * u_xx over inputs (t, x):

def heat_residual(module, coll):
    u = module(coll)
    grads = torch.autograd.grad(u, coll, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_t, u_x = grads[:, 0:1], grads[:, 1:2]
    u_xx = torch.autograd.grad(u_x, coll, grad_outputs=torch.ones_like(u_x), create_graph=True)[0][:, 1:2]
    return u_t - ALPHA * u_xx

model = PINNRegression(make_mlp(2, [32, 32], 1), heat_residual, domain=([0.0, -1.0], [1.0, 1.0]))
class PINNRegression(module, residual_fn, domain, *, noise=1.0, residual_weight=1.0, n_collocation=64, m_steps=40, lr=0.01, seed=0, name=None, device=None)[source]

Bases: NeuralGaussian

NeuralGaussian plus a PDE/ODE-residual penalty evaluated on sampled collocation points.

domain is a (low, high) pair of per-dimension box bounds for collocation sampling; residual_fn computes the physics residual (see module docstring); residual_weight scales the penalty relative to the data-fit NLL; n_collocation is how many collocation points are drawn fresh every M-step.

Parameters:

  • module (Any)

  • residual_fn (Any)

  • domain (tuple[Any, Any])

  • noise (float)

  • residual_weight (float)

  • n_collocation (int)

  • m_steps (int)

  • lr (float)

  • seed (int)

  • name (str | None)

  • device (Any)

estimator(pseudo_count=None)[source]

Return the estimator that combines weighted data fit with residual collocation penalties.

Parameters:

pseudo_count (float | None)

Return type:

PINNRegressionEstimator

dist_to_encoder()[source]

Return the neural-Gaussian encoder for (x, y) observation pairs.

Return type:

NeuralGaussianEncoder

to_dict()[source]

Serialize the module, residual function reference, domain, and PINN hyperparameters.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Rebuild a PINNRegression from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

PINNRegression

class PINNRegressionEstimator(module, residual_fn, domain, *, noise=1.0, residual_weight=1.0, n_collocation=64, m_steps=40, lr=0.01, seed=0, name=None, device=None)[source]

Bases: NeuralGaussianEstimator

EM estimator for PINNRegression: the M-step adds a residual penalty on fresh collocation points to the same weighted-NLL gradient descent NeuralGaussianEstimator runs.

Collocation sampling is deterministic given seed (a private numpy.random.RandomState, advanced once per M-step) – refitting with the same seed draws the same collocation batches.

Parameters:

  • module (Any)

  • residual_fn (Any)

  • domain (tuple[np.ndarray, np.ndarray])

  • noise (float)

  • residual_weight (float)

  • n_collocation (int)

  • m_steps (int)

  • lr (float)

  • seed (int)

  • name (str | None)

  • device (Any)

accumulator_factory()[source]

Return the neural-Gaussian accumulator factory for weighted observation pairs.

Return type:

NeuralGaussianAccumulatorFactory

estimate(nobs, suff_stat)[source]

Run the data-plus-residual M-step and return the updated PINN leaf.

Parameters:
Return type:

PINNRegression