Model composition#

In this tutorial, we’ll wire three small models together into a composed model, inspect how NEML2 resolves the connections, and evaluate the result. Then we’ll compare it to doing the same plumbing by hand.

Most constitutive theories are a chain of small maps. A Perzyna-type viscoplastic model, for instance, threads together

\[\begin{align*} \boldsymbol{\varepsilon}^e &= \boldsymbol{\varepsilon} - \boldsymbol{\varepsilon}^p, \\ \boldsymbol{\sigma} &= 3K\operatorname{vol}\boldsymbol{\varepsilon}^e + 2G\operatorname{dev}\boldsymbol{\varepsilon}^e, \\ \bar{\sigma} &= J_2(\boldsymbol{\sigma}), \\ f &= \bar{\sigma} - \sigma_y, \\ \boldsymbol{N} &= \partial f/\partial\boldsymbol{\sigma}, \\ \dot{\gamma} &= \left(\langle f\rangle / \eta\right)^n, \\ \dot{\boldsymbol{\varepsilon}}^p &= \dot{\gamma}\,\boldsymbol{N}. \end{align*}\]

Each line on the right-hand side maps to a separate Model in NEML2’s catalog, and ComposedModel glues a chosen set together. That keeps each piece testable and swappable, and the Python-side overhead of stepping through many sub-models drops out once we export the composed graph through the compilation pipeline (see Compiled models).

A worked example#

To keep the wiring visible we’ll use three small models from the catalog instead of the full plasticity stack:

(1)#\[\begin{align} \bar{a} &= I_1(\boldsymbol{a}), \\ \bar{b} &= \sqrt{\tfrac{3}{2}\,\operatorname{dev}(\boldsymbol{b}) : \operatorname{dev}(\boldsymbol{b})}, \\ \dot{\boldsymbol{b}} &= \bar{b}\,\boldsymbol{a} + \bar{a}\,\boldsymbol{b}. \end{align}\]

The first two are scalar invariants of symmetric tensors, both handled by SR2Invariant (the second one with invariant_type = VONMISES). The third is a linear combination of two SR2 tensors with scalar weights — see SR2LinearCombination.

%%writefile input.i
# Three small models from the catalog, plus a ComposedModel that wires
# them together via shared variable names.
#
#   eq1:  a_bar = I1(a)             (SR2Invariant)
#   eq2:  b_bar = vonMises(b)       (SR2Invariant)
#   eq3:  b_rate = b_bar * a + a_bar * b   (SR2LinearCombination)
#   eq:   the three glued together by ComposedModel
[Models]
  [eq1]
    type = SR2Invariant
    tensor = 'a'
    invariant = 'a_bar'
    invariant_type = I1
  []
  [eq2]
    type = SR2Invariant
    tensor = 'b'
    invariant = 'b_bar'
    invariant_type = VONMISES
  []
  [eq3]
    type = SR2LinearCombination
    from = 'a b'
    to = 'b_rate'
    # The two weights are wired to the OUTPUTS of eq1 and eq2 by name.
    # The dependency resolver will treat 'a_bar' and 'b_bar' as
    # producer/consumer links instead of free parameters.
    weights = 'b_bar a_bar'
  []
  [eq]
    type = ComposedModel
    models = 'eq1 eq2 eq3'
  []
[]

Writing input.i

The trick is in eq3’s weights = 'b_bar a_bar'. b_bar and a_bar aren’t literals or [Tensors] entries — they’re the output names of eq2 and eq1. ComposedModel notices that eq3 consumes two scalars that eq1 and eq2 produce, and wires them up.

Inspecting the wiring#

Before evaluating anything, ask neml2-inspect to print the resolved graph:

# `neml2-inspect` is a thin CLI over neml2 -- import its entry point and call
# it in-process so the resolved graph prints inline.
from neml2.cli.inspect import main as inspect_main

_ = inspect_main(["input.i", "eq"])

Model: ComposedModel

Inputs (2):
  a: type=SR2
  b: type=SR2

Outputs (1):
  b_rate: type=SR2

Parameters (1):
  eq3.offset: dtype=torch.float64, device=cpu, shape=[]

Buffers (0):

Three things to notice:

  1. Inputs collapsed to a and b. The intermediate scalars a_bar and b_bar aren’t free inputs — they’re produced internally.

  2. Outputs collapsed to b_rate. a_bar and b_bar are consumed downstream, so they don’t surface as outputs. (If we want them, add them under additional_outputs on the ComposedModel.)

  3. Parameters collapsed to eq3.offset. eq3.weight_0 and eq3.weight_1 are gone — replaced by the producer links from eq2 and eq1. Only the literal offset = 0 is still free.

Running neml2-inspect right after wiring a composed model is the fastest way to catch typos: a missed name shows up as a dangling input or missing output, instead of a shape mismatch later on.

Loading and evaluating the composed model#

The composed model loads and calls just like any other model:

import torch
import neml2
from neml2.types import SR2

torch.set_default_dtype(torch.float64)
eq = neml2.load_model("input.i", "eq")
eq

ComposedModel(
  (eq1): SR2Invariant()
  (eq2): SR2Invariant()
  (eq3): SR2LinearCombination()
)

The input_spec / output_spec properties echo what neml2-inspect showed:

list(eq.input_spec.keys()), list(eq.output_spec.keys())

(['a', 'b'], ['b_rate'])

To evaluate, call it like any other model — inputs positional, in input_spec order. A ComposedModel hands its outputs back as a tuple (in output_spec order), so we unpack the single b_rate:

a = SR2(torch.tensor([0.1, 0.05, -0.03, 0.02, 0.06, 0.03]))
b = SR2(torch.tensor([100.0, 20.0, 10.0, 5.0, -30.0, -20.0]))
(b_rate,) = eq(a, b)
b_rate

SR2(data=tensor([21.6372,  7.2186, -1.6912,  2.5274,  2.1823,  0.4912],
       grad_fn=<AddBackward0>), sub_batch_ndim=0, sub_batch_state=(), sub_batch_meta=(), k_ndim=0, k_state=(), k_pairing=())

Under the hood it ran eq1, then eq2, then eq3 (the only order that respects the dependencies), threaded the intermediate scalars into eq3’s weight slots, and returned b_rate.

If we’d rather pass and receive values by name — handy when a model has many inputs — eq.call_by_name({"a": a, "b": b}) runs the same evaluation with name-keyed dicts on both ends.

The same thing without ComposedModel#

To see what ComposedModel is doing for us, here’s the same calculation done by hand. The input file is the same three [Models] entries, but with weights = '1 1' on eq3 so weight_0 and weight_1 stay as free parameters:

%%writefile input_manual.i
# The same three models without a ComposedModel — the caller is
# responsible for evaluating them in the right order and threading
# intermediate values into eq3's weight parameters by hand.
[Models]
  [eq1]
    type = SR2Invariant
    tensor = 'a'
    invariant = 'a_bar'
    invariant_type = I1
  []
  [eq2]
    type = SR2Invariant
    tensor = 'b'
    invariant = 'b_bar'
    invariant_type = VONMISES
  []
  [eq3]
    type = SR2LinearCombination
    from = 'a b'
    to = 'b_rate'
    weights = '1 1'
  []
[]

Writing input_manual.i

import torch.nn as nn

eq1 = neml2.load_model("input_manual.i", "eq1")
eq2 = neml2.load_model("input_manual.i", "eq2")
eq3 = neml2.load_model("input_manual.i", "eq3")

# 1. Evaluate the two invariants.
a_bar = eq1(a)
b_bar = eq2(b)

# 2. Manually wire the weights of eq3 to those intermediate values.
# weight_0 is the first `from` variable (a), weight_1 is the second (b),
# so b_bar (the von Mises norm of b) goes to weight_0 and vice-versa.
eq3.weight_0 = nn.Parameter(b_bar.data)
eq3.weight_1 = nn.Parameter(a_bar.data)

# 3. Evaluate eq3 to get b_rate.
eq3(a, b)

SR2(data=tensor([21.6372,  7.2186, -1.6912,  2.5274,  2.1823,  0.4912],
       grad_fn=<AddBackward0>), sub_batch_ndim=0, sub_batch_state=(), sub_batch_meta=(), k_ndim=0, k_state=(), k_pairing=())

Same answer, but we had to:

  • pick the right evaluation order,

  • remember which weight slot maps to which invariant, and

  • copy the intermediate values into eq3’s parameters by hand.

A handful is manageable. A realistic constitutive theory with dozens of small maps isn’t. ComposedModel does this bookkeeping once at load time, then gets out of the way.

Note

This works because eq3’s weights accepts a list of names that can resolve to a parameter, a sibling model’s output, or a [Tensors] entry. See Model parameters revisited for the full story.

Where to go next#