API REFERENCE/omnibias-fields

omnibias.fields.torch.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.torch.ops.nonlinear

Module contract & mathematical conventions

Nonlinear ops: advection, material_derivative, p_laplacian, directional_derivative.

These ops are quadratic (or higher) in the field outputs, so they don't fit cleanly inside the closed-form derivative tower. The implementation folds them through the basic ops: each is a few-line composition that the kernel does not need to know about.

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.torch.ops.advectionSource ↗

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

Advection (velocity . nabla) target.

Three calling modes:

• advection(state, velocity=("u","v","w")) -- self-advection, target defaults to velocity. Returns (B, len(velocity)).
• advection(state, velocity=("u","v","w"), scalar="phi") -- scalar advection (u . nabla) phi. Returns (B,).
• advection(state, velocity=("u","v","w"), target=("a","b","c")) -- general vector-on-vector advection. Returns (B, len(target)).

The velocity-tuple length must equal the number of *spatial* axes.

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.

ReturnsTensor

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,
) -> Tensor:
    """Advection ``(velocity . nabla) target``.

    Three calling modes:

    - ``advection(state, velocity=("u","v","w"))`` -- self-advection,
      target defaults to ``velocity``. Returns ``(B, len(velocity))``.
    - ``advection(state, velocity=("u","v","w"), scalar="phi")`` --
      scalar advection ``(u . nabla) phi``. Returns ``(B,)``.
    - ``advection(state, velocity=("u","v","w"), target=("a","b","c"))``
      -- general vector-on-vector advection. Returns ``(B, len(target))``.

    The velocity-tuple length must equal the number of *spatial* axes.
    """
    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)              # (B, n_spatial)

    if scalar is not None:
        if target is not None:
            raise ValueError("advection: pass scalar OR target, not both")
        # (u . nabla) phi = sum_a u_a (d phi / dx_a).
        grad_phi = gradient(state, scalar)                 # (B, n_spatial)
        return (u_vec * grad_phi).sum(dim=-1)              # (B,)

    target = target if target is not None else velocity
    cols = []
    for n in target:
        grad_n = gradient(state, n)                        # (B, n_spatial)
        cols.append((u_vec * grad_n).sum(dim=-1))          # (B,)
    return torch.stack(cols, dim=-1)                       # (B, len(target))

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

Other public import paths

• omnibias.fields.torch.ops.advection
• omnibias.fields.torch.ops.nonlinear.advection

FUNCTION

directional_derivative

omnibias.fields.torch.ops.directional_derivativeSource ↗

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

d_v u = grad u . direction.

direction has shape (B, n_spatial) matching the 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.
directionkeyword-only	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

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: Tensor,
) -> Tensor:
    """``d_v u = grad u . direction``.

    ``direction`` has shape ``(B, n_spatial)`` matching the gradient.
    """
    g = gradient(state, name)                              # (B, n_spatial)
    if direction.shape != g.shape:
        raise ValueError(
            f"direction shape {tuple(direction.shape)} != "
            f"gradient shape {tuple(g.shape)}"
        )
    return (g * direction).sum(dim=-1)

Related APIsomnibias.fields.torch.ops.gradient↗

Other public import paths

• omnibias.fields.torch.ops.directional_derivative
• omnibias.fields.torch.ops.nonlinear.directional_derivative

FUNCTION

material_derivative

omnibias.fields.torch.ops.material_derivativeSource ↗

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

D/Dt = d/dt + (velocity . nabla).

For a vector target (default velocity itself) returns (B, C); for a scalar target returns (B,).

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.

ReturnsTensor

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.

Executed regression · dsl dt dx use coordinate spec axis names

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_dsl_dt_dx_use_coordinate_spec_axis_names():
    """``.dt`` / ``.dx`` route through ``coordinate_spec``, not hardcoded ``t``/``x``."""
    axes = ("xi", "eta", "time")
    polys = {
        "u": (Poly((1.0, 1.0)), Poly((1.0, 0.0)), Poly((1.0,))),
        "v": (Poly((0.0, 1.0)), Poly((1.0, -1.0)), Poly((1.0,))),
        "c": (Poly((1.0, 2.0)), Poly((1.0, -1.0)), Poly((0.5, 0.0, 1.0))),
    }
    f = make_field("torch", axes, polys, time_axis="time")
    rng = np.random.default_rng(0)
    nodes = rng.standard_normal((5, 3))
    st = f(torch.as_tensor(nodes, dtype=torch.float64))
    assert np.allclose(
        _np(st.c.dt),
        _np(st.ops.derivative(st, "c", axis="time")),
        rtol=0.0,
        atol=0.0,
    )
    assert np.allclose(
        _np(st.c.dx),
        _np(st.ops.derivative(st, "c", axis="xi")),
        rtol=0.0,
        atol=0.0,
    )
    assert np.allclose(
        _np(st.c.dy),
        _np(st.ops.derivative(st, "c", axis="eta")),
        rtol=0.0,
        atol=0.0,
    )
    md = st.ops.material_derivative(st, velocity=("u", "v"), scalar="c")
    assert md.shape == (5,)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_dsl_dt_dx_use_coordinate_spec_axis_names" -q

Read the implementation

def material_derivative(
    state: FieldState,
    *,
    velocity: tuple[str, ...],
    scalar: str | None = None,
) -> Tensor:
    """``D/Dt = d/dt + (velocity . nabla)``.

    For a vector target (default ``velocity`` itself) returns
    ``(B, C)``; for a scalar target returns ``(B,)``.
    """
    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)  # (B,)
        return dt_target + advection(state, velocity=velocity, scalar=scalar)
    dt_target = torch.stack(
        [derivative(state, n, axis=time_axis) for n in velocity], dim=-1,
    )                                                       # (B, C)
    return dt_target + advection(state, velocity=velocity)

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

Other public import paths

• omnibias.fields.torch.ops.material_derivative
• omnibias.fields.torch.ops.nonlinear.material_derivative

FUNCTION

p_laplacian

omnibias.fields.torch.ops.p_laplacianSource ↗

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

Delta_p u = nabla . (|nabla u|^{p-2} nabla u) of shape (B,).

Implementation follows docs/pinn-derivations.md Section 4:

Delta_p u = |g|^{p-2} Delta u + (p-2) |g|^{p-4} (g^T H g),

where g = nabla u and H is the Hessian, both restricted to spatial axes. eps regularises |g| to avoid singularity at zero 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.
pkeyword-only	floatRequired	See the operation contract and implementation below.
epskeyword-only	floatDefault: 1e-08	See the operation contract and implementation below.

ReturnsTensor

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,
) -> Tensor:
    """``Delta_p u = nabla . (|nabla u|^{p-2} nabla u)`` of shape ``(B,)``.

    Implementation follows ``docs/pinn-derivations.md`` Section 4:

        Delta_p u = |g|^{p-2} Delta u + (p-2) |g|^{p-4} (g^T H g),

    where ``g = nabla u`` and ``H`` is the Hessian, both restricted to
    spatial axes. ``eps`` regularises ``|g|`` to avoid singularity at
    zero gradient.
    """
    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)                              # (B, n_spatial)
    L = laplacian(state, name)                             # (B,)
    if abs(p - 2.0) < 1e-15:
        return L
    g2 = (g * g).sum(dim=-1)                               # (B,)
    g_norm_sq_reg = g2 + eps * eps                         # (B,)
    # Hessian restricted to spatial axes.
    H_full = hessian(state, name)                          # (B, D, D)
    spatial_idx = [
        state.coordinate_spec.axis_index(a) for a in sa
    ]
    H_sp = H_full[..., spatial_idx, :][..., :, spatial_idx]  # (B, ns, ns)
    gHg = torch.einsum("bi,bij,bj->b", g, H_sp, g)         # (B,)
    base = g_norm_sq_reg.pow((p - 2.0) / 2.0)              # (B,)
    cross = (p - 2.0) * g_norm_sq_reg.pow((p - 4.0) / 2.0) * gHg
    return base * L + cross

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

Other public import paths

• omnibias.fields.torch.ops.p_laplacian
• omnibias.fields.torch.ops.nonlinear.p_laplacian

FUNCTION

skew_symmetric_advection

omnibias.fields.torch.ops.skew_symmetric_advectionSource ↗

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

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.

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

Read the implementation

def skew_symmetric_advection(
    state: FieldState,
    *,
    velocity: tuple[str, ...],
    scalar: str | None = None,
) -> Tensor:
    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.torch.ops.divergence↗omnibias.fields.torch.ops.stack_components↗omnibias.fields.torch.ops.value↗omnibias.fields.torch.ops.advection↗

Other public import paths

• omnibias.fields.torch.ops.skew_symmetric_advection
• omnibias.fields.torch.ops.nonlinear.skew_symmetric_advection

← Package overviewRunnable learning examples →