neml2.drivers

Native [Drivers] registrations.

One file per registered driver type. Importing this package fires the @register_neml2_object decorators so the factory can resolve [Drivers] blocks from HIT input.

class neml2.drivers.TransientDriver(model, prescribed_time, time_name, forces, ics)

Bases: Driver

Drive a model over a prescribed time history.

Mirrors C++ TransientDriver. See module docstring for differences.

Parameters:

• model (Model)

• prescribed_time (Scalar)

• time_name (str)

• forces (dict[str, TensorWrapper])

• ics (dict[str, TensorWrapper])

classmethod from_hit(node, factory)

Parameters:

• node (nmhit.Node)

• factory (_NativeInputFile)

Return type:

TransientDriver

hit = <neml2.schema.HitSchema object>

result()

Flatten result_in / result_out into a input.<i>.<name> / output.<i>.<name> dict, matching the C++ ModuleDict layout that torch.jit.load(gold).named_buffers() produces. Typed wrappers are unwrapped to raw torch.Tensor at this boundary.

Return type:

dict[str, Tensor]

run()

Execute the driver. Returns True on success.

Return type:

bool

save_gold(path)

Serialize result() to a .pt file as a flat torch.save dict.

The dict carries the same keys as result(): input.<step>.<var> / output.<step>.<var>, each mapping to a detached, cloned torch.Tensor. Read it back with torch.load(path, weights_only=True).

Replaces the old TorchScript-Module format (a Module-of-Modules with register_buffer per entry, dumped via torch.jit.script(...).save(...)). The flat-dict format avoids the torch.jit.{script,load} deprecation warnings and survives round-trip with weights_only=True (which is safe for plain tensor dicts and prevents arbitrary code execution on load). TransientRegression reads both formats so existing on-disk goldens from the v2 C++ pipeline keep working — see TransientRegression.

Parameters:

path (str | Path)

Return type:

None

class neml2.drivers.TransientRegression(driver, reference, rtol=1e-05, atol=1e-08)

Bases: Driver

Run a TransientDriver and diff its result against a gold .pt file.

Parameters:

• driver (TransientDriver)

• reference (Path)

• rtol (float)

• atol (float)

classmethod from_hit(node, factory)

Parameters:

• node (nmhit.Node)

• factory (_NativeInputFile)

Return type:

TransientRegression

hit = <neml2.schema.HitSchema object>

run()

Execute the driver. Returns True on success.

Return type:

bool

class neml2.drivers.Verification(driver, refs, rtol=1e-05, atol=1e-08, time_steps=())

Bases: Driver

Run a TransientDriver and diff its result against per-variable references.

Parameters:

• driver (TransientDriver)

• refs (dict[str, TensorWrapper])

• rtol (float)

• atol (float)

• time_steps (tuple[int, ...])

classmethod from_hit(node, factory)

Parameters:

• node (nmhit.Node)

• factory (_NativeInputFile)

Return type:

Verification

hit = <neml2.schema.HitSchema object>

run()

Execute the driver. Returns True on success.

Return type:

bool