Ctrl+K

nvalchemi.models.uma.UMAWrapper#

class nvalchemi.models.uma.UMAWrapper(predict_unit, task_name='omol', train=False)[source]#

Wrapper for fairchem’s UMA (Universal Models for Atoms).

Wraps a fairchem.core.units.mlip_unit.MLIPPredictUnit — the level at which energy / forces / stress are computed by fairchem (the raw backbone only produces node embeddings). Task is fixed at construction; active_outputs reflects what that task supports.

Parameters:

  • predict_unit (fairchem.core.units.mlip_unit.MLIPPredictUnit) – Pre-loaded UMA prediction unit. Use from_checkpoint() for the typical construction path that resolves a registered checkpoint name and downloads via HuggingFace Hub.

  • task_name (str) – UMA task: one of _UMA_TASKS. Determines which per-task head in the multi-task model is used and which inputs (charge, spin) must be populated.

  • train (bool) – False (default) freezes all weights for inference (lossless for autograd forces); True keeps fairchem’s trainable/frozen split so weights stay exposed for fine-tuning.

model_config#

Task-dependent outputs + autograd + neighbor config.

Type:

ModelConfig

task_name#

The UMA task this wrapper is pinned to.

Type:

str

adapt_input(data, **kwargs)[source]#

Convert an nvalchemi AtomicData / Batch to a fairchem graph.

Tensor-native (no ASE round trip): tensors stay on data.positions.device, preserving GPU residency and autograd. edge_index is left empty (2, 0) so fairchem’s MLIPPredictUnit rebuilds the graph internally, matching the default FAIRChemCalculator path (r_edges=False), so outputs are equivalent. Charge/spin default per the ASE-calculator convention (per-system LongTensors; spin defaults to the closed-shell singlet for OMol, 0 for periodic tasks) unless the caller provides them on the batch.

Parameters:

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

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

Returns:

The fairchem graph: pos [N, 3], atomic_numbers [N], cell [B, 3, 3], pbc [B, 3], per-system charge / spin [B], and empty edge_index [2, 0].

Return type:

fairchem.core.datasets.atomic_data.AtomicData

adapt_output(raw, data=None)[source]#

Map fairchem’s prediction dict to nvalchemi’s output keys.

Parameters:

  • raw (dict) – fairchem’s prediction dict: "energy" (per-system), "forces" (per-atom [N, 3]), and optionally "stress" (per-system).

  • data (AtomicData | Batch | None, optional) – The input system the outputs were computed for. Unused here.

Returns:

The active subset of energy [B, 1], forces [N, 3], and stress [B, 3, 3].

Return type:

ModelOutputs

compute_embeddings(data, **kwargs)[source]#

Run the backbone only and attach node embeddings.

UMA/eSEN backbones return {"embedding": [N, sph, ch], "batch": [N]}; the embedding is attached as a node property so pipelines can consume it without re-running the heads.

Parameters:
Returns:

data, with node_embeddings [N, sph, ch] attached when the backbone returns an embedding.

Return type:

AtomicData | Batch

property cutoff: float#

Radial cutoff (Å) for neighbor-list construction.

distribution_spec(strategy=None)[source]#

MLIPSpec for UMA under domain decomposition.

Each rank computes a full forward over its owned + ghost atoms on plain tensors; DD happens only at the boundaries: per-block ghost-row feature refresh, owned-only + all-reduce per-system energy reduction, and forces/stress through fairchem’s autograd (ghost contributions routed to owners in consolidation). Nothing is sharded (shard_fields is empty); the spec carries the fixed-shape-caps _UMAGraphPadder.

Returns:

The memoized halo spec: boundary adapters, empty shard_fields, and a CompilePolicy carrying the graph padder.

Return type:

MLIPSpec

Parameters:

strategy (Any)

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

Shape of the per-node backbone embedding.

eSCN-MD (and eSEN) backbones produce [N, (lmax+1)^2, sphere_channels], read off the backbone’s sph_feature_size / sphere_channels attrs.

Returns:

{"node_embeddings": (sph_feature_size, sphere_channels)}.

Return type:

dict[str, tuple[int, …]]

Raises:

RuntimeError – If the predict unit’s module exposes no backbone.

extra_repr()#

Format the model config for nn.Module.__repr__.

Parameters:

self (Any)

Return type:

str

forward(data, **kwargs)[source]#

Run the UMA predict unit on data.

Pipeline: adapt_input -> MLIPPredictUnit.predict -> adapt_output. The single distribution touchpoint is ctx.maybe_pad_graph, which under compiled domain decomposition pads the fairchem graph to stable per-rank shapes (a no-op single-process). The two blocks below handle fairchem’s own compile requirements: CPU routing for the first-call MoLE merge, and forcing static shapes.

Parameters:

  • data (AtomicData | Batch) – Input structure(s); promoted to a Batch by adapt_input.

  • kwargs (Any)

Returns:

energy (per system) plus forces / stress per model_config.active_outputs.

Return type:

ModelOutputs

classmethod from_checkpoint(name_or_path, task_name='omol', device='cpu', inference_settings='default', overrides=None, train=False)[source]#

Resolve and load a UMA checkpoint.

Accepts either a registered model name ("uma-s-1p1" / "uma-s-1p2" / "uma-m-1p1") or a local filesystem path to a .pt file. The multi-task checkpoints ship all five task heads; task_name picks which one the wrapper exposes via model_config.active_outputs.

Parameters:

  • name_or_path (str | Path) – Registered model name (see fairchem.core.calculate.pretrained_mlip.available_models) or a local file path.

  • task_name (str) – One of omol, omat, oc20, odac, omc. Defaults to omol (molecular chemistry) — the most common entry point; override explicitly for crystals / catalysis.

  • device (str | torch.device) – Target device for inference. Defaults to "cpu".

  • inference_settings (InferenceSettings | str) – fairchem inference configuration. Either a preset name ("default" or "turbo") or a fairchem.core.units.mlip_unit.api.inference.InferenceSettings instance. torch.compile is reached through this argument — see the module docstring’s torch.compile section. Defaults to "default".

  • overrides (dict | None) – Optional overrides forwarded to fairchem’s inference-settings builder.

  • train (bool) – If False (default), freeze all weights for inference — lossless, since conservative forces come from autograd on positions. If True, keep fairchem’s loaded trainable/frozen split so weights remain exposed for fine-tuning. Note: the forward path goes through fairchem’s inference predict (eval mode, detached forces); gradient-based training requires a separate path through the raw model.

Returns:

A wrapper pinned to task_name over the loaded predict unit.

Return type:

UMAWrapper

Raises:

ValueError – If name_or_path is neither a registered model name nor a local file path.

On this page
Show Source