Model composition

You’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 you 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.

Listing 5 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'
  []
[]

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:

import subprocess
print(subprocess.run(
    ["neml2-inspect", "input.i", "eq"],
    capture_output=True, text=True, check=True,
).stdout)

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 you 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, pass the inputs as an {name: SR2} dict to call_by_name:

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]))
eq.call_by_name({"a": a, "b": 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.

The same thing without ComposedModel

To see what ComposedModel is doing for you, 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:

Listing 6 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'
  []
[]

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 you 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