neml2.pyzag

class neml2.pyzag.NEML2PyzagModel(sys, *args, exclude_parameters=None, include_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)

• include_parameters (list[str] | None)

Keyword Arguments:

• exclude_parameters (list of str) – NEML2 parameters to not mirror as torch parameters on the wrapper (so every other parameter is trained). Mutually exclusive with include_parameters.

• include_parameters (list of str) – the only NEML2 parameters to mirror as torch parameters on the wrapper (so every other parameter is held fixed). Mutually exclusive with exclude_parameters.

exclude_parameters and include_parameters are both optional and mutually exclusive: omit both to train every parameter, pass exclude_parameters to train all-but-some, or pass include_parameters to train only-these. Passing both raises ValueError.

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 (base_name, lag).

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

Parameters:

var (str)

Return type:

tuple[str, int]