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