NEML2PyzagModel Class Reference

Detailed Description

Wraps a NEML2 nonlinear system into a `nonlinear.NonlinearRecursiveFunction`

Args:
    sys: the NEML2 nonlinear system to wrap

Keyword Args:
    exclude_parameters (list of str): exclude these parameters from being wrapped as a pytorch parameter

Additional args and kwargs are forwarded to NonlinearRecursiveFunction (and hence torch.nn.Module) verbatim

Inheritance diagram for NEML2PyzagModel:

Public Member Functions
	__init__ (self, neml2.NonlinearSystem sys, *args, list[str] exclude_parameters=[], **kwargs)
int	lookback (self)
	lookback (self, int lookback)
int	nstate (self)
int	nforce (self)
tuple[torch.Tensor, torch.Tensor]	forward (self, torch.Tensor state, torch.Tensor forces)

Public Attributes
	sys = sys
	model = sys.model()
list	parameter_names = []
	rlayout = self.sys.blayout()
	slayout = self.sys.ulayout()
	snlayout = _extract_sublayout(self.slayout, 0, 1)
	flayout = _extract_sublayout(glayout, 0)
	fnlayout = _extract_sublayout(glayout, 0, 1)
	svars = self.slayout.vars()
	fvars = self.flayout.vars()

Protected Member Functions
	_check_model (self)
	_setup_parameters (self, list[str] exclude_parameters)
	_update_parameter_values (self)
	_setup_maps (self)
	_update_sys (self, torch.Tensor state, torch.Tensor forces)
tuple[torch.Tensor, torch.Tensor]	_adapt_for_pyzag (self, AssembledMatrix A, AssembledMatrix B, AssembledVector b)

Protected Attributes
int	_lookback = 1
list	_sn_to_g_map = [-1] * self.snlayout.nvar()

Constructor & Destructor Documentation

__init__()

__init__	(		self,
		neml2.NonlinearSystem	sys,
		*	args,
		list[str]	exclude_parameters = [],
		**	kwargs )

Member Function Documentation

_adapt_for_pyzag()

tuple[torch.Tensor, torch.Tensor] _adapt_for_pyzag ( self, AssembledMatrix A, AssembledMatrix B, AssembledVector b )

protected

Assemble the residual and Jacobians from A, B, and b

Args:
    A (AssembledMatrix): Jacobians w.r.t. unknowns
    B (AssembledMatrix): Jacobians w.r.t. given variables
    b (AssembledVector): (negative) residuals

Returns:
    tuple of neml2.Tensor: residual, Jacobian w.r.t. unknowns, Jacobian w.r.t. old state

_check_model()

_check_model ( self )

protected

Simple consistency checks, could be a debug check but we only call this once

_setup_maps()

_setup_maps ( self )

protected

Setup the maps for assembly purposes

In the C++ backend, the nonlinear system distinguishes between unknowns (u) and given variables (g).
However, pyzag expects contiguous-in-time representation of the state and forces, where unknowns and old
unknowns come from the state tensor, and forces and old forces come from the forces tensor.

Note that the given variables include old unknowns, forces, and old forces. So we need to setup maps from
the pyzag representation to the nonlinear system representation for both unknowns and given variables.

_setup_parameters()

_setup_parameters ( self, list[str] exclude_parameters )

protected

Mirror parameters of the NEML2 model with torch.nn.Parameter

Args:
    exclude_parameters (list of str): NEML2 parameters to exclude

_update_parameter_values()

_update_parameter_values ( self )

protected

Copy over new parameter values

_update_sys()

_update_sys ( self, torch.Tensor state, torch.Tensor forces )

protected

Disassemble the model input forces, old forces, and old state from
the flat tensors and update them in the nonlinear system

Args:
    state (torch.Tensor): tensor containing the model state
    forces (torch.Tensor): tensor containing the model forces

forward()

tuple[torch.Tensor, torch.Tensor] forward	(		self,
		torch.Tensor	state,
		torch.Tensor	forces )

Actually call the NEML2 model and return the residual and Jacobian

Args:
    state (torch.Tensor): tensor with the flattened state
    forces (torch.Tensor): tensor with the flattened forces

lookback() [1/2]

int lookback

(

self

)

lookback() [2/2]

lookback	(		self,
		int	lookback )

nforce()

int nforce

(

self

)

nstate()

int nstate

(

self

)

Member Data Documentation

_lookback

int _lookback = 1

protected

_sn_to_g_map

list _sn_to_g_map = [-1] * self.snlayout.nvar()

protected

flayout

flayout = _extract_sublayout(glayout, 0)

fnlayout

fnlayout = _extract_sublayout(glayout, 0, 1)

fvars

fvars = self.flayout.vars()

model

model = sys.model()

parameter_names

list parameter_names = []

rlayout

rlayout = self.sys.blayout()

slayout

slayout = self.sys.ulayout()

snlayout

snlayout = _extract_sublayout(self.slayout, 0, 1)

svars

svars = self.slayout.vars()

sys

sys = sys