nvalchemi.models.aimnet2.AIMNet2Wrapper

class nvalchemi.models.aimnet2.AIMNet2Wrapper(model, *, compile_model=False, compile_kwargs=None, train=None)

Wrapper for AIMNet2 interatomic potentials.

Energy is always computed as the primitive differentiable output via the raw AIMNet2 model. Forces and stresses are derived from energy via autograd. Partial charges and node embeddings (AIM features) are taken directly from the model outputs.

The wrapper declares an external MATRIX-format neighbor list requirement at the model’s AEV cutoff. The NeighborListHook (or the pipeline’s synthesized hook) populates neighbor_matrix on the batch before each forward pass. The wrapper converts this to AIMNet2’s internal nbmat format (with a padding row for the padding atom).

Coulomb and D3 dispersion are disabled. Use PipelineModelWrapper to compose AIMNet2 with electrostatics or dispersion models.

Parameters:

• model (nn.Module) – An AIMNet2 model (loaded from checkpoint or instantiated directly). Use from_checkpoint() for the common construction path.

• compile_model (bool, optional) – torch.compile the AIMNet2 module forward via the calculator’s kernel-aware compile path (single-process inference). Distributed compilation is a separate switch, DistributedModel(..., compile=True).

• compile_kwargs (dict[str, Any] | None, optional) – Forwarded to torch.compile when compile_model=True.

• train (bool | None, optional) – Whether AIMNet2Calculator should keep the model trainable. Defaults to the wrapped module’s current training mode.

model_config

Configuration with capability and runtime fields.

Type:

ModelConfig

model

The underlying AIMNet2 model. If you want your model to be compiled, wrap with torch.compile(model, **kwargs) before passing here.

Type:

nn.Module

adapt_input(data, **kwargs)

Build the flat input dict expected by AIMNet2.forward.

Appends a single padding atom and converts the external neighbor matrix to AIMNet2’s nbmat layout. Enables gradients on positions when an autograd output is active.

Parameters:

• data (AtomicData | Batch) – The input system; an AtomicData is promoted to a single-graph Batch. Requires neighbor_matrix (from NeighborListHook).

• **kwargs – Unused; accepted for interface compatibility.

Returns:

AIMNet2 inputs with a trailing padding atom (index N): coord [N+1, 3], numbers [N+1] (pad row = 0), mol_idx [N+1] (sorted ascending), nbmat [N+1, K] (unused slots = N), charge [n_systems], and optional cell / shifts / mult.

Return type:

dict[str, Any]

Notes

Does not call super().adapt_input(): AIMNet2 uses its own key conventions (coord / numbers / nbmat).

adapt_output(model_output, data)

Map AIMNet2 outputs to nvalchemi standard keys.

Per-atom direct outputs (charges / spin_charges) carry the padding-atom row appended by adapt_input() (plus any padding rows added under compiled DD); they are sliced back to data.num_nodes. energy (per-system) and forces (autograd over real positions) need no strip.

Parameters:

• model_output (dict[str, Any]) – Raw outputs from the AIMNet2 forward pass.

• data (AtomicData | Batch) – The input system the outputs were computed for.

Returns:

Standardized outputs: energy [n_systems, 1] plus any active forces / stress / charges / spin_charges.

Return type:

ModelOutputs

compute_embeddings(data, **kwargs)

Compute AIM-feature node embeddings and attach them to data.

Parameters:

• data (AtomicData | Batch) – The input system; an AtomicData is promoted to a single-graph Batch.

• **kwargs – Forwarded to adapt_input().

Returns:

data, with node_embeddings [N, aim_dim] written in place when the model exposes AIM features.

Return type:

AtomicData | Batch

distribution_spec(strategy=None)

MLIPSpec describing AIMNet2 under domain decomposition.

Halo-only for now (graph parallel is P1/P2, out of the essential gate); the strategy argument is accepted for the framework contract and ignored.

Returns:

The halo MLIPSpec (owned+ghost local-neighbor storage). The conv kernel and Coulomb heads refresh their ghost rows each layer, and the per-system sum counts owned atoms only. Whether the distributed forward is compiled is decided by DistributedModel(..., compile=True).

Return type:

Any

Parameters:

strategy (Any)

property embedding_shapes: dict[str, tuple[int, ...]]

AIM-feature embedding shapes produced by this model.

Returns:

Maps "node_embeddings" to its per-node feature shape (aim_dim,), read from the model’s AEV output size.

Return type:

dict[str, tuple[int, …]]

export_model(path, as_state_dict=False)

Serialize the underlying AIMNet2 model without the wrapper.

Parameters:

• path (Path) – Output path.

• as_state_dict (bool, optional) – If True, save only the state_dict; otherwise pickle the full model object. Defaults to False.

Return type:

None

extra_repr()

Format the model config for nn.Module.__repr__.

Parameters:

self (Any)

Return type:

str

forward(data, **kwargs)

Run the AIMNet2 model for the active outputs.

Pure and distribution-agnostic: computes exactly what model_config.active_outputs requests. Energy is the primitive output (summed per-system, over owned atoms plus an all-reduce under domain decomposition); forces and stresses are derived via autograd.

For stresses, the affine strain trick from prepare_strain() scales positions and cell through a displacement tensor so dE/d(displacement) gives the strain.

In a pipeline with use_autograd=True, the pipeline handles derivative computation externally — it strips forces/stresses from active_outputs so this method only computes energy.

Parameters:

• data (AtomicData | Batch) – Input batch with positions, atomic numbers, charge, and neighbor_matrix (from NeighborListHook).

• **kwargs – Forwarded to adapt_input().

Returns:

The standardized outputs for the active set.

Return type:

ModelOutputs

classmethod from_checkpoint(checkpoint_path, device='cpu', compile_model=False, **compile_kwargs)

Load an AIMNet2 model from a checkpoint and return a wrapped instance.

Uses AIMNet2Calculator to resolve and load the checkpoint, then extracts the raw nn.Module and wraps it.

Parameters:

• checkpoint_path (str | Path) – Path to an AIMNet2 checkpoint file, or a model alias recognized by AIMNet2Calculator (e.g. "aimnet2").

• device (torch.device | str, optional) – Target device. Defaults to "cpu".

• compile_model (bool, optional) – torch.compile the AIMNet2 model for single-process inference. Distributed compilation is a separate switch, DistributedModel(..., compile=True).

• **compile_kwargs – Forwarded to torch.compile when compile_model=True.

Returns:

The wrapped, fp32 model on device.

Return type:

AIMNet2Wrapper