omnibias.torch.architectures.pinn

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.architectures.pinn

Module contract & mathematical conventions

Physics-informed neural networks built from OMBU primitives.

These architectures sidestep the standard PINN pain point -- repeated torch.autograd.grad calls to compute spatial / temporal derivatives of the network output -- by structuring the network as a single hidden layer of width H and computing each derivative analytically via the activation's closed-form derivative tower:

u(x_1, ..., x_d) = c_bias + sum_{i=1..H} c_i * sigma(W_i . x + b_i)

du/dx_a = sum_i c_i * W_{i,a} * sigma'(z_i) d^2u/dx_a^2 = sum_i c_i * W_{i,a}^2 * sigma''(z_i) d^2u/dx_a dx_b = sum_i c_i * W_{i,a} * W_{i,b} * sigma''(z_i)

with z_i = W_i . x + b_i and sigma^(n) evaluated through the fast-path kernel of the chosen ActivationSpec. There is no torch.autograd.grad in the inner loop and no quotient-of-differences in the bias-collapse regime.

The PINNOMBU base class collects the chain-rule helpers; the concrete PINNHeat implements the 1D heat equation u_t = alpha * u_xx.

For *deep* networks and *arbitrary* derivative order, JetMLP replaces the hand-coded order-2 chain rule with the exact multivariate jet kernel omnibias.torch.jet_mv.mlp_jet_mv: one forward pass yields the value, gradient, Hessian and every mixed partial up to total order N with no torch.autograd.grad stacking. DeepPINNHeat is the deep counterpart of PINNHeat built on it.

*Spectral bias* -- the tendency of plain MLPs to learn low frequencies first -- is mitigated by two omnibias-native, closed-form-tower constructs that share the same exact jet readout (_JetMLPCore):

FourierFeatureMLP lifts the input through a random Fourier-feature encoding gamma(x) = [cos(B x), sin(B x)] (Tancik et al. 2020). Because cos(z) = sin(z + pi/2), the whole encoding is a *single* omnibias sin layer, so D^alpha u(x) stays exactly closed form to arbitrary order -- no autodiff for the PDE residual. A tuple of frequency_scale values gives a multi-scale encoding.
make_siren builds a SIREN (Sitzmann et al. 2020): a JetMLP with sin activations and the SIREN initialisation. Its derivative tower is exact for every order (sin^{(n)}(z) = sin(z + n pi/2)).

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

CLASS

DeepPINNHeat

omnibias.torch.architectures.DeepPINNHeatSource ↗

class DeepPINNHeat(nn.Module)

Deep 1-D heat-equation PINN u_t = alpha * u_xx on JetMLP.

The deep, arbitrary-depth counterpart of PINNHeat: the spatial and temporal derivatives come from a single closed-form order-2 multivariate jet per collocation point (no autograd-through-derivative), so the PDE residual is exact.

Parameters

hidden : int, default 32 Hidden width. depth : int, default 3 Number of hidden layers. base : str or ActivationSpec, default "tanh" Base activation with a fast path supporting orders 1 and 2. alpha : float, default 0.1 Diffusion coefficient.

Parameters

Name	Type & default	Contract
`hidden`positional-or-keyword	`int`Default: 32	int, default 32 Hidden width.
`depth`positional-or-keyword	`int`Default: 3	int, default 3 Number of hidden layers.
`base`positional-or-keyword	`str \| ActivationSpec[Tensor]`Default: 'tanh'	str or `ActivationSpec`, default `"tanh"` Base activation with a fast path supporting orders 1 and 2.
`alpha`positional-or-keyword	`float`Default: 0.1	float, default 0.1 Diffusion coefficient.

Executed regression · deeppinnheat shape mismatch raises

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_deeppinnheat_shape_mismatch_raises() -> None:
    model = DeepPINNHeat(hidden=8, depth=2)
    with pytest.raises(ValueError, match="matching shape"):
        model(torch.rand(5), torch.rand(6))

python -m pytest "packages/omnibias-torch/tests/test_jet_pinn.py::test_deeppinnheat_shape_mismatch_raises" -q

Regression example · deeppinnheat training decreases

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_deeppinnheat_training_decreases() -> None:
    torch.manual_seed(0)
    model = DeepPINNHeat(hidden=16, depth=2, base="tanh", alpha=0.1)
    optim = torch.optim.Adam(model.parameters(), lr=1e-2)
    losses: list[float] = []
    for _ in range(12):
        optim.zero_grad()
        x = torch.rand(64)
        t = torch.rand(64)
        u, res = model(x, t)
        loss = (res**2).mean() + (u[0] ** 2)
        loss.backward()
        optim.step()
        losses.append(loss.item())
    assert losses[-1] < losses[0]

python -m pytest packages/omnibias-torch/tests/test_jet_pinn.py -k test_deeppinnheat_training_decreases -q

Fields & methods

__init__method

__init__(self, hidden: int=32, depth: int=3, base: str | ActivationSpec[Tensor]='tanh', alpha: float=0.1) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`hidden`positional-or-keyword	`int`Default: 32	See the operation contract and implementation below.
`depth`positional-or-keyword	`int`Default: 3	See the operation contract and implementation below.
`base`positional-or-keyword	`str \| ActivationSpec[Tensor]`Default: 'tanh'	See the operation contract and implementation below.
`alpha`positional-or-keyword	`float`Default: 0.1	See the operation contract and implementation below.

Returns: None

Executed regression · deeppinnheat shape mismatch raises

Full file, imports & fixtures ↗

def test_deeppinnheat_shape_mismatch_raises() -> None:
    model = DeepPINNHeat(hidden=8, depth=2)
    with pytest.raises(ValueError, match="matching shape"):
        model(torch.rand(5), torch.rand(6))

python -m pytest "packages/omnibias-torch/tests/test_jet_pinn.py::test_deeppinnheat_shape_mismatch_raises" -q

Method implementation

    def __init__(
        self,
        hidden: int = 32,
        depth: int = 3,
        base: str | ActivationSpec[Tensor] = "tanh",
        alpha: float = 0.1,
    ) -> None:
        super().__init__()
        self.alpha = alpha
        self.net = JetMLP(in_dim=2, hidden=hidden, out_dim=1, depth=depth, base=base)
        self.net._check_fastpath(max_order=2)

Name	Type & default	Contract
`x`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.
`t`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

Name	Type & default	Contract
`in_dim`positional-or-keyword	`int`Required	Number of input coordinates `D`.
`num_features`positional-or-keyword	`int`Default: 64	Number of Fourier features `F` per frequency band; the encoding width is `2 * F * len(scales)`.
`hidden`positional-or-keyword	`int`Default: 64	Hidden width of the body MLP.
`out_dim`positional-or-keyword	`int`Default: 1	Number of output components (default 1).
`depth`positional-or-keyword	`int`Default: 3	Number of hidden (activated) body layers `>= 0`; `depth=0` is a pure random-Fourier-feature model with a linear readout.
`base`positional-or-keyword	`str \| ActivationSpec[Tensor]`Default: 'tanh'	Body activation (name or `ActivationSpec`) with a closed-form derivative fast path (default `"tanh"`).
`frequency_scale`keyword-only	`float \| Sequence[float]`Default: 1.0	Bandwidth(s) of the Gaussian frequency matrix. A scalar selects one band with `B ~ N(0, (2 pi * scale)^2)`; a sequence concatenates several bands for a multi-scale encoding (helps when the target mixes low and high frequencies).
`trainable_features`keyword-only	`bool`Default: False	If `True` the frequency matrix / phases are learnable parameters; otherwise they are fixed buffers (the classic random-features regime).
`seed`keyword-only	`int`Default: 0	Seed for the frequency-matrix draw.

Name	Type & default	Contract
`in_dim`positional-or-keyword	`int`Required	Number of input coordinates `D` (e.g. 2 for `(x, t)`).
`hidden`positional-or-keyword	`int`Required	Hidden width.
`out_dim`positional-or-keyword	`int`Default: 1	Number of output components (default 1).
`depth`positional-or-keyword	`int`Default: 3	Number of hidden (activated) layers `>= 1`; the readout is a pure affine map. `depth=1` recovers a single-hidden-layer network like `PINNOMBU`.
`base`positional-or-keyword	`str \| ActivationSpec[Tensor]`Default: 'tanh'	Base activation (name or `ActivationSpec`) with a closed-form derivative fast path (e.g. `"tanh"`, `"sigmoid"`, `"softplus"`).

Name	Type & default	Contract
`hidden`positional-or-keyword	`int`Default: 64	int, default 64 Number of hidden units (`H`).
`base`positional-or-keyword	`str \| ActivationSpec[Tensor]`Default: 'softplus'	str or `ActivationSpec`, default `"softplus"` Base activation. Must have a fast-path supporting orders 1 and 2.
`alpha`positional-or-keyword	`float`Default: 0.1	float, default 0.1 Diffusion coefficient.

Name	Type & default	Contract
`z`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.
`axis`positional-or-keyword	`int`Required	See the operation contract and implementation below.

Name	Type & default	Contract
`in_dim`positional-or-keyword	`int`Required	As for `JetMLP` (`depth >= 1`).
`hidden`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`out_dim`positional-or-keyword	`int`Default: 1	See the operation contract and implementation below.
`depth`positional-or-keyword	`int`Default: 3	See the operation contract and implementation below.
`omega_0`keyword-only	`float`Default: 30.0	First-layer frequency scale (default `30`, the SIREN default).
`seed`keyword-only	`int`Default: 0	Seed for the weight draw.

omnibias.torch.architectures.pinn

DeepPINNHeat​

Parameters

Fields & methods

Parameters

Parameters

FourierFeatureMLP​

Parameters

Fields & methods

Parameters

JetMLP​

Parameters

Fields & methods

Parameters

PINNHeat​

Parameters

Fields & methods

Parameters

Parameters

PINNOMBU​

Fields & methods

Parameters

Parameters

Parameters

Parameters

make_siren​

Parameters

DeepPINNHeat

FourierFeatureMLP

JetMLP

PINNHeat

PINNOMBU

make_siren