atlas.workflows package

Subpackages

Submodules

atlas.workflows.aiida_utils module

Utility functions for running several through aiida for active learning loops.

class atlas.workflows.aiida_utils.CalcType(value)

Bases: Enum

Class representing the available calculation types for DFT. This allows multiple aliases for the same calculation type.

DEFAULT = 'single_point_bulk'
RELAX_BULK = 'relaxation_bulk'
RELAX_CLUSTER = 'relaxation_cluster'
RELAX_SURFACE = 'relaxation_surface'
SP_BULK = 'single_point_bulk'
SP_CLUSTER = 'single_point_cluster'
SP_ISOLATEDATOM = 'single_point_isolatedatom'
SP_SURFACE = 'single_point_surface'
classmethod from_string(value)

Get the CalcType corresponding to a string, considering aliases.

atlas.workflows.aiida_utils.add_aiida_group_to_db(db_obj: str, group_identifier, copy=False)
atlas.workflows.aiida_utils.can_submit_calculation(code: str, limit: int, computer: aiida.orm.Computer = None) → bool

Check if a calculation can be submitted to a given computer’s queue.

This function checks if the number of jobs currently running is smaller than the limit given. If smaller, the calculation can be submitted.

In some supercomputer clusters there is a limit on the number of jobs that can be submitted at the same time, and going over this limit will result in the job not entering the queue. AiiDA will reattempt to submit after the time given by exponential_backoff_retry is elapsed and for task_maximum_attempts attempts, and if the number of attempts is exceeded, the calculation will be paused, holding the loop.

Parameters:

  • code (str) – AiiDA code label. Will be used to get the computer.

  • computer (orm.Computer) – AiiDA computer.

  • limit (int) – Maximum number of calculations that can be submitted. The limit will be set by the supercomputer’s scheduler, in the case of SLURM, it can be checked with sacctmgr show qos XXXX under the MaxSubmitPU column.

  • user (str, optional) – Username for the calculation. If not given, the default user will be used.

Returns:

Whether the calculation can be submitted or not.

Return type:

bool

atlas.workflows.aiida_utils.choose_queue(node_type: int, tot_procs: int = None)

Choose the scheduler and aiida’s computer and code options.

Parameters:

  • node_type (int) – Whether to use 12, 24 or 28 core nodes in tekla.

  • tot_procs (int) – How many cores to be used. Using more than the total number of cores in a single node will make use of more nodes.

Returns:

  • dict – The OPTIONS dict is aiida jobfile equivalent

  • str – The CODE_STRING str is aiida code identifier.

atlas.workflows.aiida_utils.choose_queue_from_struct(queue_data: dict, computer: aiida.orm.Computer)

Choose scheduler options for aiida.

Choose the scheduler and aiida’s computer and code options according to the size of the structure.

Parameters:

  • structure (pymatgen.core.structure.Structure) – Structure

  • queue_data (dict) – Dictionary specifying which queue gets assigned to every atom interval.

Returns:

  • dict – The OPTIONS dict is aiida jobfile equivalent

  • str – The CODE_STRING str is aiida code identifier.

  • int – The number of nodes that will be used.

atlas.workflows.aiida_utils.determine_vacuum_direction(structure)

A heuristic function that determines the vacuum direction by examining the spread of atomic positions in fractional coordinates.

Returns the index (0, 1, or 2) corresponding to the direction with the largest vacuum gap.

atlas.workflows.aiida_utils.gather_calc_data_from_row(target_row, curr_structure=None)
atlas.workflows.aiida_utils.generate_incar(structure, phase: str, calc_type: str, kspacing: dict = {'alpha': 0.133203528512207, 'beta-prime': 0.100530964914873, 'epsilon': 0.153309721495182, 'eta': 0.0948760981384118, 'gamma': 0.141371669411541, 'm1': 0.100530964914873, 'm2': 0.100530964914873, 'm3': 0.166504410640259, 'm4': 0.0948760981384118})

Generate an incar file using depending on the calculation type.

This incar includes a kspacing variable that depends on the phase.

Parameters:

  • structure (pymatgen.core.structure.Structure) – Current structure

  • phase (str) – Phase of the current structure

  • calc_type (str) – Calculation type, can be either ‘relax’ for relaxation or ‘sp’ for single point

Returns:

dictionary representation of the INCAR

Return type:

dict

atlas.workflows.aiida_utils.generate_kpoints_data(structure, calc_type, kspacing=None, kspacing_vec=None)
atlas.workflows.aiida_utils.generate_potential_mapping(assign_dict=None) → dict

Generate a dictionary specifying the pseudopotential mapping for vasp.

This function only assigns the default potential for every atom. The user can specify different mappings if needed.

Inputs

The function itself requires no input when called, but it will attempt to read the contents of a ‘potential_mapping’ file on the same folder where the code is being executed (CWD). There will be an example on the github repo. Also, an example follows:

# Header lines for comments that will be ignored.
# The lines following the headers should not have any comment marks, i.e. '#'.
Ag = Ag_gw
Au = Au_gw
returns:

Dictionary containing the potential assignation for each atom of the periodic table, with the shape: {‘H’: ‘H’, ‘He’: ‘He’, …}.

rtype:

dict

atlas.workflows.aiida_utils.get_or_create_containerized_code(label: str, computer: aiida.orm.Computer, image_name: str, engine_command: str, filepath_executable: str) → aiida.orm.ContainerizedCode

Return a cached ContainerizedCode, or create one if none exists.

Matches on (label, computer, image_name, engine_command, filepath_executable). prepend_text is never baked into the code; callers should pass it dynamically via builder.metadata.options.prepend_text.

Unlike get_or_create_portable_code(), no source hash is needed: a ContainerizedCode stores only the image path, not the image contents, so rebuilding the container image at the same path is automatically picked up at runtime. Old codes remain in the database for provenance but are not reused when the construction parameters change.

atlas.workflows.aiida_utils.get_or_create_portable_code(label: str, filepath_files, filepath_executable: str, prepend_text: str = '') → aiida.orm.PortableCode

Return a PortableCode for label, recreating it if source has changed.

Queries for the most recent PortableCode with the given label. If one is found, its bundled *.py files are compared (via a SHA-256 hash stored as the source_hash_v2 extra) against the current source directory. When the hash differs, i.e. the source files were edited since the code was created, a new PortableCode is created instead of reusing the stale one. This prevents situations where an edited PortableCode script (e.g. atl_process_structure.py) silently keeps running the old version.

Environment setup (module loads, venv activation, OMP_NUM_THREADS) is not baked into the PortableCode. Callers should pass it dynamically via builder.metadata.options.prepend_text instead, which AiiDA appends after the code’s (empty) prepend_text in the final submission script. This way changing the TOML settings takes effect immediately without recreating the code node.

The prepend_text parameter is accepted for backward compatibility but is ignored; callers should set it on the builder.

The old node remains in the database for provenance of past calculations but is not reused for new submissions.

atlas.workflows.aiida_utils.get_vdw_params(structure, incar)
atlas.workflows.aiida_utils.kpoint_mesh_from_density(structure, kspacing)

Return kpoint mesh (3x3) from kpoint array, intended for surfaces.

atlas.workflows.aiida_utils.parse_stress_from_structure(results_dict: dict, new_struct: Atoms) → None

Parse stress from results dictionary and add it to the structure info, updating it in-place.

Parameters:

  • results_dict (dict) – Dictionary containing the results from a VASP DFT calculation.

  • new_struct (Atoms) – Structure to update.

atlas.workflows.aiida_utils.run_dataframe_vasp_aiida_queue(initial_db: InitialDatabase | list[Atoms], config_dict: dict, log_file_path: str)
atlas.workflows.aiida_utils.run_dataframe_vasp_simulations_aiida(sel_struct_df, group_name: str, calc_type: CalcType, kspacing_dict: dict, max_batch: int, start_on: int, potential_mapping: dict, potential_family: str, initial_db: InitialDatabase, queue_dict: dict, incar_dict: dict = None, dry_run: bool = False)
atlas.workflows.aiida_utils.select_kspacing(curr_structure, incar: dict, phase: str, kspacing_dict: dict, calc_type)
atlas.workflows.aiida_utils.sort_chunk_size(chunk)
atlas.workflows.aiida_utils.submit_aiida_vasp_calculation(index, target_structure, phase, material_name, unique_id, kspacing_dict, incar_settings_dict, calc_type, queue_dict, potential_family, potential_mapping, dry_run, return_builder, group, aiida_vasp_settings=None)
atlas.workflows.aiida_utils.update_db_with_dft_results(sel_struct_db, queue)

Update the original database with the results from the calculations.

Parameters:

  • sel_struct_db (list[Atoms]) – List of structures to be updated.

  • queue (list[aiida.orm.nodes.process.workflow.workchain.WorkChainNode]) – List of AiiDA workchain nodes with the results of the calculations.

atlas.workflows.aiida_utils.validate_kspacing(kspacing_dict: dict, phases: list[str]) → None

Validate that kspacing values exist for all phases.

Checks each phase against the kspacing dict. Missing phases fall back to ATL_DEFAULT (or the legacy MDB_DEFAULT). If neither default exists and some phases are missing, raises ValueError listing the missing phases and suggesting how to fix the TOML.

Parameters:

  • kspacing_dict (dict) – The kspacing dictionary from the TOML [dft.vasp.kspacing] section. Modified in-place: if MDB_DEFAULT is present but ATL_DEFAULT is not, MDB_DEFAULT is copied to ATL_DEFAULT.

  • phases (list of str) – Phase names present in the structures to be submitted for DFT.

atlas.workflows.monitors module

Monitors for the ATLAS workflows.

atlas.workflows.monitors.output_monitor(node: CalcJobNode, transport: Transport) → str

Retrieve and inspect files in working directory of job to determine whether the job should be killed.

Parameters:

  • node – The node representing the calculation job.

  • transport – The transport that can be used to retrieve files from remote working directory.

Returns:

A string if the job should be killed, None otherwise.

atlas.workflows.simple_active_learning module

AiiDA workchain for NNP active learning loops using MD.

class atlas.workflows.simple_active_learning.SimpleActiveLearningBaseWorkChain(*args: Any, **kwargs: Any)

Bases: BaseRestartWorkChain

Base workchain for ATL active learning workflows.

This workchain is used as a base for the SimpleActiveLearningWorkChain workchain. It handles setup of the workchain and the main loop, where the active learning steps are launched. After every step, the results are checked and added to the database, and the next step is prepared. The workchain will loop until the stopping conditions are met. If the stopping conditions and the safeguard check is enabled, a MD-based safeguard mechanism will be employed to ensure the current sampler model is robust, and therefore the loop can be stopped.

It takes all the inputs of the SimpleActiveLearningWorkChain workchain, except for the inputs that are specific to the active learning loop. Additionally, the mandatory log_path and optional resume_dict can be provided. See the define method for more information on the inputs and outputs.

Note

This workchain can be restarted from a previous running workchain if its files are recovered. The optional resume_dict input is used for this. resume_dict must be a dictionary that contains the following keys:

  • last_iteration: The current iteration of the active learning loop.

  • train_db_path: The path to the last training database.

  • seed_db_path: The path to the last seed generation database.

The loop will only be restarted from the beginning of the last step.

Check SimpleActiveLearningWorkChain for information on what is done in each step.

add_dft_results_to_db()

Incorporate DFT calculation results into the training/seed generation databases.

This method updates the training and seed generation databases with DFT calculation results. If any DFT calculations have been performed, their results are appended to both the training database and the seed generation database.

check_al_loop_conditions() → bool

Evaluate conditions to determine whether to continue the active learning loop.

This method assesses multiple conditions to decide if the active learning (AL) loop should continue. It considers whether the maximum number of iterations has been reached, and the following two criteria:

  • Depletion of the seed generation database. (if True: stop)

  • No disagreement among predictions for an entire MD seed. However, this condition can optionally be overriden if the user enables the XXXXXXXXXX option, which forces the AL Loop to continue if the sampler model validation or test performance is above a certain threshold, which can be configured. (if True: stop)

The function updates the workchain context with the iteration status and generates reports based on the evaluation of these conditions.

Returns:

A boolean value indicating whether the AL loop should continue. Returns True if conditions are met for another iteration; otherwise, returns False.

Return type:

orm.Bool

Notes

The method uses self.ctx.is_finished, self.ctx.iteration, self.inputs.active_learning.max_iterations.value, self.ctx.stop_md_seed_no_disagreement.value, and self.ctx.seed_gen_db_all_structs_removed.value

check_resume_mode()
compile_safeguard_model()

Optionally pre-compile the sampler model for the safeguard node.

No-op unless the backend implements MLIPModelCompiler. Result is collected in gather_safeguard_compiled() and used by the safeguard MD calc; inference falls back to on-demand compilation otherwise.

copy_input_toml_file()

Copy the input toml file to the results directory.

classmethod define(spec)

Define the process specification.

gather_safeguard_compiled()

Store the compiled safeguard model if its compile calc succeeded.

get_al_loop_break_conditions()

Evaluate and set conditions to potentially break the active learning loop.

This function checks for specific conditions that might warrant terminating the active learning (AL) loop early:

  • Gathers stop_md_seed_no_disagreement from the outputs of the inner workchain and stores it in the workchain’s context. If this is True, the workchain will stop.

  • Checks whether all structures have been removed from the seed generation database (indicating no further candidates for evaluation). If this is True, the workchain will stop.

The results of these checks are stored in the workflow’s context.

get_data_reduction_seed()

Filter structures for data reduction mode.

In data reduction mode, we need to filter out any structures that are already in the training database from the seed database. This ensures we don’t re-select structures that have already been chosen.

The actual structure selection for training will happen in the SimpleActiveLearningWorkChain.

get_database()

Loading initial database.

get_input_report()
get_md_seed()

Selects a random subset of structures from the seed generation database to create a MD seed for the active learning loop.

This function calculates the number of structures to be included in the MD seed based on the specified fraction of the seed generation database’s length. It then randomly selects and populates the training seed with these structures.

Returns:

  • None. The function updates current_md_seed_structs with the

  • selected structures.

get_results_loop()

Attach the outputs specified in the spec from the last completed process.

get_seed_structures()

Dispatcher method to get seed structures based on the active learning mode.

For MD mode: calls get_md_seed() For data_reduction mode: calls get_data_reduction_seed()

log_atl_version()
parse_safeguard_check_results()

Check safeguard output, setting flags to stop/continue the AL Loop.

results_final()

Finalize the results at the end of the workchain.

This method prepares and returns the final training database at the conclusion of the workchain. It serializes the structures within the training database to a format compatible with AiiDA storage and subsequent processing. The serialized structures are then used to prepare the final training database, which is outputted from the workchain. This signifies the completion of the workchain and the availability of the processed training data for further use.

run_safeguard_check()
safeguard_check_md()
setup()

Call BaseRestartWorkChain setup and create input dict in self.ctx.inputs.

This self.ctx.inputs dictionary will be used by the BaseRestartWorkChain to submit the process in the internal loop.

setup_textfile_logging()
should_run_according_to_safeguard_check() → bool

Decide whether to continue the AL loop based on the safeguard check.

The loop will continue if the safeguard check has not been completed or has passed correctly in the previous attempt, and the maximum number of iterations has not been reached.

Returns:

True if the loop should continue, False otherwise.

Return type:

bool

should_run_safeguard() → bool
class atlas.workflows.simple_active_learning.SimpleActiveLearningWorkChain(*args: Any, **kwargs: Any)

Bases: WorkChain

WorkChain to run an active learning loop for a MLIP using MD simulations to generate training data.

check_extrapolation_enabled()

Check if the extrapolation check is enabled.

compile_models_for_inference()

Optionally pre-compile trained model(s) for each inference target.

No-op unless the MLIP backend implements MLIPModelCompiler (e.g. Allegro). Otherwise, for each distinct (computer, device, model) an enabled inference consumer will use, submit one compile CalcJob on that computer so the artifact matches the node’s device/toolchain. Results are collected in gather_compiled_models(). Purely an optimization: consumers fall back to on-demand compilation if an artifact is missing.

classmethod define(spec)

Specify inputs and outputs.

gather_compiled_models()

Collect compiled artifacts from the compile calcs into the cache.

gen_descriptors_and_concave_hull()
get_descriptor_results()

Get descriptor results from the calculation and store them in the context.

get_mlip_train_output()

Retrieve and process MLIP training output for model selection.

Evaluates training results to select the best model based on a weighted sum of RMSE for energy and forces. Creates a LAMMPS-compatible potential and updates context with model info.

get_test_db_results()

Retrieve and process test database evaluation results.

is_test_db_evaluation_enabled()

Check if test database evaluation is enabled.

perform_test_db_evaluation()
return_seed_dft_and_model()

Gather and output the last NNP model and DFT calculations for the current seed.

This function collects DFT calculations for the structures in the current seed, which are then returned as outputs in the workchain using the namespace dft_calcs. A check is performed to determine if the results agree using MACE models, and this check also outputted to the workchain using the namespace stop_md_seed_no_disagreement.

run_md_seed()
select_data_reduction_structures()

Select additional structures from the large database for training.

This method is only called in data reduction mode and after the first iteration. It selects structures based on the iterative_selection_method and adds them to the training database while removing them from the seed database.

send_calc_or_remove_structures()

Decide which structures to keep and send to DFT or remove from DB.

set_step_logger()
should_select_data_reduction_structures()

Check if we should select additional structures for data reduction mode.

step_setup()
train_mlip_model()

Submit MLIP training calculations using the active backend.

Generates a training database, then submits N training calculations (one per committee model). Uses PortableCode by default; falls back to a registered AiiDA Code if the code key is present in the training settings.

Module contents

Module containing utilities to interact with the workflow tool AiiDA.