API REFERENCE/omnibias-fields

omnibias.fields.jax.ops.nonlinear

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

Module contract & mathematical conventions

JAX nonlinear ops (mirrors omnibias.fields.torch.ops.nonlinear).

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

omnibias.fields.jax.ops.advectionSource ↗

advection(state: FieldState, *, velocity: tuple[str, ...], target: tuple[str, ...] | None=None, scalar: str | None=None) -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

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

ReturnsArray

Explicit validation & failure cases

• When len(velocity) != len(sa), raises ValueError(f'advection: velocity has {len(velocity)} components, but coordinate spec has {len(sa)} spatial axes ({sa!r})').
• When target is not None, raises ValueError('advection: pass scalar OR target, not both').

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[3d-<lambda>6]" -q

Read the implementation

def advection(
    state: FieldState,
    *,
    velocity: tuple[str, ...],
    target: tuple[str, ...] | None = None,
    scalar: str | None = None,
) -> Array:
    sa = state.coordinate_spec.spatial_axes
    if len(velocity) != len(sa):
        raise ValueError(
            f"advection: velocity has {len(velocity)} components, "
            f"but coordinate spec has {len(sa)} spatial axes ({sa!r})"
        )
    u_vec = stack_components(state, velocity)
    if scalar is not None:
        if target is not None:
            raise ValueError("advection: pass scalar OR target, not both")
        grad_phi = gradient(state, scalar)
        return (u_vec * grad_phi).sum(axis=-1)
    target = target if target is not None else velocity
    cols = []
    for n in target:
        grad_n = gradient(state, n)
        cols.append((u_vec * grad_n).sum(axis=-1))
    return jnp.stack(cols, axis=-1)

Related APIsomnibias.fields.jax.ops.gradient↗omnibias.fields.jax.ops.stack_components↗

Other public import paths

• omnibias.fields.jax.ops.advection
• omnibias.fields.jax.ops.nonlinear.advection

FUNCTION

directional_derivative

omnibias.fields.jax.ops.directional_derivativeSource ↗

directional_derivative(state: FieldState, name: str, *, direction: Array) -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

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.
directionkeyword-only	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When direction.shape != g.shape, raises ValueError(f'direction shape {tuple(direction.shape)} != gradient shape {tuple(g.shape)}').

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

Read the implementation

def directional_derivative(
    state: FieldState, name: str, *, direction: Array,
) -> Array:
    g = gradient(state, name)
    direction = jnp.asarray(direction)
    if direction.shape != g.shape:
        raise ValueError(
            f"direction shape {tuple(direction.shape)} != "
            f"gradient shape {tuple(g.shape)}"
        )
    return (g * direction).sum(axis=-1)

Related APIsomnibias.fields.jax.ops.gradient↗

Other public import paths

• omnibias.fields.jax.ops.directional_derivative
• omnibias.fields.jax.ops.nonlinear.directional_derivative

FUNCTION

material_derivative

omnibias.fields.jax.ops.material_derivativeSource ↗

material_derivative(state: FieldState, *, velocity: tuple[str, ...], scalar: str | None=None) -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

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

ReturnsArray

Explicit validation & failure cases

• When state.coordinate_spec.time_axis is None, raises ValueError('material_derivative requires a time axis on the coordinate spec').

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

Read the implementation

def material_derivative(
    state: FieldState,
    *,
    velocity: tuple[str, ...],
    scalar: str | None = None,
) -> Array:
    if state.coordinate_spec.time_axis is None:
        raise ValueError(
            "material_derivative requires a time axis on the coordinate spec"
        )
    time_axis = state.coordinate_spec.time_axis
    assert time_axis is not None  # guarded above
    if scalar is not None:
        dt_target = derivative(state, scalar, axis=time_axis)
        return dt_target + advection(state, velocity=velocity, scalar=scalar)
    dt_target = jnp.stack(
        [derivative(state, n, axis=time_axis) for n in velocity], axis=-1,
    )
    return dt_target + advection(state, velocity=velocity)

Related APIsomnibias.fields.jax.ops.derivative↗omnibias.fields.jax.ops.advection↗

Other public import paths

• omnibias.fields.jax.ops.material_derivative
• omnibias.fields.jax.ops.nonlinear.material_derivative

FUNCTION

p_laplacian

omnibias.fields.jax.ops.p_laplacianSource ↗

p_laplacian(state: FieldState, name: str, *, p: float, eps: float=1e-08) -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

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.
pkeyword-only	floatRequired	See the operation contract and implementation below.
epskeyword-only	floatDefault: 1e-08	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When p < 1, raises ValueError(f'p must be >= 1 for p-Laplacian, got {p}').

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

Read the implementation

def p_laplacian(
    state: FieldState, name: str, *, p: float, eps: float = 1e-8,
) -> Array:
    if p < 1:
        raise ValueError(f"p must be >= 1 for p-Laplacian, got {p}")
    sa = state.coordinate_spec.spatial_axes
    g = gradient(state, name)
    L = laplacian(state, name)
    if abs(p - 2.0) < 1e-15:
        return L
    g2 = (g * g).sum(axis=-1)
    g_norm_sq_reg = g2 + eps * eps
    H_full = hessian(state, name)
    spatial_idx = jnp.array(
        [state.coordinate_spec.axis_index(a) for a in sa]
    )
    H_sp = H_full[..., spatial_idx, :][..., :, spatial_idx]
    gHg = jnp.einsum("bi,bij,bj->b", g, H_sp, g)
    base = g_norm_sq_reg ** ((p - 2.0) / 2.0)
    cross = (p - 2.0) * g_norm_sq_reg ** ((p - 4.0) / 2.0) * gHg
    return base * L + cross

Related APIsomnibias.fields.jax.ops.gradient↗omnibias.fields.jax.ops.laplacian↗omnibias.fields.jax.ops.hessian↗

Other public import paths

• omnibias.fields.jax.ops.p_laplacian
• omnibias.fields.jax.ops.nonlinear.p_laplacian

FUNCTION

skew_symmetric_advection

omnibias.fields.jax.ops.skew_symmetric_advectionSource ↗

skew_symmetric_advection(state: FieldState, *, velocity: tuple[str, ...], scalar: str | None=None) -> Array

Energy/enstrophy-conserving (skew-symmetric) advection.

\tfrac12[(u\cdot\nabla)\phi + \nabla\cdot(u\phi)] = (u\cdot\nabla)\phi + \tfrac12(\nabla\cdot u)\phi. For a divergence-free u this equals the standard advection; the extra 0.5(div u)phi term is what makes the discrete operator skew-symmetric (conserves the quadratic invariant). With scalar=None it is the vector self-advection form.

Parameters

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

ReturnsArray

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>6]" -q

Read the implementation

def skew_symmetric_advection(
    state: FieldState,
    *,
    velocity: tuple[str, ...],
    scalar: str | None = None,
) -> Array:
    r"""Energy/enstrophy-conserving (skew-symmetric) advection.

    :math:`\tfrac12[(u\cdot\nabla)\phi + \nabla\cdot(u\phi)]
    = (u\cdot\nabla)\phi + \tfrac12(\nabla\cdot u)\phi`. For a divergence-free
    ``u`` this equals the standard advection; the extra ``0.5(div u)phi`` term is
    what makes the discrete operator skew-symmetric (conserves the quadratic
    invariant). With ``scalar=None`` it is the vector self-advection form.
    """
    div_u = divergence(state, velocity)
    if scalar is not None:
        adv = advection(state, velocity=velocity, scalar=scalar)
        return adv + 0.5 * div_u * value(state, scalar)
    adv = advection(state, velocity=velocity)
    u_i = stack_components(state, velocity)
    return adv + 0.5 * div_u[..., None] * u_i

Related APIsomnibias.fields.jax.ops.divergence↗omnibias.fields.jax.ops.stack_components↗omnibias.fields.jax.ops.value↗omnibias.fields.jax.ops.advection↗

Other public import paths

• omnibias.fields.jax.ops.skew_symmetric_advection
• omnibias.fields.jax.ops.nonlinear.skew_symmetric_advection

← Package overviewRunnable learning examples →