change name ForceFieldTaskDocument => ForceFieldStructureTaskDocument

Author: yaoyi92

src/atomate2/forcefields/jobs.py

@@ -15,7 +15,7 @@
 
 from atomate2.ase.jobs import AseRelaxMaker
 from atomate2.forcefields import MLFF, _get_formatted_ff_name
-from atomate2.forcefields.schemas import ForceFieldTaskDocument
+from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -48,7 +48,7 @@ def forcefield_job(method: Callable) -> job:
     This is a thin wrapper around :obj:`~jobflow.core.job.Job` that configures common
     settings for all forcefield jobs. For example, it ensures that large data objects
     (currently only trajectories) are all stored in the atomate2 data store.
-    It also configures the output schema to be a ForceFieldTaskDocument :obj:`.TaskDoc`.
+    It also configures the output schema to be a ForceFieldStructureTaskDocument :obj:`.TaskDoc`.
 
     Any makers that return forcefield jobs (not flows) should decorate the
     ``make`` method with @forcefield_job. For example:
@@ -73,7 +73,7 @@ def make(structure):
         A decorated version of the make function that will generate forcefield jobs.
     """
     return job(
-        method, data=_FORCEFIELD_DATA_OBJECTS, output_schema=ForceFieldTaskDocument
+        method, data=_FORCEFIELD_DATA_OBJECTS, output_schema=ForceFieldStructureTaskDocument
     )
 
 
@@ -118,7 +118,7 @@ class ForceFieldRelaxMaker(AseRelaxMaker):
     tags : list[str] or None
         A list of tags for the task.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = "Force field relax"
@@ -147,7 +147,7 @@ def __post_init__(self) -> None:
     @forcefield_job
     def make(
         self, structure: Structure, prev_dir: str | Path | None = None
-    ) -> ForceFieldTaskDocument:
+    ) -> ForceFieldStructureTaskDocument:
         """
         Perform a relaxation of a structure using a force field.
 
@@ -170,7 +170,7 @@ def make(
                 stacklevel=1,
             )
 
-        return ForceFieldTaskDocument.from_ase_compatible_result(
+        return ForceFieldStructureTaskDocument.from_ase_compatible_result(
             str(self.force_field_name),  # make mypy happy
             ase_result,
             self.steps,
@@ -212,7 +212,7 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = "Force field static"
@@ -255,7 +255,7 @@ class CHGNetRelaxMaker(ForceFieldRelaxMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.CHGNet} relax"
@@ -291,7 +291,7 @@ class CHGNetStaticMaker(ForceFieldStaticMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.CHGNet} static"
@@ -334,7 +334,7 @@ class M3GNetRelaxMaker(ForceFieldRelaxMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.M3GNet} relax"
@@ -372,7 +372,7 @@ class M3GNetStaticMaker(ForceFieldStaticMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.M3GNet} static"
@@ -415,7 +415,7 @@ class NEPRelaxMaker(ForceFieldRelaxMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.NEP} relax"
@@ -451,7 +451,7 @@ class NEPStaticMaker(ForceFieldStaticMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.NEP} static"
@@ -494,7 +494,7 @@ class NequipRelaxMaker(ForceFieldRelaxMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.Nequip} relax"
@@ -529,7 +529,7 @@ class NequipStaticMaker(ForceFieldStaticMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.Nequip} static"
@@ -576,7 +576,7 @@ class MACERelaxMaker(ForceFieldRelaxMaker):
         trained for Matbench Discovery on the MPtrj dataset available at
         https://figshare.com/articles/dataset/22715158.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.MACE_MP_0} relax"
@@ -616,7 +616,7 @@ class MACEStaticMaker(ForceFieldStaticMaker):
         trained for Matbench Discovery on the MPtrj dataset available at
         https://figshare.com/articles/dataset/22715158.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.MACE_MP_0} static"
@@ -665,7 +665,7 @@ class SevenNetRelaxMaker(ForceFieldRelaxMaker):
         trained for Matbench Discovery on the MPtrj dataset available at
         https://figshare.com/articles/dataset/22715158.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.SevenNet} relax"
@@ -707,7 +707,7 @@ class SevenNetStaticMaker(ForceFieldStaticMaker):
         trained for Matbench Discovery on the MPtrj dataset available at
         https://figshare.com/articles/dataset/22715158.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.SevenNet} static"
@@ -747,7 +747,7 @@ class GAPRelaxMaker(ForceFieldRelaxMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.GAP} relax"
@@ -783,7 +783,7 @@ class GAPStaticMaker(ForceFieldStaticMaker):
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
     task_document_kwargs : dict (deprecated)
-        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()`.
     """
 
     name: str = f"{MLFF.GAP} static"

src/atomate2/forcefields/md.py

@@ -15,7 +15,7 @@
     _DEFAULT_CALCULATOR_KWARGS,
     _FORCEFIELD_DATA_OBJECTS,
 )
-from atomate2.forcefields.schemas import ForceFieldTaskDocument
+from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -126,13 +126,13 @@ def __post_init__(self) -> None:
 
     @job(
         data=[*_FORCEFIELD_DATA_OBJECTS, "ionic_steps"],
-        output_schema=ForceFieldTaskDocument,
+        output_schema=ForceFieldStructureTaskDocument,
     )
     def make(
         self,
         structure: Structure,
         prev_dir: str | Path | None = None,
-    ) -> ForceFieldTaskDocument:
+    ) -> ForceFieldStructureTaskDocument:
         """
         Perform MD on a structure using forcefields and jobflow.
 
@@ -156,7 +156,7 @@ def make(
                 stacklevel=1,
             )
 
-        return ForceFieldTaskDocument.from_ase_compatible_result(
+        return ForceFieldStructureTaskDocument.from_ase_compatible_result(
             str(self.force_field_name),  # make mypy happy
             md_result,
             relax_cell=(self.ensemble == MDEnsemble.npt),

src/atomate2/forcefields/schemas.py

@@ -36,7 +36,7 @@ class ForcefieldObject(ValueEnum):
     TRAJECTORY = "trajectory"
 
 
-class ForceFieldTaskDocument(AseStructureTaskDoc):
+class ForceFieldStructureTaskDocument(AseStructureTaskDoc):
     """Document containing information on structure manipulation using a force field."""
 
     forcefield_name: Optional[str] = Field(
@@ -87,7 +87,7 @@ def from_ase_compatible_result(
         store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO,
         tags: list[str] | None = None,
         **task_document_kwargs,
-    ) -> ForceFieldTaskDocument:
+    ) -> ForceFieldStructureTaskDocument:
         """Create an AseTaskDoc for a task that has ASE-compatible outputs.
 
         Parameters