NEML2 2.0.0
|
Refer to Syntax Documentation for the list of available objects.
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\). \(\left[ \cdot \right]\) 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, \quad \sum_{i=1}^p \lvert x_i \rvert = m, \]
where \(\lvert x \rvert\) 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, \quad \sum_{i=1}^q \lvert y_i \rvert = n. \]
Translating the above mathematical definition into NEML2 is straightforward.
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.
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 Model
s 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:
To demonstrate the utility of the four key features of ComposedModel
, let us consider the composition of three functions \(f\), \(g\), and \(h\):
\begin{align*} y_1 &= f(x_1, x_2), \\ y_2 &= g(y_1, x_3), \\ y &= h(y_1, y_2, x_4). \end{align*}
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 Model
s, 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 \(\exists x\) such that
\[ x \in \mathcal{I}_i \wedge x \in \mathcal{O}_j. \]
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, \quad \forall i \neq 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, \quad \forall i \neq 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.
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.
Chain rule can be applied to evaluate the derivative of the forward operator with respect to the input variables, i.e.,
\begin{align*} \frac{\partial y}{\partial x_1} &= \left( \frac{\partial y}{\partial y_1} + \frac{\partial y}{\partial y_2} \frac{\partial y_2}{\partial y_1} \right) \frac{\partial y_1}{\partial x_1}, \\ \frac{\partial y}{\partial x_2} &= \left( \frac{\partial y}{\partial y_1} + \frac{\partial y}{\partial y_2} \frac{\partial y_2}{\partial y_1} \right) \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{align*}
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}, \quad \frac{\partial y_1}{\partial x_2}; \]
similarly, Model
\(g\) only defines
\[ \frac{\partial y_2}{\partial y_1}, \quad \frac{\partial y_2}{\partial x_3} \]
and Model
\(h\) only defines
\[ \frac{\partial y}{\partial y_1}, \quad \frac{\partial y}{\partial y_2}, \quad \frac{\partial y}{\partial x_4}. \]
The assembly of the partial derivatives into the total derivative \(\partial y / \partial \boldsymbol{x}\) using the chain rule is handled by NEML2. This design serves as the fundation for a modular model implementation:
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:
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.
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):
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:
LabeledAxis
An axis can be labeled recursively, e.g.,
The above example represents an axis of size 15. This axis has 4 items: a
, b
, c
, and sub
.
SR2
).Scalar
.Scalar
.LabeledAxis
. "sub" by itself represents an axis of size 7, containing 2 items: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.
/
), 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.
NEML2 implements two types of "assemblers" to assemble (or split) the implicit system given the axis layout defined by LabeledAxis:
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.
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.