API REFERENCE/omnibias-torch

omnibias.torch.architectures.hardbc

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.hardbc

Module contract & mathematical conventions

Hard-constraint boundary/initial-condition ansatz (torch).

A boundary-value problem is usually trained with a *soft* penalty: the PINN loss adds a term |u - g|^2 on the boundary and hopes the optimiser drives it to zero. The boundary condition is then only approximate, it fights the PDE residual for gradient budget, and the multi-objective loss is ill-conditioned.

The *hard-constraint* ansatz removes the boundary term entirely by baking the condition into the architecture (Lagaris et al. 1998; Sukumar & Srivastava 2022):

u(x) = g(x) + b(x)\, N(x),

where

• g (the **lift**) is a fixed smooth function that already satisfies the boundary data (g = g_\partial on \partial\Omega);
• b (the **boundary mask** / approximate distance function) vanishes *exactly* on the constrained set (b = 0 there);
• N is the free neural network (any omnibias.torch.architectures.JetMLP / ~omnibias.torch.architectures.FourierFeatureMLP / SIREN).

Then u satisfies the boundary condition for **every** value of the network parameters, so only the PDE residual remains in the loss. This pairs directly with the second-order optimiser in omnibias.torch.optim (one fewer loss term to balance).

The omnibias-native part is that the ansatz is assembled at the *jet* level. The lift g and the mask b are polynomials whose exact Taylor jets come from omnibias.torch.jet_mv.identity_jet, and the product b * N is the multivariate Leibniz rule omnibias.torch.jet_mv.jet_multiply. Hence the constrained field keeps the closed-form, arbitrary-order derivative contract: D^alpha u is exact with no torch.autograd.grad for the differential operator. HardConstraintField reuses the readout machinery of omnibias.torch.architectures.pinn._JetMLPCore.

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

AffineFactor

omnibias.torch.architectures.AffineFactorSource ↗

class AffineFactor

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

A degree-1 boundary factor scale * x[axis] + offset.

The factor vanishes on the hyperplane x[axis] = -offset / scale; a product of such factors is zero wherever *any* of them is, which is how a BoundaryMask encodes the constrained set.

Parameters

Name	Type & default	Contract
axisdataclass field	intRequired	See the operation contract and implementation below.
scaledataclass field	floatDefault: 1.0	See the operation contract and implementation below.
offsetdataclass field	floatDefault: 0.0	See the operation contract and implementation below.

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Regression example · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    bad = BoundaryMask((AffineFactor(3, 1.0, 0.0),))
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net, bad)

python -m pytest packages/omnibias-torch/tests/test_hardbc.py -k test_mask_factor_axis_out_of_range -q

Fields & methods

axisfield

axis: int

Declared int field on AffineFactor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

scalefield

scale: float = 1.0

Declared float field on AffineFactor. Default expression: 1.0.

View source ↗

offsetfield

offset: float = 0.0

Declared float field on AffineFactor. Default expression: 0.0.

View source ↗

__post_init__method

__post_init__(self) -> None

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

Returns: None

Validation & failure cases

• When self.axis < 0, raises ValueError(f'axis must be >= 0, got {self.axis}').
• When self.scale == 0.0, raises ValueError('scale must be non-zero (a constant factor is not a mask)').

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Method implementation

    def __post_init__(self) -> None:
        if self.axis < 0:
            raise ValueError(f"axis must be >= 0, got {self.axis}")
        if self.scale == 0.0:
            raise ValueError("scale must be non-zero (a constant factor is not a mask)")

View source ↗

Read the implementation

class AffineFactor:
    """A degree-1 boundary factor ``scale * x[axis] + offset``.

    The factor vanishes on the hyperplane ``x[axis] = -offset / scale``; a product
    of such factors is zero wherever *any* of them is, which is how a
    :class:`BoundaryMask` encodes the constrained set.
    """

    axis: int
    scale: float = 1.0
    offset: float = 0.0

    def __post_init__(self) -> None:
        if self.axis < 0:
            raise ValueError(f"axis must be >= 0, got {self.axis}")
        if self.scale == 0.0:
            raise ValueError("scale must be non-zero (a constant factor is not a mask)")

Other public import paths

• omnibias.torch.architectures.AffineFactor
• omnibias.torch.architectures.hardbc.AffineFactor

CLASS

AffineLift

omnibias.torch.architectures.AffineLiftSource ↗

class AffineLift

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Affine lift g(x) = W x + c matching the boundary data exactly.

weight is the (out_dim, D) matrix W and bias the (out_dim,) vector c, both stored as plain Python floats so the lift is fully device/dtype-agnostic and carries no trainable state (the boundary data is fixed). A constant lift is W = 0; a 1-D Dirichlet interpolation is a single non-zero slope.

Parameters

Name	Type & default	Contract
weightdataclass field	tuple[tuple[float, ...], ...]Required	See the operation contract and implementation below.
biasdataclass field	tuple[float, ...]Required	See the operation contract and implementation below.

Executed regression · lift dim mismatch rejected

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_lift_dim_mismatch_rejected() -> None:
    net = JetMLP(2, 8, 1, depth=2).double()
    mask = BoundaryMask((AffineFactor(0, 1.0, 0.0),))
    bad_lift = AffineLift(((0.0,), (0.0,)), (0.0, 0.0))  # in_dim 1, out_dim 2
    with pytest.raises(ValueError, match="out_dim|in_dim"):
        HardConstraintField(net, mask, bad_lift)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_lift_dim_mismatch_rejected" -q

Fields & methods

weightfield

weight: tuple[tuple[float, ...], ...]

Declared tuple[tuple[float, ...], ...] field on AffineLift. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

biasfield

bias: tuple[float, ...]

Declared tuple[float, ...] field on AffineLift. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

__post_init__method

__post_init__(self) -> None

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

Returns: None

Validation & failure cases

• When len(self.weight) != len(self.bias), raises ValueError(f'weight has {len(self.weight)} rows but bias has {len(self.bias)}').
• When any((len(row) != width for row in self.weight)), raises ValueError('weight rows must all have the same length').

Executed regression · lift dim mismatch rejected

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_lift_dim_mismatch_rejected() -> None:
    net = JetMLP(2, 8, 1, depth=2).double()
    mask = BoundaryMask((AffineFactor(0, 1.0, 0.0),))
    bad_lift = AffineLift(((0.0,), (0.0,)), (0.0, 0.0))  # in_dim 1, out_dim 2
    with pytest.raises(ValueError, match="out_dim|in_dim"):
        HardConstraintField(net, mask, bad_lift)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_lift_dim_mismatch_rejected" -q

Method implementation

    def __post_init__(self) -> None:
        if len(self.weight) != len(self.bias):
            raise ValueError(
                f"weight has {len(self.weight)} rows but bias has {len(self.bias)}"
            )
        if self.weight:
            width = len(self.weight[0])
            if any(len(row) != width for row in self.weight):
                raise ValueError("weight rows must all have the same length")

View source ↗

out_dimproperty

out_dim(self) -> int

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

Returns: int

Executed regression · lift dim mismatch rejected

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_lift_dim_mismatch_rejected() -> None:
    net = JetMLP(2, 8, 1, depth=2).double()
    mask = BoundaryMask((AffineFactor(0, 1.0, 0.0),))
    bad_lift = AffineLift(((0.0,), (0.0,)), (0.0, 0.0))  # in_dim 1, out_dim 2
    with pytest.raises(ValueError, match="out_dim|in_dim"):
        HardConstraintField(net, mask, bad_lift)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_lift_dim_mismatch_rejected" -q

Method implementation

    def out_dim(self) -> int:
        return len(self.bias)

View source ↗

in_dimproperty

in_dim(self) -> int

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

Returns: int

Executed regression · parameters are only the network

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_parameters_are_only_the_network() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=2.0)
    field_params = {id(p) for p in field.parameters()}
    net_params = {id(p) for p in net.parameters()}
    assert field_params == net_params  # mask + lift carry no trainable state

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_parameters_are_only_the_network" -q

Method implementation

    def in_dim(self) -> int:
        return len(self.weight[0]) if self.weight else 0

View source ↗

valuemethod

value(self, x: Tensor) -> Tensor

Lift value g(x), shape (..., out_dim).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · dirichlet interval endpoints exact

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_dirichlet_interval_endpoints_exact() -> None:
    net = JetMLP(1, 16, 1, depth=3, base="tanh").double()
    field = dirichlet_interval(net, 0.0, 2.0, lower_value=1.0, upper_value=-3.0)
    ends = torch.tensor([[0.0], [2.0]], dtype=torch.float64)
    u = field.value(ends)
    assert torch.allclose(u[:, 0], torch.tensor([1.0, -3.0]).double(), atol=1e-12)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_dirichlet_interval_endpoints_exact" -q

Method implementation

    def value(self, x: Tensor) -> Tensor:
        """Lift value ``g(x)``, shape ``(..., out_dim)``."""
        w = torch.tensor(self.weight, dtype=x.dtype, device=x.device)
        c = torch.tensor(self.bias, dtype=x.dtype, device=x.device)
        return x @ w.t() + c

View source ↗

point_jetmethod

point_jet(self, x0: Tensor, dim: int, order: int) -> Tensor

Exact jet of g at x0, shape (M, out_dim).

Parameters

Name	Type & default	Contract
x0positional-or-keyword	TensorRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · partials third order match autograd

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_partials_third_order_match_autograd() -> None:
    net = JetMLP(1, 14, 1, depth=2, base="tanh").double()
    field = dirichlet_interval(net, 0.0, 1.0, lower_value=1.0, upper_value=0.0)
    x = torch.tensor([[0.37]], dtype=torch.float64)
    u = _scalar_value(field)
    d3_ad = jacrev(jacrev(jacrev(u)))(x[0])  # (1,1,1)
    parts = field.partials(x, 3)
    assert torch.allclose(parts[(3,)][:, 0], d3_ad.reshape(()), atol=1e-9)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_partials_third_order_match_autograd" -q

Method implementation

    def point_jet(self, x0: Tensor, dim: int, order: int) -> Tensor:
        """Exact jet of ``g`` at ``x0``, shape ``(M, out_dim)``."""
        idj = identity_jet(x0, order)  # (M, D)
        w = torch.tensor(self.weight, dtype=x0.dtype, device=x0.device)  # (out, D)
        c = torch.tensor(self.bias, dtype=x0.dtype, device=x0.device)  # (out,)
        g = idj @ w.t()  # (M, out): row 0 = W x0, unit rows = columns of W
        m = g.shape[0]
        e0 = (torch.arange(m, device=g.device) == 0).to(g.dtype)
        return g + e0[:, None] * c

View source ↗

Read the implementation

class AffineLift:
    """Affine lift ``g(x) = W x + c`` matching the boundary data exactly.

    ``weight`` is the ``(out_dim, D)`` matrix ``W`` and ``bias`` the ``(out_dim,)``
    vector ``c``, both stored as plain Python floats so the lift is fully
    device/dtype-agnostic and carries no trainable state (the boundary data is
    fixed). A constant lift is ``W = 0``; a 1-D Dirichlet interpolation is a
    single non-zero slope.
    """

    weight: tuple[tuple[float, ...], ...]
    bias: tuple[float, ...]

    def __post_init__(self) -> None:
        if len(self.weight) != len(self.bias):
            raise ValueError(
                f"weight has {len(self.weight)} rows but bias has {len(self.bias)}"
            )
        if self.weight:
            width = len(self.weight[0])
            if any(len(row) != width for row in self.weight):
                raise ValueError("weight rows must all have the same length")

    @property
    def out_dim(self) -> int:
        return len(self.bias)

    @property
    def in_dim(self) -> int:
        return len(self.weight[0]) if self.weight else 0

    def value(self, x: Tensor) -> Tensor:
        """Lift value ``g(x)``, shape ``(..., out_dim)``."""
        w = torch.tensor(self.weight, dtype=x.dtype, device=x.device)
        c = torch.tensor(self.bias, dtype=x.dtype, device=x.device)
        return x @ w.t() + c

    def point_jet(self, x0: Tensor, dim: int, order: int) -> Tensor:
        """Exact jet of ``g`` at ``x0``, shape ``(M, out_dim)``."""
        idj = identity_jet(x0, order)  # (M, D)
        w = torch.tensor(self.weight, dtype=x0.dtype, device=x0.device)  # (out, D)
        c = torch.tensor(self.bias, dtype=x0.dtype, device=x0.device)  # (out,)
        g = idj @ w.t()  # (M, out): row 0 = W x0, unit rows = columns of W
        m = g.shape[0]
        e0 = (torch.arange(m, device=g.device) == 0).to(g.dtype)
        return g + e0[:, None] * c

Related APIsomnibias.torch.identity_jet↗

Other public import paths

• omnibias.torch.architectures.AffineLift
• omnibias.torch.architectures.hardbc.AffineLift

CLASS

BoundaryMask

omnibias.torch.architectures.BoundaryMaskSource ↗

class BoundaryMask

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Product of affine factors b(x) = \prod_k (s_k x[a_k] + o_k).

b vanishes exactly on the union of the factor hyperplanes -- the set on which the hard constraint is imposed. Both the batched value and the exact multivariate jet are provided so the mask plugs into the jet-level product.

Parameters

Name	Type & default	Contract
factorsdataclass field	tuple[AffineFactor, ...]Required	See the operation contract and implementation below.

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Regression example · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    bad = BoundaryMask((AffineFactor(3, 1.0, 0.0),))
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net, bad)

python -m pytest packages/omnibias-torch/tests/test_hardbc.py -k test_mask_factor_axis_out_of_range -q

Fields & methods

factorsfield

factors: tuple[AffineFactor, ...]

Declared tuple[AffineFactor, ...] field on BoundaryMask. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

__post_init__method

__post_init__(self) -> None

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

Returns: None

Validation & failure cases

• When not self.factors, raises ValueError('BoundaryMask needs at least one factor').

Executed regression · empty mask and zero scale rejected

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_empty_mask_and_zero_scale_rejected() -> None:
    with pytest.raises(ValueError, match="at least one factor"):
        BoundaryMask(())
    with pytest.raises(ValueError, match="non-zero"):
        AffineFactor(0, 0.0, 1.0)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_empty_mask_and_zero_scale_rejected" -q

Method implementation

    def __post_init__(self) -> None:
        if not self.factors:
            raise ValueError("BoundaryMask needs at least one factor")

View source ↗

valuemethod

value(self, x: Tensor) -> Tensor

Mask value b(x), shape (...,) for input x of shape (..., D).

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · dirichlet homogeneous has no lift

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_dirichlet_homogeneous_has_no_lift() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    field = dirichlet_interval(net, -1.0, 1.0)
    assert field.lift is None
    ends = torch.tensor([[-1.0], [1.0]], dtype=torch.float64)
    assert torch.allclose(field.value(ends), torch.zeros(2, 1, dtype=torch.float64), atol=1e-13)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_dirichlet_homogeneous_has_no_lift" -q

Method implementation

    def value(self, x: Tensor) -> Tensor:
        """Mask value ``b(x)``, shape ``(...,)`` for input ``x`` of shape ``(..., D)``."""
        f0 = self.factors[0]
        out = f0.scale * x[..., f0.axis] + f0.offset
        for f in self.factors[1:]:
            out = out * (f.scale * x[..., f.axis] + f.offset)
        return out

View source ↗

point_jetmethod

point_jet(self, x0: Tensor, dim: int, order: int) -> Tensor

Exact jet of b at x0, shape (M,) (M multi-indices).

Parameters

Name	Type & default	Contract
x0positional-or-keyword	TensorRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · value grad hessian consistent

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_value_grad_hessian_consistent() -> None:
    net = JetMLP(2, 10, 1, depth=2, base="tanh").double()
    field = homogeneous_box(net, [0.0, 0.0], [1.0, 1.0])
    x = torch.rand(5, 2, dtype=torch.float64)
    v, g, h = field.value_grad_hessian(x)
    assert torch.allclose(v, field.value(x), atol=1e-12)
    assert torch.allclose(g, field.gradient(x), atol=1e-12)
    assert torch.allclose(h, field.hessian(x), atol=1e-12)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_value_grad_hessian_consistent" -q

Method implementation

    def point_jet(self, x0: Tensor, dim: int, order: int) -> Tensor:
        """Exact jet of ``b`` at ``x0``, shape ``(M,)`` (``M`` multi-indices)."""
        idj = identity_jet(x0, order)  # (M, D)
        m = idj.shape[0]
        e0 = (torch.arange(m, device=idj.device) == 0).to(idj.dtype)
        f0 = self.factors[0]
        acc = f0.scale * idj[:, f0.axis] + f0.offset * e0
        for f in self.factors[1:]:
            fac = f.scale * idj[:, f.axis] + f.offset * e0
            acc = jet_multiply(acc, fac, dim, order)
        return acc

View source ↗

Read the implementation

class BoundaryMask:
    r"""Product of affine factors ``b(x) = \prod_k (s_k x[a_k] + o_k)``.

    ``b`` vanishes exactly on the union of the factor hyperplanes -- the set on
    which the hard constraint is imposed. Both the batched value and the exact
    multivariate jet are provided so the mask plugs into the jet-level product.
    """

    factors: tuple[AffineFactor, ...]

    def __post_init__(self) -> None:
        if not self.factors:
            raise ValueError("BoundaryMask needs at least one factor")

    def value(self, x: Tensor) -> Tensor:
        """Mask value ``b(x)``, shape ``(...,)`` for input ``x`` of shape ``(..., D)``."""
        f0 = self.factors[0]
        out = f0.scale * x[..., f0.axis] + f0.offset
        for f in self.factors[1:]:
            out = out * (f.scale * x[..., f.axis] + f.offset)
        return out

    def point_jet(self, x0: Tensor, dim: int, order: int) -> Tensor:
        """Exact jet of ``b`` at ``x0``, shape ``(M,)`` (``M`` multi-indices)."""
        idj = identity_jet(x0, order)  # (M, D)
        m = idj.shape[0]
        e0 = (torch.arange(m, device=idj.device) == 0).to(idj.dtype)
        f0 = self.factors[0]
        acc = f0.scale * idj[:, f0.axis] + f0.offset * e0
        for f in self.factors[1:]:
            fac = f.scale * idj[:, f.axis] + f.offset * e0
            acc = jet_multiply(acc, fac, dim, order)
        return acc

Related APIsomnibias.torch.identity_jet↗omnibias.torch.jet_multiply↗

Other public import paths

• omnibias.torch.architectures.BoundaryMask
• omnibias.torch.architectures.hardbc.BoundaryMask

CLASS

HardConstraintField

omnibias.torch.architectures.HardConstraintFieldSource ↗

class HardConstraintField(_JetMLPCore)

Network wrapped so a boundary/initial condition holds exactly: u = g + b N.

net is any omnibias.torch.architectures.pinn._JetMLPCore (a ~omnibias.torch.architectures.JetMLP, ~omnibias.torch.architectures.FourierFeatureMLP, or SIREN). The constrained field exposes the same exact readouts as the base network -- value, gradient, hessian, value_grad_hessian, jet, partials -- but every one is taken on u = g + b N rather than on net alone, and stays closed form because the wrapping is a jet-level add/multiply.

Use the dirichlet_interval, homogeneous_box and initial_value helpers for the common cases, or pass a custom mask / lift for full control.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	_JetMLPCoreRequired	See the operation contract and implementation below.
maskpositional-or-keyword	BoundaryMaskRequired	See the operation contract and implementation below.
liftpositional-or-keyword	AffineLift | NoneDefault: None	See the operation contract and implementation below.

Executed regression · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    bad = BoundaryMask((AffineFactor(3, 1.0, 0.0),))
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net, bad)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_mask_factor_axis_out_of_range" -q

Regression example · lift dim mismatch rejected

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_lift_dim_mismatch_rejected() -> None:
    net = JetMLP(2, 8, 1, depth=2).double()
    mask = BoundaryMask((AffineFactor(0, 1.0, 0.0),))
    bad_lift = AffineLift(((0.0,), (0.0,)), (0.0, 0.0))  # in_dim 1, out_dim 2
    with pytest.raises(ValueError, match="out_dim|in_dim"):
        HardConstraintField(net, mask, bad_lift)

python -m pytest packages/omnibias-torch/tests/test_hardbc.py -k test_lift_dim_mismatch_rejected -q

Fields & methods

netfield

net: _JetMLPCore

Declared _JetMLPCore field on HardConstraintField. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

maskfield

mask: BoundaryMask

Declared BoundaryMask field on HardConstraintField. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

liftfield

lift: AffineLift | None

Declared AffineLift | None field on HardConstraintField. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

__init__method

__init__(self, net: _JetMLPCore, mask: BoundaryMask, lift: AffineLift | None=None) -> None

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

Parameters

Name	Type & default	Contract
netpositional-or-keyword	_JetMLPCoreRequired	See the operation contract and implementation below.
maskpositional-or-keyword	BoundaryMaskRequired	See the operation contract and implementation below.
liftpositional-or-keyword	AffineLift | NoneDefault: None	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When f.axis >= self.in_dim, raises ValueError(f'mask factor axis {f.axis} out of range for in_dim={self.in_dim}').
• When lift.out_dim != self.out_dim, raises ValueError(f'lift out_dim {lift.out_dim} != network out_dim {self.out_dim}').
• When lift.in_dim != self.in_dim, raises ValueError(f'lift in_dim {lift.in_dim} != network in_dim {self.in_dim}').

Executed regression · mask factor axis out of range

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_mask_factor_axis_out_of_range() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    bad = BoundaryMask((AffineFactor(3, 1.0, 0.0),))
    with pytest.raises(ValueError, match="out of range"):
        HardConstraintField(net, bad)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_mask_factor_axis_out_of_range" -q

Method implementation

    def __init__(
        self,
        net: _JetMLPCore,
        mask: BoundaryMask,
        lift: AffineLift | None = None,
    ) -> None:
        super().__init__()
        self.net = net
        self.in_dim = int(net.in_dim)
        self.out_dim = int(net.out_dim)
        self.mask = mask
        self.lift = lift
        for f in mask.factors:
            if f.axis >= self.in_dim:
                raise ValueError(
                    f"mask factor axis {f.axis} out of range for in_dim={self.in_dim}"
                )
        if lift is not None:
            if lift.out_dim != self.out_dim:
                raise ValueError(
                    f"lift out_dim {lift.out_dim} != network out_dim {self.out_dim}"
                )
            if lift.in_dim != self.in_dim:
                raise ValueError(
                    f"lift in_dim {lift.in_dim} != network in_dim {self.in_dim}"
                )

View source ↗

valuemethod

value(self, x: Tensor) -> Tensor

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

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · dirichlet homogeneous has no lift

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_dirichlet_homogeneous_has_no_lift() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    field = dirichlet_interval(net, -1.0, 1.0)
    assert field.lift is None
    ends = torch.tensor([[-1.0], [1.0]], dtype=torch.float64)
    assert torch.allclose(field.value(ends), torch.zeros(2, 1, dtype=torch.float64), atol=1e-13)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_dirichlet_homogeneous_has_no_lift" -q

Method implementation

    def value(self, x: Tensor) -> Tensor:
        u = self.mask.value(x).unsqueeze(-1) * self.net.value(x)
        if self.lift is not None:
            u = u + self.lift.value(x)
        return u

View source ↗

in_dimattribute

in_dim = int(net.in_dim)

Public instance attribute assigned by __init__. The source assignment is in_dim = int(net.in_dim); it describes instance state, not a literal constructor default.

View source ↗

out_dimattribute

out_dim = int(net.out_dim)

Public instance attribute assigned by __init__. The source assignment is out_dim = int(net.out_dim); it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class HardConstraintField(_JetMLPCore):
    r"""Network wrapped so a boundary/initial condition holds exactly: ``u = g + b N``.

    ``net`` is any :class:`omnibias.torch.architectures.pinn._JetMLPCore` (a
    :class:`~omnibias.torch.architectures.JetMLP`,
    :class:`~omnibias.torch.architectures.FourierFeatureMLP`, or SIREN). The
    constrained field exposes the same exact readouts as the base network --
    :meth:`value`, :meth:`gradient`, :meth:`hessian`, :meth:`value_grad_hessian`,
    :meth:`jet`, :meth:`partials` -- but every one is taken on ``u = g + b N``
    rather than on ``net`` alone, and stays closed form because the wrapping is a
    jet-level add/multiply.

    Use the :func:`dirichlet_interval`, :func:`homogeneous_box` and
    :func:`initial_value` helpers for the common cases, or pass a custom
    ``mask`` / ``lift`` for full control.
    """

    net: _JetMLPCore
    mask: BoundaryMask
    lift: AffineLift | None

    def __init__(
        self,
        net: _JetMLPCore,
        mask: BoundaryMask,
        lift: AffineLift | None = None,
    ) -> None:
        super().__init__()
        self.net = net
        self.in_dim = int(net.in_dim)
        self.out_dim = int(net.out_dim)
        self.mask = mask
        self.lift = lift
        for f in mask.factors:
            if f.axis >= self.in_dim:
                raise ValueError(
                    f"mask factor axis {f.axis} out of range for in_dim={self.in_dim}"
                )
        if lift is not None:
            if lift.out_dim != self.out_dim:
                raise ValueError(
                    f"lift out_dim {lift.out_dim} != network out_dim {self.out_dim}"
                )
            if lift.in_dim != self.in_dim:
                raise ValueError(
                    f"lift in_dim {lift.in_dim} != network in_dim {self.in_dim}"
                )

    def _layer_specs(
        self,
    ) -> list[tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | None]]:
        """Delegate the fast-path check to the wrapped network's layers."""
        return self.net._layer_specs()

    def _point_jet(self, xi: Tensor, order: int) -> Tensor:
        net_jet = self.net._point_jet(xi, order)  # (M, out_dim)
        mask_jet = self.mask.point_jet(xi, self.in_dim, order)  # (M,)
        u_jet = jet_multiply(mask_jet, net_jet, self.in_dim, order)
        if self.lift is not None:
            u_jet = u_jet + self.lift.point_jet(xi, self.in_dim, order)
        return u_jet

    def value(self, x: Tensor) -> Tensor:
        u = self.mask.value(x).unsqueeze(-1) * self.net.value(x)
        if self.lift is not None:
            u = u + self.lift.value(x)
        return u

Related APIsomnibias.torch.jet_multiply↗

Other public import paths

• omnibias.torch.architectures.HardConstraintField
• omnibias.torch.architectures.hardbc.HardConstraintField

FUNCTION

dirichlet_interval

omnibias.torch.architectures.dirichlet_intervalSource ↗

dirichlet_interval(net: _JetMLPCore, lower: float, upper: float, *, lower_value: float=0.0, upper_value: float=0.0, axis: int=0) -> HardConstraintField

Hard 1-D Dirichlet BC: u(lower) = lower_value, u(upper) = upper_value.

Builds the mask b(x) = (x_axis - lower)(upper - x_axis) (zero at both ends) and the linear lift that interpolates the two endpoint values, so the field equals the prescribed data at lower and upper exactly. The same boundary profile is applied to every output component.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	_JetMLPCoreRequired	See the operation contract and implementation below.
lowerpositional-or-keyword	floatRequired	See the operation contract and implementation below.
upperpositional-or-keyword	floatRequired	See the operation contract and implementation below.
lower_valuekeyword-only	floatDefault: 0.0	See the operation contract and implementation below.
upper_valuekeyword-only	floatDefault: 0.0	See the operation contract and implementation below.
axiskeyword-only	intDefault: 0	See the operation contract and implementation below.

ReturnsHardConstraintField

Explicit validation & failure cases

• When upper <= lower, raises ValueError(f'upper ({upper}) must exceed lower ({lower})').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · dirichlet interval rejects bad bounds

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_dirichlet_interval_rejects_bad_bounds() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    with pytest.raises(ValueError, match="must exceed"):
        dirichlet_interval(net, 1.0, 0.0)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_dirichlet_interval_rejects_bad_bounds" -q

Regression example · dirichlet homogeneous has no lift

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_dirichlet_homogeneous_has_no_lift() -> None:
    net = JetMLP(1, 8, 1, depth=2).double()
    field = dirichlet_interval(net, -1.0, 1.0)
    assert field.lift is None
    ends = torch.tensor([[-1.0], [1.0]], dtype=torch.float64)
    assert torch.allclose(field.value(ends), torch.zeros(2, 1, dtype=torch.float64), atol=1e-13)

python -m pytest packages/omnibias-torch/tests/test_hardbc.py -k test_dirichlet_homogeneous_has_no_lift -q

Read the implementation

def dirichlet_interval(
    net: _JetMLPCore,
    lower: float,
    upper: float,
    *,
    lower_value: float = 0.0,
    upper_value: float = 0.0,
    axis: int = 0,
) -> HardConstraintField:
    r"""Hard 1-D Dirichlet BC: ``u(lower) = lower_value``, ``u(upper) = upper_value``.

    Builds the mask ``b(x) = (x_axis - lower)(upper - x_axis)`` (zero at both
    ends) and the linear lift that interpolates the two endpoint values, so the
    field equals the prescribed data at ``lower`` and ``upper`` exactly. The same
    boundary profile is applied to every output component.
    """
    if upper <= lower:
        raise ValueError(f"upper ({upper}) must exceed lower ({lower})")
    mask = BoundaryMask(
        (AffineFactor(axis, 1.0, -lower), AffineFactor(axis, -1.0, upper))
    )
    lift: AffineLift | None = None
    if lower_value != 0.0 or upper_value != 0.0:
        slope = (upper_value - lower_value) / (upper - lower)
        intercept = lower_value - slope * lower
        row = tuple(slope if j == axis else 0.0 for j in range(net.in_dim))
        lift = AffineLift((row,) * net.out_dim, (intercept,) * net.out_dim)
    return HardConstraintField(net, mask, lift)

Related APIsomnibias.torch.architectures.AffineFactor↗omnibias.torch.architectures.AffineLift↗omnibias.torch.architectures.BoundaryMask↗omnibias.torch.architectures.HardConstraintField↗

Other public import paths

• omnibias.torch.architectures.dirichlet_interval
• omnibias.torch.architectures.hardbc.dirichlet_interval

FUNCTION

homogeneous_box

omnibias.torch.architectures.homogeneous_boxSource ↗

homogeneous_box(net: _JetMLPCore, lows: Sequence[float], highs: Sequence[float]) -> HardConstraintField

Hard homogeneous Dirichlet BC u = 0 on the boundary of a box.

The box is \prod_i [lows_i, highs_i] and the mask is the product \prod_i (x_i - lows_i)(highs_i - x_i), which vanishes on every face. The lift is zero, so u = b N.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	_JetMLPCoreRequired	See the operation contract and implementation below.
lowspositional-or-keyword	Sequence[float]Required	See the operation contract and implementation below.
highspositional-or-keyword	Sequence[float]Required	See the operation contract and implementation below.

ReturnsHardConstraintField

Explicit validation & failure cases

• When len(lows) != net.in_dim or len(highs) != net.in_dim, raises ValueError(f'lows/highs must have length in_dim={net.in_dim}, got {len(lows)} and {len(highs)}').
• When highs[i] <= lows[i], raises ValueError(f'highs[{i}] ({highs[i]}) must exceed lows[{i}] ({lows[i]})').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · homogeneous box rejects wrong length

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_homogeneous_box_rejects_wrong_length() -> None:
    net = JetMLP(2, 8, 1, depth=2).double()
    with pytest.raises(ValueError, match="length in_dim"):
        homogeneous_box(net, [0.0], [1.0])

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_homogeneous_box_rejects_wrong_length" -q

Regression example · multi output box parity

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_multi_output_box_parity() -> None:
    tnet, jnet = _jetmlp_pair(2, 10, 3, 2, "tanh", seed=4)
    tfield = torch_hardbc.homogeneous_box(tnet, [0.0, 0.0], [1.0, 1.0])
    jfield = jax_hardbc.homogeneous_box(jnet, [0.0, 0.0], [1.0, 1.0])
    xt, xj = _xpair(seed=4, b=4, d=2)
    _assert_field_parity(tfield, jfield, xt, xj)

python -m pytest tests/test_hardbc_parity.py -k test_multi_output_box_parity -q

Read the implementation

def homogeneous_box(
    net: _JetMLPCore,
    lows: Sequence[float],
    highs: Sequence[float],
) -> HardConstraintField:
    r"""Hard homogeneous Dirichlet BC ``u = 0`` on the boundary of a box.

    The box is ``\prod_i [lows_i, highs_i]`` and the mask is the product
    ``\prod_i (x_i - lows_i)(highs_i - x_i)``, which vanishes on every face. The
    lift is zero, so ``u = b N``.
    """
    if len(lows) != net.in_dim or len(highs) != net.in_dim:
        raise ValueError(
            f"lows/highs must have length in_dim={net.in_dim}, "
            f"got {len(lows)} and {len(highs)}"
        )
    factors: list[AffineFactor] = []
    for i in range(net.in_dim):
        if highs[i] <= lows[i]:
            raise ValueError(
                f"highs[{i}] ({highs[i]}) must exceed lows[{i}] ({lows[i]})"
            )
        factors.append(AffineFactor(i, 1.0, -float(lows[i])))
        factors.append(AffineFactor(i, -1.0, float(highs[i])))
    return HardConstraintField(net, BoundaryMask(tuple(factors)), None)

Related APIsomnibias.torch.architectures.AffineFactor↗omnibias.torch.architectures.BoundaryMask↗omnibias.torch.architectures.HardConstraintField↗

Other public import paths

• omnibias.torch.architectures.homogeneous_box
• omnibias.torch.architectures.hardbc.homogeneous_box

FUNCTION

initial_value

omnibias.torch.architectures.initial_valueSource ↗

initial_value(net: _JetMLPCore, *, t_axis: int=-1, t0: float=0.0, value: float=0.0) -> HardConstraintField

Hard (constant) initial condition u(x, t0) = value.

The mask is the one-sided factor t - t0 (zero only on the initial slice, leaving later times unconstrained) and the lift is the constant value. For a non-constant initial profile, compose a custom AffineLift (or a product BoundaryMask) and build HardConstraintField directly.

Parameters

Name	Type & default	Contract
netpositional-or-keyword	_JetMLPCoreRequired	See the operation contract and implementation below.
t_axiskeyword-only	intDefault: -1	See the operation contract and implementation below.
t0keyword-only	floatDefault: 0.0	See the operation contract and implementation below.
valuekeyword-only	floatDefault: 0.0	See the operation contract and implementation below.

ReturnsHardConstraintField

Explicit validation & failure cases

• When not 0 <= ax < net.in_dim, raises ValueError(f't_axis {t_axis} out of range for in_dim={net.in_dim}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · initial value negative axis

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_initial_value_negative_axis() -> None:
    net = JetMLP(2, 8, 1, depth=2).double()
    field = initial_value(net, t_axis=-1, t0=0.5, value=0.0)
    assert field.lift is None
    ic = torch.tensor([[0.3, 0.5], [0.8, 0.5]], dtype=torch.float64)
    assert torch.allclose(field.value(ic), torch.zeros(2, 1, dtype=torch.float64), atol=1e-13)

python -m pytest "packages/omnibias-torch/tests/test_hardbc.py::test_initial_value_negative_axis" -q

Regression example · initial value parity

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_initial_value_parity() -> None:
    tnet, jnet = _jetmlp_pair(2, 10, 1, 2, "sigmoid", seed=3)
    tfield = torch_hardbc.initial_value(tnet, t_axis=1, t0=0.0, value=2.0)
    jfield = jax_hardbc.initial_value(jnet, t_axis=1, t0=0.0, value=2.0)
    xt, xj = _xpair(seed=3, b=5, d=2)
    _assert_field_parity(tfield, jfield, xt, xj)

python -m pytest tests/test_hardbc_parity.py -k test_initial_value_parity -q

Read the implementation

def initial_value(
    net: _JetMLPCore,
    *,
    t_axis: int = -1,
    t0: float = 0.0,
    value: float = 0.0,
) -> HardConstraintField:
    r"""Hard (constant) initial condition ``u(x, t0) = value``.

    The mask is the one-sided factor ``t - t0`` (zero only on the initial slice,
    leaving later times unconstrained) and the lift is the constant ``value``. For
    a non-constant initial profile, compose a custom :class:`AffineLift` (or a
    product :class:`BoundaryMask`) and build :class:`HardConstraintField` directly.
    """
    ax = t_axis if t_axis >= 0 else net.in_dim + t_axis
    if not (0 <= ax < net.in_dim):
        raise ValueError(f"t_axis {t_axis} out of range for in_dim={net.in_dim}")
    mask = BoundaryMask((AffineFactor(ax, 1.0, -t0),))
    lift: AffineLift | None = None
    if value != 0.0:
        zeros = (0.0,) * net.in_dim
        lift = AffineLift((zeros,) * net.out_dim, (float(value),) * net.out_dim)
    return HardConstraintField(net, mask, lift)

Related APIsomnibias.torch.architectures.AffineFactor↗omnibias.torch.architectures.AffineLift↗omnibias.torch.architectures.BoundaryMask↗omnibias.torch.architectures.HardConstraintField↗

Other public import paths

• omnibias.torch.architectures.initial_value
• omnibias.torch.architectures.hardbc.initial_value

← Package overviewRunnable learning examples →