Mixed partials for a heat residual
Use coordinates (t, x) and form the residual u_t - diffusivity * u_xx. A multivariate jet stores Taylor coefficients indexed by multi-index; jet_partials restores raw derivatives.
mixed_heat.py
"""Mixed Taylor coefficients for a heat-equation residual at one point."""
import torch
from omnibias.torch.jet_mv import jet_partials, mlp_jet_mv
weights = torch.tensor([[0.5, 0.7], [-0.2, 0.3]], dtype=torch.float64,
requires_grad=True)
bias = torch.tensor([0.1, -0.1], dtype=torch.float64)
readout = torch.tensor([[0.6, -0.4]], dtype=torch.float64)
layers = [(weights, bias, "tanh"), (readout, None, None)]
point = torch.tensor([0.2, 0.4], dtype=torch.float64) # coordinates (t, x)
jet = mlp_jet_mv(point, layers, order=2)
partials = jet_partials(jet, dim=2, order=2)
residual = partials[(1, 0)] - 0.1 * partials[(0, 2)]
residual.square().mean().backward()
assert partials[(1, 1)].shape == (1,)
assert weights.grad is not None and torch.isfinite(weights.grad).all()
print("Mixed partials:", sorted(partials), "finite parameter gradients: True")
The mixed-jet input is one point shaped (dim,). Here dim=2; partial (1, 0) differentiates time and (0, 2) differentiates space twice. The output has one channel. Batched evaluation needs explicit mapping over points.
This script validates a differentiable residual. A heat-equation solution also needs initial and boundary conditions, a collocation strategy, a training loop, and independent validation. The full PINN workflow explains that boundary.