neml2.pyzag

class neml2.pyzag.NEML2PyzagModel(sys, *args, exclude_parameters=None, **kwargs)

Bases: NonlinearRecursiveFunction

Wrap a native NEML2 NonlinearSystem as a pyzag NonlinearRecursiveFunction.

Parameters:

• sys (ModelNonlinearSystem) – the native NEML2 nonlinear system to wrap.

• exclude_parameters (list[str] | None)

Keyword Arguments:

exclude_parameters (list of str) – NEML2 parameters to not mirror as torch parameters on the wrapper (and therefore not optimize against).

Additional args and kwargs are forwarded to torch.nn.Module verbatim.

forward(state, forces)

Compute the residual and stacked Jacobians for one pyzag block step.

Parameters:

• state (Tensor) – flat tensor with the current+previous state, shape (n_block + lookback, ..., nstate).

• forces (Tensor) – flat tensor with the current+previous forces, shape (n_block + lookback, ..., nforce).

Return type:

tuple[Tensor, Tensor]

property lookback: int

property nforce: int

property nstate: int

neml2.pyzag.change_lag_order(var, new_order)

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)

Split a variable name into its base name and lag order.

var is either "name" (lag 0) or "name~n" (lag n). Used by the interface to identify old-state / old-force variables — pyzag wants a contiguous-in-time state and forces tensor, while the NEML2 nonlinear system keeps current values on the unknown axis and old values on the given axis with a ~1 suffix.

Parameters:

var (str)

Return type:

tuple[str, int]