diff --git a/src/atomate2/common/powerups.py b/src/atomate2/common/powerups.py index b4f8bddb2..d89330ef4 100644 --- a/src/atomate2/common/powerups.py +++ b/src/atomate2/common/powerups.py @@ -5,30 +5,43 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from jobflow import Flow, Maker + from jobflow.core.flow import Flow + from jobflow.core.maker import Maker def add_metadata_to_flow( - flow: Flow, additional_fields: dict, class_filter: Maker + flow: Flow, additional_fields: dict, class_filter: type[Maker] ) -> Flow: """ - Return the flow with additional field(metadata) to the task doc. + Add additional metadata fields to task documents in a flow. - This allows adding metadata to the task-docs, could be useful - to query results from DB. + This function updates the task document kwargs for jobs in the flow, + allowing metadata to be added for easier querying of results from a database. + + This function does not add metadata to the job themselves, only to the output + generated upon job completion. Parameters ---------- - flow: + flow : Flow + The jobflow Flow object to modify. additional_fields : dict - A dict with metadata. - class_filter: .Maker - The Maker to which additional metadata needs to be added + Dictionary containing metadata fields and their values to add to task documents. + class_filter : type[Maker] + The Maker class type to which additional metadata should be added. + Only jobs created by this Maker type will be modified. Returns ------- Flow - Flow with added metadata to the task-doc. + The modified flow with added metadata in the task documents. + + Examples + -------- + >>> from atomate2.vasp.flows.core import RelaxBandStructureMaker + >>> flow = RelaxBandStructureMaker().make(structure) + >>> metadata = {"project": "battery_materials", "version": "2.0"} + >>> flow = add_metadata_to_flow(flow, metadata, RelaxBandStructureMaker) """ flow.update_maker_kwargs( { @@ -45,28 +58,40 @@ def add_metadata_to_flow( def update_custodian_handlers( - flow: Flow, custom_handlers: tuple, class_filter: Maker + flow: Flow, custom_handlers: tuple, class_filter: type[Maker] ) -> Flow: """ - Return the flow with custom custodian handlers for VASP jobs. + Update custodian error handlers for VASP jobs in a flow. - This allows user to selectively set error correcting handlers for VASP jobs - or completely unset error handlers. + This function allows selective configuration of error-correcting handlers + for VASP jobs or complete removal of error handlers. Parameters ---------- - flow: + flow : Flow + The jobflow Flow object to modify. custom_handlers : tuple - A tuple with custodian handlers. - class_filter: .Maker - The Maker to which custom custodian handler needs to be added + Tuple of custodian handler objects to use for error correction. + Pass an empty tuple to disable all error handlers. + class_filter : type[Maker] + The Maker class type to which custom custodian handlers should be applied. + Only jobs created by this Maker type will be modified. Returns ------- Flow - Flow with modified custodian handlers. + The modified flow with updated custodian handlers. + + Examples + -------- + >>> from custodian.vasp.handlers import VaspErrorHandler, MeshSymmetryErrorHandler + >>> from atomate2.vasp.flows.core import RelaxBandStructureMaker + >>> flow = RelaxBandStructureMaker().make(structure) + >>> handlers = (VaspErrorHandler(), MeshSymmetryErrorHandler()) + >>> flow = update_custodian_handlers(flow, handlers, RelaxBandStructureMaker) """ code = class_filter.name.split(" ")[1] + flow.update_maker_kwargs( {"_set": {f"run_{code}_kwargs->handlers": custom_handlers}}, dict_mod=True, diff --git a/src/atomate2/vasp/powerups.py b/src/atomate2/vasp/powerups.py index c73813bbb..753c01b92 100644 --- a/src/atomate2/vasp/powerups.py +++ b/src/atomate2/vasp/powerups.py @@ -1,98 +1,111 @@ -"""Powerups for performing common modifications on VASP jobs and flows.""" +"""Powerups for performing common modifications on VASP jobs and flows. + +This module provides utility functions (powerups) to modify VASP computational +workflows, including updating INCAR settings, POTCAR configurations, k-points, +and custodian handlers. All powerup functions return modified copies without +altering the original objects. +""" from __future__ import annotations from copy import deepcopy -from typing import Any +from typing import Any, TypeVar -from jobflow import Flow, Job, Maker +from jobflow.core.flow import Flow +from jobflow.core.job import Job +from jobflow.core.maker import Maker from pymatgen.io.vasp import Kpoints from atomate2.common.powerups import add_metadata_to_flow as base_add_metadata_to_flow from atomate2.common.powerups import update_custodian_handlers as base_custodian_handler from atomate2.vasp.jobs.base import BaseVaspMaker +JobType = TypeVar("JobType", Job, Flow, Maker) +"""TypeVar for generic job, flow, or maker types. Used to ensure powerup functions +return the same type as their input.""" + def update_vasp_input_generators( - flow: Job | Flow | Maker, + flow: JobType, dict_mod_updates: dict[str, Any], name_filter: str | None = None, class_filter: type[Maker] | None = BaseVaspMaker, -) -> Job | Flow | Maker: - """ - Update any VaspInputGenerators or Makers in the flow. +) -> JobType: + """Update VaspInputGenerators or Makers in a job, flow, or maker. - Note, this returns a copy of the original Job/Flow/Maker. I.e., the update does not - happen in place. + This function applies modifications to VASP input generators throughout a + workflow. It creates a deep copy of the input, so the original remains unchanged. Parameters ---------- - flow : .Job or .Flow or .Maker - A job, flow or Maker. - dict_mod_updates : dict - The updates to apply. Existing keys will not be modified unless explicitly - specified in ``dict_mod_updates``. - name_filter : str or None - A filter for the name of the jobs. - class_filter : Maker or None - A filter for the VaspMaker class used to generate the flows. Note the class - filter will match any subclasses. + flow : Job or Flow or Maker + A job, flow, or maker to update. + dict_mod_updates : dict[str, Any] + Dictionary of updates to apply using arrow notation (e.g., + 'input_set_generator->user_incar_settings->ENCUT'). Existing keys are + preserved unless explicitly overridden. + name_filter : str or None, optional + Filter to apply updates only to jobs matching this name pattern. + Default is None (no filtering). + class_filter : type[Maker] or None, optional + Filter to apply updates only to makers of this class or its subclasses. + Default is BaseVaspMaker. Returns ------- Job or Flow or Maker - A copy of the input flow/job/maker modified to use the updated incar settings. + A deep copy of the input with updated VASP input generator settings. """ updated_flow = deepcopy(flow) + + tmp_dict = { + "update": {"_set": dict_mod_updates}, + "name_filter": name_filter, + "class_filter": class_filter, + "dict_mod": True, + } + if isinstance(updated_flow, Maker): - updated_flow = updated_flow.update_kwargs( - {"_set": dict_mod_updates}, - name_filter=name_filter, - class_filter=class_filter, - dict_mod=True, - ) + updated_flow = updated_flow.update_kwargs(**tmp_dict) else: - updated_flow.update_maker_kwargs( - {"_set": dict_mod_updates}, - name_filter=name_filter, - class_filter=class_filter, - dict_mod=True, - ) + updated_flow.update_maker_kwargs(**tmp_dict) + return updated_flow def update_user_incar_settings( - flow: Job | Flow | Maker, + flow: JobType, incar_updates: dict[str, Any], name_filter: str | None = None, class_filter: type[Maker] | None = BaseVaspMaker, -) -> Job | Flow | Maker: - """ - Update the user_incar_settings of any VaspInputGenerators in the flow. - - Alternatively, if a Maker is supplied, the user_incar_settings of the maker will - be updated. +) -> JobType: + """Update user INCAR settings in VaspInputGenerators. - Note, this returns a copy of the original Job/Flow/Maker. I.e., the update does not - happen in place. + Modifies the user_incar_settings attribute of VASP input generators within + jobs, flows, or makers. Creates a copy of the input. Parameters ---------- - flow : .Job or .Flow or .Maker - A job, flow or Maker. - incar_updates : dict - The updates to apply. Existing keys in user_incar_settings will not be modified - unless explicitly specified in ``incar_updates``. - name_filter : str or None - A filter for the name of the jobs. - class_filter : Maker or None - A filter for the VaspMaker class used to generate the flows. Note the class - filter will match any subclasses. + flow : Job or Flow or Maker + A job, flow, or maker to update. + incar_updates : dict[str, Any] + Dictionary mapping INCAR tags to their new values (e.g., + {'ENCUT': 520, 'EDIFF': 1e-5}). Only specified keys are modified. + name_filter : str or None, optional + Filter to apply updates only to jobs matching this name pattern. + Default is None (no filtering). + class_filter : type[Maker] or None, optional + Filter to apply updates only to makers of this class or its subclasses. + Default is BaseVaspMaker. Returns ------- Job or Flow or Maker - A copy of the input flow/job/maker modified to use the updated incar settings. + A deep copy of the input with updated INCAR settings. + + Examples + -------- + >>> flow = update_user_incar_settings(flow, {"ENCUT": 520, "ISMEAR": 0}) """ return update_vasp_input_generators( flow=flow, @@ -106,37 +119,38 @@ def update_user_incar_settings( def update_user_potcar_settings( - flow: Job | Flow | Maker, + flow: JobType, potcar_updates: dict[str, Any], name_filter: str | None = None, class_filter: type[Maker] | None = BaseVaspMaker, -) -> Job | Flow | Maker: - """ - Update the user_potcar_settings of any VaspInputGenerators in the flow. - - Alternatively, if a Maker is supplied, the user_potcar_settings of the maker will - be updated. +) -> JobType: + """Update user POTCAR settings in VaspInputGenerators. - Note, this returns a copy of the original Job/Flow/Maker. I.e., the update does not - happen in place. + Modifies the user_potcar_settings attribute of VASP input generators within + jobs, flows, or makers. Creates a copy of the input. Parameters ---------- - flow : .Job or .Flow or .Maker - A job, flow or Maker. - potcar_updates : dict - The updates to apply. Existing keys in user_potcar_settings will not be modified - unless explicitly specified in ``potcar_updates``. - name_filter : str or None - A filter for the name of the jobs. - class_filter : Maker or None - A filter for the VaspMaker class used to generate the flows. Note the class - filter will match any subclasses. + flow : Job or Flow or Maker + A job, flow, or maker to update. + potcar_updates : dict[str, Any] + Dictionary mapping element symbols to POTCAR specifications (e.g., + {'Fe': 'Fe_pv', 'O': 'O'}). Only specified elements are modified. + name_filter : str or None, optional + Filter to apply updates only to jobs matching this name pattern. + Default is None (no filtering). + class_filter : type[Maker] or None, optional + Filter to apply updates only to makers of this class or its subclasses. + Default is BaseVaspMaker. Returns ------- Job or Flow or Maker - A copy of the input flow/job/maker modified to use the updated potcar settings. + A deep copy of the input with updated POTCAR settings. + + Examples + -------- + >>> flow = update_user_potcar_settings(flow, {"Fe": "Fe_pv", "O": "O_s"}) """ return update_vasp_input_generators( flow=flow, @@ -150,36 +164,37 @@ def update_user_potcar_settings( def update_user_potcar_functional( - flow: Job | Flow | Maker, + flow: JobType, potcar_functional: str, name_filter: str | None = None, class_filter: type[Maker] | None = BaseVaspMaker, -) -> Job | Flow | Maker: - """ - Update the user_potcar_functional of any VaspInputGenerators in the flow. - - Alternatively, if a Maker is supplied, the user_potcar_functional of the maker will - be updated. +) -> JobType: + """Update POTCAR functional in VaspInputGenerators. - Note, this returns a copy of the original Job/Flow/Maker. I.e., the update does not - happen in place. + Modifies the user_potcar_functional attribute of VASP input generators within + jobs, flows, or makers. Creates a copy of the input. Parameters ---------- - flow : .Job or .Flow or .Maker - A job, flow or Maker. + flow : Job or Flow or Maker + A job, flow, or maker to update. potcar_functional : str - The new potcar functional to use. - name_filter : str or None - A filter for the name of the jobs. - class_filter : Maker or None - A filter for the VaspMaker class used to generate the flows. Note the class - filter will match any subclasses. + The POTCAR functional to use (e.g., 'PBE', 'PBE_52', 'PBE_54', 'LDA'). + name_filter : str or None, optional + Filter to apply updates only to jobs matching this name pattern. + Default is None (no filtering). + class_filter : type[Maker] or None, optional + Filter to apply updates only to makers of this class or its subclasses. + Default is BaseVaspMaker. Returns ------- Job or Flow or Maker - A copy of the input flow/job/maker modified to use the updated potcar settings. + A deep copy of the input with updated POTCAR functional. + + Examples + -------- + >>> flow = update_user_potcar_functional(flow, "PBE_54") """ return update_vasp_input_generators( flow=flow, @@ -192,38 +207,42 @@ def update_user_potcar_functional( def update_user_kpoints_settings( - flow: Job | Flow | Maker, + flow: JobType, kpoints_updates: dict[str, Any] | Kpoints, name_filter: str | None = None, class_filter: type[Maker] | None = BaseVaspMaker, -) -> Job | Flow | Maker: - """ - Update the user_kpoints_settings of any VaspInputGenerators in the flow. - - Alternatively, if a Maker is supplied, the user_kpoints_settings of the maker will - be updated. +) -> JobType: + """Update user k-points settings in VaspInputGenerators. - Note, this returns a copy of the original Job/Flow/Maker. I.e., the update does not - happen in place. + Modifies the user_kpoints_settings attribute of VASP input generators within + jobs, flows, or makers. Creates a copy of the input. Parameters ---------- - flow : .Job or .Flow or .Maker - A job, flow or Maker. - kpoints_updates : dict - The updates to apply. Can be specified as a dictionary or as a Kpoints object. - If a dictionary is supplied, existing keys in user_kpoints_settings will not be - modified unless explicitly specified in ``kpoints_updates``. - name_filter : str or None - A filter for the name of the jobs. - class_filter : Maker or None - A filter for the VaspMaker class used to generate the flows. Note the class - filter will match any subclasses. + flow : Job or Flow or Maker + A job, flow, or maker to update. + kpoints_updates : dict[str, Any] or Kpoints + K-points updates to apply. Can be either: + - A dictionary with k-points settings (e.g., {'reciprocal_density': 100}) + - A Kpoints object that replaces the entire user_kpoints_settings + name_filter : str or None, optional + Filter to apply updates only to jobs matching this name pattern. + Default is None (no filtering). + class_filter : type[Maker] or None, optional + Filter to apply updates only to makers of this class or its subclasses. + Default is BaseVaspMaker. Returns ------- Job or Flow or Maker - A copy of the input flow/job/maker modified to use the updated kpoints settings. + A deep copy of the input with updated k-points settings. + + Examples + -------- + >>> flow = update_user_kpoints_settings(flow, {"reciprocal_density": 200}) + >>> # Or with a Kpoints object + >>> kpts = Kpoints.gamma_automatic((4, 4, 4)) + >>> flow = update_user_kpoints_settings(flow, kpts) """ if isinstance(kpoints_updates, Kpoints): dict_mod_updates = { @@ -243,35 +262,38 @@ def update_user_kpoints_settings( def use_auto_ispin( - flow: Job | Flow | Maker, + flow: JobType, value: bool = True, name_filter: str | None = None, class_filter: type[Maker] | None = BaseVaspMaker, -) -> Job | Flow | Maker: - """ - Update the auto_ispin setting of any VaspInputGenerators in the flow. +) -> JobType: + """Update automatic ISPIN setting in VaspInputGenerators. - Alternatively, if a Maker is supplied, the auto_ispin of the maker will be updated. - - Note, this returns a copy of the original Job/Flow/Maker. I.e., the update does not - happen in place. + Controls whether ISPIN is automatically determined based on the magnetic + moments in the structure. Creates a copy of the input. Parameters ---------- - flow : .Job or .Flow or .Maker - A job, flow or Maker. - value : bool - The value of auto_ispin to set. - name_filter : str or None - A filter for the name of the jobs. - class_filter : Maker or None - A filter for the VaspMaker class used to generate the flows. Note the class - filter will match any subclasses. + flow : Job or Flow or Maker + A job, flow, or maker to update. + value : bool, optional + Whether to enable automatic ISPIN determination. Default is True. + name_filter : str or None, optional + Filter to apply updates only to jobs matching this name pattern. + Default is None (no filtering). + class_filter : type[Maker] or None, optional + Filter to apply updates only to makers of this class or its subclasses. + Default is BaseVaspMaker. Returns ------- Job or Flow or Maker - A copy of the input flow/job/maker but with auto_ispin set. + A deep copy of the input with updated auto_ispin setting. + + Notes + ----- + When auto_ispin is True, ISPIN=2 is used if the structure has non-zero + magnetic moments, otherwise ISPIN=1 is used. """ return update_vasp_input_generators( flow=flow, @@ -282,27 +304,34 @@ def use_auto_ispin( def add_metadata_to_flow( - flow: Flow, additional_fields: dict, class_filter: Maker = BaseVaspMaker + flow: Flow, additional_fields: dict, class_filter: type[Maker] = BaseVaspMaker ) -> Flow: - """ - Return the VASP flow with additional field(metadata) to the task doc. + """Add custom metadata fields to VASP task documents in a flow. - This allows adding metadata to the task-docs, could be useful - to query results from DB. + Adds user-defined metadata to the task documents generated by VASP jobs, + which is useful for organizing and querying results in databases. Parameters ---------- flow : Flow - The flow to which to add metadata. + The flow to which metadata will be added. additional_fields : dict - A dict with metadata. - class_filter: .BaseVaspMaker - The Maker to which additional metadata needs to be added + Dictionary of metadata fields to add to task documents. Keys are field + names, values are the metadata values. + class_filter : Maker, optional + The maker class to which metadata will be added. Only jobs created by + this maker class or its subclasses will have metadata added. + Default is BaseVaspMaker. Returns ------- Flow - Flow with added metadata to the task-doc. + A copy of the flow with metadata added to matching task documents. + + Examples + -------- + >>> metadata = {"project": "battery_materials", "batch": "exp_001"} + >>> flow = add_metadata_to_flow(flow, metadata) """ return base_add_metadata_to_flow( flow=flow, class_filter=class_filter, additional_fields=additional_fields @@ -310,26 +339,41 @@ def add_metadata_to_flow( def update_vasp_custodian_handlers( - flow: Flow, custom_handlers: tuple, class_filter: Maker = BaseVaspMaker + flow: Flow, custom_handlers: tuple, class_filter: type[Maker] = BaseVaspMaker ) -> Flow: - """ - Return the flow with custom custodian handlers for VASP jobs. + """Update custodian error handlers for VASP jobs in a flow. - This allows user to selectively set error correcting handlers for VASP jobs - or completely unset error handlers. + Replaces the default custodian error handlers with custom handlers, + allowing users to customize error handling and recovery behavior or + disable error handling entirely. Parameters ---------- - flow: + flow : Flow + The flow whose custodian handlers will be updated. custom_handlers : tuple - A tuple with custodian handlers. - class_filter: .Maker - The Maker to which custom custodian handler needs to be added + Tuple of custodian handler objects to use for error handling. + Pass an empty tuple () to disable error handling. + class_filter : Maker, optional + The maker class for which handlers will be updated. Only jobs created + by this maker class or its subclasses will have their handlers modified. + Default is BaseVaspMaker. Returns ------- Flow - Flow with modified custodian handlers. + A copy of the flow with updated custodian handlers. + + Notes + ----- + Custodian handlers are executed in the order they appear in the tuple. + Common handlers include VaspErrorHandler, MeshSymmetryErrorHandler, etc. + + Examples + -------- + >>> from custodian.vasp.handlers import VaspErrorHandler + >>> handlers = (VaspErrorHandler(),) + >>> flow = update_vasp_custodian_handlers(flow, handlers) """ return base_custodian_handler( flow=flow, custom_handlers=custom_handlers, class_filter=class_filter