pcmt

src/atomate2/ase/jobs.py

@@ -9,7 +9,6 @@
 from ase.io import Trajectory as AseTrajectory
 from emmet.core.vasp.calculation import StoreTrajectoryOption
 from jobflow import Maker, job
-from pymatgen.core import Structure
 from pymatgen.core.trajectory import Trajectory as PmgTrajectory
 
 from atomate2.ase.schemas import AseResult, AseTaskDoc
@@ -21,19 +20,20 @@
     from pathlib import Path
 
     from ase.calculators.calculator import Calculator
-    from pymatgen.core import Molecule
+    from pymatgen.core import Molecule, Structure
 
     from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc
 
 _ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory]
 
+
 @dataclass
 class AseMaker(Maker):
     """
     Define basic template of ASE-based jobs.
 
     This class defines two functions relevant attributes
-    for the ASE TaskDoc schemas, as well as two methods 
+    for the ASE TaskDoc schemas, as well as two methods
     that must be implemented in subclasses:
         1. `calculator`: the ASE .Calculator object
         2. `run_ase`: which actually makes the call to ASE.
@@ -48,18 +48,27 @@ class AseMaker(Maker):
         Keyword arguments that will get passed to the ASE calculator.
     ionic_step_data : tuple[str,...] or None
         Quantities to store in the TaskDocument ionic_steps.
-        Acceptable options are "struct_or_mol", "energy", "forces", "stress", and "magmoms"
+        Possible options are "struct_or_mol", "energy",
+        "forces", "stress", and "magmoms".
+        "structure" and "molecule" are aliases for "struct_or_mol".
     store_trajectory : emmet .StoreTrajectoryOption = "no"
         Whether to store trajectory information ("no") or complete trajectories
         ("partial" or "full", which are identical).
     tags : list[str] or None
         A list of tags for the task.
     """
+
     name: str = "ASE maker"
     calculator_kwargs: dict = field(default_factory=dict)
-    ionic_step_data : tuple[str,...] | None = ("energy","forces","magmoms","stress","mol_or_struct",)
-    store_trajectory : StoreTrajectoryOption = StoreTrajectoryOption.NO
-    tags : list[str] | None = None
+    ionic_step_data: tuple[str, ...] | None = (
+        "energy",
+        "forces",
+        "magmoms",
+        "stress",
+        "mol_or_struct",
+    )
+    store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO
+    tags: list[str] | None = None
 
     def run_ase(
         self,
@@ -117,7 +126,9 @@ class AseRelaxMaker(AseMaker):
         Keyword arguments that will get passed to the ASE calculator.
     ionic_step_data : tuple[str,...] or None
         Quantities to store in the TaskDocument ionic_steps.
-        Acceptable options are "struct_or_mol", "energy", "forces", "stress", and "magmoms"
+        Possible options are "struct_or_mol", "energy",
+        "forces", "stress", and "magmoms".
+        "structure" and "molecule" are aliases for "struct_or_mol".
     store_trajectory : emmet .StoreTrajectoryOption = "no"
         Whether to store trajectory information ("no") or complete trajectories
         ("partial" or "full", which are identical).
@@ -151,22 +162,21 @@ def make(
                 added to match the method signature of other makers.
 
         Returns
-        --------
+        -------
         AseStructureTaskDoc or AseMoleculeTaskDoc
         """
-
         return AseTaskDoc.to_mol_or_struct_metadata_doc(
             getattr(self.calculator, "name", self.calculator.__class__),
             self.run_ase(mol_or_struct, prev_dir=prev_dir),
             self.steps,
             relax_kwargs=self.relax_kwargs,
             optimizer_kwargs=self.optimizer_kwargs,
-            relax_cell = self.relax_cell,
+            relax_cell=self.relax_cell,
             fix_symmetry=self.fix_symmetry,
             symprec=self.symprec if self.fix_symmetry else None,
-            ionic_step_data = self.ionic_step_data,
-            store_trajectory = self.store_trajectory,
-            tags = self.tags,
+            ionic_step_data=self.ionic_step_data,
+            store_trajectory=self.store_trajectory,
+            tags=self.tags,
         )
 
     def run_ase(
@@ -200,6 +210,7 @@ def run_ase(
         )
         return relaxer.relax(mol_or_struct, steps=self.steps, **self.relax_kwargs)
 
+
 @dataclass
 class LennardJonesRelaxMaker(AseRelaxMaker):
     """

src/atomate2/ase/md.py

@@ -122,7 +122,9 @@ class AseMDMaker(AseMaker):
         kwargs to pass to the ASE calculator class
     ionic_step_data : tuple[str,...] or None
         Quantities to store in the TaskDocument ionic_steps.
-        Acceptable options are "struct_or_mol", "energy", "forces", "stress", and "magmoms"
+        Possible options are "struct_or_mol", "energy",
+        "forces", "stress", and "magmoms".
+        "structure" and "molecule" are aliases for "struct_or_mol".
     store_trajectory : emmet .StoreTrajectoryOption = "partial"
         Whether to store trajectory information ("no") or complete trajectories
         ("partial" or "full", which are identical).
@@ -158,8 +160,8 @@ class AseMDMaker(AseMaker):
     pressure: float | Sequence | np.ndarray | None = None
     ase_md_kwargs: dict | None = None
     calculator_kwargs: dict = field(default_factory=dict)
-    ionic_step_data : tuple[str,...] | None = None
-    store_trajectory : StoreTrajectoryOption = StoreTrajectoryOption.PARTIAL
+    ionic_step_data: tuple[str, ...] | None = None
+    store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.PARTIAL
     traj_file: str | Path | None = None
     traj_file_fmt: Literal["pmg", "ase"] = "ase"
     traj_interval: int = 1
@@ -235,7 +237,7 @@ def _get_ensemble_defaults(self) -> None:
                 10.0 * 1e-3 / units.fs,  # Same default as in VASP: 10 ps^-1
             )
 
-    @job(data=_ASE_DATA_OBJECTS + ["output.ionic_steps"])
+    @job(data=[*_ASE_DATA_OBJECTS, "output.ionic_steps"])
     def make(
         self,
         mol_or_struct: Molecule | Structure,
@@ -252,20 +254,19 @@ def make(
             A previous calculation directory to copy output files from. Unused, just
             added to match the method signature of other makers.
         """
-
         return AseTaskDoc.to_mol_or_struct_metadata_doc(
             getattr(self.calculator, "name", self.calculator.__class__),
             self.run_ase(mol_or_struct, prev_dir=prev_dir),
             steps=self.n_steps,
             relax_kwargs=None,
             optimizer_kwargs=None,
-            fix_symmetry = False,
-            symprec = None,
-            ionic_step_data = self.ionic_step_data,
-            store_trajectory = self.store_trajectory,
-            tags = self.tags,
+            fix_symmetry=False,
+            symprec=None,
+            ionic_step_data=self.ionic_step_data,
+            store_trajectory=self.store_trajectory,
+            tags=self.tags,
         )
-    
+
     def run_ase(
         self,
         mol_or_struct: Molecule | Structure,

src/atomate2/ase/schemas.py

@@ -86,7 +86,7 @@ class AseObject(ValueEnum):
 
     TRAJECTORY = "trajectory"
     IONIC_STEPS = "ionic_steps"
-    
+
 
 class AseBaseModel(BaseModel):
     """Base document class for ASE input and output."""
@@ -149,7 +149,9 @@ class OutputDoc(AseBaseModel):
         None, description="Step-by-step trajectory of the relaxation."
     )
 
-    elapsed_time: Optional[float] = Field(None, description="The time taken to run the calculation in seconds.")
+    elapsed_time: Optional[float] = Field(
+        None, description="The time taken to run the calculation in seconds."
+    )
 
     n_steps: int = Field(
         None, description="total number of steps needed in the relaxation."
@@ -348,18 +350,18 @@ def from_ase_compatible_result(
         steps: int,
         relax_kwargs: dict = None,
         optimizer_kwargs: dict = None,
-        relax_cell : bool = True,
+        relax_cell: bool = True,
         fix_symmetry: bool = False,
         symprec: float = 1e-2,
-        ionic_step_data: tuple[str,...] | None = (
+        ionic_step_data: tuple[str, ...] | None = (
             "energy",
             "forces",
             "magmoms",
             "stress",
             "mol_or_struct",
         ),
         store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO,
-        tags : list[str] | None = None,
+        tags: list[str] | None = None,
         **task_document_kwargs,
     ) -> AseTaskDoc:
         """Create an AseTaskDoc for a task that has ASE-compatible outputs.
@@ -411,7 +413,7 @@ def from_ase_compatible_result(
 
         input_doc = InputDoc(
             mol_or_struct=input_mol_or_struct,
-            relax_cell = relax_cell,
+            relax_cell=relax_cell,
             fix_symmetry=fix_symmetry,
             symprec=symprec,
             steps=steps,
@@ -457,9 +459,12 @@ def from_ase_compatible_result(
                 }
 
                 current_mol_or_struct = (
-                    trajectory[idx] if any(
-                        v in ionic_step_data for v in ("mol_or_struct", "structure","molecule")
-                    ) else None
+                    trajectory[idx]
+                    if any(
+                        v in ionic_step_data
+                        for v in ("mol_or_struct", "structure", "molecule")
+                    )
+                    else None
                 )
 
                 # include "magmoms" in :obj:`ionic_step` if the trajectory has "magmoms"
@@ -480,7 +485,7 @@ def from_ase_compatible_result(
                 )
 
                 ionic_steps.append(ionic_step)
-            
+
         objects: dict[AseObject, Any] = {}
         if store_trajectory != StoreTrajectoryOption.NO:
             # For VASP calculations, the PARTIAL trajectory option removes
@@ -510,7 +515,7 @@ def from_ase_compatible_result(
             is_force_converged=result.is_force_converged,
             energy_downhill=result.energy_downhill,
             dir_name=result.dir_name,
-            tags = tags,
+            tags=tags,
             **task_document_kwargs,
         )
 

src/atomate2/forcefields/jobs.py

@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
-import warnings
+
 from ase.io import Trajectory as AseTrajectory
 from ase.units import GPa as _GPa_to_eV_per_A3
 from jobflow import job
@@ -21,7 +22,6 @@
     from typing import Callable
 
     from ase.calculators.calculator import Calculator
-    from emmet.core.vasp.calculation import StoreTrajectoryOption
     from pymatgen.core.structure import Structure
 
 logger = logging.getLogger(__name__)
@@ -97,7 +97,9 @@ class ForceFieldRelaxMaker(AseRelaxMaker):
         Keyword arguments that will get passed to the ASE calculator.
     ionic_step_data : tuple[str,...] or None
         Quantities to store in the TaskDocument ionic_steps.
-        Acceptable options are "struct_or_mol", "energy", "forces", "stress", and "magmoms"
+        Possible options are "struct_or_mol", "energy",
+        "forces", "stress", and "magmoms".
+        "structure" and "molecule" are aliases for "struct_or_mol".
     store_trajectory : emmet .StoreTrajectoryOption = "no"
         Whether to store trajectory information ("no") or complete trajectories
         ("partial" or "full", which are identical).
@@ -140,7 +142,8 @@ def make(
             warnings.warn(
                 "`task_document_kwargs` is now deprecated, please use the top-level "
                 "attributes `ionic_step_data` and `store_trajectory`",
-                DeprecationWarning
+                category=DeprecationWarning,
+                stacklevel=1,
             )
 
         return ForceFieldTaskDocument.from_ase_compatible_result(
@@ -149,12 +152,12 @@ def make(
             self.steps,
             relax_kwargs=self.relax_kwargs,
             optimizer_kwargs=self.optimizer_kwargs,
-            relax_cell = self.relax_cell,
+            relax_cell=self.relax_cell,
             fix_symmetry=self.fix_symmetry,
             symprec=self.symprec if self.fix_symmetry else None,
             ionic_step_data=self.ionic_step_data,
             store_trajectory=self.store_trajectory,
-            tags = self.tags,
+            tags=self.tags,
             **self.task_document_kwargs,
         )
 

src/atomate2/forcefields/md.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
-import warnings
+
+from jobflow import job
 
 from atomate2.ase.md import AseMDMaker
 from atomate2.forcefields import MLFF
-from atomate2.forcefields.jobs import forcefield_job
+from atomate2.forcefields.jobs import _FORCEFIELD_DATA_OBJECTS
 from atomate2.forcefields.schemas import ForceFieldTaskDocument
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
@@ -72,7 +74,9 @@ class ForceFieldMDMaker(AseMDMaker):
         kwargs to pass to the ASE calculator class
     ionic_step_data : tuple[str,...] or None
         Quantities to store in the TaskDocument ionic_steps.
-        Acceptable options are "struct_or_mol", "energy", "forces", "stress", and "magmoms"
+        Possible options are "struct_or_mol", "energy",
+        "forces", "stress", and "magmoms".
+        "structure" and "molecule" are aliases for "struct_or_mol".
     store_trajectory : emmet .StoreTrajectoryOption = "partial"
         Whether to store trajectory information ("no") or complete trajectories
         ("partial" or "full", which are identical).
@@ -101,9 +105,12 @@ class ForceFieldMDMaker(AseMDMaker):
 
     name: str = "Forcefield MD"
     force_field_name: str = f"{MLFF.Forcefield}"
-    task_document_kwargs : dict = None
+    task_document_kwargs: dict = None
 
-    @forcefield_job
+    @job(
+        data=[*_FORCEFIELD_DATA_OBJECTS, "output.ionic_steps"],
+        output_schema=ForceFieldTaskDocument,
+    )
     def make(
         self,
         structure: Structure,
@@ -128,7 +135,8 @@ def make(
             warnings.warn(
                 "`task_document_kwargs` is now deprecated, please use the top-level "
                 "attributes `ionic_step_data` and `store_trajectory`",
-                DeprecationWarning
+                category=DeprecationWarning,
+                stacklevel=1,
             )
 
         return ForceFieldTaskDocument.from_ase_compatible_result(
@@ -138,11 +146,11 @@ def make(
             steps=self.n_steps,
             relax_kwargs=None,
             optimizer_kwargs=None,
-            fix_symmetry = False,
-            symprec = None,
-            ionic_step_data = self.ionic_step_data,
-            store_trajectory = self.store_trajectory,
-            tags = self.tags,
+            fix_symmetry=False,
+            symprec=None,
+            ionic_step_data=self.ionic_step_data,
+            store_trajectory=self.store_trajectory,
+            tags=self.tags,
             **self.task_document_kwargs,
         )