API REFERENCE/omnibias-fields

omnibias.fields.jax.ops.vector

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

Module contract & mathematical conventions

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

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

curl

omnibias.fields.jax.ops.curlSource ↗

curl(state: FieldState, names: tuple[str, ...]) -> 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.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When len(names) != len(sa), raises ValueError(f'curl: vector length {len(names)} must equal n_spatial {len(sa)} ({sa!r})').
• When len(sa) == 1, raises ValueError('curl is not defined for 1D vector fields; use derivative instead').

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

Read the implementation

def curl(state: FieldState, names: tuple[str, ...]) -> Array:
    sa = state.coordinate_spec.spatial_axes
    if len(names) != len(sa):
        raise ValueError(
            f"curl: vector length {len(names)} must equal n_spatial {len(sa)} "
            f"({sa!r})"
        )
    if len(sa) == 1:
        raise ValueError(
            "curl is not defined for 1D vector fields; use derivative instead"
        )
    if len(sa) == 2:
        u, v = names
        x_axis, y_axis = sa[0], sa[1]
        cz = derivative(state, v, axis=x_axis) - derivative(state, u, axis=y_axis)
        return cz[..., None]
    if len(sa) == 3:
        u, v, w = names
        x_axis, y_axis, z_axis = sa[0], sa[1], sa[2]
        cx = derivative(state, w, axis=y_axis) - derivative(state, v, axis=z_axis)
        cy = derivative(state, u, axis=z_axis) - derivative(state, w, axis=x_axis)
        cz = derivative(state, v, axis=x_axis) - derivative(state, u, axis=y_axis)
        return jnp.stack([cx, cy, cz], axis=-1)
    raise ValueError(
        f"curl supported for 2D / 3D vector fields, got {len(sa)} spatial axes"
    )

Related APIsomnibias.fields.jax.ops.derivative↗

Other public import paths

• omnibias.fields.jax.ops.curl
• omnibias.fields.jax.ops.vector.curl

FUNCTION

curl_of_curl

omnibias.fields.jax.ops.curl_of_curlSource ↗

curl_of_curl(state: FieldState, names: tuple[str, ...]) -> Array

\nabla\times(\nabla\times u) for 2D / 3D vector fields.

Expanded directly from second partials (not via the identity), so a test can independently assert curl_of_curl == gradient_of_divergence - vector_laplacian.

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.

ReturnsArray

Explicit validation & failure cases

• When len(names) != len(sa), raises ValueError(f'curl_of_curl: vector length {len(names)} must equal n_spatial {len(sa)} ({sa!r})').

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

Read the implementation

def curl_of_curl(state: FieldState, names: tuple[str, ...]) -> Array:
    r""":math:`\nabla\times(\nabla\times u)` for 2D / 3D vector fields.

    Expanded directly from second partials (not via the identity), so a test can
    independently assert ``curl_of_curl == gradient_of_divergence - vector_laplacian``.
    """
    sa = state.coordinate_spec.spatial_axes
    if len(names) != len(sa):
        raise ValueError(
            f"curl_of_curl: vector length {len(names)} must equal n_spatial "
            f"{len(sa)} ({sa!r})"
        )
    if len(sa) == 2:
        u, v = names
        x, y = sa
        s_y = mixed_partial(state, v, (x, y), (1, 1)) - derivative(state, u, axis=y, order=2)
        s_x = derivative(state, v, axis=x, order=2) - mixed_partial(state, u, (x, y), (1, 1))
        return jnp.stack([s_y, -s_x], axis=-1)
    if len(sa) == 3:
        u, v, w = names
        x, y, z = sa
        ccx = (
            mixed_partial(state, v, (x, y), (1, 1)) + mixed_partial(state, w, (x, z), (1, 1))
            - derivative(state, u, axis=y, order=2) - derivative(state, u, axis=z, order=2)
        )
        ccy = (
            mixed_partial(state, w, (y, z), (1, 1)) + mixed_partial(state, u, (x, y), (1, 1))
            - derivative(state, v, axis=z, order=2) - derivative(state, v, axis=x, order=2)
        )
        ccz = (
            mixed_partial(state, u, (x, z), (1, 1)) + mixed_partial(state, v, (y, z), (1, 1))
            - derivative(state, w, axis=x, order=2) - derivative(state, w, axis=y, order=2)
        )
        return jnp.stack([ccx, ccy, ccz], axis=-1)
    raise ValueError(
        f"curl_of_curl supported for 2D / 3D vector fields, got {len(sa)} spatial axes"
    )

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

Other public import paths

• omnibias.fields.jax.ops.curl_of_curl
• omnibias.fields.jax.ops.vector.curl_of_curl

FUNCTION

deformation_gradient

omnibias.fields.jax.ops.deformation_gradientSource ↗

deformation_gradient(state: FieldState, names: tuple[str, ...]) -> 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.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When len(names) != len(sa), raises ValueError(f'deformation_gradient: vector length {len(names)} must equal n_spatial {len(sa)} ({sa!r})').

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

Read the implementation

def deformation_gradient(state: FieldState, names: tuple[str, ...]) -> Array:
    sa = state.coordinate_spec.spatial_axes
    if len(names) != len(sa):
        raise ValueError(
            f"deformation_gradient: vector length {len(names)} must equal "
            f"n_spatial {len(sa)} ({sa!r})"
        )
    return spatial_jacobian(state, names)

Related APIsomnibias.fields.jax.ops.spatial_jacobian↗

Other public import paths

• omnibias.fields.jax.ops.deformation_gradient
• omnibias.fields.jax.ops.vector.deformation_gradient

FUNCTION

gradient_of_divergence

omnibias.fields.jax.ops.gradient_of_divergenceSource ↗

gradient_of_divergence(state: FieldState, names: tuple[str, ...]) -> Array

\nabla(\nabla\cdot u), shape (B, d).

Component i is \sum_j \partial^2_{ij} u_j.

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.

ReturnsArray

Explicit validation & failure cases

• When len(names) != len(sa), raises ValueError(f'gradient_of_divergence: vector length {len(names)} must equal n_spatial {len(sa)} ({sa!r})').

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

Read the implementation

def gradient_of_divergence(state: FieldState, names: tuple[str, ...]) -> Array:
    r""":math:`\nabla(\nabla\cdot u)`, shape ``(B, d)``.

    Component ``i`` is :math:`\sum_j \partial^2_{ij} u_j`.
    """
    sa = state.coordinate_spec.spatial_axes
    if len(names) != len(sa):
        raise ValueError(
            f"gradient_of_divergence: vector length {len(names)} must equal "
            f"n_spatial {len(sa)} ({sa!r})"
        )
    cols = []
    for i in sa:
        acc = None
        for n, j in zip(names, sa, strict=False):
            term = mixed_partial(state, n, (i, j), (1, 1))
            acc = term if acc is None else acc + term
        cols.append(acc)
    return jnp.stack(cols, axis=-1)

Related APIsomnibias.fields.jax.ops.mixed_partial↗

Other public import paths

• omnibias.fields.jax.ops.gradient_of_divergence
• omnibias.fields.jax.ops.vector.gradient_of_divergence

FUNCTION

rate_of_rotation_tensor

omnibias.fields.jax.ops.rate_of_rotation_tensorSource ↗

rate_of_rotation_tensor(state: FieldState, names: tuple[str, ...]) -> Array

Antisymmetric velocity-gradient (spin) tensor W = \tfrac12(J - J^T).

Complement of strain_rate; together J = strain_rate + rate_of_rotation_tensor.

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.

ReturnsArray

Explicit validation & failure cases

• When J.shape[-2] != J.shape[-1], raises ValueError(f'rate_of_rotation_tensor requires square Jacobian; got shape {tuple(J.shape)}').

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

Read the implementation

def rate_of_rotation_tensor(state: FieldState, names: tuple[str, ...]) -> Array:
    r"""Antisymmetric velocity-gradient (spin) tensor :math:`W = \tfrac12(J - J^T)`.

    Complement of :func:`strain_rate`; together ``J = strain_rate + rate_of_rotation_tensor``.
    """
    J = deformation_gradient(state, names)
    if J.shape[-2] != J.shape[-1]:
        raise ValueError(
            f"rate_of_rotation_tensor requires square Jacobian; got shape {tuple(J.shape)}"
        )
    return 0.5 * (J - jnp.swapaxes(J, -1, -2))

Related APIsomnibias.fields.jax.ops.deformation_gradient↗

Other public import paths

• omnibias.fields.jax.ops.rate_of_rotation_tensor
• omnibias.fields.jax.ops.vector.rate_of_rotation_tensor

FUNCTION

spatial_jacobian

omnibias.fields.jax.ops.spatial_jacobianSource ↗

spatial_jacobian(state: FieldState, names: tuple[str, ...]) -> 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.
namespositional-or-keyword	tuple[str, ...]Required	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>2]" -q

Read the implementation

def spatial_jacobian(state: FieldState, names: tuple[str, ...]) -> Array:
    rows = [gradient(state, n) for n in names]
    return jnp.stack(rows, axis=-2)

Related APIsomnibias.fields.jax.ops.gradient↗

Other public import paths

• omnibias.fields.jax.ops.spatial_jacobian
• omnibias.fields.jax.ops.vector.spatial_jacobian

FUNCTION

strain_rate

omnibias.fields.jax.ops.strain_rateSource ↗

strain_rate(state: FieldState, names: tuple[str, ...]) -> 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.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When J.shape[-2] != J.shape[-1], raises ValueError(f'strain_rate requires square Jacobian; got shape {tuple(J.shape)}').

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

Read the implementation

def strain_rate(state: FieldState, names: tuple[str, ...]) -> Array:
    J = deformation_gradient(state, names)
    if J.shape[-2] != J.shape[-1]:
        raise ValueError(
            f"strain_rate requires square Jacobian; got shape {tuple(J.shape)}"
        )
    return 0.5 * (J + jnp.swapaxes(J, -1, -2))

Related APIsomnibias.fields.jax.ops.deformation_gradient↗

Other public import paths

• omnibias.fields.jax.ops.strain_rate
• omnibias.fields.jax.ops.vector.strain_rate

FUNCTION

vorticity

omnibias.fields.jax.ops.vorticitySource ↗

vorticity(state: FieldState, names: tuple[str, ...]) -> 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.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def vorticity(state: FieldState, names: tuple[str, ...]) -> Array:
    return curl(state, names)

Related APIsomnibias.fields.jax.ops.curl↗

Other public import paths

• omnibias.fields.jax.ops.vorticity
• omnibias.fields.jax.ops.vector.vorticity

← Package overviewRunnable learning examples →