atlas.active_learning package

Subpackages

Submodules

atlas.active_learning.active_learning_utils module

General utility functions for the active learning workflows.

atlas.active_learning.active_learning_utils.aiida_serialized_ase_dict_to_atoms(struct_dict: dict) → Atoms

Convert a serialized Atoms dictionary to an Atoms object.

atlas.active_learning.active_learning_utils.aiida_wait_submit(builder, computer: orm.Computer, calc_count: int = 0, code: orm.Code | str = None)
atlas.active_learning.active_learning_utils.calculate_fps_scores_descriptor(selected_uuids: list[str], descriptor_dict: dict, n_to_select: int) → dict[str, dict]

Calculates FPS scores for candidate structures by finding points farthest from an ALREADY SELECTED subset of structures.

atlas.active_learning.active_learning_utils.check_atl_ids(atoms_list: list[Atoms])

Checks for ‘atl_id’ key in info dicts and reports missing or repeated IDs.

atlas.active_learning.active_learning_utils.check_atom_in_domain(concave_hull: aiida.orm.ArrayData, descriptors: aiida.orm.ArrayData) → tuple[ndarray, ndarray, ndarray]
atlas.active_learning.active_learning_utils.check_md_seed_agreement(return_list_path: str | None, md_structs_in_domain: bool | None) → aiida.orm.Bool

Check if all predictions agree for current seed.

Parameters:

return_list_path (str | None) – Path pointing to the file that contains all calculations for predictions where the models disagreed.

Returns:

True if all the predictions have agreed for the current MD seed on the current AL iteration. False if there is no agreement on on all structures.

Return type:

orm.Bool

atlas.active_learning.active_learning_utils.convert_database_to_ase_atoms(database: list, deserialize: bool = False) → list[Atoms]

Converts a struture list/array into containing both dicts and ase.Atoms into a list containing only ase.Atoms.

atlas.active_learning.active_learning_utils.create_mace_lammps_model(model_file: aiida.orm.SinglefileData)

Create a LAMMPS potential from a MACE model.

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.active_learning_utils.filter_dft_calcs_threshold(dft_calc_list: list, threshold_E_meV: float, threshold_F_meV: float, workchain=None) → list

Filter DFT calculations based on energy and force thresholds.

Returns a list of serialized ASE Atoms objects that have forces and energy below the specified thresholds.

atlas.active_learning.active_learning_utils.flush_system_memory()

Forces Python to sweep unreferenced objects and instructs the Linux C-allocator to return freed memory blocks to the OS. Use this only after deleting massive arrays or dataframes.

atlas.active_learning.active_learning_utils.gather_dft_calcs_mlip(dft_calc_list: list, results_dir: str, workchain=None) → aiida.orm.List

Collect and preprocess MLIP evaluation results for active learning input.

atlas.active_learning.active_learning_utils.gather_dft_calcs_vasp(dft_calc_list: list) → aiida.orm.List

Collect and preprocess VASP DFT calculation results for active learning input.

This function takes a list of DFT calculation nodes, extracts the calculation results, and processes these results into a format suitable for active learning input. Specifically, it converts VASP runs into ASE Atoms objects and collects additional calculation data like forces. It also augments the Atoms objects with metadata necessary for the active learning workflow. Failed calculations are skipped ensuring that only successfully completed CalcJobs are included. The function returns a list of serialized ASE Atoms objects, ready for inclusion in the active learning database.

Parameters:

dft_calc_list (list) – A list of identifiers for completed DFT calculation nodes.

Returns:

An AiiDA orm.List object containing serialized ASE Atoms objects, each representing a completed DFT calculation augmented with necessary metadata and calculation results.

Return type:

orm.List

Notes

  • The ASE Atoms objects are serialized to ensure compatibility with AiiDA’s data

storage and manipulation frameworks. - Extra care is taken to include forces (and optionally, stress) in the Atoms objects, as these are critical for many active learning applications but are not included by default in the extxyz format’s Properties tag. - Skips any DFT calculations that encountered errors.

atlas.active_learning.active_learning_utils.gen_xyz_file_from_traj(struct_list)

Generate a temporary xyz file from a list of structures.

atlas.active_learning.active_learning_utils.generate_descriptors(database: list[Atoms] | ndarray, descriptor_type: str, descriptor_settings: dict, model_path: str | Path | None = None, outer_average_mace: bool = False, verbose: bool = False) → tuple[dict, ndarray, list[str]]

Wrapper function to generate descriptors for a given database.

Allows for the generation of descriptors using different methods (e.g., MACE, SOAP) based on the descriptor_type parameter.

Parameters:

  • database (list[Atoms] | np.ndarray) – List or array of structures for which to generate descriptors.

  • descriptor_type (str) – Type of descriptor to generate. Options are ‘soap’ or ‘mace’.

  • device (str, optional) – Compute device, by default ‘cpu’

  • dtype (str, optional) – Floating point number precision, by default ‘float32’

  • model_path (str, optional) – For MLIP based descriptors, the pretrained model path, by default None

  • descriptor_settings (dict, optional) – Descriptor settings dictionary.

Returns:

A tuple containing a dictionary of descriptors and a numpy array of vstacked descriptors.

Return type:

tuple[dict, np.ndarray, list[str]]

atlas.active_learning.active_learning_utils.generate_descriptors_mace(model_path: str, database, descriptor_settings: dict, outer_average: bool = False, verbose: bool = False) → tuple[dict, ndarray, list[str]]
atlas.active_learning.active_learning_utils.generate_descriptors_soap(database: Atoms | list[Atoms], descriptor_settings: dict, verbose: bool = False) → tuple[dict, ndarray, list[str]]
atlas.active_learning.active_learning_utils.generate_model_name()

Generate a unique NNP model name combining random words and a number.

This function creates a unique model name by concatenating randomly selected adjective, noun, and verb, followed by a random number. This combination ensures the generation of distinctive and memorable names suitable for labeling models in simulations or learning tasks.

Returns:

A string consisting of a random adjective, noun, and verb followed by a hyphen and a random number between 1 and 99, forming a unique model name.

Return type:

str

atlas.active_learning.active_learning_utils.get_dft_calc_builder_mlip_list(struct_list: list, dft_settings: dict, container_settings: dict)

Get an MLIP evaluation builder for a given structure list.

atlas.active_learning.active_learning_utils.get_dft_calc_builder_vasp(struct, row, calc_idx: int, group, dft_settings: dict)

Generate a aiida-vasp calculation builder for a given structure and row.

atlas.active_learning.active_learning_utils.get_final_db_path(result_dir_path, final_db_name, node)

Get the path to the final database file.

atlas.active_learning.active_learning_utils.get_model_energies_std(energies_dict: dict) → ndarray

Get the standard deviation of the energies for each structure in the dict.

atlas.active_learning.active_learning_utils.get_model_energies_variance(energies_dict: dict) → ndarray

Get the variance of the energies for each structure in the dict.

atlas.active_learning.active_learning_utils.get_model_forces_std(forces_dict: dict) → ndarray

Get the standard deviation of the forces for each structure in the dict.

atlas.active_learning.active_learning_utils.get_model_forces_variance(forces_dict: dict) → ndarray

Get the variance of the forces for each structure in the dict.

atlas.active_learning.active_learning_utils.get_outliers_from_calc_list(curr_struct_res, result_list, outlier_list)
atlas.active_learning.active_learning_utils.get_results_dir_path(result_dir_path, node, check_temp_dir=True)

Get the path to the results directory.

atlas.active_learning.active_learning_utils.get_species_from_database(database: list[Atoms] | Atoms) → list[str]

Get the list of species from the database of structures.

Parameters:

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

Returns:

List of unique species in the database.

Return type:

list[str]

atlas.active_learning.active_learning_utils.get_total_num_frames(len_traj, md_tstep_duration_ps, frame_interval)

Compute the number of frames to get from the trajectory using user input.

atlas.active_learning.active_learning_utils.iqr_outlier_check(res_list: list) → ndarray

Identifies outliers in a list of E/F values using the interquartile range (IQR).

Parameters:

res_list (list) – A list of numerical values to check for outliers.

Returns:

An array where outliers are replaced with NaN and non-outliers are retained.

Return type:

np.ndarray

Notes

The function calculates the 30th and 70th percentiles of the input list to determine the interquartile range (IQR). Values outside the range [Q1 - 1.5 * IQR, Q2 + 1.5 * IQR] are considered outliers and replaced with NaN.

atlas.active_learning.active_learning_utils.load_database(path: str) → list[Atoms]

Load an extended xyz file (database) from a given path as a list of ASE Atoms.

Automatically standardizes legacy mdb_* info keys to atl_* and ensures every structure has an atl_id UUID.

atlas.active_learning.active_learning_utils.manual_progress_display(dyn)
atlas.active_learning.active_learning_utils.md_apply_temperature_ramp(dyn, total_steps, T_start, T_end, T_list)

Function to compute the temperature ramp during ASE MD simulations.

Parameters:

  • step (int) – Current step in the MD simulation.

  • total_steps (int) – Total number of steps in the MD simulation.

  • T_start (float) – Initial temperature of the MD simulation.

  • T_end (float) – Final temperature of the MD simulation.

Returns:

Temperature to set for the current step in the MD simulation.

Return type:

float

atlas.active_learning.active_learning_utils.md_coexistence_final_step_log(dyn, T_list)

Function to compute the temperature ramp during ASE MD simulations.

Parameters:

  • step (int) – Current step in the MD simulation.

  • total_steps (int) – Total number of steps in the MD simulation.

  • T_start (float) – Initial temperature of the MD simulation.

  • T_end (float) – Final temperature of the MD simulation.

Returns:

Temperature to set for the current step in the MD simulation.

Return type:

float

atlas.active_learning.active_learning_utils.md_save_gen_structs(dyn, struct_list)
atlas.active_learning.active_learning_utils.md_stop_explode_filter(dyn, max_T, max_T_multiplier, T_list, remove_positive_E)
atlas.active_learning.active_learning_utils.md_write_frame_traj(dyn, traj, stage_name: str = None)

Function to write frames to a trajectory during ASE MD simulations.

Parameters:

  • dyn (ASE MD object) – ASE MD object used to run the MD simulation.

  • traj (TrajectoryWriter) – ASE trajectory object to store the MD simulation.

  • stage_name (str) – Name of the MD stage.

atlas.active_learning.active_learning_utils.model_res_dict_to_arr(res_dict: dict, dict_type: str) → ndarray

Convert a dictionary of model results to a numpy array.

Parameters:

  • res_dict (dict) – Dictionary containing the model results.

  • dict_type (str) – Type of the dictionary. Either “energy” or “forces”.

Returns:

Numpy array containing the model results.

Return type:

np.ndarray

atlas.active_learning.active_learning_utils.plot_concave_hull(concave_hull: ndarray, point_inside: ndarray, point_outside: ndarray, latent_space: ndarray, filename: str = 'concave_hull.png')
atlas.active_learning.active_learning_utils.prepare_output_dataframe(md_seed_results_df)

Prepare the output dataframe for the active learning workflow.

atlas.active_learning.active_learning_utils.prepare_output_final_training_db(training_db_path)

Convert the training database to a orm.SinglefileData object.

atlas.active_learning.active_learning_utils.prepare_test_set(test_db_path: str, test_db_frac: float, training_db: Atoms | list[Atoms])

Prepare a test set from the training database based on the provided settings.

A test set is prepared either by reading a user provided file or by randomly selecting structures from the training database (Dt). Structures selected from Dt are then removed in order to avoid data leakage.

Parameters:

  • test_db_path (str) – Path to the user provided test set file. If None, random selection is used.

  • test_db_frac (float) – Fraction of structures to select from the training database for the test set.

  • training_db (Atoms | list[Atoms]) – List of ASE Atoms objects representing the training database.

Returns:

  • orm.SinglefileData – A SinglefileData object containing the test set structures.

  • list[Atoms] – The obtained test database.

  • list[Atoms] – The updated training database with test set structures removed.

atlas.active_learning.active_learning_utils.process_call_root(process)

Show root process of the call stack for the given process.

atlas.active_learning.active_learning_utils.read_toml_settings(settings_file: str | Path) → dict

Read a TOML file containing settings for the active learning workflow.

atlas.active_learning.active_learning_utils.remove_isolated_atoms(train_db, E_dft_list_per_at, E_nn_list_per_at, F_dft_list_per_at, F_nn_list_per_at, E_diff_list_meV, F_diff_list_meV)

Remove isolated atoms from the training database and associated lists. This function identifies isolated atoms in the training database and removes them from the database and associated lists. It returns the updated training database and lists without the isolated atoms.

atlas.active_learning.active_learning_utils.remove_structs_from_seed_gen_db(seed_gen_path: aiida.orm.Str, delete_indices: list) → aiida.orm.List

Remove specified structures from a seed generation database based on UUIDs.

This function iterates over a list of UUIDs (delete_indices) and removes the corresponding structures from a seed generation database. The database is accessed via the seed_gen_db object, which is loaded from the seed_gen_path using ASE. Each element of the list is an ase.Atoms object with an unique identifier (atl_id/aiida_uuid) in the info attribute. The function writes the modified list back into seed_gen_path after the specified ones have been removed. No list is returned into the workchain to avoid having to serialize the atoms list.

Parameters:

  • seed_gen (orm.Str | str) – The path to the seed generation database.

  • delete_indices (list) – A list of UUIDs (strings) identifying the structures to be removed from the seed generation database.

atlas.active_learning.active_learning_utils.return_code_from_settings(current_settings: dict, code_settings: dict, workchain: aiida.orm.Node, num_threads: int, executable_name: str, code_path: str, portable_code_label: str, builder) → tuple[aiida.orm.Code, str]

Return a (code, prepend_text) pair for the given settings.

The code is cached via get_or_create_portable_code or get_or_create_containerized_code so it is reused across calls with identical parameters. prepend_text is returned separately so the caller can pass it dynamically via builder.metadata.options.prepend_text.

atlas.active_learning.active_learning_utils.run_mace_md_ase(init_conf: Atoms, md_params: dict, T_start: float, traj_obj: TrajectoryWriter | None, prepend_path: str | Path = '.', explode_filter_dict: dict = None, mode='normal', md_struct_list: list = None, enable_cueq: bool = False, model_name: str = None, stage_name: str = None)

Run MD simulations using ASE and MACE.

Parameters:

  • init_conf (Atoms) – Initial structure to use for the MD simulation.

  • md_params (dict) – Dictionary containing the MD parameters.

  • T_start (float) – Initial temperature of the MD simulation.

  • traj_obj (ASE trajectory object) – ASE trajectory object to store the MD simulation.

  • prepend_path (str, optional) – Path to prepend to the model path, by default None

  • explode_filter (bool, optional) – Whether to apply the MD explode filter.

  • mode (str) – Operation mode for this function. One of ‘normal’ or ‘init_db’.

  • enable_cueq (bool, optional) – Whether to enable the CUEQ mode for the MD simulation. If True, the MD simulation will be run in CUEQ mode. Default is False.

  • model_name (str) – Name of the model to use. If None, ‘curr_model.model’ is used. Default is None.

  • stage_name (str) – Name of the MD stage. Default is None.

atlas.active_learning.active_learning_utils.sampler_populate_E_and_F_list(structure_list: list[Atoms], model_file: aiida.orm.SinglefileData, backend_name: str = 'mace')
atlas.active_learning.active_learning_utils.save_descriptors(save_path: str | Path, descriptor_dict: dict, save_format: str = 'pkl')
atlas.active_learning.active_learning_utils.select_dft_structures(struct_arr, frame_interval)

Select DFT structures using the interval given as an input of the workchain.

Parameters:

  • struct_arr (np.array) – Array containing all possible structures to compute.

  • frame_interval (orm.Int) – Integer representing the interval between structures to keep.

Returns:

Array containing only the selected structures.

Return type:

np.array

atlas.active_learning.active_learning_utils.select_md_frames_to_keep(frame_interval: int, md_tstep_duration_ps: float, traj, steps_E_F_arr: array, forces: array)

Select MD frames to keep using the frame interval and total number of frames.

atlas.active_learning.active_learning_utils.select_structures_data_reduction(database: list[Atoms], n_structures: int, selection_method: str, descriptor_settings: dict = None, model_files: list = None, **kwargs) → list[Atoms]

Select structures from database using the specified method for data reduction.

Parameters:

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

  • n_structures (int) – Number of structures to select.

  • selection_method (str) – Selection method: ‘random’, ‘lowest_energy’, ‘fps’, ‘uncertainty’.

  • descriptor_settings (dict, optional) – Settings for descriptor calculation (needed for fps and uncertainty).

  • model_files (list, optional) – List of model files for uncertainty calculation.

  • **kwargs – Additional arguments for specific selection methods.

Returns:

List of selected structures.

Return type:

list[Atoms]

atlas.active_learning.active_learning_utils.select_structures_fps(candidate_db: list[Atoms], selected_db: list[Atoms], n_structures: int, descriptor_settings: dict, outer_average_mace: bool = False, model_path: str = None) → list[Atoms]

Select n structures using Farthest Point Sampling.

atlas.active_learning.active_learning_utils.select_structures_lowest_energy(database: list[Atoms], n_structures: int) → list[Atoms]

Select the n structures with lowest energy from the database.

Parameters:

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

  • n_structures (int) – Number of structures to select.

Returns:

List of selected structures sorted by energy.

Return type:

list[Atoms]

atlas.active_learning.active_learning_utils.select_structures_random(database: list[Atoms], n_structures: int) → list[Atoms]

Randomly select n structures from the database.

Parameters:

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

  • n_structures (int) – Number of structures to select.

Returns:

List of selected structures.

Return type:

list[Atoms]

atlas.active_learning.active_learning_utils.select_structures_uncertainty(database: list[Atoms], n_structures: int, model_files: list, descriptor_settings: dict) → list[Atoms]

Select n structures with highest uncertainty based on model committee disagreement.

Parameters:

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

  • n_structures (int) – Number of structures to select.

  • model_files (list) – List of model file paths for the committee.

  • descriptor_settings (dict) – Settings for model evaluation.

Returns:

List of selected structures with highest uncertainty.

Return type:

list[Atoms]

atlas.active_learning.active_learning_utils.serialize_ase(curr_s: dict | Atoms) → dict

Serialize an ASE Atoms object to a dictionary.

atlas.active_learning.active_learning_utils.simplify_forces_struct(forces: ndarray)
atlas.active_learning.active_learning_utils.standardize_atoms_info(atoms_list: list[Atoms]) → list[Atoms]

Ensure every ASE structure has standardized atl_* info keys.

Migrates legacy mdb_* keys to their atl_* equivalents and ensures every structure carries an atl_id UUID. Legacy keys are removed after migration. Called automatically by load_database() so that databases generated by older versions of the codebase are transparently upgraded.

Parameters:

atoms_list (list of ase.Atoms) – The list of ASE atoms objects to process.

Returns:

The updated list with standardized keys in info.

Return type:

list of ase.Atoms

atlas.active_learning.active_learning_utils.suppress_stdout()

Temporarily redirects standard output to the void.

atlas.active_learning.active_learning_utils.update_mace_train_settings_dict(settings_dict: dict, train_data_path: str, curr_model: str, curr_iter: int, db_size: int, containerized: aiida.orm.Bool = False)

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

atlas.active_learning.active_learning_utils.vasprun_add_info_dict(vasprun_dict: dict, calc_info_dict: dict) → dict

Add calculation information to the vasprun dictionary.

atlas.active_learning.active_learning_utils.write_gathered_dft_calcs_to_file(dft_calc_list: aiida.orm.List, results_dir: str, workchain=None) → tuple[Path, Path]

atlas.active_learning.conversion module

Utility functions for unit conversion and to convert data to MACE and N2P2.

class atlas.active_learning.conversion.Units(value)

Bases: Enum

Conversion factors for units used in the conversion functions.

Ang2Bohr: float = 1.8897261246257702
Bohr2Ang: float = 0.529177210903
Eh2eV: float = 27.211386245988
eV2Eh: float = 0.03674932217565499
kB: float = 8.314
atlas.active_learning.conversion.atl_database_to_mace_train(atl_database: InitialDatabase, path: str | Path, skip_dipole=True, skip_stress=True, skip_free_energy=False)

Converts an initial database to an extended XYZ format file suitable for MACE training.

Parameters:

  • atl_database (atl_indb.InitialDatabase) – The initial database containing molecular structures.

  • path (Union[str, pathlib.Path]) – The file path where the extxyz file will be written.

  • skip_dipole (bool, optional) – If True, dipole information is not written to the file (default is True).

  • skip_stress (bool, optional) – If True, stress information is not written to the file (default is True).

  • skip_free_energy (bool, optional) – If False, free energy information is included in the file (default is False).

atlas.active_learning.conversion.gather_calc_data_from_node(node, units='atomic')

Get data from a vasp calulation node to be used in a training set.

atlas.active_learning.conversion.gen_mace_train_aiida(aiida_group_list: list, filter_dict: dict, path: str = None, remove_dipole=False, remove_stress=False)

Generate a MACE training file from a list of AiiDA groups.

atlas.active_learning.conversion.gen_mace_train_structure_list(path: str | pathlib.Path, structure_list: list | orm.List, skip_dipole=True, skip_stress=True, skip_free_energy=False)

Generate a MACE training file from a list of structures.

atlas.active_learning.conversion.gen_n2p2_train_aiida(aiida_group_list: list, filter_dict: dict, path: str = None)

Generate a N2P2 training file from a list of AiiDA groups.

atlas.active_learning.conversion.get_struct_type(vasprun, dft_calc_node)

Return the structure type of a calculation from calc settings.

atlas.active_learning.mace_tools_aiida module

AiiDA plugin for MACE calculations.

MACE-specific CalcJobs and Parsers have been extracted to atlas.active_learning.backends.mace.calcjobs as part of the MLIP-agnostic refactoring. They are re-exported here for backward compatibility and to preserve existing AiiDA entry points.

class atlas.active_learning.mace_tools_aiida.EvalTestDatabaseCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

CalcJob to evaluate the test database using the sampler model (M0).

This calculation predicts all energies and forces for the test database, computing the MAE and RMSE values for both properties at each step of the active learning iteration, saving the results in the workchain and preparing a report figure showing the evolution of the error over time.

Parameters:

  • spec (aiida.engine.processes.ports.PortNamespace) – The process specification to define the inputs, outputs, and exit codes.

  • Inputs

  • ------

  • sampler_model (aiida.orm.SinglefileData) – File containing the MACE model to use for evaluation.

  • current_iteration (aiida.orm.Int) – Current iteration number.

  • settings_file_path (aiida.orm.Str) – Path to the ATL settings file in the toml format.

  • test_database (aiida.orm.SinglefileData) – File containing the structures for testing the model.

  • test_db_eval_results (aiida.orm.Dict) – Dictionary containing the evaluation results up until now.

Outputs

rmse_eaiida.orm.Float

Root mean square error for energy predictions.

rmse_faiida.orm.Float

Root mean square error for force predictions.

mae_eaiida.orm.Float

Mean absolute error for energy predictions.

mae_faiida.orm.Float

Mean absolute error for force predictions.

eval_plotaiida.orm.SinglefileData

File containing the evaluation plot.

test_db_eval_resultsaiida.orm.SinglefileData

Dictionary containing the updated evaluation results.

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.mace_tools_aiida.EvalTestDatabaseCalculationParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for a descriptor and extrapolation gathering job.

parse(**kwargs)

Parses the temporarily retrieved files. Outputs are stored in AiiDA SinglefileData objects.

parse(**kwargs)

Parse the retrieved files of the calculation job.

Returns:

  • descriptor_max (aiida.orm.SinglefileData) – File containing the maximum values for the descriptors.

  • descriptor_min (aiida.orm.SinglefileData) – File containing the minimum values for the descriptors.

  • concave_hull (aiida.orm.SinglefileData, optional) – File containing the concave hull of the latent space as an array.

  • latent_space (aiida.orm.SinglefileData, optional) – File containing the latent space represented as an array.

class atlas.active_learning.mace_tools_aiida.GetDescriptorsCombinedCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

CalcJob to gather the descriptors for the training database of an AL Loop.

This calculation job computes the descriptors for all the configurations in the training database. Additionally, further extrapolation metrics are computed depending on the the selected extrapolation type. With min-max extrapolation enabled, the the minimum and maximum range for all descriptors in the training database is computed. With advanced extrapolation, the ranges plus the concave hull of the latent space for all configurations in the training database, are provided, along with a plot showing the configuration distribution in the latent space and the concave hull.

Parameters:

  • spec (aiida.engine.processes.ports.PortNamespace) – The process specification to define the inputs, outputs, and exit codes.

  • Inputs

  • ------

  • commitee_models (PortNamespace) – A namespace to hold an arbitrary number of committee MACE potentials.

  • settings_file_path (orm.Str) – Path to the ATL settings file in the .toml format.

  • training_database_path (orm.Str) – Path to the configurations to evaluate, provided in the extxyz format.

  • autoencoder_model (orm.SinglefileData, optional) – File containing the autoencoder model. If not provided, a new autoencoder is trained is computed.

  • latent_space (orm.ArrayData, optional) – File containing the latent space represented as an array. If not provided, the latent space is computed.

  • Outputs

  • -------

  • descriptor_max (orm.ArrayData) – File containing the maximum values for the descriptors.

  • descriptor_min (orm.ArrayData) – File containing the minimum values for the descriptors.

  • latent_space – File containing the latent space as an array.

  • concave_hull (orm.ArrayData, optional) – File containing the concave hull of the latent space as an array.

  • extrapolation_plot (atl_img.ImagePNGData, optional) – Figure showing the extrapolation for the current database.

  • autoencoder_model – File containing the autoencoder model.

  • Codes (Exit)

  • ----------

  • 420 (ERROR_OUT_OF_VRAM) – CUDA out of GPU memory.

  • 421 (ERROR_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.mace_tools_aiida.GetDescriptorsCombinedParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for a descriptor and extrapolation gathering job.

parse(**kwargs)

Parses the temporarily retrieved files. Outputs are stored in AiiDA SinglefileData objects.

parse(**kwargs)

Parse the retrieved files of the calculation job.

Returns:

  • descriptor_max (aiida.orm.SinglefileData) – File containing the maximum values for the descriptors.

  • descriptor_min (aiida.orm.SinglefileData) – File containing the minimum values for the descriptors.

  • concave_hull (aiida.orm.SinglefileData, optional) – File containing the concave hull of the latent space as an array.

  • latent_space (aiida.orm.SinglefileData, optional) – File containing the latent space represented as an array.

class atlas.active_learning.mace_tools_aiida.ProcessMDSeedStructCalculation(*args: Any, **kwargs: Any)

Bases: CalcJob

Launch a calculation to process a structure in an AL Loop step.

This CalcJob will process the output of a structure calculation by doing a MD simulation using the user provided settings, followed by checking all frames for extrapolation using several possible methods.

Parameters:

  • md_structure (orm.SinglefileData) – File containing the structure to be used for the MD, in the extxyz format.

  • commitee_models (PortNamespace) – A namespace to hold an arbitrary number of committee MACE potentials.

  • autoencoder_model (orm.SinglefileData, optional) – File containing the autoencoder model.

  • m_rmse_e (orm.Float) – Validation RMSE of the best model for the energy, in meV / atom.

  • m_rmse_f (orm.Float) – Validation RMSE of the best model for the forces, in meV / A.

  • concave_hull (orm.ArrayData, optional) – Array containing the concave hull to be used for the extrapolation check.

  • desc_max_arr (orm.ArrayData) – Array containing the maximum values for the descriptors.

  • desc_min_arr (orm.ArrayData) – Array containing the minimum values for the descriptors.

  • settings_file_pth (orm.Str) – Path to the ATL settings file in the .toml format.

  • Outputs

  • -------

  • extrapolating_structures (orm.SinglefileData) – File containing all structures that were found to be extrapolating. Uses the extxyz format.

  • extrapolation_plot (atl_img.ImagePNGData) – File containing a visualization of the extrapolation check and latent space boundaries

  • Codes (Exit)

  • ----------

  • 420 (ERROR_INVALID_OUTPUT) – Structure could not be processed.

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.mace_tools_aiida.ProcessMDSeedStructCalculationParser(*args: Any, **kwargs: Any)

Bases: Parser

Parser for the retrieved files from an active learning MD sampling job.

parse(**kwargs)

Parse the retrieved files of the calculation job.

atlas.active_learning.report_utils module

Utils for generating reports.

atlas.active_learning.report_utils.gather_stdout_and_scheduler(calcjob: CalcJobNode) → tuple[str, str]
atlas.active_learning.report_utils.gen_al_loop_report(loop_id: int | str = None, log_path: str = None, get_error_plot: bool = False, device='cpu', model_path: str | Path = None, database_path: str | Path = None, threshold_E_meV: float = None, threshold_F_meV: float = None, remove_outliers: bool = False, title: str = None, get_latent_space: bool = False, autoencoder_path: str = None, limit_num_steps: int = None, enable_cueq: bool = False)
atlas.active_learning.report_utils.gen_batch_report(training_db_path: str, threshold_meV: float = 100, remove_outliers: bool = False, root_path: str | Path = '.', device_str='cuda')
atlas.active_learning.report_utils.gen_init_db_report(train_db_path, threshold_E: float = None, threshold_F: float = None, remove_outliers: bool = False, color_type: str = None, per_atom: bool = False)
atlas.active_learning.report_utils.gen_performance_report(al_loop_pk: int | list[int], output_filename: str = None)
atlas.active_learning.report_utils.generate_error_plot(al_loop_node: list[Node], device_str: str = 'cpu', ax=None, database_path: str = None, model_path: str = None, threshold_E_meV: float = None, threshold_F_meV: float = None, remove_outliers: bool = False, enable_cueq: bool = False)
atlas.active_learning.report_utils.generate_latent_space_evol(al_loop_node: Node | list[Node], device_str: str, ax, model_path: str | Path | SinglefileData, database_path: str, autoencoder_model: Node | Path | str, autoencoder_path: str = None, databases: list[str] = None, ignore_latent_spaces=False)

Generate latent space evolution plot for the active learning loop.

Parameters:

  • al_loop_node (orm.Node | list[orm.Node]) – The AiiDA node or list of nodes representing the active learning loop. If a list is provided, it should contain nodes that have been executed in the active learning loop.

  • device_str (str) – The device string to use for the model, i.e., ‘cpu’ or ‘cuda’.

  • ax (_type_) – The matplotlib axes to plot the latent space evolution on.

  • model_path (str | Path | orm.SinglefileData) – The path to the MLIP model file or a SinglefileData node containing the model.

  • database_path (str) – The path to the database file containing structures for which to compute

  • autoencoder_model (orm.Node | Path | str) – The path to the autoencoder model file or a SinglefileData node

  • autoencoder_path (str, optional) – Loaded autoencoder model, by default None

  • databases (list[str], optional) – Loaded training data, by default None

  • ignore_latent_spaces (bool, optional) – Whether to ignore latent spaces, by default False

atlas.active_learning.report_utils.get_al_loop_performance(al_loop_pk: int | list[int]) → dict
atlas.active_learning.report_utils.get_al_step_performance(al_step: WorkChainNode) → dict

Get performance information about an AL Step WorkChain.

Calculates the total duration of the step and the elapsed time spent within each major stage (Training, Descriptors, MD, DFT) by finding the time between the start of the first calculation of that type and the end of the last calculation of that type within the step. Also collects individual calculation durations.

Parameters:

al_step – The AL Step WorkChain node.

Returns:

A dictionary with performance information.

atlas.active_learning.report_utils.get_latent_spaces_database_files(device_str, model_path, autoencoder_path, databases, autoencoder_file_path, autoencoder, latent_spaces, descr_dict)
atlas.active_learning.report_utils.get_latent_spaces_workchain(database_path: str | Path, model_path: str | Path, device_str: str, latent_spaces: list, autoencoder_model: str | Path | Autoencoder = None, latent_spaces_list: list = None)
atlas.active_learning.report_utils.get_loop_report(loop_id=None, log_path=None, title=None, model_acc_multiplier=None)

Parses an active learning loop report to extract summary statistics.

Parameters:

  • loop_id – Identifier for the loop (currently not used for parsing).

  • log_path – Filesystem path to the log file. Used if report_str is not provided.

  • title – Optional title for the run. If None, attempts to extract from report.

  • model_acc_multiplier – Multiplier for model accuracy (can be None).

Returns:

  • A tuple containing

  • - title (str) (The title of the run.)

  • - al_loop_node (int) (The active learning loop node ID.)

  • - ini_db_size (int) (The initial database size.)

  • - model_acc_multiplier (float or None) (The model accuracy multiplier.)

  • - stats_dict (dict) (A dictionary with statistics for each iteration.) –

    Example: {
    1: {

    “it_idx”: 1, “mace_e”: 0.778, “mace_f”: 8.204, “train_db_size”: 3706, “seed_gen_db_sizes”: 3689, “finished”: True

    }, …

    }

atlas.active_learning.report_utils.get_mlip_eval_results(al_loop_node: list[Node], database_path: str | Path, device_str: str = 'cpu', model_path: str | Path = None, folder_path: str | Path = None, enable_cueq: bool = False)
atlas.active_learning.report_utils.get_ncores_from_calcjob(calcjob: CalcJobNode) → int

Get the number of cores used by a CalcJob.

Parameters:

calcjob (orm.CalcJobNode) – The CalcJob node to analyze.

Returns:

Number of cores used by the CalcJob.

Return type:

int

atlas.active_learning.report_utils.get_runtime_from_calcjob(calcjob: CalcJobNode) → timedelta

Get the total runtime of a CalcJob.

Allows to get information from different schedulers, including SGE and SLURM. For SLURM, see: https://slurm.schedmd.com/sacct.html Does not take into account time in the queue.

Parameters:

calcjob (orm.CalcJobNode) – CalcJob node to analyze.

Returns:

Total runtime of the CalcJob.

Return type:

datetime.timedelta

atlas.active_learning.report_utils.mark_and_remove_outliers(E_sorted_diff_list_meV: list, F_sorted_diff_list_meV: list, E_diff_list_meV: list, F_diff_list_meV: list, E_nn_list_per_at: list, F_nn_list_per_at: list, E_dft_list_per_at: list, F_dft_list_per_at: list, threshold_E_meV: float, threshold_F_meV: float, remove_outliers: bool, train_db: list)
atlas.active_learning.report_utils.output_al_stats_dict(stats_dict: dict, title: str = None)

Displays a dictionary of active learning statistics in a formatted table.

Parameters:

stats_dict (dict) –

A dictionary where keys are iteration numbers (or similar identifiers) and values are dictionaries containing statistics for that iteration. Example: {

0: {‘it_idx’: 0, ‘mace_e’: None, …}, 1: {‘it_idx’: 1, ‘mace_e’: 0.778, …}, …

}

atlas.active_learning.report_utils.plot_al_loop_report(limit_num_steps: bool, seed_gen_db_sizes: list[int], train_db_sizes: list[int], mace_e: list[float], mace_f: list[float], it_idx: list[int], ax, model_acc_multiplier: float = None)
atlas.active_learning.report_utils.plot_performance_stacked_bar(all_performance_data: list[dict], output_filename: str = 'al_loop_performance.png')

Generates a stacked bar chart of stage elapsed times per AL step.

Parameters:

  • all_performance_data (List[Dict]) – A list containing performance dictionaries for each AL step.

  • output_filename (Optional[str], optional) – The filename to save the plot. If None, the plot is displayed instead of being saved. Defaults to “al_loop_performance.png”.

atlas.active_learning.report_utils.print_performance_report(all_performance_data: list)

Displays AL performance data in a formatted table using the rich library.

This function takes a list of dictionaries, where each dictionary contains performance data for a step in an Active Learning loop. It then prints two tables: one detailing each step’s performance, and another summarizing the total times and core hours for each stage and the overall loop.

Parameters:

all_performance_data (list) –

A list of dictionaries, where each dictionary represents a step’s performance data. Expected structure for each element: {

”step”: {“pk”: str, “step_duration”: datetime.timedelta}, “stages”: {

”training”: {

“total_elapsed_time”: datetime.timedelta, “cores”: float | int

}, “descriptors”: {

”total_elapsed_time”: datetime.timedelta, “cores”: float | int

}, “md”: {

”total_elapsed_time”: datetime.timedelta, “cores”: float | int

}, “dft”: {

”total_elapsed_time”: datetime.timedelta, “cores”: float | int

},

}

}

atlas.active_learning.report_utils.simplify_timedelta_str(tdelta: timedelta) → str

Simplifies and converts a timedelta object to a string format.

Removes the milisecond resolution.

Parameters:

tdelta (datetime.timedelta) – The timedelta object to convert.

Returns:

A simplified string representation of the timedelta in the format’ "X day, %HH:%MM:%SS".

Return type:

str

Module contents

ATLAS active learning functions and utilities.