From 7af3143fba456dc3aaf6e5ccce712ed13d659569 Mon Sep 17 00:00:00 2001
From: Yi Yao <yaoyi92@gmail.com>
Date: Sat, 15 Feb 2025 22:05:44 -0600
Subject: [PATCH 1/5] output ForceFieldStructureTaskDocument or
 ForceFieldMoleculeTaskDocument based on the input type of mol_or_struct.

change name ForceFieldTaskDocument => ForceFieldStructureTaskDocument

output ForceFieldStructureTaskDocument or ForceFieldMoleculeTaskDocument based on type of mol_or_struct

update ForceFieldTaskDocument => ForceFieldStructureTaskDocument in the tests

import Union from typing

include Union in forcefield/md.py

take the suggestions from the formatter

take ruff's suggestions

try again with ruff format

ruff format again

try again ruff

ruff again

fix the mypy error

Take inputs of both Molecule and Structure

update docstring

add molecule test for forcefield
---
 src/atomate2/ase/schemas.py         |  21 -----
 src/atomate2/forcefields/jobs.py    |  57 +++++++------
 src/atomate2/forcefields/md.py      |  17 ++--
 src/atomate2/forcefields/schemas.py | 125 ++++++++++++++++++++++++++--
 tests/conftest.py                   |   7 +-
 tests/forcefields/test_jobs.py      |  47 +++++++----
 tests/test_data/molecules/water.xyz |   5 ++
 7 files changed, 200 insertions(+), 79 deletions(-)
 create mode 100644 tests/test_data/molecules/water.xyz

diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py
index 35d5319f4e..e7ed33bb3e 100644
--- a/src/atomate2/ase/schemas.py
+++ b/src/atomate2/ase/schemas.py
@@ -232,27 +232,6 @@ class AseStructureTaskDoc(StructureMetadata):
 
     tags: list[str] | None = 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):
     """Document containing information on molecule manipulation using ASE."""
diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py
index 63bdf23b45..8643fe75c2 100644
--- a/src/atomate2/forcefields/jobs.py
+++ b/src/atomate2/forcefields/jobs.py
@@ -15,7 +15,11 @@
 
 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 (
+    ForceFieldMoleculeTaskDocument,
+    ForceFieldStructureTaskDocument,
+    ForceFieldTaskDocument,
+)
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
@@ -23,7 +27,7 @@
     from pathlib import Path
 
     from ase.calculators.calculator import Calculator
-    from pymatgen.core.structure import Structure
+    from pymatgen.core.structure import Molecule, Structure
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +54,8 @@ 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:
@@ -74,9 +79,7 @@ def make(structure):
     callable
         A decorated version of the make function that will generate forcefield jobs.
     """
-    return job(
-        method, data=_FORCEFIELD_DATA_OBJECTS, output_schema=ForceFieldTaskDocument
-    )
+    return job(method, data=_FORCEFIELD_DATA_OBJECTS)
 
 
 @dataclass
@@ -120,7 +123,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"
@@ -148,15 +151,15 @@ def __post_init__(self) -> None:
 
     @forcefield_job
     def make(
-        self, structure: Structure, prev_dir: str | Path | None = None
-    ) -> ForceFieldTaskDocument:
+        self, structure: Molecule | Structure, prev_dir: str | Path | None = None
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """
         Perform a relaxation of a structure using a force field.
 
         Parameters
         ----------
-        structure: .Structure
-            pymatgen structure.
+        structure: .Structure or Molecule
+            pymatgen structure or molecule.
         prev_dir : str or Path or None
             A previous calculation directory to copy output files from. Unused, just
                 added to match the method signature of other makers.
@@ -172,7 +175,7 @@ def make(
                 stacklevel=1,
             )
 
-        return ForceFieldTaskDocument.from_ase_compatible_result(
+        return ForceFieldTaskDocument.from_ase_compatible_result_forcefield(
             str(self.force_field_name),  # make mypy happy
             ase_result,
             self.steps,
@@ -214,7 +217,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"
@@ -257,7 +260,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"
@@ -293,7 +296,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"
@@ -336,7 +339,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"
@@ -374,7 +377,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"
@@ -417,7 +420,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"
@@ -453,7 +456,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"
@@ -496,7 +499,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"
@@ -531,7 +534,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"
@@ -578,7 +581,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"
@@ -618,7 +621,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"
@@ -667,7 +670,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"
@@ -709,7 +712,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"
@@ -749,7 +752,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"
@@ -785,7 +788,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"
diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py
index c46330cd52..ef6d0f3692 100644
--- a/src/atomate2/forcefields/md.py
+++ b/src/atomate2/forcefields/md.py
@@ -15,14 +15,18 @@
     _DEFAULT_CALCULATOR_KWARGS,
     _FORCEFIELD_DATA_OBJECTS,
 )
-from atomate2.forcefields.schemas import ForceFieldTaskDocument
+from atomate2.forcefields.schemas import (
+    ForceFieldMoleculeTaskDocument,
+    ForceFieldStructureTaskDocument,
+    ForceFieldTaskDocument,
+)
 from atomate2.forcefields.utils import ase_calculator, revert_default_dtype
 
 if TYPE_CHECKING:
     from pathlib import Path
 
     from ase.calculators.calculator import Calculator
-    from pymatgen.core.structure import Structure
+    from pymatgen.core.structure import Molecule, Structure
 
 
 @dataclass
@@ -126,19 +130,18 @@ def __post_init__(self) -> None:
 
     @job(
         data=[*_FORCEFIELD_DATA_OBJECTS, "ionic_steps"],
-        output_schema=ForceFieldTaskDocument,
     )
     def make(
         self,
-        structure: Structure,
+        structure: Molecule | Structure,
         prev_dir: str | Path | None = None,
-    ) -> ForceFieldTaskDocument:
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """
         Perform MD on a structure using forcefields and jobflow.
 
         Parameters
         ----------
-        structure: .Structure
+        structure: .Structure or Molecule
             pymatgen structure.
         prev_dir : str or Path or None
             A previous calculation directory to copy output files from. Unused, just
@@ -156,7 +159,7 @@ def make(
                 stacklevel=1,
             )
 
-        return ForceFieldTaskDocument.from_ase_compatible_result(
+        return ForceFieldTaskDocument.from_ase_compatible_result_forcefield(
             str(self.force_field_name),  # make mypy happy
             md_result,
             relax_cell=(self.ensemble == MDEnsemble.npt),
diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py
index 14f737ec09..21e14442f0 100644
--- a/src/atomate2/forcefields/schemas.py
+++ b/src/atomate2/forcefields/schemas.py
@@ -8,9 +8,16 @@
 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 Molecule, Structure
 
-from atomate2.ase.schemas import AseObject, AseResult, AseStructureTaskDoc, AseTaskDoc
+from atomate2.ase.schemas import (
+    AseMoleculeTaskDoc,
+    AseObject,
+    AseResult,
+    AseStructureTaskDoc,
+    AseTaskDoc,
+    _task_doc_translation_keys,
+)
 from atomate2.forcefields import MLFF
 
 
@@ -36,7 +43,81 @@ 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(
+        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 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(
@@ -68,7 +149,7 @@ class ForceFieldTaskDocument(AseStructureTaskDoc):
     )
 
     @classmethod
-    def from_ase_compatible_result(
+    def from_ase_compatible_result_forcefield(
         cls,
         ase_calculator_name: str,
         result: AseResult,
@@ -87,8 +168,8 @@ def from_ase_compatible_result(
         store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO,
         tags: list[str] | None = None,
         **task_document_kwargs,
-    ) -> ForceFieldTaskDocument:
-        """Create an AseTaskDoc for a task that has ASE-compatible outputs.
+    ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
+        """Create an ForceField output for a task that has ASE-compatible outputs.
 
         Parameters
         ----------
@@ -155,6 +236,38 @@ 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
+    ) -> ForceFieldStructureTaskDocument | 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:`.ForceFieldStructureTaskDocument()`
+            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 = ForceFieldStructureTaskDocument
+            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) -> Optional[dict[AseObject, Any]]:
         """Alias `objects` attr for backwards compatibility."""
diff --git a/tests/conftest.py b/tests/conftest.py
index 7def0a6906..20bcacf683 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -12,7 +12,7 @@
 from jobflow.settings import JobflowSettings
 from maggma.stores import MemoryStore
 from monty.serialization import loadfn
-from pymatgen.core import Structure
+from pymatgen.core import Molecule, Structure
 
 from atomate2.utils.log import initialize_logger
 
@@ -117,6 +117,11 @@ def ba_ti_o3_structure(test_dir):
     return Structure.from_file(test_dir / "structures" / "BaTiO3.cif")
 
 
+@pytest.fixture
+def water_molecule(test_dir):
+    return Molecule.from_file(test_dir / "molecules" / "water.xyz")
+
+
 @pytest.fixture(autouse=True)
 def mock_jobflow_settings(memory_jobstore):
     """Mock the jobflow settings to use our specific jobstore (with data store)."""
diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py
index dc68c160ee..69353c641d 100644
--- a/tests/forcefields/test_jobs.py
+++ b/tests/forcefields/test_jobs.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pytest
 from jobflow import run_locally
-from pymatgen.core import Structure
+from pymatgen.core import Molecule, Structure
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 from pytest import approx, importorskip
 
@@ -24,7 +24,7 @@
     NequipRelaxMaker,
     NequipStaticMaker,
 )
-from atomate2.forcefields.schemas import ForceFieldTaskDocument
+from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument
 
 
 def test_maker_initialization():
@@ -53,7 +53,7 @@ def test_chgnet_static_maker(si_structure):
 
     # validate job outputs
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.output.energy == approx(-10.6275062, rel=1e-4)
     assert output1.output.ionic_steps[-1].magmoms is None
     assert output1.output.n_steps == 1
@@ -114,7 +114,7 @@ def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool):
 
     # validate job outputs
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     if relax_cell:
         assert not output1.is_force_converged
         assert output1.output.n_steps == max_step + 2
@@ -146,7 +146,7 @@ def test_m3gnet_static_maker(si_structure):
 
     # validate job outputs
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.output.energy == approx(-10.8, abs=0.2)
     assert output1.output.n_steps == 1
 
@@ -173,7 +173,7 @@ def test_m3gnet_relax_maker(si_structure):
 
     # validate job outputs
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.is_force_converged
     assert output1.output.energy == approx(-10.8, abs=0.2)
     assert output1.output.n_steps == 24
@@ -207,7 +207,7 @@ def test_mace_static_maker(si_structure: Structure, test_dir: Path, model):
 
     # validation the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.output.energy == approx(-0.068231, rel=1e-4)
     assert output1.output.n_steps == 1
     assert output1.forcefield_version == get_imported_version("mace-torch")
@@ -293,7 +293,7 @@ def test_mace_relax_maker(
     # validating the outputs of the job
     output1 = responses[job.uuid][1].output
     assert output1.is_force_converged
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
 
     si_atoms = si_structure.to_ase_atoms()
     symmetry_ops_init = check_symmetry(si_atoms, symprec=1.0e-3)
@@ -320,9 +320,7 @@ def test_mace_relax_maker(
         assert output1.output.n_steps == 7
 
 
-def test_mace_mpa_0_relax_maker(
-    si_structure: Structure,
-):
+def test_mace_mpa_0_relax_maker(si_structure: Structure, water_molecule: Molecule):
     job = ForceFieldRelaxMaker(
         force_field_name="MACE_MPA_0",
         steps=25,
@@ -334,12 +332,27 @@ def test_mace_mpa_0_relax_maker(
     # validating the outputs of the job
     output = responses[job.uuid][1].output
 
+    job_mol = ForceFieldRelaxMaker(
+        force_field_name="MACE_MPA_0",
+        steps=25,
+        relax_kwargs={"fmax": 0.005},
+    ).make(water_molecule)
+    # run the flow or job and ensure that it finished running successfully
+    responses_mol = run_locally(job_mol, ensure_success=True)
+
+    # validating the outputs of the job
+    output_mol = responses_mol[job_mol.uuid][1].output
+
     assert output.ase_calculator_name == "MLFF.MACE_MPA_0"
     assert output.output.energy == pytest.approx(-10.829493522644043)
     assert output.output.structure.volume == pytest.approx(40.87471552602735)
     assert len(output.output.ionic_steps) == 4
     assert output.structure.volume == output.output.structure.volume
 
+    assert output_mol.ase_calculator_name == "MLFF.MACE_MPA_0"
+    assert output_mol.output.energy == pytest.approx(-13.786081314086914)
+    assert len(output_mol.output.ionic_steps) == 20
+
 
 def test_gap_static_maker(si_structure: Structure, test_dir):
     importorskip("quippy")
@@ -360,7 +373,7 @@ def test_gap_static_maker(si_structure: Structure, test_dir):
 
     # validation the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.output.energy == approx(-10.8523, rel=1e-4)
     assert output1.output.n_steps == 1
     assert output1.forcefield_version == get_imported_version("quippy-ase")
@@ -394,7 +407,7 @@ def test_gap_relax_maker(si_structure: Structure, test_dir: Path, relax_cell: bo
 
     # validating the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     if relax_cell:
         assert not output1.is_force_converged
         assert output1.output.energy == approx(-13.08492, rel=1e-2)
@@ -428,7 +441,7 @@ def test_nep_static_maker(al2_au_structure: Structure, test_dir: Path):
 
     # validation the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.output.energy == approx(-47.65972, rel=1e-4)
     assert output1.output.n_steps == 1
 
@@ -470,7 +483,7 @@ def test_nep_relax_maker(
 
     # validate the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     if relax_cell:
         assert output1.output.energy == approx(-47.6727, rel=1e-3)
         assert output1.output.n_steps == 3
@@ -505,7 +518,7 @@ def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path):
 
     # validation the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     assert output1.output.energy == approx(-44.40017, rel=1e-4)
     assert output1.output.n_steps == 1
     assert output1.forcefield_version == get_imported_version("nequip")
@@ -544,7 +557,7 @@ def test_nequip_relax_maker(
 
     # validation the outputs of the job
     output1 = responses[job.uuid][1].output
-    assert isinstance(output1, ForceFieldTaskDocument)
+    assert isinstance(output1, ForceFieldStructureTaskDocument)
     if relax_cell:
         assert output1.output.energy == approx(-44.407, rel=1e-3)
         assert output1.output.n_steps == 5
diff --git a/tests/test_data/molecules/water.xyz b/tests/test_data/molecules/water.xyz
new file mode 100644
index 0000000000..c067ca4f1f
--- /dev/null
+++ b/tests/test_data/molecules/water.xyz
@@ -0,0 +1,5 @@
+3
+Water molecule
+O          0.00000        0.00000        0.11779
+H          0.00000        0.75545       -0.47116
+H          0.00000       -0.75545       -0.47116

From 18534e05cfdec96e165ae6ec6a42526ce6a9ec68 Mon Sep 17 00:00:00 2001
From: Yi Yao <yaoyi92@gmail.com>
Date: Wed, 26 Feb 2025 11:25:29 -0600
Subject: [PATCH 2/5] add from_ase_task_doc func back in ase/schemas.py

---
 src/atomate2/ase/schemas.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py
index e7ed33bb3e..35d5319f4e 100644
--- a/src/atomate2/ase/schemas.py
+++ b/src/atomate2/ase/schemas.py
@@ -232,6 +232,27 @@ class AseStructureTaskDoc(StructureMetadata):
 
     tags: list[str] | None = 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):
     """Document containing information on molecule manipulation using ASE."""

From 07ffd13ec8d9a2d386baaf7da58b24a2c13b7ae3 Mon Sep 17 00:00:00 2001
From: Yi Yao <yaoyi92@gmail.com>
Date: Fri, 7 Mar 2025 11:18:21 -0600
Subject: [PATCH 3/5] add or ForceFieldMoleculeTaskDocument in the function
 descriptions

---
 src/atomate2/forcefields/jobs.py | 51 +++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py
index 8643fe75c2..9ceda4aba0 100644
--- a/src/atomate2/forcefields/jobs.py
+++ b/src/atomate2/forcefields/jobs.py
@@ -55,7 +55,8 @@ def forcefield_job(method: Callable) -> job:
     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
-    ForceFieldStructureTaskDocument :obj:`.TaskDoc`.
+    ForceFieldStructureTaskDocument :obj:`.TaskDoc`. or
+    ForceFieldMoleculeTaskDocument :obj:`.TaskDoc`.
 
     Any makers that return forcefield jobs (not flows) should decorate the
     ``make`` method with @forcefield_job. For example:
@@ -123,7 +124,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = "Force field relax"
@@ -217,7 +219,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = "Force field static"
@@ -260,7 +263,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.CHGNet} relax"
@@ -296,7 +300,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.CHGNet} static"
@@ -339,7 +344,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.M3GNet} relax"
@@ -377,7 +383,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.M3GNet} static"
@@ -420,7 +427,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.NEP} relax"
@@ -456,7 +464,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.NEP} static"
@@ -499,7 +508,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.Nequip} relax"
@@ -534,7 +544,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.Nequip} static"
@@ -581,7 +592,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.MACE_MP_0} relax"
@@ -621,7 +633,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.MACE_MP_0} static"
@@ -670,7 +683,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.SevenNet} relax"
@@ -712,7 +726,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.SevenNet} static"
@@ -752,7 +767,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.GAP} relax"
@@ -788,7 +804,8 @@ 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:`.ForceFieldStructureTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = f"{MLFF.GAP} static"

From c040c8e81587ba6d15643028346ead6c3c551f3c Mon Sep 17 00:00:00 2001
From: Yi Yao <yaoyi92@gmail.com>
Date: Fri, 18 Apr 2025 11:05:11 -0500
Subject: [PATCH 4/5] update ForceFieldTaskDocument to
 ForceFieldStructureTaskDocument

---
 tests/forcefields/test_jobs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py
index 69353c641d..cddc0d2e25 100644
--- a/tests/forcefields/test_jobs.py
+++ b/tests/forcefields/test_jobs.py
@@ -647,7 +647,7 @@ def test_matpes_relax_makers(
     resp = run_locally(job)
     output = resp[job.uuid][1].output
 
-    assert isinstance(output, ForceFieldTaskDocument)
+    assert isinstance(output, ForceFieldStructureTaskDocument)
 
     ref = refs[ref_func]
     assert output.output.energy == approx(ref["energy"])

From 95fd127586e63938a721c7e7e29cb1e668514dfc Mon Sep 17 00:00:00 2001
From: Yi Yao <yaoyi92@gmail.com>
Date: Mon, 21 Apr 2025 13:53:34 -0500
Subject: [PATCH 5/5] pick the correct mdareporter version

---
 pyproject.toml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index f44a6d74b4..31f4ec403c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -121,7 +121,9 @@ strict = [
 strict-openff = [
     "mdanalysis==2.9.0",
     "monty==2025.3.3",
-    "openmm-mdanalysis-reporter==0.1.0",
+    #"openmm-mdanalysis-reporter==0.1.0",
+    "mdareporter @ git+https://github.com/sef43/openmm-mdanalysis-reporter.git@86e8bdffb63bbe9b13430ba97e2a67b85c996048",
+    "h5py==3.13",
     "openmm==8.1.1",
     "pymatgen==2025.4.20", # TODO: open ff is extremely sensitive to pymatgen version
     "mdanalysis==2.9.0"