nvalchemi.models.pipeline.PipelineModelWrapper

class nvalchemi.models.pipeline.PipelineModelWrapper(groups, additive_keys=None)

Compose multiple models via a grouped pipeline.

Models are organized into PipelineGroup instances, where each group has a derivative computation strategy. Within a group, steps execute in order so that upstream outputs can wire into downstream inputs. The pipeline sums outputs across groups using sum_outputs().

The pipeline’s default model_config.active_outputs is synthesized as the union of all sub-model model_config.active_outputs sets at construction time, so it honestly reflects what the sub-models are configured to produce. The user can then expand or narrow it.

Parameters:

• groups (list[PipelineGroup]) – Ordered list of groups. Groups execute in declaration order.

• additive_keys (set[str], optional) – Keys whose values are summed across groups. Defaults to {"energy", "forces", "stress"}.

model_config

Mutable configuration controlling what the pipeline computes.

Type:

ModelConfig

compute_embeddings(data, **kwargs)

Compute embeddings is not meaningful for pipeline models. Call compute_embeddings on individual sub-models instead.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

AtomicData | Batch

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 model is not implemented for pipeline models. Export individual sub-models instead.

Parameters:

• path (Path)

• as_state_dict (bool)

Return type:

None

extra_repr()

Show pipeline structure: groups, steps, wire mappings, and autograd strategy.

Return type:

str

forward(data, **kwargs)

Run all sub-models and accumulate outputs.

For groups with use_autograd=True, sub-models produce energies only. The group sums them and calls the derivative function (default or user-provided) to compute forces, stresses, and any other requested derivatives from the summed energy.

What gets computed is driven by self.model_config.active_outputs.

Parameters:

• data (AtomicData | Batch) – Input batch.

• kwargs (Any)

Returns:

Combined outputs across all groups.

Return type:

ModelOutputs

classmethod load(path, models, derivative_fns=None)

Load a pipeline from a file saved with save().

Models must be provided in the same order they appear in the saved config (flattened across groups). The topology (groups, wire mappings, autograd flags) is restored from the file.

Parameters:

• path (str | Path) – Path to a file created by save().

• models (list[BaseModelMixin]) – Pre-constructed model instances, one per step in the original pipeline (flattened across groups, in order).

• derivative_fns (dict[int, DerivativeFn] | None, optional) – Mapping from group index to custom derivative function. Required for groups that were saved with has_derivative_fn=True.

Return type:

PipelineModelWrapper

Raises:

ValueError – If the number of models doesn’t match the saved config, or if a group requires a derivative_fn that wasn’t provided.

make_neighbor_hooks(max_neighbors=None)

Return a single NeighborListHook for the composite neighbor config.

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[NeighborListHook]

save(path)

Save the full pipeline (topology + model weights) to a file.

The saved file contains:

• "config" — pipeline topology (groups, wire mappings, autograd flags, additive keys).

• "state_dict" — model weights for all sub-models.

• "active_outputs" — current model_config.active_outputs.

Custom derivative_fn callables are not serialized. When loading a pipeline that used a custom function, pass it again via load().

Parameters:

path (str | Path) – Destination file path.

Return type:

None