output ForceFieldStructureTaskDocument or ForceFieldMoleculeTaskDocument based on type of mol_or_struct

Author: yaoyi92

src/atomate2/ase/schemas.py

@@ -233,26 +233,6 @@ class AseStructureTaskDoc(StructureMetadata):
 
     tags: Optional[list[str]] = Field(None, description="List of tags for the task.")
 
-    @classmethod
-    def from_ase_task_doc(
-        cls, ase_task_doc: AseTaskDoc, **task_document_kwargs
-    ) -> AseStructureTaskDoc:
-        """Create an AseStructureTaskDoc for a task that has ASE-compatible outputs.
-
-        Parameters
-        ----------
-        ase_task_doc : AseTaskDoc
-            Task doc for the calculation
-        task_document_kwargs : dict
-            Additional keyword args passed to :obj:`.AseStructureTaskDoc()`.
-        """
-        task_document_kwargs.update(
-            {k: getattr(ase_task_doc, k) for k in _task_doc_translation_keys},
-            structure=ase_task_doc.mol_or_struct,
-        )
-        return cls.from_structure(
-            meta_structure=ase_task_doc.mol_or_struct, **task_document_kwargs
-        )
 
 
 class AseMoleculeTaskDoc(MoleculeMetadata):

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 ForceFieldStructureTaskDocument
+from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument, ForceFieldTaskDocument
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -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=ForceFieldStructureTaskDocument
+        method, data=_FORCEFIELD_DATA_OBJECTS
     )
 
 
@@ -147,7 +147,7 @@ def __post_init__(self) -> None:
     @forcefield_job
     def make(
         self, structure: Structure, prev_dir: str | Path | None = None
-    ) -> ForceFieldStructureTaskDocument:
+    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
         """
         Perform a relaxation of a structure using a force field.
 
@@ -170,7 +170,7 @@ def make(
                 stacklevel=1,
             )
 
-        return ForceFieldStructureTaskDocument.from_ase_compatible_result(
+        return ForceFieldTaskDocument.from_ase_compatible_result(
             str(self.force_field_name),  # make mypy happy
             ase_result,
             self.steps,

src/atomate2/forcefields/md.py

@@ -15,7 +15,7 @@
     _DEFAULT_CALCULATOR_KWARGS,
     _FORCEFIELD_DATA_OBJECTS,
 )
-from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument
+from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument, ForceFieldTaskDocument
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -126,13 +126,12 @@ def __post_init__(self) -> None:
 
     @job(
         data=[*_FORCEFIELD_DATA_OBJECTS, "ionic_steps"],
-        output_schema=ForceFieldStructureTaskDocument,
     )
     def make(
         self,
         structure: Structure,
         prev_dir: str | Path | None = None,
-    ) -> ForceFieldStructureTaskDocument:
+    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
         """
         Perform MD on a structure using forcefields and jobflow.
 
@@ -156,7 +155,7 @@ def make(
                 stacklevel=1,
             )
 
-        return ForceFieldStructureTaskDocument.from_ase_compatible_result(
+        return ForceFieldTaskDocument.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

@@ -8,9 +8,10 @@
 from emmet.core.vasp.calculation import StoreTrajectoryOption
 from monty.dev import deprecated
 from pydantic import Field
-from pymatgen.core import Structure
+from pymatgen.core import Structure, Molecule
 
-from atomate2.ase.schemas import AseObject, AseResult, AseStructureTaskDoc, AseTaskDoc
+from atomate2.ase.schemas import AseObject, AseResult, AseMoleculeTaskDoc, AseStructureTaskDoc, AseTaskDoc
+from atomate2.ase.schemas import _task_doc_translation_keys
 from atomate2.forcefields import MLFF
 
 
@@ -67,6 +68,79 @@ class ForceFieldStructureTaskDocument(AseStructureTaskDoc):
         ),
     )
 
+
+    @property
+    def forcefield_objects(self) -> Optional[dict[AseObject, Any]]:
+        """Alias `objects` attr for backwards compatibility."""
+        return self.objects
+
+class ForceFieldMoleculeTaskDocument(AseMoleculeTaskDoc):
+    """Document containing information on structure manipulation using a force field."""
+
+    forcefield_name: Optional[str] = Field(
+        None,
+        description="name of the interatomic potential used for relaxation.",
+    )
+
+    forcefield_version: Optional[str] = Field(
+        "Unknown",
+        description="version of the interatomic potential used for relaxation.",
+    )
+
+    dir_name: Optional[str] = Field(
+        None, description="Directory where the force field calculations are performed."
+    )
+
+    included_objects: Optional[list[AseObject]] = Field(
+        None, description="list of forcefield objects included with this task document"
+    )
+    objects: Optional[dict[AseObject, Any]] = Field(
+        None, description="Forcefield objects associated with this task"
+    )
+
+    is_force_converged: Optional[bool] = Field(
+        None,
+        description=(
+            "Whether the calculation is converged with respect to interatomic forces."
+        ),
+    )
+
+    @property
+    def forcefield_objects(self) -> Optional[dict[AseObject, Any]]:
+        """Alias `objects` attr for backwards compatibility."""
+        return self.objects
+
+class ForceFieldTaskDocument(AseTaskDoc):
+    """Document containing information on structure manipulation using a force field."""
+
+    forcefield_name: Optional[str] = Field(
+        None,
+        description="name of the interatomic potential used for relaxation.",
+    )
+
+    forcefield_version: Optional[str] = Field(
+        "Unknown",
+        description="version of the interatomic potential used for relaxation.",
+    )
+
+    dir_name: Optional[str] = Field(
+        None, description="Directory where the force field calculations are performed."
+    )
+
+    included_objects: Optional[list[AseObject]] = Field(
+        None, description="list of forcefield objects included with this task document"
+    )
+    objects: Optional[dict[AseObject, Any]] = Field(
+        None, description="Forcefield objects associated with this task"
+    )
+
+    is_force_converged: Optional[bool] = Field(
+        None,
+        description=(
+            "Whether the calculation is converged with respect to interatomic forces."
+        ),
+    )
+
     @classmethod
     def from_ase_compatible_result(
         cls,
@@ -87,7 +161,7 @@ def from_ase_compatible_result(
         store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO,
         tags: list[str] | None = None,
         **task_document_kwargs,
-    ) -> ForceFieldStructureTaskDocument:
+    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
         """Create an AseTaskDoc for a task that has ASE-compatible outputs.
 
         Parameters
@@ -155,6 +229,35 @@ def from_ase_compatible_result(
 
         return cls.from_ase_task_doc(ase_task_doc, **ff_kwargs)
 
+    @classmethod
+    def from_ase_task_doc(
+        cls, ase_task_doc: AseTaskDoc, **task_document_kwargs
+    ) -> Union[ForceFieldStructureTaskDocument, ForceFieldMoleculeTaskDocument]:
+        """Create an AseStructureTaskDoc for a task that has ASE-compatible outputs.
+
+        Parameters
+        ----------
+        ase_task_doc : AseTaskDoc
+            Task doc for the calculation
+        task_document_kwargs : dict
+            Additional keyword args passed to :obj:`.AseStructureTaskDoc()`.
+        """
+        task_document_kwargs.update(
+            {k: getattr(ase_task_doc, k) for k in _task_doc_translation_keys},
+        )
+        if isinstance(ase_task_doc.mol_or_struct, Structure):
+            meta_class = ForceFieldStructureTaskDocument
+            k = "structure"
+            if relax_cell := getattr(ase_task_doc, "relax_cell", None):
+                kwargs.update({"relax_cell": relax_cell})
+            task_document_kwargs.update(structure=ase_task_doc.mol_or_struct)
+        elif isinstance(ase_task_doc.mol_or_struct, Molecule):
+            meta_class = ForceFieldMoleculeTaskDocument
+            k = "molecule"
+            task_document_kwargs.update(molecule=ase_task_doc.mol_or_struct)
+        task_document_kwargs.update({k: ase_task_doc.mol_or_struct, f"meta_{k}": ase_task_doc.mol_or_struct})
+        return getattr(meta_class, f"from_{k}")(**task_document_kwargs)
+
     @property
     def forcefield_objects(self) -> Optional[dict[AseObject, Any]]:
         """Alias `objects` attr for backwards compatibility."""