Update doc: adding metadata to flow

docs/user/fireworks.md

@@ -16,6 +16,7 @@ submit an MgO band structure workflow using FireWorks:
 ```python
 from fireworks import LaunchPad
 from atomate2.vasp.flows.core import RelaxBandStructureMaker
+from atomate2.vasp.powerups import add_metadata_to_flow
 from jobflow.managers.fireworks import flow_to_workflow
 from pymatgen.core import Structure
 
@@ -29,6 +30,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 = add_metadata_to_flow(
+    flow=bandstructure_flow,
+    additional_fields={"mp_id": "mp-190"},
+)
+
 # convert the flow to a fireworks WorkFlow object
 wf = flow_to_workflow(bandstructure_flow)
 

src/atomate2/common/powerups.py

@@ -0,0 +1,74 @@
+"""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": {
+                f"task_document_kwargs->additional_fields->{field}": value
+                for field, value in additional_fields.items()
+            }
+        },
+        dict_mod=True,
+        class_filter=class_filter,
+    )
+
+    return flow
+
+
+def update_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.
+    """
+    code = class_filter.name.split(" ")[1]
+    flow.update_maker_kwargs(
+        {"_set": {f"run_{code}_kwargs->handlers": custom_handlers}},
+        dict_mod=True,
+        class_filter=class_filter,
+    )
+
+    return flow

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
+    )

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
+    )

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

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_vasp_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