neml2.pyzag#

Adapter exposing a NEML2 nonlinear system to the pyzag time-integration library.

NEML2PyzagModel wraps a NEML2 ModelNonlinearSystem as a pyzag.nonlinear.NonlinearFunctionOperatorFactory – assembling the per-chunk residual and bidiagonal Jacobian from NEML2’s equation-systems layer and mirroring the model’s HIT parameters as torch.nn.Parameter s for gradient-based calibration. The backing block-operator implementations live in neml2.pyzag.operators.

class neml2.pyzag.NEML2BlockJacobian(diag_am, sub_am, layout, _reversed=False)[source]#

Bases: BlockJacobian

Per-chunk bidiagonal Jacobian wrapping neml2 AssembledMatrix diag/sub.

Parameters:

diag_am (AssembledMatrix)
sub_am (AssembledMatrix)
_reversed (bool)

adjoint_system(inverse_operator)[source]#

Build the transposed, time-reversed bidiagonal operator for the adjoint sweep.

Uses the transposed diagonal/subdiagonal blocks in reverse time order. Must be called on the Jacobian returned by as_adjoint_walk().

as_adjoint_walk()[source]#

Return a view of this Jacobian marked for the reverse-time adjoint walk.

Return type:: NEML2BlockJacobian

property batch_size: int#: Size of the plain (non-dynamic) batch axis.

property block_size: int#: Total per-step degrees of freedom of the state layout.

couple_prev_chunk(a_first)[source]#

Propagate the adjoint across the chunk boundary via the subdiagonal coupling.

Applies the transpose of this chunk’s boundary subdiagonal block to the first adjoint variable a_first, giving the contribution into the previous chunk’s terminal adjoint.

Parameters:: a_first (BlockVector)
Return type:: NEML2BlockVector

property device: device#: Device of the backing tensors.

property dtype: dtype#: Dtype of the backing tensors.

forward_system(inverse_operator)[source]#

Build the forward bidiagonal chunk operator solved during the state sweep.

inverse_operator is the factorization pyzag requests. Must be called on a forward-walk Jacobian, not one from as_adjoint_walk().

property nblk_steps: int#: Number of time steps (blocks) in this chunk.

solve_terminal_adjoint(g_terminal)[source]#

Seed the adjoint recursion from the terminal cost gradient g_terminal.

Solves -A_last^T x = g_terminal on the last step’s transposed diagonal.

Parameters:: g_terminal (Tensor)
Return type:: NEML2BlockVector

class neml2.pyzag.NEML2BlockVector(raw_tensors, layout, intmd_dims=None)[source]#

Bases: BlockVector

Block vector held as per-group torch tensors with neml2 layout metadata.

Per-group torch tensors plus explicit intmd_dims mirror the dense backend pattern and let __setitem__ / clone work in place. A neml2 AssembledVector is materialized only at the neml2 boundary (via to_av()), where the linear solvers and matmul consume it.

Parameters:

raw_tensors (list[torch.Tensor])
layout (AxisLayout)
intmd_dims (list[int] | None)

property batch_size: int#: Size of the plain (non-dynamic) batch axis.

property block_size: int#: Total per-block degrees of freedom summed across all groups (incl. intmd dims).

classmethod cat(vectors, dim=0)[source]#

Concatenate matching-layout block vectors group-by-group along dim.

Parameters:

vectors (Sequence[BlockVector])
dim (int)

Return type:

NEML2BlockVector

clone()[source]#

Deep copy, cloning every backing tensor.

Return type:: NEML2BlockVector

property device: device#: Device of the backing tensors.

property dtype: dtype#: Dtype of the backing tensors.

flatten()[source]#

Concatenate all groups into one (batch, nblk * dofs) tensor.

Each group is transposed to put the batch axis first, then flattened and concatenated along the feature axis – the flat form pyzag’s PCR path expects.

Return type:: Tensor

flip(dim)[source]#

Reverse every group along dim (used to walk time backward in the adjoint).

Parameters:: dim (int)
Return type:: NEML2BlockVector

classmethod from_av(av)[source]#

Construct from a neml2 AssembledVector.

Parameters:: av (AssembledVector)
Return type:: NEML2BlockVector

property nblk: int#: Number of blocks along the dynamic (time) axis.

norm(dim=-1)[source]#

Combined L2 norm over the whole multi-group state, per block and batch.

Parameters:: dim (int)
Return type:: Tensor

scale_batches(factor)[source]#

Scale each plain-batch entry by the matching entry of factor.

Parameters:: factor (Tensor)
Return type:: NEML2BlockVector

to_av()[source]#

Materialize as a neml2 AssembledVector.

Return type:: AssembledVector

where(mask, other)[source]#

Batchwise select: keep self where mask is true, else other.

Parameters:

mask (Tensor)
other (BlockVector)

Return type:

NEML2BlockVector

classmethod zeros_like(other)[source]#

A zero vector with the same layout, shapes, dtype, and device as other.

Parameters:: other (BlockVector)
Return type:: NEML2BlockVector

class neml2.pyzag.NEML2PyzagModel(sys, *args, exclude_parameters=None, include_parameters=None, **kwargs)[source]#

Bases: Module, NonlinearFunctionOperatorFactory

Adapt a NEML2 ModelNonlinearSystem to pyzag’s time-integration library.

Implements pyzag’s NonlinearFunctionOperatorFactory ABC. The public name stays NEML2PyzagModel for user-facing continuity even though the base protocol is now the factory (pyzag 2.0) rather than the old NonlinearRecursiveFunction.

Parameters:

sys (ModelNonlinearSystem) – the native NEML2 nonlinear system to wrap.
exclude_parameters (list[str] | None)
include_parameters (list[str] | None)

Keyword Arguments:

exclude_parameters (list of str) – NEML2 parameters to not mirror as torch parameters. Mutually exclusive with include_parameters.
include_parameters (list of str) – the only NEML2 parameters to mirror as torch parameters. Mutually exclusive with exclude_parameters.

Construction yields a py-eager model. To accelerate the residual with in-process torch.compile (py-jit), opt in explicitly after construction with neml2.compile() – e.g. neml2.compile(model).

assemble_forces(forces_dict, dynamic_dim=2)[source]#

Build a flat forces tensor from a per-variable forces dict.

Derived fresh from the pristine force baselines (see assemble_state()); idempotent and order-independent.

Parameters:: dynamic_dim (int)
Return type:: Tensor

assemble_state(ic_dict, dynamic_dim=1)[source]#

Build a flat initial-state tensor from a per-variable IC dict.

The runtime state / old-state / residual layouts are derived fresh from the pristine baselines captured in _setup_maps(), so this is idempotent and independent of prior assemble_state() / assemble_forces() calls.

Parameters:: dynamic_dim (int)
Return type:: Tensor

evaluate_raw(x_full, forces)[source]#

Return (r, NEML2BlockJacobian) for one chunk’s state and forces.

pyzag passes the driving forces as a list (one tensor per *forces argument to pyzag.nonlinear.solve(), each chunk-sliced). They are concatenated along the feature axis into the single flat forces tensor the system consumes, in the order given – which must match the force (given-variable) layout, exactly as assemble_forces() produces it.

Parameters:

x_full (Tensor)
forces (Sequence[Tensor])

Return type:

tuple[Tensor, NEML2BlockJacobian]

property lookback: int#: Number of previous time steps the residual depends on (always 1 for NEML2).

make_operator(prev_solution, forces, inverse_operator)[source]#

Build the pyzag chunk operator for one solve window.

Parameters:

prev_solution (Tensor) – the converged state at the last step of the previous chunk.
forces (Sequence[Tensor]) – the driving forces for this chunk (one tensor per *forces arg).
inverse_operator – the bidiagonal factorization class/instance pyzag uses to invert the chunk system (Thomas, PCR, …).

Returns:

A pyzag.nonlinear.ChunkOp bound to this factory.

Return type:

ChunkOp

property nforce: int#: Total flattened size of the force (given/driving) layout.

property nstate: int#: Total flattened size of the state (unknown) layout.

property wrapper#: The NEML2Wrapper for the current state layout.

class neml2.pyzag.NEML2SolvableBlockOperator(am)[source]#

Bases: SolvableBlockOperator

Block operator backed by a neml2 AssembledMatrix.

Parameters:: am (AssembledMatrix)

property batch_size: int#: Size of the plain (non-dynamic) batch axis.

clone()[source]#

Deep copy the operator, cloning every backing tensor (cache dropped).

Return type:: NEML2SolvableBlockOperator

property device: device#: Device of the backing tensors.

property dtype: dtype#: Dtype of the backing tensors.

classmethod factored(am)[source]#

Construct the operator and eagerly cache its factorization.

For a single-group diagonal this factors all block-axis entries at once (one batched lu_factor); the cache is carried through __getitem__() so a Thomas sweep reuses it per block via lu_solve instead of re-factoring. Multi-group (Schur) diagonals are not cached here – they delegate to SchurComplement per solve.

Parameters:: am (AssembledMatrix)
Return type:: NEML2SolvableBlockOperator

matvec(x)[source]#

Grain-diagonal matrix-vector product self @ x.

Delegates to per_instance_matvec() (a BLOCK output stays per-site; only a DENSE output consuming a BLOCK input aggregates the site axis). Raw tensors appear only at the final NEML2BlockVector hand-off.

Parameters:: x (BlockVector)
Return type:: NEML2BlockVector

property nblk: int#: Number of blocks along the dynamic (time) axis.

pad_front(n=1)[source]#

Return a copy with n zero blocks prepended along the dynamic axis.

Used to align a subdiagonal operator with the diagonal so block k of the padded result is the coupling into step k.

Parameters:: n (int)
Return type:: NEML2SolvableBlockOperator

pcr_finalize(state)[source]#

Finish the PCR sweep, returning the reduced (operator, rhs) pair.

Parameters:: state (PCRState)

pcr_init(B, v)[source]#

Seed a parallel-cyclic-reduction sweep (single-group dense only).

Delegates to pyzag’s dense backend. Raises NotImplementedError for multi-group / BLOCK layouts – use the Thomas factorization for those.

Parameters:

B (BlockOperator)
v (BlockVector)

Return type:

PCRState

pcr_reduce_level(state, level)[source]#

Advance the PCR sweep by one stride-doubling level (single-group dense).

Parameters:

state (PCRState)
level (int)

solve(rhs)[source]#

Solve self @ x = rhs for the diagonal block.

Single-group layouts use the cached batched LU; a two-group BLOCK+DENSE split delegates to SchurComplement with a CachingLU primary solver. Three or more groups are unsupported.

Parameters:: rhs (BlockVector)
Return type:: NEML2BlockVector

t_matvec(x)[source]#

Transposed grain-diagonal matrix-vector product self.T @ x (see matvec()).

Parameters:: x (BlockVector)
Return type:: NEML2BlockVector

class neml2.pyzag.NEML2Wrapper(layout)[source]#

Bases: object

Convert between pyzag flat torch tensors and neml2-backed block objects.

Parameters:: layout (AxisLayout)

unwrap_vector(bv)[source]#

NEML2BlockVector -> flat torch (..., nflat).

Parameters:: bv (BlockVector)
Return type:: Tensor

wrap_jacobian(diag, sub)[source]#

Wrap diag / sub AssembledMatrix blocks into an NEML2BlockJacobian.

Parameters:

diag (AssembledMatrix)
sub (AssembledMatrix)

Return type:

NEML2BlockJacobian

wrap_vector(raw)[source]#

Flat torch (..., nflat) -> NEML2BlockVector.

Parameters:: raw (Tensor)
Return type:: NEML2BlockVector

neml2.pyzag.change_lag_order(var, new_order)[source]#

Re-tag a variable name to a different lag order. Inverse of lag_order().

Parameters:

var (str)
new_order (int)

Return type:

str

neml2.pyzag.lag_order(var)[source]#

Split a variable name into (base_name, lag).

var is either "name" (lag 0) or "name~n" (lag n).

Parameters:: var (str)
Return type:: tuple[str, int]