From 34923cd9c04ca4b393a7489c5000e9c71060e984 Mon Sep 17 00:00:00 2001 From: anaik Date: Fri, 8 Dec 2023 16:58:39 +0100 Subject: [PATCH 1/5] Update doc: adding metadata to flow --- docs/user/fireworks.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/user/fireworks.md b/docs/user/fireworks.md index defad2ca5a..47a6b0dc4b 100644 --- a/docs/user/fireworks.md +++ b/docs/user/fireworks.md @@ -29,6 +29,14 @@ mgo_structure = Structure( # make a band structure flow to optimise the structure and obtain the band structure bandstructure_flow = RelaxBandStructureMaker().make(mgo_structure) +# (Optional) add metadata to the flow task document. +# Could be useful to filter specific results from the database. +# For e.g., adding material project ID for the compound, use following lines +bandstructure_flow.update_maker_kwargs( + {"_set": {"task_document_kwargs->additional_fields": {"mp_id": "mp-xxx"}}}, + dict_mod=True, +) + # convert the flow to a fireworks WorkFlow object wf = flow_to_workflow(bandstructure_flow) From ad2b3e545b0b0c9f7f1ae1e6f36029c36b1bea05 Mon Sep 17 00:00:00 2001 From: anaik Date: Mon, 8 Jan 2024 15:45:16 +0100 Subject: [PATCH 2/5] add powerups module to utils --- docs/user/fireworks.md | 9 +++-- src/atomate2/utils/powerups.py | 68 ++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 3 deletions(-) create mode 100644 src/atomate2/utils/powerups.py diff --git a/docs/user/fireworks.md b/docs/user/fireworks.md index 47a6b0dc4b..fd4278abfe 100644 --- a/docs/user/fireworks.md +++ b/docs/user/fireworks.md @@ -15,7 +15,9 @@ submit an MgO band structure workflow using FireWorks: ```python from fireworks import LaunchPad +from atomate2.vasp.jobs.base import BaseVaspMaker from atomate2.vasp.flows.core import RelaxBandStructureMaker +from atomate2.utils.powerups import add_metadata_to_flow from jobflow.managers.fireworks import flow_to_workflow from pymatgen.core import Structure @@ -32,9 +34,10 @@ bandstructure_flow = RelaxBandStructureMaker().make(mgo_structure) # (Optional) add metadata to the flow task document. # Could be useful to filter specific results from the database. # For e.g., adding material project ID for the compound, use following lines -bandstructure_flow.update_maker_kwargs( - {"_set": {"task_document_kwargs->additional_fields": {"mp_id": "mp-xxx"}}}, - dict_mod=True, +bandstructure_flow = add_metadata_to_flow( + flow=bandstructure_flow, + class_filter=BaseVaspMaker, + additional_fields={"mp_id": "mp-190"}, ) # convert the flow to a fireworks WorkFlow object diff --git a/src/atomate2/utils/powerups.py b/src/atomate2/utils/powerups.py new file mode 100644 index 0000000000..7953c17df5 --- /dev/null +++ b/src/atomate2/utils/powerups.py @@ -0,0 +1,68 @@ +"""Utilities for modifying workflow.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from jobflow import Flow, Maker + + +def add_metadata_to_flow(flow, additional_fields: dict, class_filter: Maker) -> Flow: + """ + Return the flow with additional field(metadata) to the task doc. + + This allows adding metadata to the task-docs, could be useful + to query results from DB. + + Parameters + ---------- + flow: + additional_fields : dict + A dict with metadata. + class_filter: .Maker + The Maker to which additional metadata needs to be added + + Returns + ------- + Flow + Flow with added metadata to the task-doc. + """ + flow.update_maker_kwargs( + {"_set": {"task_document_kwargs->additional_fields": additional_fields}}, + dict_mod=True, + class_filter=class_filter, + ) + + return flow + + +def update_vasp_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: Maker +) -> Flow: + """ + Return the flow with custom custodian handlers for VASP jobs. + + This allows user to selectively set error correcting handlers for VASP jobs + or completely unset error handlers. + + Parameters + ---------- + flow: + custom_handlers : tuple + A tuple with custodian handlers. + class_filter: .Maker + The Maker to which custom custodian handler needs to be added + + Returns + ------- + Flow + Flow with modified custodian handlers. + """ + flow.update_maker_kwargs( + {"_set": {"run_vasp_kwargs->handlers": custom_handlers}}, + dict_mod=True, + class_filter=class_filter, + ) + + return flow From 40195a5beabeca5d2d4d58bbab2436ff08cbe7ef Mon Sep 17 00:00:00 2001 From: anaik Date: Wed, 10 Jan 2024 10:26:21 +0100 Subject: [PATCH 3/5] address reivew comments --- docs/user/fireworks.md | 2 +- src/atomate2/{utils => common}/powerups.py | 26 +++++++++++++++------- 2 files changed, 19 insertions(+), 9 deletions(-) rename src/atomate2/{utils => common}/powerups.py (65%) diff --git a/docs/user/fireworks.md b/docs/user/fireworks.md index fd4278abfe..c01c1a74f7 100644 --- a/docs/user/fireworks.md +++ b/docs/user/fireworks.md @@ -17,7 +17,7 @@ submit an MgO band structure workflow using FireWorks: from fireworks import LaunchPad from atomate2.vasp.jobs.base import BaseVaspMaker from atomate2.vasp.flows.core import RelaxBandStructureMaker -from atomate2.utils.powerups import add_metadata_to_flow +from atomate2.common.powerups import add_metadata_to_flow from jobflow.managers.fireworks import flow_to_workflow from pymatgen.core import Structure diff --git a/src/atomate2/utils/powerups.py b/src/atomate2/common/powerups.py similarity index 65% rename from src/atomate2/utils/powerups.py rename to src/atomate2/common/powerups.py index 7953c17df5..9fb371a249 100644 --- a/src/atomate2/utils/powerups.py +++ b/src/atomate2/common/powerups.py @@ -5,10 +5,14 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from jobflow import Flow, Maker + from jobflow import Flow + from atomate2.vasp.jobs.base import BaseVaspMaker -def add_metadata_to_flow(flow, additional_fields: dict, class_filter: Maker) -> Flow: + +def add_metadata_to_flow( + flow, additional_fields: dict, class_filter: BaseVaspMaker +) -> Flow: """ Return the flow with additional field(metadata) to the task doc. @@ -20,7 +24,7 @@ def add_metadata_to_flow(flow, additional_fields: dict, class_filter: Maker) -> flow: additional_fields : dict A dict with metadata. - class_filter: .Maker + class_filter: .BaseVaspMaker The Maker to which additional metadata needs to be added Returns @@ -29,7 +33,12 @@ def add_metadata_to_flow(flow, additional_fields: dict, class_filter: Maker) -> Flow with added metadata to the task-doc. """ flow.update_maker_kwargs( - {"_set": {"task_document_kwargs->additional_fields": additional_fields}}, + { + "_set": { + f"task_document_kwargs->additional_fields->{field}": value + for field, value in additional_fields.items() + } + }, dict_mod=True, class_filter=class_filter, ) @@ -37,8 +46,8 @@ def add_metadata_to_flow(flow, additional_fields: dict, class_filter: Maker) -> return flow -def update_vasp_custodian_handlers( - flow: Flow, custom_handlers: tuple, class_filter: Maker +def update_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: BaseVaspMaker ) -> Flow: """ Return the flow with custom custodian handlers for VASP jobs. @@ -51,7 +60,7 @@ def update_vasp_custodian_handlers( flow: custom_handlers : tuple A tuple with custodian handlers. - class_filter: .Maker + class_filter: .BaseVaspMaker The Maker to which custom custodian handler needs to be added Returns @@ -59,8 +68,9 @@ def update_vasp_custodian_handlers( Flow Flow with modified custodian handlers. """ + code = class_filter.name.split(" ")[1] flow.update_maker_kwargs( - {"_set": {"run_vasp_kwargs->handlers": custom_handlers}}, + {"_set": {f"run_{code}_kwargs->handlers": custom_handlers}}, dict_mod=True, class_filter=class_filter, ) From 6b9582ebda0279a85683ac467da7681bdda1af73 Mon Sep 17 00:00:00 2001 From: anaik Date: Wed, 10 Jan 2024 12:32:17 +0100 Subject: [PATCH 4/5] add metadata/custodian powerups for vasp/cp2k and tests, update correponding tutorial --- docs/user/fireworks.md | 4 +-- src/atomate2/common/powerups.py | 14 +++------ src/atomate2/cp2k/powerups.py | 56 +++++++++++++++++++++++++++++++++ src/atomate2/vasp/powerups.py | 56 +++++++++++++++++++++++++++++++++ tests/cp2k/test_powerups.py | 40 +++++++++++++++++++++++ tests/vasp/test_powerups.py | 34 ++++++++++++++++++++ 6 files changed, 192 insertions(+), 12 deletions(-) diff --git a/docs/user/fireworks.md b/docs/user/fireworks.md index c01c1a74f7..59defe3595 100644 --- a/docs/user/fireworks.md +++ b/docs/user/fireworks.md @@ -15,9 +15,8 @@ submit an MgO band structure workflow using FireWorks: ```python from fireworks import LaunchPad -from atomate2.vasp.jobs.base import BaseVaspMaker from atomate2.vasp.flows.core import RelaxBandStructureMaker -from atomate2.common.powerups import add_metadata_to_flow +from atomate2.vasp.powerups import add_metadata_to_flow from jobflow.managers.fireworks import flow_to_workflow from pymatgen.core import Structure @@ -36,7 +35,6 @@ bandstructure_flow = RelaxBandStructureMaker().make(mgo_structure) # For e.g., adding material project ID for the compound, use following lines bandstructure_flow = add_metadata_to_flow( flow=bandstructure_flow, - class_filter=BaseVaspMaker, additional_fields={"mp_id": "mp-190"}, ) diff --git a/src/atomate2/common/powerups.py b/src/atomate2/common/powerups.py index 9fb371a249..26c330db77 100644 --- a/src/atomate2/common/powerups.py +++ b/src/atomate2/common/powerups.py @@ -5,14 +5,10 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from jobflow import Flow + from jobflow import Flow, Maker - from atomate2.vasp.jobs.base import BaseVaspMaker - -def add_metadata_to_flow( - flow, additional_fields: dict, class_filter: BaseVaspMaker -) -> Flow: +def add_metadata_to_flow(flow, additional_fields: dict, class_filter: Maker) -> Flow: """ Return the flow with additional field(metadata) to the task doc. @@ -24,7 +20,7 @@ def add_metadata_to_flow( flow: additional_fields : dict A dict with metadata. - class_filter: .BaseVaspMaker + class_filter: .Maker The Maker to which additional metadata needs to be added Returns @@ -47,7 +43,7 @@ def add_metadata_to_flow( def update_custodian_handlers( - flow: Flow, custom_handlers: tuple, class_filter: BaseVaspMaker + flow: Flow, custom_handlers: tuple, class_filter: Maker ) -> Flow: """ Return the flow with custom custodian handlers for VASP jobs. @@ -60,7 +56,7 @@ def update_custodian_handlers( flow: custom_handlers : tuple A tuple with custodian handlers. - class_filter: .BaseVaspMaker + class_filter: .Maker The Maker to which custom custodian handler needs to be added Returns diff --git a/src/atomate2/cp2k/powerups.py b/src/atomate2/cp2k/powerups.py index 35e696d781..69fde55e6a 100644 --- a/src/atomate2/cp2k/powerups.py +++ b/src/atomate2/cp2k/powerups.py @@ -8,6 +8,8 @@ from jobflow import Flow, Job, 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.cp2k.jobs.base import BaseCp2kMaker @@ -138,3 +140,57 @@ def update_user_kpoints_settings( dict_mod=True, ) return updated_flow + + +def add_metadata_to_flow( + flow, additional_fields: dict, class_filter: Maker = BaseCp2kMaker +) -> Flow: + """ + Return the Cp2k flow with additional field(metadata) to the task doc. + + This allows adding metadata to the task-docs, could be useful + to query results from DB. + + Parameters + ---------- + flow: + additional_fields : dict + A dict with metadata. + class_filter: .BaseCp2kMaker + The Maker to which additional metadata needs to be added + + Returns + ------- + Flow + Flow with added metadata to the task-doc. + """ + return base_add_metadata_to_flow( + flow=flow, class_filter=class_filter, additional_fields=additional_fields + ) + + +def update_cp2k_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: Maker = BaseCp2kMaker +) -> Flow: + """ + Return the flow with custom custodian handlers for Cp2k jobs. + + This allows user to selectively set error correcting handlers for Cp2k jobs + or completely unset error handlers. + + Parameters + ---------- + flow: + custom_handlers : tuple + A tuple with custodian handlers. + class_filter: .BaseCp2kMaker + The Maker to which custom custodian handler needs to be added + + Returns + ------- + Flow + Flow with modified custodian handlers. + """ + return base_custodian_handler( + flow=flow, custom_handlers=custom_handlers, class_filter=class_filter + ) diff --git a/src/atomate2/vasp/powerups.py b/src/atomate2/vasp/powerups.py index a050654628..810300f820 100644 --- a/src/atomate2/vasp/powerups.py +++ b/src/atomate2/vasp/powerups.py @@ -8,6 +8,8 @@ from jobflow import Flow, Job, 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 @@ -277,3 +279,57 @@ def use_auto_ispin( name_filter=name_filter, class_filter=class_filter, ) + + +def add_metadata_to_flow( + flow, additional_fields: dict, class_filter: Maker = BaseVaspMaker +) -> Flow: + """ + Return the VASP flow with additional field(metadata) to the task doc. + + This allows adding metadata to the task-docs, could be useful + to query results from DB. + + Parameters + ---------- + flow: + additional_fields : dict + A dict with metadata. + class_filter: .BaseVaspMaker + The Maker to which additional metadata needs to be added + + Returns + ------- + Flow + Flow with added metadata to the task-doc. + """ + return base_add_metadata_to_flow( + flow=flow, class_filter=class_filter, additional_fields=additional_fields + ) + + +def update_vasp_custodian_handlers( + flow: Flow, custom_handlers: tuple, class_filter: Maker = BaseVaspMaker +) -> Flow: + """ + Return the flow with custom custodian handlers for VASP jobs. + + This allows user to selectively set error correcting handlers for VASP jobs + or completely unset error handlers. + + Parameters + ---------- + flow: + custom_handlers : tuple + A tuple with custodian handlers. + class_filter: .Maker + The Maker to which custom custodian handler needs to be added + + Returns + ------- + Flow + Flow with modified custodian handlers. + """ + return base_custodian_handler( + flow=flow, custom_handlers=custom_handlers, class_filter=class_filter + ) diff --git a/tests/cp2k/test_powerups.py b/tests/cp2k/test_powerups.py index e3b74bb819..587bff8e8d 100644 --- a/tests/cp2k/test_powerups.py +++ b/tests/cp2k/test_powerups.py @@ -60,3 +60,43 @@ def test_update_user_settings(powerup, attribute, settings): getattr(flow.jobs[1].function.__self__.input_set_generator, attribute) != settings ) + + +@pytest.mark.parametrize( + "powerup,settings", + [ + ("add_metadata_to_flow", {"mp-id": "mp-xxx"}), + ("add_metadata_to_flow", {"mp-id": "mp-170", "composition": "NaCl"}), + ], +) +def test_add_metadata_to_flow(powerup, settings): + from atomate2.cp2k import powerups + from atomate2.cp2k.flows.core import DoubleRelaxMaker + + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert ( + flow.jobs[0].function.__self__.task_document_kwargs["additional_fields"] + == settings + ) + + +@pytest.mark.parametrize( + "powerup, settings", + [("update_cp2k_custodian_handlers", ())], +) +def test_update_cp2k_custodian_handlers(powerup, settings): + from atomate2.cp2k import powerups + from atomate2.cp2k.flows.core import DoubleRelaxMaker + + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert flow.jobs[0].function.__self__.run_cp2k_kwargs["handlers"] == settings diff --git a/tests/vasp/test_powerups.py b/tests/vasp/test_powerups.py index 6b1164d2a3..90c4e74a3b 100644 --- a/tests/vasp/test_powerups.py +++ b/tests/vasp/test_powerups.py @@ -63,3 +63,37 @@ def test_update_user_settings(powerup, attribute, settings): getattr(flow.jobs[1].function.__self__.input_set_generator, attribute) != settings ) + + +@pytest.mark.parametrize( + "powerup,settings", + [ + ("add_metadata_to_flow", {"mp-id": "mp-xxx"}), + ("add_metadata_to_flow", {"mp-id": "mp-161", "composition": "NaCl"}), + ], +) +def test_add_metadata_to_flow(powerup, settings): + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert ( + flow.jobs[0].function.__self__.task_document_kwargs["additional_fields"] + == settings + ) + + +@pytest.mark.parametrize( + "powerup, settings", + [("update_vasp_custodian_handlers", ())], +) +def test_update_cp2k_custodian_handlers(powerup, settings): + powerup_func = getattr(powerups, powerup) + + # test flow + drm = DoubleRelaxMaker() + flow = drm.make(1) + flow = powerup_func(flow, settings) + assert flow.jobs[0].function.__self__.run_vasp_kwargs["handlers"] == settings From 89f610d58182fae86e0d96766f546cd3585a63a3 Mon Sep 17 00:00:00 2001 From: anaik Date: Wed, 10 Jan 2024 12:34:10 +0100 Subject: [PATCH 5/5] fix test name --- tests/vasp/test_powerups.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/vasp/test_powerups.py b/tests/vasp/test_powerups.py index 90c4e74a3b..6d61dfa7e3 100644 --- a/tests/vasp/test_powerups.py +++ b/tests/vasp/test_powerups.py @@ -89,7 +89,7 @@ def test_add_metadata_to_flow(powerup, settings): "powerup, settings", [("update_vasp_custodian_handlers", ())], ) -def test_update_cp2k_custodian_handlers(powerup, settings): +def test_update_vasp_custodian_handlers(powerup, settings): powerup_func = getattr(powerups, powerup) # test flow