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