simplify classes

Author: esoteric-ephemera

src/atomate2/forcefields/jobs.py

@@ -10,7 +10,7 @@
 from jobflow import job
 
 from atomate2.ase.jobs import AseRelaxMaker
-from atomate2.forcefields.schemas import BaseForceFieldTaskDocument
+from atomate2.forcefields.schemas import ForceFieldTaskDocument
 from atomate2.forcefields.utils import _FORCEFIELD_DATA_OBJECTS, MLFF, ForceFieldMixin
 
 if TYPE_CHECKING:
@@ -19,10 +19,7 @@
 
     from pymatgen.core.structure import Molecule, Structure
 
-    from atomate2.forcefields.schemas import (
-        ForceFieldMoleculeTaskDocument,
-        ForceFieldTaskDocument,
-    )
+    from atomate2.forcefields.schemas import ForceFieldMoleculeTaskDocument
 
 logger = logging.getLogger(__name__)
 
@@ -144,7 +141,7 @@ def make(
                 stacklevel=1,
             )
 
-        return BaseForceFieldTaskDocument.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

@@ -9,18 +9,15 @@
 from jobflow import job
 
 from atomate2.ase.md import AseMDMaker, MDEnsemble
-from atomate2.forcefields.schemas import BaseForceFieldTaskDocument
+from atomate2.forcefields.schemas import ForceFieldTaskDocument
 from atomate2.forcefields.utils import _FORCEFIELD_DATA_OBJECTS, ForceFieldMixin
 
 if TYPE_CHECKING:
     from pathlib import Path
 
     from pymatgen.core.structure import Molecule, Structure
 
-    from atomate2.forcefields.schemas import (
-        ForceFieldMoleculeTaskDocument,
-        ForceFieldTaskDocument,
-    )
+    from atomate2.forcefields.schemas import ForceFieldMoleculeTaskDocument
 
 
 @dataclass
@@ -137,7 +134,7 @@ def make(
                 stacklevel=1,
             )
 
-        return BaseForceFieldTaskDocument.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

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from emmet.core.types.enums import StoreTrajectoryOption
 from pydantic import BaseModel, Field
-from pymatgen.core import Molecule, Structure
+from pymatgen.core import Molecule
 
 from atomate2.ase.schemas import (
     AseMoleculeTaskDoc,
@@ -18,6 +18,9 @@
 )
 from atomate2.forcefields import MLFF
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 
 class ForceFieldMeta(BaseModel):
     """Add metadata to forcefield output documents."""
@@ -56,44 +59,33 @@ def forcefield_objects(self) -> dict[AseObject, Any] | None:
         return self.objects
 
 
-class ForceFieldTaskDocument(AseStructureTaskDoc, ForceFieldMeta):
-    """Document containing information on structure manipulation using a force field."""
-
-
 class ForceFieldMoleculeTaskDocument(AseMoleculeTaskDoc, ForceFieldMeta):
-    """Document containing information on structure manipulation using a force field."""
-
-
-class BaseForceFieldTaskDocument(AseTaskDoc):
-    """Document containing information on structure manipulation using a force field."""
-
-    forcefield_name: str | None = Field(
-        None,
-        description="name of the interatomic potential used for relaxation.",
-    )
+    """Document containing information on molecule manipulation using a force field."""
 
-    forcefield_version: str | None = Field(
-        "Unknown",
-        description="version of the interatomic potential used for relaxation.",
-    )
+    @classmethod
+    def from_ase_task_doc(
+        cls, ase_task_doc: AseTaskDoc, **task_document_kwargs
+    ) -> Self:
+        """Create a ForceFieldMoleculeTaskDocument from an AseTaskDoc.
 
-    dir_name: str | None = Field(
-        None, description="Directory where the force field calculations are performed."
-    )
+        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_molecule(
+            meta_molecule=ase_task_doc.mol_or_struct, **task_document_kwargs
+        )
 
-    included_objects: list[AseObject] | None = Field(
-        None, description="list of forcefield objects included with this task document"
-    )
-    objects: dict[AseObject, Any] | None = Field(
-        None, description="Forcefield objects associated with this task"
-    )
 
-    is_force_converged: bool | None = Field(
-        None,
-        description=(
-            "Whether the calculation is converged with respect to interatomic forces."
-        ),
-    )
+class ForceFieldTaskDocument(AseStructureTaskDoc, ForceFieldMeta):
+    """Document containing information on atomistic manipulation using a force field."""
 
     @classmethod
     def from_ase_compatible_result(  # type: ignore[override]
@@ -115,8 +107,8 @@ def from_ase_compatible_result(  # type: ignore[override]
         store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO,
         tags: list[str] | None = None,
         **task_document_kwargs,
-    ) -> ForceFieldTaskDocument | ForceFieldMoleculeTaskDocument:
-        """Create an ForceField output for a task that has ASE-compatible outputs.
+    ) -> Self | ForceFieldMoleculeTaskDocument:
+        """Create forcefield output for a task that has ASE-compatible outputs.
 
         Parameters
         ----------
@@ -184,41 +176,8 @@ def from_ase_compatible_result(  # type: ignore[override]
 
             ff_kwargs["forcefield_version"] = importlib.metadata.version(pkg_name)
 
-        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
-    ) -> ForceFieldTaskDocument | ForceFieldMoleculeTaskDocument:
-        """Create an ForceField output 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:`.ForceFieldTaskDocument()`
-            or `.ForceFieldMoleculeTaskDocument()`.
-        """
-        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 = ForceFieldTaskDocument
-            k = "structure"
-            if relax_cell := getattr(ase_task_doc, "relax_cell", None):
-                task_document_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) -> dict[AseObject, Any] | None:
-        """Alias `objects` attr for backwards compatibility."""
-        return self.objects
+        return (
+            ForceFieldMoleculeTaskDocument
+            if isinstance(result.final_mol_or_struct, Molecule)
+            else cls
+        ).from_ase_task_doc(ase_task_doc, **ff_kwargs)