nvalchemi.models.base.BaseModelMixin

class nvalchemi.models.base.BaseModelMixin

Abstract mixin providing a standardized interface for model wrappers.

All external MLIP wrappers should inherit from this mixin (alongside nn.Module) to ensure a consistent interface for dynamics engines, composition pipelines, and downstream tooling.

Concrete implementations must provide:

• model_config attribute — a ModelConfig instance set in __init__.

• embedding_shapes property — expected shapes of computed embeddings.

• compute_embeddings() — compute and attach embeddings to the input data structure.

The mixin provides default implementations of:

• input_data() — set of required input keys derived from the model config.

• output_data() — set of active outputs intersected with supported outputs (warns on unsupported requests).

• adapt_input() — enable gradients on required tensors and collect input dict.

• adapt_output() — map raw model output to ModelOutputs ordered dict.

adapt_input(data, **kwargs)

Adapt framework batch data to external model input format.

The base implementation enables requires_grad on tensors that need gradients (determined by model_config.autograd_inputs and model_config.gradient_keys), then collects all keys declared by input_data() into a dict.

Subclasses should call super().adapt_input(data) and then add or transform entries as needed for their underlying model.

Parameters:

• data (AtomicData | Batch | AtomsLike) – Framework data structure.

• kwargs (Any)

Returns:

Input in the format expected by the external model.

Return type:

dict[str, Any]

adapt_output(model_output, data)

Adapt external model output to ModelOutputs format.

Returns an OrderedDict keyed by output_data() entries, populated from model_output where keys match.

Note

Returned tensors may still be attached to the autograd computation graph (e.g. energies from autograd-force models like MACE). This is intentional — the model does not know whether the caller needs the graph (e.g. pipeline shared-autograd groups). Callers that do not need the graph are responsible for detaching.

Parameters:

• model_output (Any) – Raw output from the external model.

• data (AtomicData | Batch) – Original input data (may be needed for context/metadata).

Returns:

OrderedDict with expected output keys and their values (or None if not present). Tensors may be graph-attached.

Return type:

ModelOutputs

add_output_head(prefix)

Add an output head to the model.

Parameters:

prefix (str) – Prefix for the output head

Return type:

None

abstractmethod compute_embeddings(data, **kwargs)

Compute embeddings at different levels of a batch of atomic graphs.

Parameters:

• data (AtomicData | Batch) – Input atomic data containing positions, atomic numbers, etc.

• kwargs (Any)

Returns:

Data structure with embeddings attached in-place.

Return type:

AtomicData | Batch

Raises:

NotImplementedError – If the model does not support embeddings computation

direct_derivative_keys()

Return output keys this model computes analytically in forward().

When this model participates in a pipeline autograd group, the pipeline strips "forces" and "stress" from sub-model active_outputs so it can compute them via autograd on the summed energy. Keys returned by this method are kept in active_outputs — the pipeline collects them from the model output and sums them with the autograd-derived derivatives.

Override this in models that produce analytical forces or stress alongside an energy that carries autograd information (e.g. Ewald/PME with hybrid_forces=True).

Returns:

Keys (e.g. {"forces", "stress"}) that the model produces analytically and should be summed with autograd derivatives. Default: empty set (all derivatives come from autograd).

Return type:

set[str]

abstract 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)

Export the current model without the BaseModelMixin interface.

Parameters:

• path (Path)

• as_state_dict (bool)

Return type:

None

input_data()

Return the set of required input keys.

Base implementation derives keys from the model config: {positions, atomic_numbers} plus neighbor-list keys (from neighbor_config), pbc (if needs_pbc), and any extra keys in model_config.required_inputs.

Optional inputs (model_config.optional_inputs) are handled separately in adapt_input() and are NOT included here.

Returns:

Set of required input keys.

Return type:

set[str]

make_neighbor_hooks(max_neighbors=None)

Return a list of NeighborListHook instances for this model’s neighbor configuration.

Returns an empty list if the model does not require a neighbor list. Defers the import to avoid circular imports.

Parameters:

max_neighbors (int | None, optional) – Maximum neighbors per atom for MATRIX format. When None (default), auto-estimated from the cutoff at first use.

Return type:

list

output_data()

Return the set of keys the model will compute this run.

Intersects active_outputs with outputs. Warns if any active keys are not supported by the model.

Returns:

Set of output keys that are both active and supported.

Return type:

set[str]

set_config(key, value)

Set a mutable field on model_config.

Convenience method equivalent to self.model_config.<key> = value with validation that the field exists and is mutable.

Parameters:

• key (str) – Name of a mutable ModelConfig field (e.g. "active_outputs", "gradient_keys").

• value (Any) – New value for the field.

Raises:

AttributeError – If key is not a field on ModelConfig.

Return type:

None