nvalchemi.models.mace.MACEWrapper

class nvalchemi.models.mace.MACEWrapper(model)

Wrapper for any MACE model implementing the BaseModelMixin interface.

Accepts any MACE model variant (MACE, ScaleShiftMACE, cuEq-converted models, torch.compile-d models, etc.). The wrapper handles:

• One-hot node_attrs encoding via a pre-built GPU lookup table (no CPU round-trip per step).

• Gradient enabling on positions for conservative force / stress computation.

• PBC via both neighbor_list_shifts (integer image indices) and pre-computed shifts (physical Å vectors from neighbor_list_shifts @ cell) passed to MACE. shifts is always required; neighbor_list_shifts is additionally consumed when compute_displacement=True (stress path).

Parameters:

model (nn.Module) – An instantiated MACE model. Any subclass of mace.modules.MACE is accepted.

model

The underlying MACE model.

Type:

nn.Module

model_config

Mutable configuration controlling which outputs are computed.

Type:

ModelConfig

adapt_input(data, **kwargs)

Build the input dict expected by MACE.forward.

Handles AtomicData -> Batch promotion, node_attrs encoding, gradient enabling on positions, transposing edge_index from nvalchemi’s [E, 2] to MACE’s [2, E] convention, zero-filling of neighbor_list_shifts / cell for non-PBC systems, and pre-computation of physical shifts vectors from neighbor_list_shifts @ cell.

Expects COO neighbor data (neighbor_list, optionally neighbor_list_shifts) to be present on the batch. When used in a PipelineModelWrapper, the pipeline handles format conversion and cutoff filtering before calling this model.

Note

This method does not call super().adapt_input() because Batch does not implement model_dump(), which the base implementation requires. Gradient enabling on positions is handled manually here instead.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

dict[str, Any]

adapt_output(raw_output, data)

Map MACE output keys to nvalchemi standard keys.

MACE uses "energy" / "stress" / "hessian"; nvalchemi expects "energy" / "stress" / "hessian". Renaming happens before calling super() so the base auto-mapper sees the canonical key names.

Parameters:

• raw_output (dict[str, Any])

• data (AtomicData | Batch)

Return type:

OrderedDict[str, Float[Tensor, ‘B 1’] | Float[Tensor, ‘V 3’] | Float[Tensor, ‘V 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3’] | None]

compute_embeddings(data, **kwargs)

Compute node and graph embeddings without forces or stresses.

Writes node_embeddings (shape [N, hidden_dim]) and graph_embeddings (shape [B, hidden_dim], sum-pooled over atoms) into data in-place and returns it. Does not mutate model_config.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

AtomicData | Batch

property cutoff: float

Interaction cutoff in Angstroms, read from model.r_max.

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

Retrieves the expected shapes of the node, edge, and graph embeddings.

export_model(path, as_state_dict=False)

Serialize the underlying MACE model without the wrapper.

The exported file can be reloaded as a plain MACE nn.Module and used with the standard MACE / ASE interface.

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 MACE model and return the output.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

OrderedDict[str, Float[Tensor, ‘B 1’] | Float[Tensor, ‘V 3’] | Float[Tensor, ‘V 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3’] | None]

classmethod from_checkpoint(checkpoint_path, device=torch.device('cpu'), enable_cueq=False, dtype=None, compile_model=False, **compile_kwargs)

Load a MACE model from a checkpoint and return a MACEWrapper.

Accepts local file paths or named MACE-MP foundation-model checkpoints (e.g. "medium-0b2"), which are downloaded automatically to the MACE cache directory.

Operations are applied in this order:

1. Load — torch.load the checkpoint to the specified device.

2. dtype — cast model weights to the requested dtype.

3. cuEq — convert to cuEquivariance format for GPU speedup.

4. compile — torch.compile; freezes parameters and sets eval mode. The model is inference-only after this step.

For best GPU throughput, use device=torch.device("cuda"), enable_cueq=True, dtype=torch.float32, and compile_model=True. Example:

model = MACEWrapper.from_checkpoint(
    "medium-mpa-0",
    device=torch.device("cuda"),
    dtype=torch.float32,
    enable_cueq=True,
    compile_model=True,
)

Parameters:

• checkpoint_path (Path | str) – Local path to a .pt file, or a named checkpoint string such as "medium-0b2".

• device (torch.device, optional) – Target device. Defaults to CPU.

• enable_cueq (bool, optional) – Convert to cuEquivariance format for GPU speedup. Defaults to False. Requires the cuequivariance package.

• dtype (torch.dtype | None, optional) – If set, cast model weights to this dtype before cuEq conversion.

• compile_model (bool, optional) – Apply torch.compile. Sets eval mode and freezes parameters; the model is inference-only after this step.

• **compile_kwargs – Forwarded to torch.compile.

Return type:

MACEWrapper

Raises:

ImportError – If mace-torch is not installed, or if enable_cueq=True and cuequivariance is not installed.