atlas.active_learning.backends package

MLIP backend registry for model-agnostic active learning.

Subpackages

Submodules

atlas.active_learning.backends._base module

Protocol definitions for MLIP backend interfaces.

class atlas.active_learning.backends._base.MLIPCalculatorFactory(*args, **kwargs)

Bases: Protocol

Protocol for creating ASE calculators from trained MLIP models.

create_calculator(model_path: str | Path, device: str = 'cpu', dtype: str = 'float32', **kwargs) → Calculator

Create an ASE Calculator from a model file on disk.

Parameters:

  • model_path (str | Path) – Path to the model file.

  • device (str) – Device for inference (‘cpu’ or ‘cuda’).

  • dtype (str) – Data type for inference.

Returns:

An ASE-compatible calculator.

Return type:

Calculator

property lammps_pair_style: str

LAMMPS pair_style string for this backend (e.g. ‘mace’, ‘allegro’).

property model_file_extension: str

File extension for model files (e.g. ‘.model’ for MACE).

class atlas.active_learning.backends._base.MLIPCommitteeEvaluator(*args, **kwargs)

Bases: Protocol

Protocol for committee-based uncertainty quantification.

Backends that support training multiple models with different seeds implement this protocol to evaluate structures with the committee and compute disagreement statistics.

evaluate_committee(structures: list[Atoms], model_files: list[str | Path], device: str = 'cpu', dtype: str = 'float32', **kwargs) → dict[str, dict[str, list]]

Evaluate structures with multiple committee models.

Parameters:

  • structures (list[Atoms]) – Structures to evaluate.

  • model_files (list[str | Path]) – Paths to committee model files.

  • device (str) – Device for inference.

  • dtype (str) – Data type for inference.

Returns:

Dictionary keyed by model name, each containing ‘REF_energy’ (meV/atom) and ‘REF_forces’ (meV/A) lists.

Return type:

dict

property supports_committee_training: bool

Whether this backend supports training multiple committee models.

class atlas.active_learning.backends._base.MLIPDescriptorProvider(*args, **kwargs)

Bases: Protocol

Protocol for models that can provide structural descriptors.

Not all MLIP backends support descriptor extraction. Backends that do not can omit this protocol, and the workflow will fall back to SOAP descriptors.

generate_descriptors(database: list[Atoms], model_path: str | Path | None, settings: dict, **kwargs) → tuple[dict, ndarray, list[str]]

Generate structural descriptors for a list of structures.

Parameters:

  • database (list[Atoms]) – List of ASE Atoms objects.

  • model_path (str | Path | None) – Path to the model file for descriptor extraction.

  • settings (dict) – Descriptor generation settings.

Returns:

(descriptor_dict, descriptor_array, uuid_list) where descriptor_dict is keyed by atl_id, descriptor_array is a vertically stacked numpy array, and uuid_list contains the UUIDs of newly assigned structures.

Return type:

tuple

class atlas.active_learning.backends._base.MLIPModelCompiler(*args, **kwargs)

Bases: Protocol

Protocol for backends whose models must be compiled before inference.

Some backends (e.g. Allegro/NequIP) cannot run inference from the raw trained checkpoint. It must first be compiled (nequip-compile) into a device-specific artifact. Compilation is expensive and toolchain-sensitive, so it is best done once, on the computer where inference will run, and the compiled artifact reused there.

Backends that need no compilation (e.g. MACE) simply omit this protocol; callers gate on isinstance(backend, MLIPModelCompiler) and, when absent, ship the raw model unchanged.

compile_model(model_path: str | Path, device: str = 'cpu', mode: str = 'aotinductor', target: str = 'ase') → Path | None

Compile a trained model into an inference-ready artifact.

Runs on the compute node where the artifact will be used, so the compiled result matches that node’s device/toolchain.

Parameters:

  • model_path (str | Path) – Path to the trained model/checkpoint.

  • device (str) – Target device for the compiled model (‘cpu’ or ‘cuda’). The artifact is device-specific.

  • mode (str) – Compilation backend (e.g. ‘aotinductor’).

  • target (str) – Inference target the artifact is compiled for (e.g. ‘ase’).

Returns:

Path to the compiled artifact, or None on failure.

Return type:

Path | None

property compiled_model_extension: str

File extension of the compiled artifact (e.g. ‘.nequip.pt2’).

class atlas.active_learning.backends._base.MLIPTrainer(*args, **kwargs)

Bases: Protocol

Protocol for MLIP training backends.

Backends that support model training implement this protocol to provide the AiiDA CalcJob entry point, training data preparation, builder configuration, and model selection logic.

property calcjob_entry_point: str

AiiDA entry point string for the training CalcJob.

create_lammps_potential(model_file: Any) → Any | None

Convert a trained model to LAMMPS-compatible format.

Parameters:

model_file (SinglefileData) – The trained model file.

Returns:

LAMMPS potential file, or None if not supported.

Return type:

SinglefileData | None

parse_training_results(results_dir: Path) → dict

Extract model file path and metrics from training output.

Called by the remote training script after training completes. Each backend knows its own output format.

Parameters:

results_dir (Path) – Directory containing training outputs.

Returns:

Dictionary with keys ‘model_file’ (str), ‘rmse_e’ (float, meV/atom), ‘rmse_f’ (float, meV/A), and optionally ‘train_log’ (str).

Return type:

dict

property parser_entry_point: str

AiiDA entry point string for the training parser.

prepare_builder(builder: Any, settings_dict: dict, train_data_path: str, model_name: str, iteration: int, db_size: int, containerized: bool = False) → Any

Populate the AiiDA CalcJob builder with model-specific inputs.

Parameters:

  • builder (ProcessBuilder) – The AiiDA CalcJob builder to populate.

  • settings_dict (dict) – Model-specific training settings from the TOML config.

  • train_data_path (str) – Path to the training data file.

  • model_name (str) – Unique name for this training run.

  • iteration (int) – Current AL loop iteration number.

  • db_size (int) – Number of structures in the training database.

  • containerized (bool) – Whether to use containerized execution.

Returns:

The populated builder, ready for submission.

Return type:

ProcessBuilder

prepare_training_data(path: str | Path, structure_list: list[Atoms], **kwargs) → Path

Convert structure list to model-specific training format.

Parameters:

  • path (str | Path) – Output file path for the training data.

  • structure_list (list[Atoms]) – List of ASE Atoms objects with DFT reference data.

Returns:

Path to the written training data file.

Return type:

Path

run_training(config_path: str | Path) → None

Execute model training from a config file.

Called by the remote training script on the HPC node. Each backend dispatches to its own training CLI.

Parameters:

config_path (str | Path) – Path to the backend-specific training config file.

select_best_model(training_results: list, force_weight: float = 0.1) → tuple[str, Any, float, float, list[tuple[str, str]]]

Select the best model from N trained models.

Parameters:

  • training_results (list) – List of completed AiiDA CalcJobNode results.

  • force_weight (float) – Weight for forces in the RMSE-based selection criterion.

Returns:

Tuple of (best_model_name, best_model_file, rmse_e, rmse_f, committee_models_tupl_name_uuid). The last element is a list of (name, uuid) tuples for the non-best models.

Return type:

tuple

Module contents

MLIP backend registry for model-agnostic active learning.

This module provides a registry pattern for MLIP backends. Each backend (MACE, NequIP, Allegro, DeepMD, etc.) registers itself and can then be looked up by name from the TOML configuration.

Usage

from atlas.active_learning.backends import get_backend

backend = get_backend(‘mace’) calculator = backend.create_calculator(model_path, device=’cuda’)

class atlas.active_learning.backends.MLIPCalculatorFactory(*args, **kwargs)

Bases: Protocol

Protocol for creating ASE calculators from trained MLIP models.

create_calculator(model_path: str | Path, device: str = 'cpu', dtype: str = 'float32', **kwargs) → Calculator

Create an ASE Calculator from a model file on disk.

Parameters:

  • model_path (str | Path) – Path to the model file.

  • device (str) – Device for inference (‘cpu’ or ‘cuda’).

  • dtype (str) – Data type for inference.

Returns:

An ASE-compatible calculator.

Return type:

Calculator

property lammps_pair_style: str

LAMMPS pair_style string for this backend (e.g. ‘mace’, ‘allegro’).

property model_file_extension: str

File extension for model files (e.g. ‘.model’ for MACE).

class atlas.active_learning.backends.MLIPCommitteeEvaluator(*args, **kwargs)

Bases: Protocol

Protocol for committee-based uncertainty quantification.

Backends that support training multiple models with different seeds implement this protocol to evaluate structures with the committee and compute disagreement statistics.

evaluate_committee(structures: list[Atoms], model_files: list[str | Path], device: str = 'cpu', dtype: str = 'float32', **kwargs) → dict[str, dict[str, list]]

Evaluate structures with multiple committee models.

Parameters:

  • structures (list[Atoms]) – Structures to evaluate.

  • model_files (list[str | Path]) – Paths to committee model files.

  • device (str) – Device for inference.

  • dtype (str) – Data type for inference.

Returns:

Dictionary keyed by model name, each containing ‘REF_energy’ (meV/atom) and ‘REF_forces’ (meV/A) lists.

Return type:

dict

property supports_committee_training: bool

Whether this backend supports training multiple committee models.

class atlas.active_learning.backends.MLIPDescriptorProvider(*args, **kwargs)

Bases: Protocol

Protocol for models that can provide structural descriptors.

Not all MLIP backends support descriptor extraction. Backends that do not can omit this protocol, and the workflow will fall back to SOAP descriptors.

generate_descriptors(database: list[Atoms], model_path: str | Path | None, settings: dict, **kwargs) → tuple[dict, ndarray, list[str]]

Generate structural descriptors for a list of structures.

Parameters:

  • database (list[Atoms]) – List of ASE Atoms objects.

  • model_path (str | Path | None) – Path to the model file for descriptor extraction.

  • settings (dict) – Descriptor generation settings.

Returns:

(descriptor_dict, descriptor_array, uuid_list) where descriptor_dict is keyed by atl_id, descriptor_array is a vertically stacked numpy array, and uuid_list contains the UUIDs of newly assigned structures.

Return type:

tuple

class atlas.active_learning.backends.MLIPModelCompiler(*args, **kwargs)

Bases: Protocol

Protocol for backends whose models must be compiled before inference.

Some backends (e.g. Allegro/NequIP) cannot run inference from the raw trained checkpoint. It must first be compiled (nequip-compile) into a device-specific artifact. Compilation is expensive and toolchain-sensitive, so it is best done once, on the computer where inference will run, and the compiled artifact reused there.

Backends that need no compilation (e.g. MACE) simply omit this protocol; callers gate on isinstance(backend, MLIPModelCompiler) and, when absent, ship the raw model unchanged.

compile_model(model_path: str | Path, device: str = 'cpu', mode: str = 'aotinductor', target: str = 'ase') → Path | None

Compile a trained model into an inference-ready artifact.

Runs on the compute node where the artifact will be used, so the compiled result matches that node’s device/toolchain.

Parameters:

  • model_path (str | Path) – Path to the trained model/checkpoint.

  • device (str) – Target device for the compiled model (‘cpu’ or ‘cuda’). The artifact is device-specific.

  • mode (str) – Compilation backend (e.g. ‘aotinductor’).

  • target (str) – Inference target the artifact is compiled for (e.g. ‘ase’).

Returns:

Path to the compiled artifact, or None on failure.

Return type:

Path | None

property compiled_model_extension: str

File extension of the compiled artifact (e.g. ‘.nequip.pt2’).

class atlas.active_learning.backends.MLIPTrainer(*args, **kwargs)

Bases: Protocol

Protocol for MLIP training backends.

Backends that support model training implement this protocol to provide the AiiDA CalcJob entry point, training data preparation, builder configuration, and model selection logic.

property calcjob_entry_point: str

AiiDA entry point string for the training CalcJob.

create_lammps_potential(model_file: Any) → Any | None

Convert a trained model to LAMMPS-compatible format.

Parameters:

model_file (SinglefileData) – The trained model file.

Returns:

LAMMPS potential file, or None if not supported.

Return type:

SinglefileData | None

parse_training_results(results_dir: Path) → dict

Extract model file path and metrics from training output.

Called by the remote training script after training completes. Each backend knows its own output format.

Parameters:

results_dir (Path) – Directory containing training outputs.

Returns:

Dictionary with keys ‘model_file’ (str), ‘rmse_e’ (float, meV/atom), ‘rmse_f’ (float, meV/A), and optionally ‘train_log’ (str).

Return type:

dict

property parser_entry_point: str

AiiDA entry point string for the training parser.

prepare_builder(builder: Any, settings_dict: dict, train_data_path: str, model_name: str, iteration: int, db_size: int, containerized: bool = False) → Any

Populate the AiiDA CalcJob builder with model-specific inputs.

Parameters:

  • builder (ProcessBuilder) – The AiiDA CalcJob builder to populate.

  • settings_dict (dict) – Model-specific training settings from the TOML config.

  • train_data_path (str) – Path to the training data file.

  • model_name (str) – Unique name for this training run.

  • iteration (int) – Current AL loop iteration number.

  • db_size (int) – Number of structures in the training database.

  • containerized (bool) – Whether to use containerized execution.

Returns:

The populated builder, ready for submission.

Return type:

ProcessBuilder

prepare_training_data(path: str | Path, structure_list: list[Atoms], **kwargs) → Path

Convert structure list to model-specific training format.

Parameters:

  • path (str | Path) – Output file path for the training data.

  • structure_list (list[Atoms]) – List of ASE Atoms objects with DFT reference data.

Returns:

Path to the written training data file.

Return type:

Path

run_training(config_path: str | Path) → None

Execute model training from a config file.

Called by the remote training script on the HPC node. Each backend dispatches to its own training CLI.

Parameters:

config_path (str | Path) – Path to the backend-specific training config file.

select_best_model(training_results: list, force_weight: float = 0.1) → tuple[str, Any, float, float, list[tuple[str, str]]]

Select the best model from N trained models.

Parameters:

  • training_results (list) – List of completed AiiDA CalcJobNode results.

  • force_weight (float) – Weight for forces in the RMSE-based selection criterion.

Returns:

Tuple of (best_model_name, best_model_file, rmse_e, rmse_f, committee_models_tupl_name_uuid). The last element is a list of (name, uuid) tuples for the non-best models.

Return type:

tuple

atlas.active_learning.backends.find_inference_model(directory, stem: str, backend)

Return the model file to hand to backend.create_calculator.

Prefers a pre-compiled inference artifact (shipped by the optional compile-once step, e.g. {stem}.nequip.pt2) when present in directory, otherwise falls back to the raw model {stem}{backend.model_file_extension}.

Backends without compilation (no MLIPModelCompiler) always resolve to the raw model, so behaviour is unchanged for them.

atlas.active_learning.backends.get_backend(name: str)

Instantiate and return a backend by name.

Parameters:

name (str) – Registered backend name.

Raises:

ValueError – If the backend name is not registered.

atlas.active_learning.backends.list_backends() → list[str]

Return the list of registered backend names.

atlas.active_learning.backends.list_inference_models(directory, backend) → list

Return all model files in directory to use for inference.

Prefers pre-compiled artifacts (*.nequip.pt2) when any are present, otherwise the raw models (*{backend.model_file_extension}). Used by the committee evaluation, which loads every committee member. Falls back to the raw extension for backends without compilation.

atlas.active_learning.backends.model_file_ext(filename: str, default: str = '.model') → str

Return a model file’s extension, honouring compound nequip suffixes.

Path('m.nequip.zip').suffix is only '.zip'; this returns the full '.nequip.zip' so calcjobs name copied model files consistently with what the remote scripts (which use backend.model_file_extension) expect.

atlas.active_learning.backends.model_file_stem(filename: str) → str

Return a model filename with its (possibly compound) extension removed.

atlas.active_learning.backends.register_backend(name: str)

Decorator to register an MLIP backend class.

Parameters:

name (str) – Short name for the backend (e.g. ‘mace’, ‘nequip’, ‘allegro’).