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.

The problem declaration identifies the equation, coordinates, fields, residual components, boundary/initial constraints, identifiability conditions, and whether constraints are enforced by observations or by the module architecture. Observation-enforced problems require actual labeled constraint rows. Residual-only optimization is accepted only when the declaration states that a constrained architecture enforces the named conditions; an unconstrained residual fit is never represented as a solved problem.

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

problem = PINNProblemSpec(
    equation="u_t - alpha*u_xx = 0",
    coordinates=("t", "x"),
    fields=("u",),
    residual_components=("heat",),
    constraints=("u(0, x) = initial(x)",),
    identifiability_conditions=("initial condition fixes the parabolic solution",),
    constraint_enforcement="observations",
    completeness_basis="initial-value problem on the declared box",
)
model = PINNRegression(
    make_mlp(2, [32, 32], 1),
    heat_residual,
    domain=([0.0, -1.0], [1.0, 1.0]),
    problem=problem,
)
class PINNProblemSpec(equation, coordinates, fields, residual_components, constraints, identifiability_conditions, constraint_enforcement, completeness_basis)[source]

Bases: object

Explicit mathematical and enforcement schema for a PINN fit.

Parameters:
to_dict()[source]

Return the JSON-compatible problem declaration.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Validate and restore a problem declaration.

Parameters:

payload (Any)

Return type:

PINNProblemSpec

class PINNConstraintReceipt(status, equation, constraint_enforcement, constraints, identifiability_conditions, completeness_basis, labeled_constraints, residual_loss_initial, residual_loss_final)[source]

Bases: object

Durable declaration and measured outcome of a constraint-aware PINN fit.

Parameters:

  • status (str)

  • equation (str)

  • constraint_enforcement (str)

  • constraints (tuple[str, ...])

  • identifiability_conditions (tuple[str, ...])

  • completeness_basis (str)

  • labeled_constraints (int)

  • residual_loss_initial (float)

  • residual_loss_final (float)

to_dict()[source]

Return a JSON-compatible receipt.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Restore a receipt produced by to_dict().

Parameters:

payload (Any)

Return type:

PINNConstraintReceipt

class PINNRegression(module, residual_fn, domain, *, problem, noise=1.0, residual_weight=1.0, n_collocation=64, m_steps=40, lr=0.01, seed=0, name=None, device=None, constraint_receipt=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])

  • problem (PINNProblemSpec)

  • noise (float)

  • residual_weight (float)

  • n_collocation (int)

  • m_steps (int)

  • lr (float)

  • seed (int)

  • name (str | None)

  • device (Any)

  • constraint_receipt (PINNConstraintReceipt | None)

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, *, problem, 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])

  • problem (PINNProblemSpec)

  • 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