API REFERENCE/omnibias-fields

omnibias.fields.torch.ops.conservation

0.2.0rc1Release candidatePython >=3.10Apache-2.0

Structured scalar, vector, and tensor fields with derivative caches, calculus operators, and coordinate conventions.

import omnibias.fields.torch.ops.conservation

Module contract & mathematical conventions

Torch conservation / flux / wave ops (twin of omnibias.fields.jax.ops.conservation).

Closed-form compositions of the existing partial-derivative reducers; algebraically corresponding with the jax twin by construction (the math lives in the shared sigma tower, not here).

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

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 →

FUNCTION

advection_diffusion_residual

omnibias.fields.torch.ops.advection_diffusion_residualSource ↗

advection_diffusion_residual(state: FieldState, *, scalar: str, velocity: tuple[str, ...], diffusivity: float | str=1.0, source: str | float | None=None) -> Tensor

\partial_t c + (u\cdot\nabla)c - \nabla\cdot(D\nabla c) - s.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
scalarkeyword-only	strRequired	See the operation contract and implementation below.
velocitykeyword-only	tuple[str, ...]Required	See the operation contract and implementation below.
diffusivitykeyword-only	float | strDefault: 1.0	See the operation contract and implementation below.
sourcekeyword-only	str | float | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When ta is None, raises ValueError('advection_diffusion_residual requires a time axis on the spec').

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

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[st-<lambda>2]" -q

Read the implementation

def advection_diffusion_residual(
    state: FieldState,
    *,
    scalar: str,
    velocity: tuple[str, ...],
    diffusivity: float | str = 1.0,
    source: str | float | None = None,
) -> Tensor:
    r""":math:`\partial_t c + (u\cdot\nabla)c - \nabla\cdot(D\nabla c) - s`."""
    ta = state.coordinate_spec.time_axis
    if ta is None:
        raise ValueError("advection_diffusion_residual requires a time axis on the spec")
    dt_c = derivative(state, scalar, axis=ta, order=1)
    adv = advection(state, velocity=velocity, scalar=scalar)
    diff = variable_coefficient_diffusion(state, scalar, diffusivity=diffusivity)
    res = dt_c + adv - diff
    if source is not None:
        s = value(state, source) if isinstance(source, str) else float(source)
        res = res - s
    return res

Related APIsomnibias.fields.torch.ops.derivative↗omnibias.fields.torch.ops.value↗omnibias.fields.torch.ops.variable_coefficient_diffusion↗omnibias.fields.torch.ops.advection↗

Other public import paths

• omnibias.fields.torch.ops.advection_diffusion_residual
• omnibias.fields.torch.ops.conservation.advection_diffusion_residual

FUNCTION

conservation_residual

omnibias.fields.torch.ops.conservation_residualSource ↗

conservation_residual(state: FieldState, *, density: str, flux: tuple[str, ...], source: str | float | None=None) -> Tensor

Local conservation law residual \partial_t \rho + \nabla\cdot F - s.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
densitykeyword-only	strRequired	See the operation contract and implementation below.
fluxkeyword-only	tuple[str, ...]Required	See the operation contract and implementation below.
sourcekeyword-only	str | float | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When ta is None, raises ValueError('conservation_residual requires a time axis on the spec').

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

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[st-<lambda>1]" -q

Read the implementation

def conservation_residual(
    state: FieldState,
    *,
    density: str,
    flux: tuple[str, ...],
    source: str | float | None = None,
) -> Tensor:
    r"""Local conservation law residual :math:`\partial_t \rho + \nabla\cdot F - s`."""
    ta = state.coordinate_spec.time_axis
    if ta is None:
        raise ValueError("conservation_residual requires a time axis on the spec")
    res = derivative(state, density, axis=ta, order=1) + divergence(state, flux)
    if source is not None:
        s = value(state, source) if isinstance(source, str) else float(source)
        res = res - s
    return res

Related APIsomnibias.fields.torch.ops.derivative↗omnibias.fields.torch.ops.divergence↗omnibias.fields.torch.ops.value↗

Other public import paths

• omnibias.fields.torch.ops.conservation_residual
• omnibias.fields.torch.ops.conservation.conservation_residual

FUNCTION

dalembertian

omnibias.fields.torch.ops.dalembertianSource ↗

dalembertian(state: FieldState, name: str, *, c: float=1.0, signature: str='mostly_plus') -> Tensor

d'Alembert / wave operator \Box u.

signature="mostly_plus" (the (-,+,+,+) metric) gives \Box u = \Delta u - c^{-2}\,\partial_{tt} u, so the wave equation is \Box u = 0 \iff \partial_{tt} u = c^2 \Delta u. "mostly_minus" flips the overall sign. Requires a time axis.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
ckeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
signaturekeyword-only	strDefault: 'mostly_plus'	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When ta is None, raises ValueError('dalembertian requires a time axis on the spec').
• When c == 0.0, raises ValueError('wave speed c must be non-zero').

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

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[st-<lambda>0]" -q

Read the implementation

def dalembertian(
    state: FieldState, name: str, *, c: float = 1.0, signature: str = "mostly_plus",
) -> Tensor:
    r"""d'Alembert / wave operator :math:`\Box u`.

    ``signature="mostly_plus"`` (the :math:`(-,+,+,+)` metric) gives
    :math:`\Box u = \Delta u - c^{-2}\,\partial_{tt} u`, so the wave equation is
    :math:`\Box u = 0 \iff \partial_{tt} u = c^2 \Delta u`. ``"mostly_minus"``
    flips the overall sign. Requires a time axis.
    """
    ta = state.coordinate_spec.time_axis
    if ta is None:
        raise ValueError("dalembertian requires a time axis on the spec")
    if c == 0.0:
        raise ValueError("wave speed c must be non-zero")
    # axes=None -> spatial Laplacian via the fast closed-form path when available.
    spatial_lap = laplacian(state, name)
    dtt = derivative(state, name, axis=ta, order=2)
    inv_c2 = 1.0 / (c * c)
    if signature == "mostly_plus":
        return spatial_lap - inv_c2 * dtt
    if signature == "mostly_minus":
        return inv_c2 * dtt - spatial_lap
    raise ValueError(
        f"signature must be 'mostly_plus' or 'mostly_minus', got {signature!r}"
    )

Related APIsomnibias.fields.torch.ops.derivative↗omnibias.fields.torch.ops.laplacian↗

Other public import paths

• omnibias.fields.torch.ops.dalembertian
• omnibias.fields.torch.ops.conservation.dalembertian

FUNCTION

diffusive_flux

omnibias.fields.torch.ops.diffusive_fluxSource ↗

diffusive_flux(state: FieldState, name: str, *, diffusivity: float | str=1.0) -> Tensor

Fickian flux F = -D\,\nabla u, shape (B, n_spatial).

diffusivity is a constant or the name of a scalar field component.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
diffusivitykeyword-only	float | strDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>4]" -q

Read the implementation

def diffusive_flux(
    state: FieldState, name: str, *, diffusivity: float | str = 1.0,
) -> Tensor:
    r"""Fickian flux :math:`F = -D\,\nabla u`, shape ``(B, n_spatial)``.

    ``diffusivity`` is a constant or the name of a scalar field component.
    """
    g = gradient(state, name, axes=_spatial(state))
    if isinstance(diffusivity, str):
        d = value(state, diffusivity)[..., None]
        return -d * g
    return -float(diffusivity) * g

Related APIsomnibias.fields.torch.ops.gradient↗omnibias.fields.torch.ops.value↗

Other public import paths

• omnibias.fields.torch.ops.diffusive_flux
• omnibias.fields.torch.ops.conservation.diffusive_flux

FUNCTION

flux_divergence

omnibias.fields.torch.ops.flux_divergenceSource ↗

flux_divergence(state: FieldState, names: tuple[str, ...]) -> Tensor

Divergence of a named flux field \nabla\cdot F.

Conservation-law spelling of omnibias.fields.torch.ops.basic.divergence (the flux components are model outputs / named fields).

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def flux_divergence(state: FieldState, names: tuple[str, ...]) -> Tensor:
    r"""Divergence of a named flux field :math:`\nabla\cdot F`.

    Conservation-law spelling of :func:`omnibias.fields.torch.ops.basic.divergence`
    (the flux components are model outputs / named fields).
    """
    return divergence(state, names)

Related APIsomnibias.fields.torch.ops.divergence↗

Other public import paths

• omnibias.fields.torch.ops.flux_divergence
• omnibias.fields.torch.ops.conservation.flux_divergence

FUNCTION

grad_squared_norm

omnibias.fields.torch.ops.grad_squared_normSource ↗

grad_squared_norm(state: FieldState, name: str) -> Tensor

|\nabla u|^2 = \sum_i (\partial_i u)^2 (spatial gradient).

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>3]" -q

Read the implementation

def grad_squared_norm(state: FieldState, name: str) -> Tensor:
    r""":math:`|\nabla u|^2 = \sum_i (\partial_i u)^2` (spatial gradient)."""
    g = gradient(state, name, axes=_spatial(state))
    return (g * g).sum(dim=-1)

Related APIsomnibias.fields.torch.ops.gradient↗

Other public import paths

• omnibias.fields.torch.ops.grad_squared_norm
• omnibias.fields.torch.ops.conservation.grad_squared_norm

FUNCTION

gradient_of_composition

omnibias.fields.torch.ops.gradient_of_compositionSource ↗

gradient_of_composition(state: FieldState, name: str, fprime: Callable[[Tensor], Tensor]) -> Tensor

\nabla f(u) = f'(u)\,\nabla u. fprime is applied to the value.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
fprimepositional-or-keyword	Callable[[Tensor], Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>7]" -q

Read the implementation

def gradient_of_composition(
    state: FieldState, name: str, fprime: Callable[[Tensor], Tensor],
) -> Tensor:
    r""":math:`\nabla f(u) = f'(u)\,\nabla u`. ``fprime`` is applied to the value."""
    u = value(state, name)
    g = gradient(state, name, axes=_spatial(state))
    return fprime(u)[..., None] * g

Related APIsomnibias.fields.torch.ops.gradient↗omnibias.fields.torch.ops.value↗

Other public import paths

• omnibias.fields.torch.ops.gradient_of_composition
• omnibias.fields.torch.ops.conservation.gradient_of_composition

FUNCTION

laplacian_of_composition

omnibias.fields.torch.ops.laplacian_of_compositionSource ↗

laplacian_of_composition(state: FieldState, name: str, fprime: Callable[[Tensor], Tensor], fsecond: Callable[[Tensor], Tensor]) -> Tensor

\Delta f(u) = f'(u)\,\Delta u + f''(u)\,|\nabla u|^2.

The chain rule for the Laplacian of a pointwise nonlinearity (Cahn-Hilliard / Allen-Cahn style). fprime / fsecond are callables applied to u.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
fprimepositional-or-keyword	Callable[[Tensor], Tensor]Required	See the operation contract and implementation below.
fsecondpositional-or-keyword	Callable[[Tensor], Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>8]" -q

Read the implementation

def laplacian_of_composition(
    state: FieldState,
    name: str,
    fprime: Callable[[Tensor], Tensor],
    fsecond: Callable[[Tensor], Tensor],
) -> Tensor:
    r""":math:`\Delta f(u) = f'(u)\,\Delta u + f''(u)\,|\nabla u|^2`.

    The chain rule for the Laplacian of a pointwise nonlinearity (Cahn-Hilliard /
    Allen-Cahn style). ``fprime`` / ``fsecond`` are callables applied to ``u``.
    """
    u = value(state, name)
    lap = laplacian(state, name, axes=_spatial(state))
    gsq = grad_squared_norm(state, name)
    return fprime(u) * lap + fsecond(u) * gsq

Related APIsomnibias.fields.torch.ops.laplacian↗omnibias.fields.torch.ops.value↗omnibias.fields.torch.ops.grad_squared_norm↗

Other public import paths

• omnibias.fields.torch.ops.laplacian_of_composition
• omnibias.fields.torch.ops.conservation.laplacian_of_composition

FUNCTION

variable_coefficient_diffusion

omnibias.fields.torch.ops.variable_coefficient_diffusionSource ↗

variable_coefficient_diffusion(state: FieldState, name: str, *, diffusivity: float | str=1.0) -> Tensor

\nabla\cdot(D\,\nabla u).

For a constant D this is D * laplacian(u); for D a scalar field component it expands to grad(D).grad(u) + D * laplacian(u).

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
diffusivitykeyword-only	float | strDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · cross backend bit parity

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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>5]" -q

Read the implementation

def variable_coefficient_diffusion(
    state: FieldState, name: str, *, diffusivity: float | str = 1.0,
) -> Tensor:
    r""":math:`\nabla\cdot(D\,\nabla u)`.

    For a constant ``D`` this is ``D * laplacian(u)``; for ``D`` a scalar field
    component it expands to ``grad(D).grad(u) + D * laplacian(u)``.
    """
    sa = _spatial(state)
    lap = laplacian(state, name, axes=sa)
    if isinstance(diffusivity, str):
        grad_d = gradient(state, diffusivity, axes=sa)
        grad_u = gradient(state, name, axes=sa)
        d_val = value(state, diffusivity)
        return (grad_d * grad_u).sum(dim=-1) + d_val * lap
    return float(diffusivity) * lap

Related APIsomnibias.fields.torch.ops.gradient↗omnibias.fields.torch.ops.laplacian↗omnibias.fields.torch.ops.value↗

Other public import paths

• omnibias.fields.torch.ops.variable_coefficient_diffusion
• omnibias.fields.torch.ops.conservation.variable_coefficient_diffusion

FUNCTION

wave_operator

omnibias.fields.torch.ops.conservation.wave_operatorSource ↗

wave_operator(state: FieldState, name: str, *, c: float=1.0, signature: str='mostly_plus') -> Tensor

Alias for dalembertian.

Parameters

Name	Type & default	Contract
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
ckeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
signaturekeyword-only	strDefault: 'mostly_plus'	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def wave_operator(
    state: FieldState, name: str, *, c: float = 1.0, signature: str = "mostly_plus",
) -> Tensor:
    """Alias for :func:`dalembertian`."""
    return dalembertian(state, name, c=c, signature=signature)

Related APIsomnibias.fields.torch.ops.dalembertian↗

← Package overviewRunnable learning examples →