The forward operator#

The previous tutorial declared the projectile model’s inputs, outputs, and parameters. This one fills in the math: a forward() method that turns a velocity into an acceleration.

The forward method#

forward() is the function NEML2 calls to evaluate our model. The positional arguments are the inputs in the order we declared them in the schema — each one is already a typed wrapper (Vec, SR2, …) the framework built for us. Return the outputs in schema order.

There’s also a v=None keyword in the signature — that’s the chain-rule hook for first derivatives, covered later on this page. Until we need derivatives, we can leave it alone; calls without v get back just the outputs.

Implementation#

Recall the equation for the projectile model:

\[ \boldsymbol{a} = \boldsymbol{g} - \mu \boldsymbol{v}. \]

self.g is the gravity vector (a buffer, read directly), and the velocity comes in as the first positional argument. The drag coefficient mu is a parameter, so it is read through self._get_param("mu", promoted_params, Scalar) rather than self.mu — more on that below. Here is the complete projectile.py; the forward method is the new part, dissected line by line below:

%%writefile projectile.py
"""Custom ``ProjectileAcceleration`` model — the running example for the
``extension`` tutorial chain. Mirrors the C++ tutorial's projectile in
the Python-native model surface.

The equation is

    a = g - mu * v

where ``v`` is the projectile velocity (input), ``a`` is the
acceleration (output), ``g`` is the gravitational acceleration vector
(buffer; constant), and ``mu`` is the scalar dynamic viscosity
(parameter; calibratable).
"""

from __future__ import annotations

from neml2.factory import register_neml2_object
from neml2.models.chain_rule import ChainRuleDict
from neml2.models.model import Model
from neml2.schema import HitSchema, buffer, input, output, parameter
from neml2.types import Scalar, Vec


@register_neml2_object("ProjectileAcceleration")
class ProjectileAcceleration(Model):
    """Newton's second law for a projectile in a viscous medium:
    ``a = g - mu * v``.
    """

    hit = HitSchema(
        input("velocity", Vec, "Velocity of the projectile", attr="_v_name"),
        output("acceleration", Vec, "Acceleration of the projectile"),
        buffer(
            "gravitational_acceleration",
            Vec,
            "Gravity vector",
            attr="g",
            default=Vec.fill(0.0, -9.81, 0.0),
        ),
        parameter("dynamic_viscosity", Scalar, "Dynamic viscosity", attr="mu"),
    )

    _v_name: str
    g: Vec
    mu: Scalar

    def forward(  # type: ignore[override]
        self,
        v_in: Vec,
        *promoted_params,
        v: ChainRuleDict | None = None,
    ):
        # Read the drag coefficient through ``_get_param`` rather than
        # ``self.mu``. ``_get_param`` resolves a static slot from ``self`` or
        # a promoted runtime input from ``*promoted_params``, so the same forward
        # keeps working after ``mu`` is promoted (neml2-compile -p); a bare
        # ``self.mu`` would be rejected by the parameter-attribute guard.
        mu = self._get_param("mu", promoted_params, Scalar)

        # Compute the value: a = g - mu * v. ``self.g`` is a buffer (not a
        # parameter), so reading it directly is fine.
        a = self.g - mu * v_in

        # Pure forward: return the typed output and stop.
        if v is None:
            return a

        # First-order chain rule: ∂a / ∂v_in = -mu * I. The closure
        # captures the local ``mu`` and receives an incoming tangent V
        # (a ``Vec`` shaped like the input), returns the contribution
        # to ∂(acceleration)/∂(seed-leaf).
        actions = {self._v_name: lambda V: -mu * V}

        # ``apply_chain_rule`` returns the v_out dict; pair it with the
        # value so the caller can unpack ``(a, v_out)``.
        return a, self.apply_chain_rule(v, "acceleration", actions, output=a)

Writing projectile.py

The first line reads the parameter: mu = self._get_param("mu", promoted_params, Scalar). Always read a parameter this way inside forward — never as self.mu. A bare self.mu is rejected by a runtime guard, because it bypasses _get_param’s static-or-promoted dispatch: the moment mu is promoted to a runtime input (neml2-compile -p) the static nn.Parameter no longer exists and the attribute read breaks. _get_param works for both static and promoted parameters — it pulls the value from self when static and from the *promoted_params pack when promoted — so the leaf stays promotion-compatible. (Buffers like self.g are not parameters, so reading them directly is fine.) That is also why the signature is def forward(self, v_in, *promoted_params, v=None): the *promoted_params pack is where promoted parameters arrive.

The next line is the physics: Vec - Scalar * Vec gives back a Vec, batched or not. If v is None (the usual case) the method returns and we’re done.

The else branch is the chain-rule hook. actions maps each input variable to a small function that takes an incoming tangent (something the same shape as that input) and returns its contribution to the output’s tangent. For this model the math is simple: \(\partial \boldsymbol{a}/\partial \boldsymbol{v} = -\mu I\), so the closure is lambda V: -mu * V, capturing the local mu we read at the top. apply_chain_rule then sums the contribution against any tangents the caller seeded on v, without ever building the full Jacobian matrix in memory.

Evaluation#

That’s the whole model. Load it the same way we’d load any built-in type — neml2.load_model finds it through the factory as long as the module that registers it has been imported.

The input file from the previous tutorial wires the model into HIT:

%%writefile input.i
[Models]
  [accel]
    type = ProjectileAcceleration
    velocity = 'v'
    acceleration = 'a'
    dynamic_viscosity = '0.001'
  []
[]

Writing input.i

Import the module so the class registers, then load and call:

import sys, os

sys.path.insert(0, os.getcwd())

import projectile  # registers ProjectileAcceleration with the native factory
import neml2

model = neml2.load_model("input.i", "accel")
model

ProjectileAcceleration()

import torch
from neml2.types import Vec

vel = Vec.fill(10.0, 2.0, 0.0)
accel = model(vel)
accel

Vec(data=tensor([-0.0100, -9.8120,  0.0000], dtype=torch.float64,
       grad_fn=<SubBackward0>), sub_batch_ndim=0, sub_batch_state=(), sub_batch_meta=(), k_ndim=0, k_state=(), k_pairing=())

The result is \(\boldsymbol{g} - \mu \boldsymbol{v} = (0, -9.81, 0) - 0.001 \cdot (10, 2, 0) = (-0.01, -9.812, 0)\) — the typed wrapper preserves the Vec shape on the way out.

First derivatives#

Because forward() implements the v branch, the same model can hand back directional derivatives with no extra wiring. Seed a tangent on the velocity input and the JVP comes back through v_out:

# Seed the identity on Vec to read off the full Jacobian column ∂a/∂v.
seed = Vec(torch.eye(3, dtype=torch.float64))
accel, v_out = model(vel, v={"v": {"velocity_leaf": seed}})
v_out["a"]["velocity_leaf"].data

tensor([[-0.0010, -0.0000, -0.0000],
        [-0.0000, -0.0010, -0.0000],
        [-0.0000, -0.0000, -0.0010]], dtype=torch.float64,
       grad_fn=<MulBackward0>)

This matches the analytical \(\partial \boldsymbol{a}/\partial \boldsymbol{v} = -\mu I = -0.001\, I_3\) exactly.

Driving the model from a unit-test input#

ModelUnitTest is the usual way to pin a custom model’s behavior in CI. It loads the model, runs it on inputs we supply, checks the outputs, and cross-checks the derivatives against PyTorch’s autograd:

%%writefile unit_test.i
[Tensors]
  [v_in]
    type = Python
    expr = 'Vec.fill(10.0, 2.0, 0.0)'
  []
  [a_expected]
    type = Python
    expr = 'Vec.fill(-0.01, -9.812, 0.0)'
  []
[]

[Models]
  [accel]
    type = ProjectileAcceleration
    velocity = 'v'
    acceleration = 'a'
    dynamic_viscosity = '0.001'
  []
[]

[Drivers]
  [unit]
    type = ModelUnitTest
    model = 'accel'
    input_Vec_names = 'v'
    input_Vec_values = 'v_in'
    output_Vec_names = 'a'
    output_Vec_values = 'a_expected'
  []
[]

Writing unit_test.i

from neml2.drivers.ModelUnitTest import ModelUnitTest

report = ModelUnitTest.from_file("unit_test.i").run()
print(f"value checks: {report.value_checks}, JVP checks: {report.jvp_checks}")

value checks: 1, JVP checks: 1

If both counters are positive (and the cell didn’t raise), every value and every JVP matched. A zero on either side means that check was skipped, not that it failed silently.

Where to go next#