Skip to main content
API REFERENCE/omnibias-fields

omnibias.fields.jax.ops.norms

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.norms
Module contract & mathematical conventions

Inner products and L2 / Sobolev norms as standalone ops (jax).

Bit-identical twin of omnibias.fields.torch.ops.norms.

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

inner_product​

omnibias.fields.jax.ops.inner_productSource ↗
inner_product(state: FieldState, name_a: str, name_b: str, *, rule: QuadratureSpec, weight: str | None=None) -> Array

Weighted real inner product \langle a, b\rangle_w = \int_\Omega w\,a\,b\,dx.

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
name_apositional-or-keywordstrRequiredSee the operation contract and implementation below.
name_bpositional-or-keywordstrRequiredSee the operation contract and implementation below.
rulekeyword-onlyQuadratureSpecRequiredSee the operation contract and implementation below.
weightkeyword-onlystr | NoneDefault: NoneSee the operation contract and implementation below.
ReturnsArray
Executed regression · inner product 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_inner_product_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))
    t = _np(ts.ops.inner_product(ts, "u", "v", rule=rule))
    j = _np(js.ops.inner_product(js, "u", "v", rule=rule))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_inner_product_cross_backend" -q
Read the implementation
def inner_product(
    state: FieldState,
    name_a: str,
    name_b: str,
    *,
    rule: QuadratureSpec,
    weight: str | None = None,
) -> Array:
    r"""Weighted real inner product :math:`\langle a, b\rangle_w
    = \int_\Omega w\,a\,b\,dx`."""
    a = value(state, name_a)
    b = value(state, name_b)
    integrand = a * b
    if weight is not None:
        integrand = integrand * value(state, weight)
    return _integrate_values(integrand, rule)
Related APIsomnibias.fields.jax.ops.value↗
Other public import paths
  • omnibias.fields.jax.ops.inner_product
  • omnibias.fields.jax.ops.norms.inner_product
FUNCTION

l2_norm​

omnibias.fields.jax.ops.l2_normSource ↗
l2_norm(state: FieldState, name: str, *, rule: QuadratureSpec) -> Array

\lVert u\rVert_{L^2} = \sqrt{\int_\Omega u^2\,dx}.

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
rulekeyword-onlyQuadratureSpecRequiredSee the operation contract and implementation below.
ReturnsArray
Read the implementation
def l2_norm(state: FieldState, name: str, *, rule: QuadratureSpec) -> Array:
    r""":math:`\lVert u\rVert_{L^2} = \sqrt{\int_\Omega u^2\,dx}`."""
    sq = _integrate_values(value(state, name) ** 2, rule)
    return jnp.sqrt(sq)
Related APIsomnibias.fields.jax.ops.value↗
Other public import paths
  • omnibias.fields.jax.ops.l2_norm
  • omnibias.fields.jax.ops.norms.l2_norm
FUNCTION

sobolev_norm​

omnibias.fields.jax.ops.sobolev_normSource ↗
sobolev_norm(state: FieldState, name: str, *, rule: QuadratureSpec, k: int=1, weights: tuple[float, ...] | None=None) -> Array

Sobolev norm \lVert u\rVert_{H^k} for k in {0, 1, 2}.

See omnibias.fields.torch.ops.norms.sobolev_norm for the definition.

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
rulekeyword-onlyQuadratureSpecRequiredSee the operation contract and implementation below.
kkeyword-onlyintDefault: 1See the operation contract and implementation below.
weightskeyword-onlytuple[float, ...] | NoneDefault: NoneSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When k < 0, raises ValueError(f'k must be >= 0, got {k}').
  • When k > 2, raises NotImplementedError(f'sobolev_norm currently supports k in {{0, 1, 2}}; got k={k}. Higher orders need the full multi-index expansion.').
  • When len(coeffs) != k + 1, raises ValueError(f'weights must have length k+1 = {k + 1}, got {len(coeffs)}').

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

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 sobolev_norm(
    state: FieldState,
    name: str,
    *,
    rule: QuadratureSpec,
    k: int = 1,
    weights: tuple[float, ...] | None = None,
) -> Array:
    r"""Sobolev norm :math:`\lVert u\rVert_{H^k}` for ``k`` in ``{0, 1, 2}``.

    See :func:`omnibias.fields.torch.ops.norms.sobolev_norm` for the definition.
    """
    if k < 0:
        raise ValueError(f"k must be >= 0, got {k}")
    if k > 2:
        raise NotImplementedError(
            "sobolev_norm currently supports k in {0, 1, 2}; "
            f"got k={k}. Higher orders need the full multi-index expansion."
        )
    coeffs = (1.0,) * (k + 1) if weights is None else weights
    if len(coeffs) != k + 1:
        raise ValueError(f"weights must have length k+1 = {k + 1}, got {len(coeffs)}")

    total = coeffs[0] * _integrate_values(value(state, name) ** 2, rule)
    if k >= 1:
        g = gradient(state, name)
        total = total + coeffs[1] * _integrate_values((g ** 2).sum(axis=-1), rule)
    if k >= 2:
        h = spatial_hessian(state, name)
        total = total + coeffs[2] * _integrate_values((h ** 2).sum(axis=(-2, -1)), rule)
    return jnp.sqrt(total)
Related APIsomnibias.fields.jax.ops.gradient↗omnibias.fields.jax.ops.value↗omnibias.fields.jax.ops.spatial_hessian↗
Other public import paths
  • omnibias.fields.jax.ops.sobolev_norm
  • omnibias.fields.jax.ops.norms.sobolev_norm
← Package overviewRunnable learning examples →