omnibias.fields.weak.jax

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.weak.jax

Module contract & mathematical conventions

Weak-form VPINN assembly (JAX twin of torch; theory 02-04).

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 →

CLASS

ResidualTerm

omnibias.fields.weak.jax.ResidualTermSource ↗

class ResidualTerm

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

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
`name`dataclass field	`str`Required	See the operation contract and implementation below.
`path`dataclass field	`PathName`Required	See the operation contract and implementation below.
`value`dataclass field	`Array`Required	See the operation contract and implementation below.

Fields & methods

namefield

name: str

Declared str field on ResidualTerm. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

pathfield

path: PathName

Declared PathName field on ResidualTerm. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

valuefield

value: Array

Declared Array field on ResidualTerm. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

Read the implementation

class ResidualTerm:
    name: str
    path: PathName
    value: Array

FUNCTION

weak_loss

omnibias.fields.weak.jax.weak_lossSource ↗

weak_loss(field: Callable[[Array], Array] | tuple[float, ...], space: TestFunctionSpace, *, operator: WeakForm, source: Callable[[Array], Array] | tuple[float, ...] | None=None, quadrature: QuadratureSpec | None=None, include_boundary_bound: bool=True, deriv_bound: float=1.0) -> 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
`field`positional-or-keyword	`Callable[[Array], Array] \| tuple[float, ...]`Required	See the operation contract and implementation below.
`space`positional-or-keyword	`TestFunctionSpace`Required	See the operation contract and implementation below.
`operator`keyword-only	`WeakForm`Required	See the operation contract and implementation below.
`source`keyword-only	`Callable[[Array], Array] \| tuple[float, ...] \| None`Default: None	See the operation contract and implementation below.
`quadrature`keyword-only	`QuadratureSpec \| None`Default: None	See the operation contract and implementation below.
`include_boundary_bound`keyword-only	`bool`Default: True	See the operation contract and implementation below.
`deriv_bound`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def weak_loss(
    field: Callable[[Array], Array] | tuple[float, ...],
    space: TestFunctionSpace,
    *,
    operator: WeakForm,
    source: Callable[[Array], Array] | tuple[float, ...] | None = None,
    quadrature: QuadratureSpec | None = None,
    include_boundary_bound: bool = True,
    deriv_bound: float = 1.0,
) -> Array:
    residual, _terms = weak_residual(
        field, space, operator=operator, source=source, quadrature=quadrature
    )
    loss = jnp.sum(residual * residual)
    if include_boundary_bound:
        bound = boundary_bound(space, deriv_bound=deriv_bound)
        loss = loss + jnp.asarray(float(bound.hi), dtype=loss.dtype)
    return loss

FUNCTION

weak_residual

omnibias.fields.weak.jax.weak_residualSource ↗

weak_residual(field: Callable[[Array], Array] | tuple[float, ...], space: TestFunctionSpace, *, operator: WeakForm, source: Callable[[Array], Array] | tuple[float, ...] | None=None, quadrature: QuadratureSpec | None=None) -> tuple[Array, tuple[ResidualTerm, ...]]

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
`field`positional-or-keyword	`Callable[[Array], Array] \| tuple[float, ...]`Required	See the operation contract and implementation below.
`space`positional-or-keyword	`TestFunctionSpace`Required	See the operation contract and implementation below.
`operator`keyword-only	`WeakForm`Required	See the operation contract and implementation below.
`source`keyword-only	`Callable[[Array], Array] \| tuple[float, ...] \| None`Default: None	See the operation contract and implementation below.
`quadrature`keyword-only	`QuadratureSpec \| None`Default: None	See the operation contract and implementation below.

Returnstuple[Array, tuple[ResidualTerm, ...]]

Explicit validation & failure cases

When space.window is None, raises ValueError('weak_residual needs a box window').

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

Executed regression · g5 torch jax exact residual 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_g5_torch_jax_exact_residual_parity() -> None:
    import jax
    import jax.numpy as jnp
    from omnibias.fields.weak.jax import weak_residual as jax_weak_residual

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    space = _space()
    op = WeakForm(diffusion=(1.0,), source=(2.0,))
    field = (0.0, 1.0, -1.0)
    t_res, t_terms = weak_residual(field, space, operator=op)
    j_res, j_terms = jax_weak_residual(field, space, operator=op)
    t = t_res.detach().cpu().numpy()
    j = np.asarray(j_res)
    assert np.allclose(t, j, rtol=0, atol=1e-15)
    assert all(term.path == "exact" for term in t_terms)
    assert all(term.path == "exact" for term in j_terms)
    del jnp

python -m pytest "packages/omnibias-fields/tests/test_weak.py::test_g5_torch_jax_exact_residual_parity" -q

Read the implementation

def weak_residual(
    field: Callable[[Array], Array] | tuple[float, ...],
    space: TestFunctionSpace,
    *,
    operator: WeakForm,
    source: Callable[[Array], Array] | tuple[float, ...] | None = None,
    quadrature: QuadratureSpec | None = None,
) -> tuple[Array, tuple[ResidualTerm, ...]]:
    if space.window is None:
        raise ValueError("weak_residual needs a box window")
    src = operator.source if source is None else source
    values: list[Array] = []
    terms: list[ResidualTerm] = []
    for i in range(space.size):
        stiffness, stiff_path = _stiffness_term(field, space, i, operator.diffusion, quadrature)
        load, load_path = _load_term(src, space, i, quadrature)
        r = stiffness - load
        values.append(r)
        terms.append(ResidualTerm(f"stiffness[{i}]", stiff_path, stiffness))
        terms.append(ResidualTerm(f"load[{i}]", load_path, load))
    return jnp.stack(values), tuple(terms)

← Package overview Runnable learning examples →

omnibias.fields.weak.jax

ResidualTerm​

Parameters

Fields & methods

weak_loss​

Parameters

weak_residual​

Parameters

ResidualTerm

weak_loss

weak_residual