Model

Table of Contents

• Model definition
• Model composition
  • A symbolic example
  • Automatic dependency registration
  • Automatic input/output identification
  • Automatic dependency resolution
  • Automatic chain rule
• Automatic differentiation
• Model assembly
  • Axis labeling
  • Tensor assembly

Refer to Syntax Documentation for the list of available objects.

Model definition

A NEML2 model is a function (in the context of mathematics)

$$f:{\mathbb{R}}^m\to {\mathbb{R}}^n$$

mapping from the input space ${\mathbb{R}}^m$ of dimension $m$ to the output space ${\mathbb{R}}^n$ of dimension $n$. $[\cdot ]$ be the flatten-concatenation operator, the input vector is the concatenation of $p$ flattened variables, i.e.,

$$x={\left[x_i\right]}_{i=1}^p\in {\mathbb{R}}^m,\sum _{i=1}^p|x_i|=m,$$

where $|x|$ denotes the modulus of flattened variable $x$. Similarly, the output vector is the concatenation of $q$ flattened variables, i.e.,

$$y={\left[y_i\right]}_{i=1}^q\in {\mathbb{R}}^n,\sum _{i=1}^q|y_i|=n.$$

Translating the above mathematical definition into NEML2 is straightforward.

• A model following this definition derives from Model.
• declare_input_variable declares an input variable $x_i$ in the input space ${\mathbb{R}}^m$.
• declare_output_variable declares an output variable $y_i$ in the output space ${\mathbb{R}}^n$.
• set_value is a method defining the forward operator $f$ itself.

Both declare_input_variable and declare_output_variable are templated on the variable type – recall that only a variable of the NEML2 primitive tensor type can be registered. Furthermore, both methods return a Variable<T> & used for retrieving and setting the variable value inside the forward operator, i.e. set_value. Note that the reference returned by declare_input_variable is writable, while the reference returned by declare_output_variable is read-only.

Model composition

Quoting Wikipedia:

In mathematics, function composition is an operation $\circ$ that takes two functions $f$ and $g$, and produces a function $h=g\circ f$ such that $h(x)=g(f(x))$.

Since NEML2 Model is a function (in the context of mathematics) at its core, it should be possible, in theory, to compose different NEML2 Models into a new NEML2 Model. The ComposedModel is precisely for that purpose.

Similar to the statement "a composed function is a function" in the context of mathematics, in NEML2, the equivalent statement "a `ComposedModel` is a `Model`" also holds. In addition, the ComposedModel provides four key features to help simplify the composition and reduces computational cost:

• Automatic dependency registration
• Automatic input/output identification
• Automatic dependency resolution
• Automatic chain rule

A symbolic example

To demonstrate the utility of the four key features of ComposedModel, let us consider the composition of three functions $f$, $g$, and $h$:

$$\begin{array}{rl}{y}_1& =f(x_1,x_2),\\ y_2& =g(y_1,x_3),\\ y& =h(y_1,y_2,x_4).\end{array}$$

Automatic dependency registration

It is obvious to us that the function $h$ depends on functions $f$ and $g$ because the input of $h$ depends on the outputs of $f$ and $g$. Such dependency is automatically identified and registered while composing a ComposedModel in NEML2. This procedure is called "automatic dependency registration".

In order to identify dependencies among different Models, we keep track of the set of consumed variables, ${\mathcal{I}}_i$, and a set of provided variables, ${\mathcal{O}}_i$, for each Model $f_i$. When a set of models (functions) are composed together, Model $f_i$ is said to depend on $f_j$ if and only if $\mathrm{\exists }x$ such that

$$x\in {\mathcal{I}}_i\wedge x\in {\mathcal{O}}_j.$$

Automatic input/output identification

The only possible composition $r$ of these three functions is

$$y=r(x_1,x_2,x_3,x_4):=h(f(x_1,x_2),g(f(x_1,x_2),x_3),x_4).$$

The input variables of the composed function $r$ are $[x_1,x_2,x_3,x_4]$ (or their flattened concatenation), and the output variable of the composed function is simply $y$. The input/output variables are automatically identified while composing a ComposedModel in NEML2. This procedure is referred to as "automatic input/output identification".

In a ComposedModel, a leaf model is a model which does not depend on any other model, and a root model is a model which is not depent upon by any other model. A ComposedModel may have multiple leaf models and multiple root models. An input variable is said to be a root input variable if it is not consumed by any other model, i.e. $x\in {\mathcal{I}}_i$ is a root input variable if and only if

$$x\notin {\mathcal{O}}_j,\mathrm{\forall }i\ne j.$$

Similarly, an output variable is said to be a leaf output variable if it is not provided by any other model, i.e. $x\in {\mathcal{O}}_i$ is a leaf output variable if an only if

$$x\notin {\mathcal{I}}_j,\mathrm{\forall }i\ne j.$$

The input variables of a ComposedModel is the union of the set of all the root input variables, and the output variables of a ComposedModel is the set of all the leaf output variables.

Automatic dependency resolution

To evaluate the forward operator of the composed model $r$, one has to first evaluate model $f$, then model $g$, and finally model $h$. The process of sorting out such evaluation order is called "dependency resolution".

While it is possible to sort the evaluation order "by hand" for this simple example composition, it is generally not a trivial task for practical compositions with more involved dependencies. To that end, NEML2 uses topological sort to sort the model evaluation order, such that by the time each model is evaluated, all of its dependent models have already been evaluated.

Automatic chain rule

Chain rule can be applied to evaluate the derivative of the forward operator with respect to the input variables, i.e.,

$$\begin{array}{rl}\frac{\partial y}{\partial x_1}& =(\frac{\partial y}{\partial y_1}+\frac{\partial y}{\partial y_2}\frac{\partial y_2}{\partial y_1})\frac{\partial y_1}{\partial x_1},\\ \frac{\partial y}{\partial x_2}& =(\frac{\partial y}{\partial y_1}+\frac{\partial y}{\partial y_2}\frac{\partial y_2}{\partial y_1})\frac{\partial y_1}{\partial x_2},\\ \frac{\partial y}{\partial x_3}& =\frac{\partial y}{\partial y_2}\frac{\partial y_2}{\partial x_3},\\ \frac{\partial y}{\partial x_4}& =\frac{\partial y}{\partial x_4}.\end{array}$$

Spelling out this chain rule can be cumbersome and error-prone, especially for more complicated model compositions. The evaluation of the chain rule is automated in NEML2, and the user is only responsible for implementing the partial derivatives of each model. For example, in the implementation of Model $f$, the user only need to define the partial derivatives

$$\frac{\partial y_1}{\partial x_1},\frac{\partial y_1}{\partial x_2};$$

similarly, Model $g$ only defines

$$\frac{\partial y_2}{\partial y_1},\frac{\partial y_2}{\partial x_3}$$

and Model $h$ only defines

$$\frac{\partial y}{\partial y_1},\frac{\partial y}{\partial y_2},\frac{\partial y}{\partial x_4}.$$

The assembly of the partial derivatives into the total derivative $\partial y/\partial \mathit{x}$ using the chain rule is handled by NEML2. This design serves as the fundation for a modular model implementation:

• Each model does not need to know its composition with others.
• The same model partial derivatives can be reused in any composition.

Automatic differentiation

Deriving and implementing derivatives of the forward operator can be cumbersome from times to times. NEML2 offers the option to use automatic differentiation (AD) to obtain derivatives. To enable automatic differentiation, one simply need to override the neml2::Model::request_AD method and specify which derivatives should be computed using AD:

void
MyModel::request_AD()
{
  std::vector<const VariableBase *> inputs = {&foo, &bar, &baz, &T};
 
  // First derivatives
  foo_dot.request_AD(inputs);
  bar_dot.request_AD(inputs);
  baz_dot.request_AD(inputs);
 
  // Second derivatives
  foo_dot.request_AD(inputs, inputs);
  bar_dot.request_AD(inputs, inputs);
  baz_dot.request_AD(inputs, inputs);
}

Note

Each model can use a mix of hand-coded derivatives and AD derivatives. However, an error will be raised if hand-coded derivatives are provided for those marked by neml2::Variable::request_AD.

Since a composed model uses chain rule to efficiently evaluate the total derivatives, automatic differentiation is disabled for ComposedModel. However, each of the child model can still use AD to calculate the partial derivatives of its own forward operator. Moreover, AD and non-AD models can be composed together.

Model assembly

NEML2 stores each variable in contiguous memory, but does not guarantee contiguity across variables. This choice is made to allow for massive asynchronous evaluation (with the help of lazy tensors) and to reduce memory consumption (since variables can have different number of batch shapes). However, this choice is not ideal for a family of nonlinear material models whose constitutive updates require solving one (or more) implicit system of equations. To address such issue, NEML2 offers two mechanisms to facilitate the creation of the implicit system (e.g., its residual and Jacobian):

• Axis labeling for setting up the layout of the implicit system
• Tensor assembly for assembling and disassembling the implicit system

Axis labeling

NEML2 provides a data structure named LabeledAxis to create a contiguous layout for scattered input/output variables. Typically, each model contains an input axis for input variables and an output axis for output variables.

The LabeledAxis contains all information regarding how the variables of interest should be contiguously laid out. In other words, the labeled axis maintains the mapping between variables and their contiguous slice along an axis. The following naming convention is used:

• Item: A labelable slice of data
• Variable: An item that is also of a NEML2 primitive tensor type
• Sub-axis: An item of type LabeledAxis

An axis can be labeled recursively, e.g.,

     0 1 2 3 4 5     6     7 8 9 10 11 12   13   14
/// |-----------| |-----| |              | |  | |  |
///       a          b    |              | |  | |  |
/// |-------------------| |--------------| |--| |--|
///          sub                  a          b    c

The above example represents an axis of size 15. This axis has 4 items: a, b, c, and sub.

• "a" is a variable of storage size 6 (possibly of type SR2).
• "b" is a variable of type Scalar.
• "c" is a variable of type Scalar.
• "sub" is a sub-axis of type LabeledAxis. "sub" by itself represents an axis of size 7, containing 2 items:
  • "a" is a variable of storage size 6.
  • "b" is a variable of type Scalar.

Duplicate labels are not allowed on the same level of the axis, e.g. "a", "b", "c", and "sub" share the same level and so must be different. However, items on different levels of an axis can share the same label, e.g., "a" on the sub-axis "sub" has the same label as "a" on the main axis. In NEML2 convention, item names are always fully qualified, and a sub-axis is prefixed with a left slash, e.g. item "b" on the sub-axis "sub" can be denoted as "sub/b" on the main axis.

Note

A label cannot contain: white spaces, quotes, left slash (/), or new line.

Due to performance considerations, a LabeledAxis can only be modified, e.g., adding/removing variables and sub-axis, at the time a model is constructed. After the model construction phase, the LabeledAxis associated with that model can no longer be modified over the entire course of the simulation.

Refer to the documentation for a complete list of APIs for creating and modifying a LabeledAxis.

Tensor assembly

NEML2 implements two types of "assemblers" to assemble (or split) the implicit system given the axis layout defined by LabeledAxis:

• VectorAssembler: Assemble a map of vectors into a single vector (neml2::VectorAssembler::assemble_by_variable), or split a single vector into a map of vectors (neml2::VectorAssembler::split_by_variable).
• MatrixAssembler: Assemble a map of map of matrices into a single matrix (neml2::MatrixAssembler::assemble_by_variable), or split a single matrix into a map of map of matrices (neml2::MatrixAssembler::split_by_variable).

The assemble_by_variable methods take a map (1D map for the vector assembler and 2D map for the matrix assembler) as input argument. The keys of the map are variable names.

Note

Variable values not provided by the map are filled with zeros.

The VectorAssembler is useful for working with the residual and solution vectors of the implicit system, and the MatrixAssembler is primarily used to work with the Jacobian matrix of the implicit system.

In addition to the assemble_by_variable and split_by_variable methods, the assemblers also provide a third method called split_by_subaxis. The split_by_subaxis method is similar to split_by_variable, but it splits the tensor by subaxes instead of variables.