omnibias.fields.locus.torch

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.locus.torch

Module contract & mathematical conventions

PyTorch Newton projection on an equality locus (theory 01-09).

newton_project is Gauss-Newton in the forward pass. Its gradient uses the implicit function theorem (dx* / d theta = - DF^+ dF/d theta), not the unrolled iteration, so memory does not scale with iteration count.

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

AnsatzSolutionField

omnibias.fields.locus.AnsatzSolutionFieldSource ↗

class AnsatzSolutionField(nn.Module)

Level-2: equality inside a named tanh-class. Not a general PDE solver.

Parameters

Name	Type & default	Contract
`pde`positional-or-keyword	`PDESpec`Required	See the operation contract and implementation below.
`ansatz`positional-or-keyword	`TravellingWaveAnsatz`Required	See the operation contract and implementation below.

Executed regression · g5 ansatz rejects wrong

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_ansatz_rejects_wrong() -> None:
    pde = classical_pdes()["burgers"]
    wrong = published_ansatz("kdv")
    with pytest.raises(ValueError, match="symbolic verification"):
        AnsatzSolutionField(pde, wrong)
    good = AnsatzSolutionField(pde, published_ansatz("burgers"))
    cert = good.certificate()
    assert cert["verified"] is True
    assert cert["level3_general_solver"] is False

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g5_ansatz_rejects_wrong" -q

Fields & methods

__init__method

__init__(self, pde: PDESpec, ansatz: TravellingWaveAnsatz) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`pde`positional-or-keyword	`PDESpec`Required	See the operation contract and implementation below.
`ansatz`positional-or-keyword	`TravellingWaveAnsatz`Required	See the operation contract and implementation below.

Returns: None

Validation & failure cases

When not verify_exact(pde, ansatz), raises ValueError(f'ansatz {ansatz.kind!r} failed symbolic verification for {pde.name!r}; no claim is made outside a named ansatz class').

Executed regression · g5 ansatz rejects wrong

Full file, imports & fixtures ↗

def test_g5_ansatz_rejects_wrong() -> None:
    pde = classical_pdes()["burgers"]
    wrong = published_ansatz("kdv")
    with pytest.raises(ValueError, match="symbolic verification"):
        AnsatzSolutionField(pde, wrong)
    good = AnsatzSolutionField(pde, published_ansatz("burgers"))
    cert = good.certificate()
    assert cert["verified"] is True
    assert cert["level3_general_solver"] is False

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g5_ansatz_rejects_wrong" -q

Method implementation

    def __init__(self, pde: PDESpec, ansatz: TravellingWaveAnsatz) -> None:
        super().__init__()
        if not verify_exact(pde, ansatz):
            raise ValueError(
                f"ansatz {ansatz.kind!r} failed symbolic verification for {pde.name!r}; "
                "no claim is made outside a named ansatz class"
            )
        self.pde = pde
        self.ansatz = ansatz

View source ↗

forwardmethod

forward(self, x: Tensor, t: Tensor) -> Tensor

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.
`t`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

Returns: Tensor

Method implementation

    def forward(self, x: Tensor, t: Tensor) -> Tensor:
        from omnibias.core.tanh_method import evaluate_ansatz

        xs = x.detach().reshape(-1).tolist()
        ts = t.detach().reshape(-1).tolist()
        vals = [
            evaluate_ansatz(self.ansatz, float(xi), float(ti))
            for xi, ti in zip(xs, ts, strict=True)
        ]
        return torch.tensor(vals, dtype=x.dtype, device=x.device).reshape(x.shape)

View source ↗

certificatemethod

certificate(self) -> dict[str, object]

See the class contract above and the exact method implementation below.

Returns: dict[str, object]

Executed regression · g5 ansatz rejects wrong

Full file, imports & fixtures ↗

def test_g5_ansatz_rejects_wrong() -> None:
    pde = classical_pdes()["burgers"]
    wrong = published_ansatz("kdv")
    with pytest.raises(ValueError, match="symbolic verification"):
        AnsatzSolutionField(pde, wrong)
    good = AnsatzSolutionField(pde, published_ansatz("burgers"))
    cert = good.certificate()
    assert cert["verified"] is True
    assert cert["level3_general_solver"] is False

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g5_ansatz_rejects_wrong" -q

Method implementation

    def certificate(self) -> dict[str, object]:
        return {
            "class": self.pde.name,
            "verified": True,
            "level3_general_solver": False,
        }

View source ↗

pdeattribute

pde: PDESpec = pde

Public instance attribute assigned by __init__. The source assignment is pde = pde; it describes instance state, not a literal constructor default.

View source ↗

ansatzattribute

ansatz: TravellingWaveAnsatz = ansatz

Public instance attribute assigned by __init__. The source assignment is ansatz = ansatz; it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class AnsatzSolutionField(nn.Module):
    """Level-2: equality inside a named tanh-class. Not a general PDE solver."""

    def __init__(self, pde: PDESpec, ansatz: TravellingWaveAnsatz) -> None:
        super().__init__()
        if not verify_exact(pde, ansatz):
            raise ValueError(
                f"ansatz {ansatz.kind!r} failed symbolic verification for {pde.name!r}; "
                "no claim is made outside a named ansatz class"
            )
        self.pde = pde
        self.ansatz = ansatz

    def forward(self, x: Tensor, t: Tensor) -> Tensor:
        from omnibias.core.tanh_method import evaluate_ansatz

        xs = x.detach().reshape(-1).tolist()
        ts = t.detach().reshape(-1).tolist()
        vals = [
            evaluate_ansatz(self.ansatz, float(xi), float(ti))
            for xi, ti in zip(xs, ts, strict=True)
        ]
        return torch.tensor(vals, dtype=x.dtype, device=x.device).reshape(x.shape)

    def certificate(self) -> dict[str, object]:
        return {
            "class": self.pde.name,
            "verified": True,
            "level3_general_solver": False,
        }

Other public import paths

omnibias.fields.locus.AnsatzSolutionField
omnibias.fields.locus.torch.AnsatzSolutionField

CLASS

EqualityLocusLayer

omnibias.fields.locus.EqualityLocusLayerSource ↗

class EqualityLocusLayer(nn.Module)

Gauss-Newton projection with IFT backward. Always returns branch / condition.

Parameters

Name	Type & default	Contract
`system`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`branch`keyword-only	`int \| Literal['nearest', 'all']`Default: 'nearest'	See the operation contract and implementation below.
`max_iter`keyword-only	`int`Default: 20	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.
`require_transversal`keyword-only	`bool`Default: True	See the operation contract and implementation below.
`dtype`keyword-only	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

Executed regression · g1 forward residual

Full file, imports & fixtures ↗

def test_g1_forward_residual() -> None:
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    layer = EqualityLocusLayer(sys, dtype=torch.float64)
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out = layer(x0)
    assert bool(out.converged.detach())
    nrm = abs(residual(sys, tuple(float(v) for v in out.point.detach().tolist()))[0])
    assert nrm <= 1e-12
    assert out.branch.numel() >= 1
    assert float(out.condition.detach()) >= 1.0

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g1_forward_residual" -q

Regression example · g3 degeneracy refusal

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_g3_degeneracy_refusal() -> None:
    torch.set_default_dtype(torch.float64)
    sys = EqualitySystem(
        (
            UnitTerm(1, 1.0, (1.0, 0.0), 0.0),
            UnitTerm(1, 1.0, (1.0, 0.0), 0.0),
        )
    )
    layer = EqualityLocusLayer(sys, require_transversal=True, dtype=torch.float64)
    out = layer(torch.tensor([0.2, 0.3], dtype=torch.float64))
    assert bool(out.converged.detach()) is False
    assert float(out.condition.detach()) > 1e6

python -m pytest packages/omnibias-fields/tests/test_locus_layer.py -k test_g3_degeneracy_refusal -q

Fields & methods

__init__method

__init__(self, system: EqualitySystem, *, branch: int | Literal['nearest', 'all']='nearest', max_iter: int=20, tol: float=1e-12, require_transversal: bool=True, dtype: torch.dtype | None=None) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`system`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`branch`keyword-only	`int \| Literal['nearest', 'all']`Default: 'nearest'	See the operation contract and implementation below.
`max_iter`keyword-only	`int`Default: 20	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.
`require_transversal`keyword-only	`bool`Default: True	See the operation contract and implementation below.
`dtype`keyword-only	`torch.dtype \| None`Default: None	See the operation contract and implementation below.

Returns: None

Executed regression · g1 forward residual

Full file, imports & fixtures ↗

def test_g1_forward_residual() -> None:
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    layer = EqualityLocusLayer(sys, dtype=torch.float64)
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out = layer(x0)
    assert bool(out.converged.detach())
    nrm = abs(residual(sys, tuple(float(v) for v in out.point.detach().tolist()))[0])
    assert nrm <= 1e-12
    assert out.branch.numel() >= 1
    assert float(out.condition.detach()) >= 1.0

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g1_forward_residual" -q

Method implementation

    def __init__(
        self,
        system: EqualitySystem,
        *,
        branch: int | Literal["nearest", "all"] = "nearest",
        max_iter: int = 20,
        tol: float = 1e-12,
        require_transversal: bool = True,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.system = system
        self.branch_mode = branch
        self.max_iter = int(max_iter)
        self.tol = float(tol)
        self.require_transversal = bool(require_transversal)
        self._dtype = torch.get_default_dtype() if dtype is None else dtype

View source ↗

forwardmethod

forward(self, x0: Tensor) -> LocusOutput

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`x0`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

Returns: LocusOutput

Executed regression · g1 forward residual

Full file, imports & fixtures ↗

def test_g1_forward_residual() -> None:
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    layer = EqualityLocusLayer(sys, dtype=torch.float64)
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out = layer(x0)
    assert bool(out.converged.detach())
    nrm = abs(residual(sys, tuple(float(v) for v in out.point.detach().tolist()))[0])
    assert nrm <= 1e-12
    assert out.branch.numel() >= 1
    assert float(out.condition.detach()) >= 1.0

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g1_forward_residual" -q

Method implementation

    def forward(self, x0: Tensor) -> LocusOutput:
        result = newton_result(self.system, x0, max_iter=self.max_iter, tol=self.tol)
        point = newton_project(self.system, x0, max_iter=self.max_iter, tol=self.tol)
        sigma_min = float(result.condition)
        cond = 1.0 / max(sigma_min, 1e-30)
        ok = bool(result.converged)
        if self.require_transversal:
            ok = ok and bool(result.transversal)
        return LocusOutput(
            point=point,
            branch=torch.tensor(result.branch, dtype=torch.int64, device=x0.device),
            condition=torch.tensor(cond, dtype=x0.dtype, device=x0.device),
            converged=torch.tensor(ok, dtype=torch.bool, device=x0.device),
        )

View source ↗

systemattribute

system: EqualitySystem = system

Public instance attribute assigned by __init__. The source assignment is system = system; it describes instance state, not a literal constructor default.

View source ↗

branch_modeattribute

branch_mode: int | Literal['nearest', 'all'] = branch

Public instance attribute assigned by __init__. The source assignment is branch_mode = branch; it describes instance state, not a literal constructor default.

View source ↗

max_iterattribute

max_iter = int(max_iter)

Public instance attribute assigned by __init__. The source assignment is max_iter = int(max_iter); it describes instance state, not a literal constructor default.

View source ↗

tolattribute

tol = float(tol)

Public instance attribute assigned by __init__. The source assignment is tol = float(tol); it describes instance state, not a literal constructor default.

View source ↗

require_transversalattribute

require_transversal = bool(require_transversal)

Public instance attribute assigned by __init__. The source assignment is require_transversal = bool(require_transversal); it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class EqualityLocusLayer(nn.Module):
    """Gauss-Newton projection with IFT backward. Always returns branch / condition."""

    def __init__(
        self,
        system: EqualitySystem,
        *,
        branch: int | Literal["nearest", "all"] = "nearest",
        max_iter: int = 20,
        tol: float = 1e-12,
        require_transversal: bool = True,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.system = system
        self.branch_mode = branch
        self.max_iter = int(max_iter)
        self.tol = float(tol)
        self.require_transversal = bool(require_transversal)
        self._dtype = torch.get_default_dtype() if dtype is None else dtype

    def forward(self, x0: Tensor) -> LocusOutput:
        result = newton_result(self.system, x0, max_iter=self.max_iter, tol=self.tol)
        point = newton_project(self.system, x0, max_iter=self.max_iter, tol=self.tol)
        sigma_min = float(result.condition)
        cond = 1.0 / max(sigma_min, 1e-30)
        ok = bool(result.converged)
        if self.require_transversal:
            ok = ok and bool(result.transversal)
        return LocusOutput(
            point=point,
            branch=torch.tensor(result.branch, dtype=torch.int64, device=x0.device),
            condition=torch.tensor(cond, dtype=x0.dtype, device=x0.device),
            converged=torch.tensor(ok, dtype=torch.bool, device=x0.device),
        )

Other public import paths

omnibias.fields.locus.EqualityLocusLayer
omnibias.fields.locus.torch.EqualityLocusLayer

CLASS

LocusOutput

omnibias.fields.locus.LocusOutputSource ↗

class LocusOutput

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.

Required diagnostics. The locus is a constraint manifold, not a PDE solver.

Parameters

Name	Type & default	Contract
`point`dataclass field	`Tensor`Required	See the operation contract and implementation below.
`branch`dataclass field	`Tensor`Required	See the operation contract and implementation below.
`condition`dataclass field	`Tensor`Required	See the operation contract and implementation below.
`converged`dataclass field	`Tensor`Required	See the operation contract and implementation below.

Fields & methods

pointfield

point: Tensor

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

View source ↗

branchfield

branch: Tensor

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

View source ↗

conditionfield

condition: Tensor

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

View source ↗

convergedfield

converged: Tensor

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

View source ↗

Read the implementation

class LocusOutput:
    """Required diagnostics. The locus is a constraint manifold, not a PDE solver."""

    point: Tensor
    branch: Tensor
    condition: Tensor
    converged: Tensor

Other public import paths

omnibias.fields.locus.LocusOutput
omnibias.fields.locus.torch.LocusOutput

FUNCTION

jacobian

omnibias.fields.locus.torch.jacobianSource ↗

jacobian(sys: EqualitySystem, x: Tensor) -> Tensor

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
`sys`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`x`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def jacobian(sys: EqualitySystem, x: Tensor) -> Tensor:
    return torch.tensor(jacobian_core(sys, _as_tuple(x)), dtype=x.dtype, device=x.device)

FUNCTION

locus_tangent_tensor

omnibias.fields.locus.torch.locus_tangent_tensorSource ↗

locus_tangent_tensor(sys: EqualitySystem, x: Tensor) -> Tensor

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
`sys`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`x`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def locus_tangent_tensor(sys: EqualitySystem, x: Tensor) -> Tensor:
    from omnibias.core.locus import locus_tangent as _lt

    basis = _lt(sys, _as_tuple(x))
    if not basis:
        return torch.zeros((0, sys.dim), dtype=x.dtype, device=x.device)
    return torch.tensor(basis, dtype=x.dtype, device=x.device)

FUNCTION

newton_project

omnibias.fields.locus.torch.newton_projectSource ↗

newton_project(sys: EqualitySystem, x0: Tensor, *, weights: Tensor | None=None, max_iter: int=20, tol: float=1e-12) -> Tensor

Gauss-Newton projection. Differentiable through the IFT, not the loop.

Parameters

Name	Type & default	Contract
`sys`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`x0`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.
`weights`keyword-only	`Tensor \| None`Default: None	See the operation contract and implementation below.
`max_iter`keyword-only	`int`Default: 20	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.

ReturnsTensor

Executed regression · g1 forward residual

Full file, imports & fixtures ↗

def test_g1_forward_residual() -> None:
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    layer = EqualityLocusLayer(sys, dtype=torch.float64)
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out = layer(x0)
    assert bool(out.converged.detach())
    nrm = abs(residual(sys, tuple(float(v) for v in out.point.detach().tolist()))[0])
    assert nrm <= 1e-12
    assert out.branch.numel() >= 1
    assert float(out.condition.detach()) >= 1.0

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g1_forward_residual" -q

Regression example · g5 ift matches unrolled

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_ift_matches_unrolled() -> None:
    sys = _case_b()
    x0 = torch.tensor([0.0, 0.20], dtype=torch.float64)
    w = torch.tensor([1.0, -2.0], dtype=torch.float64, requires_grad=True)
    x_ift = newton_project(sys, x0, weights=w, max_iter=8, tol=1e-14)
    loss = x_ift[1]
    loss.backward()
    g_ift = w.grad.detach().clone()

    w2 = torch.tensor([1.0, -2.0], dtype=torch.float64, requires_grad=True)
    x_un = newton_project_unrolled(sys, x0, weights=w2, max_iter=8, tol=1e-14)
    x_un[1].backward()
    g_un = w2.grad.detach().clone()
    rel = torch.norm(g_ift - g_un) / torch.norm(g_un).clamp_min(1e-30)
    assert float(rel.detach()) <= 1e-8

python -m pytest packages/omnibias-fields/tests/test_locus.py -k test_g5_ift_matches_unrolled -q

Read the implementation

def newton_project(
    sys: EqualitySystem,
    x0: Tensor,
    *,
    weights: Tensor | None = None,
    max_iter: int = 20,
    tol: float = 1e-12,
) -> Tensor:
    """Gauss-Newton projection. Differentiable through the IFT, not the loop."""
    if weights is None:
        w = torch.tensor(
            [t.weight for t in sys.terms],
            dtype=x0.dtype,
            device=x0.device,
        )
    else:
        w = weights
    return cast(Tensor, _NewtonIFT.apply(x0, w, sys, max_iter, tol))

FUNCTION

newton_project_unrolled

omnibias.fields.locus.torch.newton_project_unrolledSource ↗

newton_project_unrolled(sys: EqualitySystem, x0: Tensor, *, weights: Tensor, max_iter: int=20, tol: float=1e-12) -> Tensor

Unrolled Gauss-Newton kept on the autodiff graph (G5 comparison only).

Implemented for codimension 1 (the IFT-vs-unrolled gate). Higher codimension should use newton_project.

Parameters

Name	Type & default	Contract
`sys`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`x0`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.
`weights`keyword-only	`Tensor`Required	See the operation contract and implementation below.
`max_iter`keyword-only	`int`Default: 20	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

When sys.codimension != 1, raises ValueError('newton_project_unrolled is the G5 probe for m=1').

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

Executed regression · g5 ift matches unrolled

Full file, imports & fixtures ↗

def test_g5_ift_matches_unrolled() -> None:
    sys = _case_b()
    x0 = torch.tensor([0.0, 0.20], dtype=torch.float64)
    w = torch.tensor([1.0, -2.0], dtype=torch.float64, requires_grad=True)
    x_ift = newton_project(sys, x0, weights=w, max_iter=8, tol=1e-14)
    loss = x_ift[1]
    loss.backward()
    g_ift = w.grad.detach().clone()

    w2 = torch.tensor([1.0, -2.0], dtype=torch.float64, requires_grad=True)
    x_un = newton_project_unrolled(sys, x0, weights=w2, max_iter=8, tol=1e-14)
    x_un[1].backward()
    g_un = w2.grad.detach().clone()
    rel = torch.norm(g_ift - g_un) / torch.norm(g_un).clamp_min(1e-30)
    assert float(rel.detach()) <= 1e-8

python -m pytest "packages/omnibias-fields/tests/test_locus.py::test_g5_ift_matches_unrolled" -q

Read the implementation

def newton_project_unrolled(
    sys: EqualitySystem,
    x0: Tensor,
    *,
    weights: Tensor,
    max_iter: int = 20,
    tol: float = 1e-12,
) -> Tensor:
    """Unrolled Gauss-Newton kept on the autodiff graph (G5 comparison only).

    Implemented for codimension 1 (the IFT-vs-unrolled gate). Higher
    codimension should use :func:`newton_project`.
    """
    if sys.codimension != 1:
        raise ValueError("newton_project_unrolled is the G5 probe for m=1")
    from omnibias.torch.activations import get_activation

    act = get_activation(sys.base)
    x = x0.clone()
    normals = [
        torch.tensor(t.normal, dtype=x.dtype, device=x.device) for t in sys.terms
    ]
    for _ in range(max_iter):
        zs = []
        for t, nrm in zip(sys.terms, normals, strict=True):
            zs.append((x * nrm).sum() + t.bias)
        sigmas = [cast(Tensor, act.fastpath(z, t.order)) for t, z in zip(sys.terms, zs, strict=True)]
        sigmas_p = [
            cast(Tensor, act.fastpath(z, t.order + 1)) for t, z in zip(sys.terms, zs, strict=True)
        ]
        f = weights[0] * sigmas[0] - weights[1] * sigmas[1]
        if float(f.detach().abs()) <= tol:
            return x
        df = weights[0] * sigmas_p[0] * normals[0] - weights[1] * sigmas_p[1] * normals[1]
        gram = (df * df).sum()
        x = x - df * (f / gram)
    return x

FUNCTION

newton_result

omnibias.fields.locus.torch.newton_resultSource ↗

newton_result(sys: EqualitySystem, x0: Tensor, *, max_iter: int=20, tol: float=1e-12) -> NewtonResult

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
`sys`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`x0`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.
`max_iter`keyword-only	`int`Default: 20	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.

ReturnsNewtonResult

Executed regression · g1 forward residual

Full file, imports & fixtures ↗

def test_g1_forward_residual() -> None:
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    layer = EqualityLocusLayer(sys, dtype=torch.float64)
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out = layer(x0)
    assert bool(out.converged.detach())
    nrm = abs(residual(sys, tuple(float(v) for v in out.point.detach().tolist()))[0])
    assert nrm <= 1e-12
    assert out.branch.numel() >= 1
    assert float(out.condition.detach()) >= 1.0

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g1_forward_residual" -q

Read the implementation

def newton_result(
    sys: EqualitySystem, x0: Tensor, *, max_iter: int = 20, tol: float = 1e-12
) -> NewtonResult:
    return newton_project_core(sys, _as_tuple(x0), max_iter=max_iter, tol=tol)

FUNCTION

residual

omnibias.fields.locus.torch.residualSource ↗

residual(sys: EqualitySystem, x: Tensor) -> Tensor

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
`sys`positional-or-keyword	`EqualitySystem`Required	See the operation contract and implementation below.
`x`positional-or-keyword	`Tensor`Required	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def residual(sys: EqualitySystem, x: Tensor) -> Tensor:
    return torch.tensor(residual_core(sys, _as_tuple(x)), dtype=x.dtype, device=x.device)

← Package overview Runnable learning examples →

omnibias.fields.locus.torch

AnsatzSolutionField​

Parameters

Fields & methods

Parameters

Parameters

EqualityLocusLayer​

Parameters

Fields & methods

Parameters

Parameters

LocusOutput​

Parameters

Fields & methods

jacobian​

Parameters

locus_tangent_tensor​

Parameters

newton_project​

Parameters

newton_project_unrolled​

Parameters

newton_result​

Parameters

residual​

Parameters

AnsatzSolutionField

EqualityLocusLayer

LocusOutput

jacobian

locus_tangent_tensor

newton_project

newton_project_unrolled

newton_result

residual