API REFERENCE/omnibias-fields

omnibias.fields.torch.ops.high_order

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

Module contract & mathematical conventions

High-order ops: biharmonic, polylaplacian, hessian, jacobian.

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

biharmonic

omnibias.fields.torch.ops.biharmonicSource ↗

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

Delta^2 f_name of shape (B,).

Closed-form on one-layer omnibias fields (one extra fast-path call at order 4); spectral fields get a diagonal-multiplier shortcut.

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

Read the implementation

def biharmonic(state: FieldState, name: str) -> Tensor:
    """``Delta^2 f_name`` of shape ``(B,)``.

    Closed-form on one-layer omnibias fields (one extra fast-path call
    at order 4); spectral fields get a diagonal-multiplier shortcut.
    """
    if _is_one_layer(state):
        sigma_4 = _sigma_of_order(state, 4)
        return state.field.polylaplacian(sigma_4, name, k=2)
    if _is_spectral(state):
        return state.field.biharmonic(state, name)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        return state.field.biharmonic(state, name)  # type: ignore[attr-defined]
    if _is_jet_mlp(state):
        return state.field.biharmonic(state, name)  # type: ignore[attr-defined]
    # Generic fallback: chain through laplacian-of-laplacian.
    raise NotImplementedError(
        f"biharmonic not implemented for field type {type(state.field).__name__}"
    )

Other public import paths

• omnibias.fields.torch.ops.biharmonic
• omnibias.fields.torch.ops.high_order.biharmonic

FUNCTION

gradient_of_derivative

omnibias.fields.torch.ops.gradient_of_derivativeSource ↗

gradient_of_derivative(state: FieldState, name: str, *, axis: int | str) -> Tensor

Helper: grad (d f / dx_axis). Used to assemble the Hessian for field types without a closed-form full Hessian.

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.
axiskeyword-only	int | strRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · sobolev norm cross backend

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_sobolev_norm_cross_backend(make_torch_field, make_jax_field):  # type: ignore[no-untyped-def]
    rule = gauss_legendre(BOUNDS, (8, 8))
    ts = make_torch_field()(torch.as_tensor(rule.nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(rule.nodes, dtype=jnp.float64))
    for k in (0, 1, 2):
        t = _np(ts.u.sobolev_norm(rule=rule, k=k))
        j = _np(js.u.sobolev_norm(rule=rule, k=k))
        assert np.allclose(t, j, rtol=1e-12, atol=1e-12), k

python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_sobolev_norm_cross_backend" -q

Read the implementation

def gradient_of_derivative(
    state: FieldState, name: str, *, axis: int | str,
) -> Tensor:
    """Helper: ``grad (d f / dx_axis)``. Used to assemble the Hessian for
    field types without a closed-form full Hessian."""
    a = _resolve_axis(state, axis)
    D = state.coordinate_spec.ndim
    cols = []
    for j in range(D):
        if a == j:
            cols.append(derivative(state, name, axis=a, order=2))
        else:
            from omnibias.fields.torch.ops.basic import mixed_partial
            cols.append(mixed_partial(state, name, (a, j), (1, 1)))
    return torch.stack(cols, dim=-1)

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

Other public import paths

• omnibias.fields.torch.ops.gradient_of_derivative
• omnibias.fields.torch.ops.high_order.gradient_of_derivative

FUNCTION

hessian

omnibias.fields.torch.ops.hessianSource ↗

hessian(state: FieldState, name: str, *, axes: tuple[int | str, ...] | None=None) -> Tensor

Hessian over selected axes, shape (B, len(axes), len(axes)).

axes=None keeps the historical behaviour: all coordinate axes, including time. Use spatial_hessian for the common spatial-only Hessian.

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.
axeskeyword-only	tuple[int | str, ...] | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · sobolev norm cross backend

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_sobolev_norm_cross_backend(make_torch_field, make_jax_field):  # type: ignore[no-untyped-def]
    rule = gauss_legendre(BOUNDS, (8, 8))
    ts = make_torch_field()(torch.as_tensor(rule.nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(rule.nodes, dtype=jnp.float64))
    for k in (0, 1, 2):
        t = _np(ts.u.sobolev_norm(rule=rule, k=k))
        j = _np(js.u.sobolev_norm(rule=rule, k=k))
        assert np.allclose(t, j, rtol=1e-12, atol=1e-12), k

python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_sobolev_norm_cross_backend" -q

Read the implementation

def hessian(
    state: FieldState,
    name: str,
    *,
    axes: tuple[int | str, ...] | None = None,
) -> Tensor:
    """Hessian over selected axes, shape ``(B, len(axes), len(axes))``.

    ``axes=None`` keeps the historical behaviour: all coordinate axes,
    including time. Use :func:`spatial_hessian` for the common spatial-only
    Hessian.
    """
    axis_idx = (
        tuple(range(state.coordinate_spec.ndim))
        if axes is None
        else tuple(_resolve_axis(state, a) for a in axes)
    )
    if _is_one_layer(state):
        sigma_pp = _sigma_of_order(state, 2)
        full = state.field.hessian_full(sigma_pp, name)
        if axis_idx == tuple(range(full.shape[-1])):
            return full
        idx = list(axis_idx)
        return full[..., idx, :][..., :, idx]
    if _is_jet_mlp(state):
        # The order-2 block of a single jet is the whole Hessian.
        full = state.field.hessian_full(state, name)  # type: ignore[attr-defined]
        if axis_idx == tuple(range(full.shape[-1])):
            return full
        idx = list(axis_idx)
        return full[..., idx, :][..., :, idx]
    # Generic fallback: build by taking first-derivative-of-gradient.
    rows = []
    for i in axis_idx:
        # d/dx_i grad f -> the i-th column is grad of (d f / dx_i).
        gi = gradient_of_derivative(state, name, axis=i)
        rows.append(gi[..., list(axis_idx)])
    return torch.stack(rows, dim=-2)  # (B, D, D); rows index outer axis

Related APIsomnibias.fields.torch.ops.gradient_of_derivative↗

Other public import paths

• omnibias.fields.torch.ops.hessian
• omnibias.fields.torch.ops.high_order.hessian

FUNCTION

jacobian

omnibias.fields.torch.ops.jacobianSource ↗

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

Jacobian of the vector (f_{names[0]}, ...), shape (B, C, D).

J[b, c, d] = d f_{names[c]} / dx_d over *all* axes (including the time axis), since the user can subset afterwards.

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 jacobian(state: FieldState, names: tuple[str, ...]) -> Tensor:
    """Jacobian of the vector ``(f_{names[0]}, ...)``, shape ``(B, C, D)``.

    ``J[b, c, d] = d f_{names[c]} / dx_d`` over *all* axes (including
    the time axis), since the user can subset afterwards.
    """
    if _is_one_layer(state):
        sigma_p = _sigma_of_order(state, 1)
        rows = [
            state.field.gradient_full(sigma_p, n) for n in names
        ]
        return torch.stack(rows, dim=-2)  # (B, C, D)
    if _is_jet_mlp(state):
        rows = [state.field.gradient_full(state, n) for n in names]  # type: ignore[attr-defined]
        return torch.stack(rows, dim=-2)  # (B, C, D)
    rows = []
    for n in names:
        rows.append(gradient(state, n, axes=tuple(range(state.coordinate_spec.ndim))))
    return torch.stack(rows, dim=-2)

Related APIsomnibias.fields.torch.ops.gradient↗

Other public import paths

• omnibias.fields.torch.ops.jacobian
• omnibias.fields.torch.ops.high_order.jacobian

FUNCTION

polylaplacian

omnibias.fields.torch.ops.polylaplacianSource ↗

polylaplacian(state: FieldState, name: str, *, k: int) -> Tensor

Evaluate Delta^k of the named scalar field; the result has shape (B,).

On a supported one-layer ridge field, the implementation evaluates the activation derivative at order 2k and contracts it with powers of weight norms and readout coefficients. The contraction is not the cost of the entire calculation: activation-derivative work depends on order, and the field representation and projection depend on dimension. Spectral, Chebyshev, and jet-MLP fields take different dispatch paths with different costs.

k must be at least one. k=1 delegates to the ordinary Laplacian. Unsupported field representations raise NotImplementedError. Use the corresponding field representation's derivative-order and numerical-domain contract.

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

ReturnsTensor

Explicit validation & failure cases

• When k < 1, raises ValueError(f'polylaplacian k must be >= 1, got {k}').

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

Read the implementation

def polylaplacian(state: FieldState, name: str, *, k: int) -> Tensor:
    """``Delta^k f_name`` of shape ``(B,)``.

    Memory and time both ``O(B * H)``, *independent of* ``k``, because
    the omnibias closed-form expresses ``Delta^k`` as one ``sigma^{(2k)}``
    fast-path call plus a per-mode reduction.
    """
    if k < 1:
        raise ValueError(f"polylaplacian k must be >= 1, got {k}")
    if k == 1:
        from omnibias.fields.torch.ops.basic import laplacian
        return laplacian(state, name)
    if _is_one_layer(state):
        sigma_2k = _sigma_of_order(state, 2 * k)
        return state.field.polylaplacian(sigma_2k, name, k=k)
    if _is_spectral(state):
        return state.field.polylaplacian(state, name, k=k)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        # Composition fallback: Delta^k = Delta(Delta^{k-1}).
        return state.field.polylaplacian(state, name, k=k)  # type: ignore[attr-defined]
    if _is_jet_mlp(state):
        # Multinomial expansion of (sum_i d_i^2)^k over one order-2k jet.
        return state.field.polylaplacian(state, name, k=k)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"polylaplacian not implemented for field type {type(state.field).__name__}"
    )

Other public import paths

• omnibias.fields.torch.ops.polylaplacian
• omnibias.fields.torch.ops.high_order.polylaplacian

FUNCTION

spatial_hessian

omnibias.fields.torch.ops.spatial_hessianSource ↗

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

Spatial-only Hessian, excluding the 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.

ReturnsTensor

Executed regression · sobolev norm cross backend

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_sobolev_norm_cross_backend(make_torch_field, make_jax_field):  # type: ignore[no-untyped-def]
    rule = gauss_legendre(BOUNDS, (8, 8))
    ts = make_torch_field()(torch.as_tensor(rule.nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(rule.nodes, dtype=jnp.float64))
    for k in (0, 1, 2):
        t = _np(ts.u.sobolev_norm(rule=rule, k=k))
        j = _np(js.u.sobolev_norm(rule=rule, k=k))
        assert np.allclose(t, j, rtol=1e-12, atol=1e-12), k

python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_sobolev_norm_cross_backend" -q

Read the implementation

def spatial_hessian(state: FieldState, name: str) -> Tensor:
    """Spatial-only Hessian, excluding the time axis."""
    return hessian(state, name, axes=tuple(state.coordinate_spec.spatial_axes))

Related APIsomnibias.fields.torch.ops.hessian↗

Other public import paths

• omnibias.fields.torch.ops.spatial_hessian
• omnibias.fields.torch.ops.high_order.spatial_hessian

FUNCTION

vector_biharmonic

omnibias.fields.torch.ops.vector_biharmonicSource ↗

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

Biharmonic per component, shape (B, len(names)).

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 vector_biharmonic(state: FieldState, names: tuple[str, ...]) -> Tensor:
    """Biharmonic per component, shape ``(B, len(names))``."""
    cols = [biharmonic(state, n) for n in names]
    return torch.stack(cols, dim=-1)

Related APIsomnibias.fields.torch.ops.biharmonic↗

Other public import paths

• omnibias.fields.torch.ops.vector_biharmonic
• omnibias.fields.torch.ops.high_order.vector_biharmonic

FUNCTION

vector_hessian

omnibias.fields.torch.ops.vector_hessianSource ↗

vector_hessian(state: FieldState, names: tuple[str, ...], *, axes: tuple[int | str, ...] | None=None) -> Tensor

Stack component Hessians, shape (B, C, D, D) or (B, C, A, A).

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.
axeskeyword-only	tuple[int | str, ...] | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def vector_hessian(
    state: FieldState,
    names: tuple[str, ...],
    *,
    axes: tuple[int | str, ...] | None = None,
) -> Tensor:
    """Stack component Hessians, shape ``(B, C, D, D)`` or ``(B, C, A, A)``."""
    rows = [hessian(state, n, axes=axes) for n in names]
    return torch.stack(rows, dim=-3)

Related APIsomnibias.fields.torch.ops.hessian↗

Other public import paths

• omnibias.fields.torch.ops.vector_hessian
• omnibias.fields.torch.ops.high_order.vector_hessian

FUNCTION

vector_laplacian

omnibias.fields.torch.ops.vector_laplacianSource ↗

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

Laplacian applied per component, shape (B, len(names)).

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

Executed regression · curl of curl identity

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_curl_of_curl_identity(which):
    """curl(curl u) == grad(div u) - vector_laplacian(u) in 2D and 3D."""
    ndim = 2 if which == "2d" else 3
    names = ("u", "v") if which == "2d" else ("u", "v", "w")
    nodes = _nodes(ndim)
    st = _state("torch", which, nodes)
    cc = _np(st.ops.curl_of_curl(st, names))
    gd = _np(st.ops.gradient_of_divergence(st, names))
    vl = _np(st.ops.vector_laplacian(st, names))
    assert np.allclose(cc, gd - vl, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_curl_of_curl_identity[2d]" -q

Read the implementation

def vector_laplacian(state: FieldState, names: tuple[str, ...]) -> Tensor:
    """Laplacian applied per component, shape ``(B, len(names))``."""
    from omnibias.fields.torch.ops.basic import laplacian
    cols = [laplacian(state, n) for n in names]
    return torch.stack(cols, dim=-1)

Other public import paths

• omnibias.fields.torch.ops.vector_laplacian
• omnibias.fields.torch.ops.high_order.vector_laplacian

FUNCTION

vector_polylaplacian

omnibias.fields.torch.ops.vector_polylaplacianSource ↗

vector_polylaplacian(state: FieldState, names: tuple[str, ...], *, k: int) -> Tensor

Polylaplacian applied per component, shape (B, len(names)).

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

ReturnsTensor

Read the implementation

def vector_polylaplacian(state: FieldState, names: tuple[str, ...], *, k: int) -> Tensor:
    """Polylaplacian applied per component, shape ``(B, len(names))``."""
    cols = [polylaplacian(state, n, k=k) for n in names]
    return torch.stack(cols, dim=-1)

Related APIsomnibias.fields.torch.ops.polylaplacian↗

Other public import paths

• omnibias.fields.torch.ops.vector_polylaplacian
• omnibias.fields.torch.ops.high_order.vector_polylaplacian

← Package overviewRunnable learning examples →