NEML2 2.1.0
Loading...
Searching...
No Matches
Public Member Functions | Public Attributes | Static Public Attributes | Protected Member Functions | Protected Attributes | List of all members
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:
Inheritance graph
[legend]

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)
 __init__ (self, neml2.core.NonlinearSystem sys, *args, list[str] exclude_parameters=list(), **kwargs)
tuple[torch.Tensor, torch.Tensor] forward (self, torch.Tensor state, torch.Tensor forces)
int nforce (self)
int nstate (self)

Public Attributes

 sys = sys
 model = sys.model()
list parameter_names = []
 smap = self.sys.umap()
 slayout = self.sys.ulayout()
list fmap = [v for v in self.sys.gmap() if VariableName(v).start_with(FORCES)]
list flayout

Static Public Attributes

str FORCES = 'forces'
str OLD_FORCES = 'old_forces'
str OLD_STATE = 'old_state'
str PARAMETERS = 'parameters'
str RESIDUAL = 'residual'
str STATE = 'state'
list subaxis_names = ['state', 'old_state', 'forces', 'old_forces', 'residual', 'parameters']

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[neml2.Tensor, neml2.Tensor, neml2.Tensor] _assemble (self, list[neml2.Tensor] A, list[neml2.Tensor] B, list[neml2.Tensor] b)
tuple[torch.Tensor, torch.Tensor] _adapt_for_pyzag (self, neml2.Tensor r, neml2.Tensor J, neml2.Tensor Jn)
tuple[torch.Tensor, torch.Tensor] _adapt_for_pyzag (self, neml2.tensors.Tensor r, neml2.tensors.Tensor J, neml2.tensors.Tensor Jn)
tuple[neml2.tensors.Tensor, neml2.tensors.Tensor, neml2.tensors.Tensor] _assemble (self, list[neml2.tensors.Tensor] A, list[neml2.tensors.Tensor] B, list[neml2.tensors.Tensor] b)
 _check_model (self)
 _setup_maps (self)
 _setup_parameters (self, list[str] exclude_parameters)
 _update_parameter_values (self)
 _update_sys (self, torch.Tensor state, torch.Tensor forces)

Protected Attributes

int _lookback = 1
 _sn_to_g_map = _index_map_from_to(snmap, self.sys.gmap())
 _f_to_g_map = _index_map_from_to(self.fmap, self.sys.gmap())
 _fn_to_g_map = _index_map_from_to(fn_map, self.sys.gmap())
 _g_to_sn_map = _index_map_from_to(self.sys.gmap(), snmap)

Constructor & Destructor Documentation

◆ __init__() [1/2]

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

◆ __init__() [2/2]

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

Member Function Documentation

◆ _adapt_for_pyzag() [1/2]

tuple[torch.Tensor, torch.Tensor] _adapt_for_pyzag ( self,
neml2.Tensor r,
neml2.Tensor J,
neml2.Tensor Jn )
protected 
Adapt the residual and Jacobians for pyzag

pyzag has additional requirements on residual and Jacobians:
  1. The residual and Jacobians should have the same batch shape
  2. The Jacobians should be square

The second requirement has been taken care of by the previous _assemble call.

Args:
    r (neml2.Tensor): residual
    J (neml2.Tensor): Jacobian
    Jn (neml2.Tensor): Jacobian for the old state

Returns:
    tuple of torch.Tensor: residual, Jacobian

◆ _adapt_for_pyzag() [2/2]

tuple[torch.Tensor, torch.Tensor] _adapt_for_pyzag ( self,
neml2.tensors.Tensor r,
neml2.tensors.Tensor J,
neml2.tensors.Tensor Jn )
protected 
Adapt the residual and Jacobians for pyzag

        pyzag has additional requirements on residual and Jacobians:
          1. The residual and Jacobians should have the same batch shape
          2. The Jacobians should be square

        The second requirement has been taken care of by the previous _assemble call.

        Args:
            r (neml2.Tensor): residual
            J (neml2.Tensor): Jacobian
            Jn (neml2.Tensor): Jacobian for the old state

        Returns:
            tuple of torch.Tensor: residual, Jacobian

◆ _assemble() [1/2]

tuple[neml2.Tensor, neml2.Tensor, neml2.Tensor] _assemble ( self,
list[neml2.Tensor] A,
list[neml2.Tensor] B,
list[neml2.Tensor] b )
protected 
Assemble the residual and Jacobians from A, B, and b

Args:
    A (list of neml2.Tensor): list of Jacobians w.r.t. unknowns
    B (list of neml2.Tensor): list of Jacobians w.r.t. given variables
    b (list of neml2.Tensor): list of (negative) residuals

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

◆ _assemble() [2/2]

tuple[neml2.tensors.Tensor, neml2.tensors.Tensor, neml2.tensors.Tensor] _assemble ( self,
list[neml2.tensors.Tensor] A,
list[neml2.tensors.Tensor] B,
list[neml2.tensors.Tensor] b )
protected 
Assemble the residual and Jacobians from A, B, and b

        Args:
            A (list of neml2.Tensor): list of Jacobians w.r.t. unknowns
            B (list of neml2.Tensor): list of Jacobians w.r.t. given variables
            b (list of neml2.Tensor): list of (negative) residuals

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

◆ _check_model() [1/2]

_check_model ( self)
protected 
Simple consistency checks, could be a debug check but we only call this once

◆ _check_model() [2/2]

_check_model ( self)
protected 
Simple consistency checks, could be a debug check but we only call this once

◆ _setup_maps() [1/2]

_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_maps() [2/2]

_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() [1/2]

_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

◆ _setup_parameters() [2/2]

_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() [1/2]

_update_parameter_values ( self)
protected 
Copy over new parameter values

◆ _update_parameter_values() [2/2]

_update_parameter_values ( self)
protected 
Copy over new parameter values

◆ _update_sys() [1/2]

_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

◆ _update_sys() [2/2]

_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() [1/2]

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

◆ forward() [2/2]

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() [1/2]

int nforce ( self)

◆ nforce() [2/2]

int nforce ( self)

◆ nstate() [1/2]

int nstate ( self)

◆ nstate() [2/2]

int nstate ( self)

Member Data Documentation

◆ _f_to_g_map

_f_to_g_map = _index_map_from_to(self.fmap, self.sys.gmap())
protected

◆ _fn_to_g_map

_fn_to_g_map = _index_map_from_to(fn_map, self.sys.gmap())
protected

◆ _g_to_sn_map

_g_to_sn_map = _index_map_from_to(self.sys.gmap(), snmap)
protected

◆ _lookback

int _lookback = 1
protected

◆ _sn_to_g_map

_sn_to_g_map = _index_map_from_to(snmap, self.sys.gmap())
protected

◆ flayout

list flayout
Initial value:
= [
self.sys.glayout()[i]
for i, v in enumerate(self.sys.gmap())
if VariableName(v).start_with(FORCES)
]

◆ fmap

list fmap = [v for v in self.sys.gmap() if VariableName(v).start_with(FORCES)]

◆ FORCES

str FORCES = 'forces'
static

◆ model

model = sys.model()

◆ OLD_FORCES

str OLD_FORCES = 'old_forces'
static

◆ OLD_STATE

str OLD_STATE = 'old_state'
static

◆ parameter_names

list parameter_names = []

◆ PARAMETERS

str PARAMETERS = 'parameters'
static

◆ RESIDUAL

str RESIDUAL = 'residual'
static

◆ slayout

slayout = self.sys.ulayout()

◆ smap

smap = self.sys.umap()

◆ STATE

str STATE = 'state'
static

◆ subaxis_names

list subaxis_names = ['state', 'old_state', 'forces', 'old_forces', 'residual', 'parameters']
static

◆ sys

sys = sys