nvalchemi.models.mace.MACEWrapper

class nvalchemi.models.mace.MACEWrapper(model, *, reconstruction_spec=None)

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.

• reconstruction_spec (BaseSpec | None)

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.

Parameters:

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

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

Returns:

MACE inputs: positions, node_attrs, batch, ptr, edge_index [2, E], shifts [E, 3], cell [B, 3, 3].

Return type:

dict[str, Any]

Notes

Does not call super().adapt_input() (Batch has no model_dump); gradient enabling on positions is handled here.

adapt_output(raw_output, data)

Map MACE raw outputs to nvalchemi standard keys.

Normalizes energy shape, forwards forces / stress / hessian when present, and exposes MACE’s node_energy as atomic_energies, then delegates to the base auto-mapper.

Parameters:

• raw_output (dict[str, Any]) – The dict returned by MACE.forward.

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

Returns:

The standardized outputs (subset of energy, forces, stress, hessian, atomic_energies).

Return type:

ModelOutputs

checkpoint_spec()

Return the factory spec used to reconstruct this wrapper, if known.

Wrappers created by from_checkpoint() store a callable spec for that factory so strategy checkpoints can rebuild optimized MACE models without introspecting the transformed inner MACE module constructor. Wrappers around arbitrary live modules return None and use the generic constructor-introspection fallback.

Return type:

BaseSpec | None

compute_embeddings(data, **kwargs)

Compute node and graph embeddings without forces or stresses.

Parameters:

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

• **kwargs – Forwarded to adapt_input().

Returns:

data, with node_embeddings [N, hidden_dim] and graph_embeddings [B, hidden_dim] (sum-pooled) written in place. model_config is not mutated.

Return type:

AtomicData | Batch

property cutoff: float

Interaction cutoff in Angstroms, read from model.r_max.

distribution_spec(strategy=None)

MLIPSpec for MACE under domain decomposition.

MACE uses the MPNN halo spec: every message-passing layer scatters over edges into node_feats (halo rows kept in sync), and a final per-graph scatter over node energies produces total energy (halo rows dropped, then all-reduced across ranks). For cueq-converted checkpoints the fused conv_tp kernel hides that gather/scatter, so the spec installs a mode-dependent conv adapter for the DD scope (_cueq_conv_unfuse_adapters): under eager DD it unfuses to the external gather + scatter plain MACE uses (halo handlers fire); under compiled DD it keeps the conv fused for memory parity with single-GPU and relies on the refresh adapter’s scatter_to_owners for halo correctness.

Memoized on first access. The per-checkpoint additions over the base spec are: the message-passing halo refresh (neighbor_refresh_adapters discovers the concrete InteractionBlocks and declares their per-node output halo-corrected under compile; NVALCHEMI_MACE_NO_REFRESH=1 drops it, debug only) and, for cueq, the conv unfuse adapters.

Parameters:

strategy (Any)

Return type:

Any

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 for the active outputs.

Pure and distribution-agnostic: computes exactly what model_config.active_outputs requests. Forces/stress run MACE’s internal autograd; an atomic_energies-only request runs an energy-only forward returning per-atom energy.

Parameters:

• data (AtomicData | Batch) – The input system (with neighbor data attached).

• **kwargs – Forwarded to adapt_input().

Returns:

The standardized outputs for the active set.

Return type:

ModelOutputs

classmethod from_checkpoint(checkpoint_path, device=torch.device('cpu'), enable_cueq=False, dtype=None, compile_model=False, atomic_energies=None, atomic_energies_path=None, **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.

• atomic_energies (Mapping[int | str, float] | None, optional) – Per-element E0 overrides keyed by atomic number.

• atomic_energies_path (Path | str | None, optional) – JSON file containing per-element E0 overrides keyed by atomic number.

• **compile_kwargs – Forwarded to torch.compile.

Return type:

MACEWrapper

Raises:

• ImportError – If mace-torch is not installed, or if enable_cueq=True and either cuequivariance or the cuequivariance-ops-torch CUDA kernels (the cu12 / cu13 dependency groups) are missing.

• ValueError – If enable_cueq=True and device is not a CUDA device.

modify_ema_methods()

Restore cuEquivariance methods discarded by EMA model copying.

torch.optim.swa_utils.AveragedModel deep-copies its source model. cuEquivariance fused convolution modules attach their specialized forward method at runtime, and that instance method is not retained by the copy. Reapply MACE’s fusion wrapper only when it is missing.

Return type:

None