omnibias.fields.jax.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.jax.ops.conservation

Module contract & mathematical conventions

JAX conservation / flux / wave ops (mirrors omnibias.fields.torch.ops.conservation).

These are closed-form compositions of the existing partial-derivative reducers (no autodiff on the closed-form path): every term below is a single analytic combination of value / gradient / laplacian / mixed_partial and therefore stays algebraically corresponding with the torch twin by construction.

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.jax.ops.advection_diffusion_residualSource ↗

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

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

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`scalar`keyword-only	`str`Required	See the operation contract and implementation below.
`velocity`keyword-only	`tuple[str, ...]`Required	See the operation contract and implementation below.
`diffusivity`keyword-only	`float \| str`Default: 1.0	See the operation contract and implementation below.
`source`keyword-only	`str \| float \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

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,
) -> Array:
    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

Other public import paths

omnibias.fields.jax.ops.advection_diffusion_residual
omnibias.fields.jax.ops.conservation.advection_diffusion_residual

FUNCTION

conservation_residual

omnibias.fields.jax.ops.conservation_residualSource ↗

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

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

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`density`keyword-only	`str`Required	See the operation contract and implementation below.
`flux`keyword-only	`tuple[str, ...]`Required	See the operation contract and implementation below.
`source`keyword-only	`str \| float \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

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

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,
) -> Array:
    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

Other public import paths

omnibias.fields.jax.ops.conservation_residual
omnibias.fields.jax.ops.conservation.conservation_residual

FUNCTION

dalembertian

omnibias.fields.jax.ops.dalembertianSource ↗

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

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
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`c`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.
`signature`keyword-only	`str`Default: 'mostly_plus'	See the operation contract and implementation below.

ReturnsArray

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

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",
) -> Array:
    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}"
    )

Other public import paths

omnibias.fields.jax.ops.dalembertian
omnibias.fields.jax.ops.conservation.dalembertian

FUNCTION

diffusive_flux

omnibias.fields.jax.ops.diffusive_fluxSource ↗

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

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
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`diffusivity`keyword-only	`float \| str`Default: 1.0	See the operation contract and implementation below.

ReturnsArray

Executed regression · cross backend bit parity

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,
) -> Array:
    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

Other public import paths

omnibias.fields.jax.ops.diffusive_flux
omnibias.fields.jax.ops.conservation.diffusive_flux

FUNCTION

flux_divergence

omnibias.fields.jax.ops.flux_divergenceSource ↗

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

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

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

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`names`positional-or-keyword	`tuple[str, ...]`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

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

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

Other public import paths

omnibias.fields.jax.ops.flux_divergence
omnibias.fields.jax.ops.conservation.flux_divergence

FUNCTION

grad_squared_norm

omnibias.fields.jax.ops.grad_squared_normSource ↗

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

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

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.

ReturnsArray

Executed regression · cross backend bit parity

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) -> Array:
    r""":math:`|\nabla u|^2 = \sum_i (\partial_i u)^2` (spatial gradient)."""
    g = gradient(state, name, axes=_spatial(state))
    return (g * g).sum(axis=-1)

Other public import paths

omnibias.fields.jax.ops.grad_squared_norm
omnibias.fields.jax.ops.conservation.grad_squared_norm

FUNCTION

gradient_of_composition

omnibias.fields.jax.ops.gradient_of_compositionSource ↗

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

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

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`fprime`positional-or-keyword	`Callable[[Array], Array]`Required	See the operation contract and implementation below.

ReturnsArray

Executed regression · cross backend bit parity

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[[Array], Array],
) -> Array:
    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

Other public import paths

omnibias.fields.jax.ops.gradient_of_composition
omnibias.fields.jax.ops.conservation.gradient_of_composition

FUNCTION

laplacian_of_composition

omnibias.fields.jax.ops.laplacian_of_compositionSource ↗

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

\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
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`fprime`positional-or-keyword	`Callable[[Array], Array]`Required	See the operation contract and implementation below.
`fsecond`positional-or-keyword	`Callable[[Array], Array]`Required	See the operation contract and implementation below.

ReturnsArray

Executed regression · cross backend bit parity

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[[Array], Array],
    fsecond: Callable[[Array], Array],
) -> Array:
    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

Other public import paths

omnibias.fields.jax.ops.laplacian_of_composition
omnibias.fields.jax.ops.conservation.laplacian_of_composition

FUNCTION

variable_coefficient_diffusion

omnibias.fields.jax.ops.variable_coefficient_diffusionSource ↗

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

\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
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`diffusivity`keyword-only	`float \| str`Default: 1.0	See the operation contract and implementation below.

ReturnsArray

Executed regression · cross backend bit parity

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,
) -> Array:
    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(axis=-1) + d_val * lap
    return float(diffusivity) * lap

Other public import paths

omnibias.fields.jax.ops.variable_coefficient_diffusion
omnibias.fields.jax.ops.conservation.variable_coefficient_diffusion

FUNCTION

wave_operator

omnibias.fields.jax.ops.wave_operatorSource ↗

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

Alias for dalembertian.

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`name`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`c`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.
`signature`keyword-only	`str`Default: 'mostly_plus'	See the operation contract and implementation below.

ReturnsArray

Read the implementation

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

Other public import paths

omnibias.fields.jax.ops.wave_operator
omnibias.fields.jax.ops.conservation.wave_operator

← Package overview Runnable learning examples →

omnibias.fields.jax.ops.conservation

advection_diffusion_residual​

Parameters

conservation_residual​

Parameters

dalembertian​

Parameters

diffusive_flux​

Parameters

flux_divergence​

Parameters

grad_squared_norm​

Parameters

gradient_of_composition​

Parameters

laplacian_of_composition​

Parameters

variable_coefficient_diffusion​

Parameters

wave_operator​

Parameters

advection_diffusion_residual

conservation_residual

dalembertian

diffusive_flux

flux_divergence

grad_squared_norm

gradient_of_composition

laplacian_of_composition

variable_coefficient_diffusion

wave_operator