atlas.active_learning.backends.mace package

Concrete MACE implementation of the MLIP backend protocols.

Submodules

atlas.active_learning.backends.mace.calcjobs module

MACE-specific AiiDA CalcJob and Parser classes.

These classes were extracted from mace_tools_aiida.py as part of the MLIP-agnostic refactoring. The original module re-exports them for backward compatibility and to preserve AiiDA entry points.

class atlas.active_learning.backends.mace.calcjobs.CheckMACECommiteeResultsCalculationParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for processing the retrieved files from a MACE committee results job.

parse(**kwargs)

Parse the retrieved files of the calculation job.

class atlas.active_learning.backends.mace.calcjobs.CheckMACECommitteeResultsCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

CalcJob to check the E and F of structures using a committee of MACE models.

Inputs

commitee_modelsPortNamespace

A namespace to hold an arbitrary number of committee MACE potentials.

mace_settings_dictaiida.orm.Dict

Dictionary containing MACE settings.

configurations_to_evaluateaiida.orm.orm.SinglefileData

Path to the configurations to evaluate in extxyz format.

Outputs

energy_result_dictaiida.orm.Dict

Dictionary of values for the energy prediction.

forces_result_dictaiida.orm.Dict

Dictionary of arrays of values for the force prediction.

num_threadsaiida.orm.Int

Number of OpenMP threads to use for the evaluation.

Exit Codes

420ERROR_OUT_OF_VRAM

CUDA out of GPU memory.

421ERROR_OUTPUT_NOT_FOUND

Missing output file.

classmethod define(spec)

prepare_for_submission(folder)

Write the input files that are required for the code to run.

Parameters:

folder – an Folder to temporarily write files on disk

Returns:

CalcInfo instance

class atlas.active_learning.backends.mace.calcjobs.EvaluateMACEConfigsCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

CalcJob to evaluate E and F of structures using a MACE model.

classmethod define(spec)

prepare_for_submission(folder)

Write the input files that are required for the code to run.

Parameters:

folder – an Folder to temporarily write files on disk

Returns:

CalcInfo instance

class atlas.active_learning.backends.mace.calcjobs.EvaluateMACEConfigsCalculationParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for MACE E and F evaluation calculation jobs.

parse(**kwargs)

Parse the retrieved files of the calculation job.

class atlas.active_learning.backends.mace.calcjobs.GetMACEDescriptorsCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

Calculation to obtain descriptors for a structure database from MACE.

classmethod define(spec)

prepare_for_submission(folder)

Write the input files that are required for the code to run.

Parameters:

folder – a Folder to temporarily write files on disk

Returns:

CalcInfo instance

class atlas.active_learning.backends.mace.calcjobs.GetMACEDescriptorsCalculationParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for the retrieved files from a MACE descriptors job.

parse(**kwargs)

Parse the retrieved files of the calculation job.

class atlas.active_learning.backends.mace.calcjobs.LAMMPSMACERawParser(*args: Any, **kwargs: Any)

Bases: Parser

Base parser for LAMMPS output.

parse(**kwargs)

Parse the output files stored in the retrieved output node.

class atlas.active_learning.backends.mace.calcjobs.RunMDCalculationGPULAMMPS(*args: Any, **kwargs: Any)

Bases: LammpsRawCalculation

aiida-lammps raw calculation modified to run on GPU using Kokkos.

This CalcJob is backend-agnostic. The LAMMPS pair_style (mace, allegro, etc.) is determined by the input script, not this class.

prepare_for_submission(folder: aiida.common.folders.Folder) → aiida.common.datastructures.CalcInfo

Prepare the calculation for submission.

Parameters:

folder – A temporary folder on the local file system.

Returns:

A aiida.common.datastructures.CalcInfo instance.

atlas.active_learning.backends.mace.calcjobs.RunMDCalculationGPULAMMPSMACE

alias of RunMDCalculationGPULAMMPS

class atlas.active_learning.backends.mace.calcjobs.TrainMACEModelCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

Implementation of a CalcJob to perform a MACE training using a settings dir.

Inputs

mace_settings_dictorm.Dict

Dictionary containing MACE settings.

mace_train_file_pathorm.Str

Local machine path to the structures to evaluate in extxyz format.

test_fileorm.SinglefileData

Local machine path to the structures for testing in extxyz format.

mace_train_file_pathorm.Str

Path to the configurations to evaluate in extxyz format.

model_nameorm.Str

Name given to the model.

use_containerorm.Bool

Use code in container. Default is False. Will be set automatically by the code if the containerized mode is enabled.

Outputs

model_fileorm.SinglefileData

Trained MACE model.

train_fileorm.SinglefileData

Log file containing training information.

m_rmse_eorm.Float

Validation RMSE for the energy, in meV / atom.

m_rmse_form.Float

Validation RMSE for the forces, in meV / A.

Exit Codes

420ERROR_INVALID_OUTPUT

Training calculation could not run.

classmethod define(spec)

Define the input and output specifications for the CalcJob.

prepare_for_submission(folder)

Write the input files that are required for the code to run.

Parameters:

folder – an Folder to temporarily write files on disk

Returns:

CalcInfo instance

class atlas.active_learning.backends.mace.calcjobs.TrainMACEModelCalculationParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for the retrieved files from a MACE training calculation job.

parse(**kwargs)

Parse the retrieved files of the calculation job.

atlas.active_learning.backends.mace.calcjobs.prepare_cli_args_mace(params_list: list, settings_dict: dict, use_container: bool = False)

Prepare the command line arguments for the MACE calculation.

atlas.active_learning.backends.mace.calculator module

MACE ASE calculator creation.

Extracted from atlas.active_learning.active_learning_utils as part of the MLIP-agnostic refactoring.

atlas.active_learning.backends.mace.calculator.create_mace_calculator(model_path: str | Path, device: str = 'cpu', dtype: str = 'float32', **kwargs) → Calculator

Create an ASE Calculator from a trained MACE model.

Parameters:

• model_path (str | Path) – Path to the MACE .model file, or a foundation model identifier (e.g. "mace:mp-small").

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

• dtype (str) – Data type for inference.

• **kwargs – Additional keyword arguments passed to MACECalculator.

Returns:

An ASE-compatible MACE calculator.

Return type:

Calculator

atlas.active_learning.backends.mace.descriptors module

MACE descriptor generation.

Extracted from atlas.active_learning.active_learning_utils as part of the MLIP-agnostic refactoring.

atlas.active_learning.backends.mace.descriptors.generate_descriptors_mace(model_path: str, database: list[Atoms], descriptor_settings: dict, outer_average: bool = False, verbose: bool = False) → tuple[dict, ndarray, list[str]]

Generate per-structure MACE descriptors for a database.

Parameters:

• model_path (str) – Path to a trained MACE model file, or a foundation model identifier (e.g. "mace:mp-small", "mace:off-medium").

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

• descriptor_settings (dict) – Settings dict containing device and dtype keys.

• outer_average (bool) – If True, average atom-level descriptors into a single structure-level vector (analogous to SOAP outer averaging).

• verbose (bool) – If True, show a progress bar.

Returns:

• descriptor_dict (dict) – Mapping {atl_id: {'descriptors': [...], 'latent_space': []}}

• descriptor_arr (np.ndarray) – Vertically stacked descriptor array.

• uuid_list (list[str]) – UUIDs assigned to structures that lacked an atl_id.

atlas.active_learning.backends.mace.training module

MACE training utilities.

Extracted from atlas.active_learning.active_learning_utils as part of the MLIP-agnostic refactoring.

atlas.active_learning.backends.mace.training.create_mace_lammps_model_impl(model_file: SinglefileData)

Create a LAMMPS potential from a MACE model (inner logic, no @calcfunction).

The @calcfunction-decorated version lives in active_learning_utils.create_mace_lammps_model for AiiDA provenance.

Parameters:

model_file (orm.SinglefileData) – A MACE model file to convert to a LAMMPS potential.

Returns:

A LAMMPS potential file generated from the MACE model.

Return type:

orm.SinglefileData

atlas.active_learning.backends.mace.training.update_mace_train_settings_dict(settings_dict: dict, train_data_path: str, curr_model: str, curr_iter: int, db_size: int, containerized: Bool = False)

Update the MACE training settings dictionary with the new database path.

Module contents

MACE backend for ATLAS active learning.

This module provides the concrete MACE implementation of the MLIP backend protocols. It wraps existing MACE-specific functions from the codebase, delegating to them without duplicating logic.

class atlas.active_learning.backends.mace.MACEBackend

Bases: object

MACE backend implementing all four MLIP protocols.

This backend delegates to the existing MACE-specific functions in active_learning_utils, conversion, and mace_tools_aiida.

property calcjob_entry_point: str

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

create_lammps_potential(model_file) → object | None

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

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

property lammps_pair_style: str

property model_file_extension: str

parse_training_results(results_dir: Path) → dict

property parser_entry_point: str

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

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

run_training(config_path: str | Path) → None

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

property supports_committee_training: bool