Merge branch 'elastic' of github.com:mjwen/atomate2 into elastic

docs/user/codes/vasp.md

@@ -285,6 +285,40 @@ run_locally(lobster, create_folders=True, store=SETTINGS.JOB_STORE)
 It is, however,  computationally very beneficial to define two different types of job scripts for the VASP and Lobster runs, as VASP and Lobster runs are parallelized differently (MPI vs. OpenMP).
 [FireWorks](https://github.com/materialsproject/fireworks) allows to run the VASP and Lobster jobs with different job scripts. Please check out the [jobflow documentation on FireWorks](https://materialsproject.github.io/jobflow/tutorials/8-fireworks.html#setting-the-manager-configs) for more information.
 
+Specifically, you might want to change the `_fworker` for the LOBSTER runs and define a separate `lobster` worker within FireWorks:
+
+```python
+from fireworks import LaunchPad
+from jobflow.managers.fireworks import flow_to_workflow
+from pymatgen.core.structure import Structure
+
+from atomate2.vasp.flows.lobster import VaspLobsterMaker
+from atomate2.vasp.powerups import update_user_incar_settings
+
+structure = Structure(
+    lattice=[[0, 2.13, 2.13], [2.13, 0, 2.13], [2.13, 2.13, 0]],
+    species=["Mg", "O"],
+    coords=[[0, 0, 0], [0.5, 0.5, 0.5]],
+)
+
+lobster = VaspLobsterMaker().make(structure)
+lobster = update_user_incar_settings(lobster, {"NPAR": 4})
+
+# update the fireworker of the Lobster jobs
+for job, _ in lobster.iterflow():
+    config = {"manager_config": {"_fworker": "worker"}}
+    if "get_lobster" in job.name:
+        config["response_manager_config"] = {"_fworker": "lobster"}
+    job.update_config(config)
+
+# convert the flow to a fireworks WorkFlow object
+wf = flow_to_workflow(lobster)
+
+# submit the workflow to the FireWorks launchpad
+lpad = LaunchPad.auto_load()
+lpad.add_wf(wf)
+```
+
 Outputs from the automatic analysis with LobsterPy can easily be extracted from the database and also plotted:
 
 ```python

pyproject.toml

@@ -187,5 +187,7 @@ isort.known-first-party = ["atomate2"]
 [tool.ruff.per-file-ignores]
 "__init__.py" = ["F401"]
 "**/tests/*" = ["D"]
-# things inside TYPE_CHECKING aren't available at runtime and so can't be used by pydantic models
-"**/schemas/*" = ["TCH"]
+# flake8-type-checking (TCH): things inside TYPE_CHECKING aren't available
+#     at runtime and so can't be used by pydantic models
+# flake8-future-annotations (FA): future annotations only work in pydantic models in python 3.10+
+"**/schemas/*" = ["FA", "TCH"]

src/atomate2/amset/jobs.py

@@ -102,8 +102,7 @@ def make(
                 transport_data = loadfn(next(Path().glob("transport_*.json")))
                 converged = check_converged(transport_data, loadfn(prev_transport_file))
 
-        if "include_mesh" not in self.task_document_kwargs:
-            self.task_document_kwargs["include_mesh"] = converged is not False
+        self.task_document_kwargs.setdefault("include_mesh", converged is not False)
 
         # parse amset outputs
         task_doc = AmsetTaskDocument.from_directory(

src/atomate2/common/schemas/cclib.py

@@ -295,7 +295,7 @@ def cclib_calculate(
             f"A cube file must be provided for {method}. Returning None."
         )
     if method in ["ddec6", "hirshfeld"] and not proatom_dir:
-        if "PROATOM_DIR" not in os.environ:
+        if os.getenv("PROATOM_DIR") is None:
             raise OSError("PROATOM_DIR environment variable not set. Returning None.")
         proatom_dir = os.path.expandvars(os.environ["PROATOM_DIR"])
     if proatom_dir and not os.path.exists(proatom_dir):

src/atomate2/cp2k/jobs/base.py

@@ -147,9 +147,8 @@ def make(self, structure: Structure, prev_cp2k_dir: str | Path | None = None):
             structure = transmuter.transformed_structures[-1].final_structure
 
             # to avoid MongoDB errors, ":" is automatically converted to "."
-            if "transformations:json" not in self.write_additional_data:
-                tjson = transmuter.transformed_structures[-1]
-                self.write_additional_data["transformations:json"] = tjson
+            t_json = transmuter.transformed_structures[-1]
+            self.write_additional_data.setdefault("transformations:json", t_json)
 
         # copy previous inputs
         from_prev = prev_cp2k_dir is not None

src/atomate2/cp2k/jobs/core.py

@@ -331,16 +331,11 @@ def make(
         """
         self.input_set_generator.mode = mode
 
-        if "parse_dos" not in self.task_document_kwargs:
-            # parse DOS only for uniform band structure
-            self.task_document_kwargs["parse_dos"] = mode == "uniform"
-
-        if "parse_bandstructure" not in self.task_document_kwargs:
-            self.task_document_kwargs["parse_bandstructure"] = mode
-
+        # parse DOS only for uniform band structure
+        self.task_document_kwargs.setdefault("parse_dos", mode == "uniform")
+        self.task_document_kwargs.setdefault("parse_bandstructure", mode)
         # copy previous inputs
-        if "additional_cp2k_files" not in self.copy_cp2k_kwargs:
-            self.copy_cp2k_kwargs["additional_cp2k_files"] = ("wfn",)
+        self.copy_cp2k_kwargs.setdefault("additional_cp2k_files", ("wfn",))
 
         return super().make.original(self, structure, prev_cp2k_dir)
 
@@ -413,9 +408,8 @@ def make(
         structure = transmuter.transformed_structures[-1].final_structure
 
         # to avoid MongoDB errors, ":" is automatically converted to "."
-        if "transformations:json" not in self.write_additional_data:
-            tjson = transmuter.transformed_structures[-1]
-            self.write_additional_data["transformations:json"] = tjson
+        tjson = transmuter.transformed_structures[-1]
+        self.write_additional_data.setdefault("transformations:json", tjson)
 
         return super().make.original(self, structure, prev_cp2k_dir)
 

src/atomate2/cp2k/schemas/task.py

@@ -372,7 +372,7 @@ def from_directory(
         cp2k_objects = all_cp2k_objects[-1]
         included_objects = None
         if cp2k_objects:
-            included_objects = list(cp2k_objects.keys())
+            included_objects = list(cp2k_objects)
 
         if isinstance(calcs_reversed[-1].output.structure, Structure):
             attr = "from_structure"

src/atomate2/forcefields/schemas.py

@@ -209,7 +209,7 @@ def from_ase_compatible_result(
                 )
 
             # otherwise do not include "magmoms" in :obj:`cur_ionic_step`
-            elif "magmoms" not in trajectory.keys():
+            elif "magmoms" not in trajectory:
                 cur_ionic_step = IonicStep(
                     energy=cur_energy,
                     forces=cur_forces,

src/atomate2/vasp/jobs/base.py

@@ -136,8 +136,7 @@ def make(self, structure: Structure, prev_vasp_dir: str | Path | None = None):
         if prev_vasp_dir is not None:
             copy_vasp_outputs(prev_vasp_dir, **self.copy_vasp_kwargs)
 
-        if "from_prev" not in self.write_input_set_kwargs:
-            self.write_input_set_kwargs["from_prev"] = from_prev
+        self.write_input_set_kwargs.setdefault("from_prev", from_prev)
 
         # write vasp input files
         write_vasp_input_set(
@@ -173,16 +172,14 @@ def get_vasp_task_document(
     **kwargs,
 ):
     """Get VASP Task Document using atomate2 settings."""
-    if "store_additional_json" not in kwargs:
-        kwargs["store_additional_json"] = SETTINGS.VASP_STORE_ADDITIONAL_JSON
+    kwargs.setdefault("store_additional_json", SETTINGS.VASP_STORE_ADDITIONAL_JSON)
 
-    if "volume_change_warning_tol" not in kwargs:
-        kwargs["volume_change_warning_tol"] = SETTINGS.VASP_VOLUME_CHANGE_WARNING_TOL
+    kwargs.setdefault(
+        "volume_change_warning_tol", SETTINGS.VASP_VOLUME_CHANGE_WARNING_TOL
+    )
 
-    if "run_bader" not in kwargs:
-        kwargs["run_bader"] = SETTINGS.VASP_RUN_BADER and _BADER_EXE_EXISTS
+    kwargs.setdefault("run_bader", SETTINGS.VASP_RUN_BADER and _BADER_EXE_EXISTS)
 
-    if "store_volumetric_data" not in kwargs:
-        kwargs["store_volumetric_data"] = SETTINGS.VASP_STORE_VOLUMETRIC_DATA
+    kwargs.setdefault("store_volumetric_data", SETTINGS.VASP_STORE_VOLUMETRIC_DATA)
 
     return TaskDoc.from_directory(path, **kwargs)

src/atomate2/vasp/jobs/core.py

@@ -213,16 +213,11 @@ def make(
         """
         self.input_set_generator.mode = mode
 
-        if "parse_dos" not in self.task_document_kwargs:
-            # parse DOS only for uniform band structure
-            self.task_document_kwargs["parse_dos"] = mode == "uniform"
-
-        if "parse_bandstructure" not in self.task_document_kwargs:
-            self.task_document_kwargs["parse_bandstructure"] = mode
-
+        # parse DOS only for uniform band structure
+        self.task_document_kwargs.setdefault("parse_dos", mode == "uniform")
+        self.task_document_kwargs.setdefault("parse_bandstructure", mode)
         # copy previous inputs
-        if "additional_vasp_files" not in self.copy_vasp_kwargs:
-            self.copy_vasp_kwargs["additional_vasp_files"] = ("CHGCAR",)
+        self.copy_vasp_kwargs.setdefault("additional_vasp_files", ("CHGCAR",))
 
         return super().make.original(self, structure, prev_vasp_dir)
 
@@ -401,20 +396,15 @@ def make(
             )
             mode = "uniform"
 
-        if "parse_dos" not in self.task_document_kwargs:
-            # parse DOS only for uniform band structure
-            self.task_document_kwargs["parse_dos"] = "uniform" in mode
+        # parse DOS only for uniform band structure
+        self.task_document_kwargs.setdefault("parse_dos", "uniform" in mode)
 
-        if "parse_bandstructure" not in self.task_document_kwargs:
-            parse_bandstructure = "uniform" if mode == "gap" else mode
-            self.task_document_kwargs["parse_bandstructure"] = parse_bandstructure
+        parse_bandstructure = "uniform" if mode == "gap" else mode
+        self.task_document_kwargs.setdefault("parse_bandstructure", parse_bandstructure)
 
         # copy previous inputs
-        if (
-            prev_vasp_dir is not None
-            and "additional_vasp_files" not in self.copy_vasp_kwargs
-        ):
-            self.copy_vasp_kwargs["additional_vasp_files"] = ("CHGCAR",)
+        if prev_vasp_dir is not None:
+            self.copy_vasp_kwargs.setdefault("additional_vasp_files", ("CHGCAR",))
 
         return super().make.original(self, structure, prev_vasp_dir)
 
@@ -530,9 +520,8 @@ def make(
         structure = transmuter.transformed_structures[-1].final_structure
 
         # to avoid MongoDB errors, ":" is automatically converted to "."
-        if "transformations:json" not in self.write_additional_data:
-            tjson = transmuter.transformed_structures[-1]
-            self.write_additional_data["transformations:json"] = tjson
+        tjson = transmuter.transformed_structures[-1]
+        self.write_additional_data.setdefault("transformations:json", tjson)
 
         return super().make.original(self, structure, prev_vasp_dir)
 

src/atomate2/vasp/jobs/phonons.py

@@ -79,14 +79,11 @@ def get_supercell_size(
     **kwargs:
         Additional parameters that can be set.
     """
-    if "min_atoms" not in kwargs:
-        kwargs["min_atoms"] = None
-    if "force_diagonal" not in kwargs:
-        kwargs["force_diagonal"] = False
+    kwargs.setdefault("min_atoms", None)
+    kwargs.setdefault("force_diagonal", False)
 
     if not prefer_90_degrees:
-        if "max_atoms" not in kwargs:
-            kwargs["max_atoms"] = None
+        kwargs.setdefault("max_atoms", None)
         transformation = CubicSupercellTransformation(
             min_length=min_length,
             min_atoms=kwargs["min_atoms"],
@@ -97,9 +94,8 @@ def get_supercell_size(
         transformation.apply_transformation(structure=structure)
 
     else:
-        max_atoms = 1000 if "max_atoms" not in kwargs else kwargs["max_atoms"]
-        if "angle_tolerance" not in kwargs:
-            kwargs["angle_tolerance"] = 1e-2
+        max_atoms = kwargs.get("max_atoms", 1000)
+        kwargs.setdefault("angle_tolerance", 1e-2)
         try:
             transformation = CubicSupercellTransformation(
                 min_length=min_length,
@@ -112,8 +108,7 @@ def get_supercell_size(
             transformation.apply_transformation(structure=structure)
 
         except AttributeError:
-            if "max_atoms" not in kwargs:
-                kwargs["max_atoms"] = None
+            kwargs.setdefault("max_atoms", None)
 
             transformation = CubicSupercellTransformation(
                 min_length=min_length,

src/atomate2/vasp/run.py

@@ -132,8 +132,7 @@ def run_vasp(
     split_vasp_cmd = shlex.split(vasp_cmd)
     split_vasp_gamma_cmd = shlex.split(vasp_gamma_cmd)
 
-    if "auto_npar" not in vasp_job_kwargs:
-        vasp_job_kwargs["auto_npar"] = False
+    vasp_job_kwargs.setdefault("auto_npar", False)
 
     vasp_job_kwargs.update({"gamma_vasp_cmd": split_vasp_gamma_cmd})
 

src/atomate2/vasp/sets/base.py

@@ -183,11 +183,7 @@ def is_valid(self) -> bool:
 
         if self.incar.get("LHFCALC", False) is True and self.incar.get(
             "ALGO", "Normal"
-        ) not in [
-            "Normal",
-            "All",
-            "Damped",
-        ]:
+        ) not in ["Normal", "All", "Damped"]:
             warnings.warn(
                 "Hybrid functionals only support Algo = All, Damped, or Normal.",
                 BadInputSetWarning,
@@ -332,7 +328,7 @@ def __post_init__(self):
             if self.vdw not in vdw_par:
                 raise KeyError(
                     "Invalid or unsupported van-der-Waals functional. Supported "
-                    f"functionals are {vdw_par.keys()}"
+                    f"functionals are {list(vdw_par)}"
                 )
             self.config_dict["INCAR"].update(vdw_par[self.vdw])
 
@@ -642,7 +638,7 @@ def _get_incar(
 
         # apply previous incar settings, be careful not to override user_incar_settings
         # also skip LDAU/MAGMOM as structure may have changed.
-        skip = list(self.user_incar_settings.keys())
+        skip = list(self.user_incar_settings)
         skip += ["MAGMOM", "NUPDOWN", "LDAUU", "LDAUL", "LDAUJ", "LMAXMIX"]
         _apply_incar_updates(incar, previous_incar, skip=skip)
 
@@ -938,7 +934,7 @@ def _set_u_params(incar, incar_settings, structure):
     has_u = incar_settings.get("LDAU", False) and sum(incar["LDAUU"]) > 0
 
     if not has_u:
-        for key in list(incar.keys()):
+        for key in list(incar):
             if key.startswith("LDAU"):
                 del incar[key]
 
@@ -948,7 +944,7 @@ def _set_u_params(incar, incar_settings, structure):
     # investigation it was determined that this would lead to a significant difference
     # between SCF -> NonSCF even without Hubbard U enabled. Thanks to Andrew Rosen for
     # investigating and reporting.
-    if "LMAXMIX" not in incar_settings.keys():
+    if "LMAXMIX" not in incar_settings:
         # contains f-electrons
         if any(el.Z > 56 for el in structure.composition):
             incar["LMAXMIX"] = 6

src/atomate2/vasp/sets/core.py

@@ -926,7 +926,7 @@ def _get_ensemble_defaults(structure: Structure, ensemble: str) -> dict[str, Any
         try:
             return defaults[ensemble.lower()]  # type: ignore
         except KeyError as err:
-            supported = tuple(defaults.keys())
+            supported = tuple(defaults)
             raise ValueError(
                 f"Expect `ensemble` to be one of {supported}; got {ensemble}."
             ) from err

tests/vasp/flows/test_defect.py

@@ -183,7 +183,7 @@ def _check_plnr_locpot(name):
         plnr_locpot = SETTINGS.JOB_STORE.query_one({"output.task_label": name})[
             "output"
         ]["calcs_reversed"][0]["output"]["locpot"]
-        assert set(plnr_locpot.keys()) == {"0", "1", "2"}
+        assert set(plnr_locpot) == {"0", "1", "2"}
 
     for k in ref_paths:
         _check_plnr_locpot(k)