Merge pull request #1127 from yaoyi92/mol_for_forcefield

Molecule for forcefield

Author: esoteric-ephemera

.gitignore

@@ -79,4 +79,5 @@ docs/reference/atomate2.*
 .ipynb_checkpoints
 .aider*
 
-tests/test_data/forcefields/deepmd_graph.pb
+# deepmd-kit files
+**/*.pb

src/atomate2/forcefields/jobs.py

@@ -17,7 +17,9 @@
     from collections.abc import Callable
     from pathlib import Path
 
-    from pymatgen.core.structure import Structure
+    from pymatgen.core.structure import Molecule, Structure
+
+    from atomate2.forcefields.schemas import ForceFieldMoleculeTaskDocument
 
 logger = logging.getLogger(__name__)
 
@@ -29,7 +31,9 @@ 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
+    ForceFieldTaskDocument :obj:`.TaskDoc`. or
+    ForceFieldMoleculeTaskDocument :obj:`.TaskDoc`.
 
     Any makers that return forcefield jobs (not flows) should decorate the
     ``make`` method with @forcefield_job. For example:
@@ -53,9 +57,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
@@ -99,7 +101,8 @@ class ForceFieldRelaxMaker(ForceFieldMixin, 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:`.ForceFieldTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = "Force field relax"
@@ -115,15 +118,15 @@ class ForceFieldRelaxMaker(ForceFieldMixin, AseRelaxMaker):
 
     @forcefield_job
     def make(
-        self, structure: Structure, prev_dir: str | Path | None = None
-    ) -> ForceFieldTaskDocument:
+        self, structure: Molecule | Structure, prev_dir: str | Path | None = None
+    ) -> ForceFieldTaskDocument | 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,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:`.ForceFieldTaskDocument()`.
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()` or
+          :obj: `ForceFieldMoleculeTaskDocument`.
     """
 
     name: str = "Force field static"

src/atomate2/forcefields/md.py

@@ -15,7 +15,9 @@
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from pymatgen.core.structure import Structure
+    from pymatgen.core.structure import Molecule, Structure
+
+    from atomate2.forcefields.schemas import ForceFieldMoleculeTaskDocument
 
 
 @dataclass
@@ -104,19 +106,18 @@ class ForceFieldMDMaker(ForceFieldMixin, AseMDMaker):
 
     @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:
+    ) -> ForceFieldTaskDocument | 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

src/atomate2/forcefields/schemas.py

@@ -2,17 +2,28 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from emmet.core.types.enums import StoreTrajectoryOption
-from pydantic import Field
-
-from atomate2.ase.schemas import AseObject, AseResult, AseStructureTaskDoc, AseTaskDoc
+from pydantic import BaseModel, Field
+from pymatgen.core import Molecule
+
+from atomate2.ase.schemas import (
+    AseMoleculeTaskDoc,
+    AseObject,
+    AseResult,
+    AseStructureTaskDoc,
+    AseTaskDoc,
+    _task_doc_translation_keys,
+)
 from atomate2.forcefields import MLFF
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 
-class ForceFieldTaskDocument(AseStructureTaskDoc):
-    """Document containing information on structure manipulation using a force field."""
+class ForceFieldMeta(BaseModel):
+    """Add metadata to forcefield output documents."""
 
     forcefield_name: str | None = Field(
         None,
@@ -42,6 +53,40 @@ class ForceFieldTaskDocument(AseStructureTaskDoc):
         ),
     )
 
+    @property
+    def forcefield_objects(self) -> dict[AseObject, Any] | None:
+        """Alias `objects` attr for backwards compatibility."""
+        return self.objects
+
+
+class ForceFieldMoleculeTaskDocument(AseMoleculeTaskDoc, ForceFieldMeta):
+    """Document containing information on molecule manipulation using a force field."""
+
+    @classmethod
+    def from_ase_task_doc(
+        cls, ase_task_doc: AseTaskDoc, **task_document_kwargs
+    ) -> Self:
+        """Create a ForceFieldMoleculeTaskDocument from an AseTaskDoc.
+
+        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
+        )
+
+
+class ForceFieldTaskDocument(AseStructureTaskDoc, ForceFieldMeta):
+    """Document containing information on atomistic manipulation using a force field."""
+
     @classmethod
     def from_ase_compatible_result(
         cls,
@@ -62,8 +107,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.
+    ) -> Self | ForceFieldMoleculeTaskDocument:
+        """Create forcefield output for a task that has ASE-compatible outputs.
 
         Parameters
         ----------
@@ -131,9 +176,8 @@ def from_ase_compatible_result(
 
             ff_kwargs["forcefield_version"] = importlib.metadata.version(pkg_name)
 
-        return cls.from_ase_task_doc(ase_task_doc, **ff_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)

tests/forcefields/conftest.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
 import hashlib
+import tempfile
 import urllib.request
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import pytest
 import torch
 from emmet.core.utils import get_hash_blocked
 
 if TYPE_CHECKING:
-    from pathlib import Path
     from typing import Any
 
 
@@ -21,15 +22,17 @@ def pytest_runtest_setup(item: Any) -> None:
 
 
 @pytest.fixture(scope="session", autouse=True)
-def download_deepmd_pretrained_model(test_dir: Path) -> None:
+def get_deepmd_pretrained_model_path(test_dir: Path) -> Path:
     # Download DeepMD pretrained model from GitHub
     file_url = "https://raw.github.com/sliutheorygroup/UniPero/main/model/graph.pb"
-    local_path = test_dir / "forcefields" / "deepmd_graph.pb"
+    local_path = tempfile.NamedTemporaryFile(suffix=".pb")  # noqa : SIM115
     ref_md5 = "2814ae7f2eb1c605dd78f2964187de40"
-    _, http_message = urllib.request.urlretrieve(file_url, local_path)  # noqa: S310
+    _, http_message = urllib.request.urlretrieve(file_url, local_path.name)  # noqa: S310
     if "Content-Type: text/html" in http_message:
         raise RuntimeError(f"Failed to download from: {file_url}")
 
     # Check MD5 to ensure file integrity
-    if (file_md5 := get_hash_blocked(local_path, hasher=hashlib.md5())) != ref_md5:
+    if (file_md5 := get_hash_blocked(local_path.name, hasher=hashlib.md5())) != ref_md5:
         raise RuntimeError(f"MD5 mismatch: {file_md5} != {ref_md5}")
+    yield Path(local_path.name)
+    local_path.close()

tests/forcefields/test_jobs.py

@@ -4,12 +4,15 @@
 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
 
 from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker
-from atomate2.forcefields.schemas import ForceFieldTaskDocument
+from atomate2.forcefields.schemas import (
+    ForceFieldMoleculeTaskDocument,
+    ForceFieldTaskDocument,
+)
 
 
 def test_maker_initialization():
@@ -299,9 +302,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, test_dir: Path, tmp_dir):
     job = ForceFieldRelaxMaker(
         force_field_name="MACE_MPA_0",
         steps=25,
@@ -313,12 +314,29 @@ def test_mace_mpa_0_relax_maker(
     # validating the outputs of the job
     output = responses[job.uuid][1].output
 
+    water_molecule = Molecule.from_file(test_dir / "molecules" / "water.xyz.gz")
+    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 isinstance(output_mol, ForceFieldMoleculeTaskDocument)
+
     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")
@@ -522,14 +540,16 @@ def test_nequip_relax_maker(
     assert final_spg_num == 99
 
 
-def test_deepmd_static_maker(sr_ti_o3_structure: Structure, test_dir: Path):
+def test_deepmd_static_maker(
+    sr_ti_o3_structure: Structure, test_dir: Path, get_deepmd_pretrained_model_path
+):
     importorskip("deepmd")
 
     # generate job
     job = ForceFieldStaticMaker(
         force_field_name="DeepMD",
         ionic_step_data=("structure", "energy"),
-        calculator_kwargs={"model": test_dir / "forcefields" / "deepmd_graph.pb"},
+        calculator_kwargs={"model": get_deepmd_pretrained_model_path},
     ).make(sr_ti_o3_structure)
 
     # run the flow or job and ensure that it finished running successfully
@@ -552,6 +572,7 @@ def test_deepmd_relax_maker(
     test_dir: Path,
     relax_cell: bool,
     fix_symmetry: bool,
+    get_deepmd_pretrained_model_path: Path,
 ):
     importorskip("deepmd")
     # translate one atom to ensure a small number of relaxation steps are taken
@@ -563,7 +584,7 @@ def test_deepmd_relax_maker(
         optimizer_kwargs={"optimizer": "BFGSLineSearch"},
         relax_cell=relax_cell,
         fix_symmetry=fix_symmetry,
-        calculator_kwargs={"model": test_dir / "forcefields" / "deepmd_graph.pb"},
+        calculator_kwargs={"model": get_deepmd_pretrained_model_path},
     ).make(sr_ti_o3_structure)
 
     # run the flow or job and ensure that it finished running successfully

tests/forcefields/test_md.py

@@ -49,6 +49,7 @@ def test_ml_ff_md_maker(
     al2_au_structure,
     test_dir,
     clean_dir,
+    get_deepmd_pretrained_model_path,
 ):
     if ff_name in map(MLFF, ("Forcefield", "MACE")):
         return  # nothing to test here, MLFF.Forcefield is just a generic placeholder
@@ -103,7 +104,7 @@ def test_ml_ff_md_maker(
         }
         unit_cell_structure = sr_ti_o3_structure.copy()
     elif ff_name == MLFF.DeepMD:
-        calculator_kwargs = {"model": test_dir / "forcefields" / "deepmd_graph.pb"}
+        calculator_kwargs = {"model": get_deepmd_pretrained_model_path}
         unit_cell_structure = sr_ti_o3_structure.copy()
 
     structure = unit_cell_structure.to_conventional() * (2, 2, 2)

tests/forcefields/test_phonon.py

@@ -55,7 +55,7 @@ def test_supercell_orthorhombic(clean_dir, si_structure: Structure):
 
 
 def test_phonon_maker_initialization_with_all_mlff(
-    si_structure: Structure, test_dir: Path
+    si_structure: Structure, test_dir: Path, get_deepmd_pretrained_model_path: Path
 ):
     """Test PhononMaker can be initialized with all MLFF static and relax makers."""
 
@@ -74,7 +74,7 @@ def test_phonon_maker_initialization_with_all_mlff(
         calc_kwargs = {
             MLFF.Nequip: {"model_path": f"{chk_pt_dir}/nequip/nequip_ff_sr_ti_o3.pth"},
             MLFF.NEP: {"model_filename": f"{test_dir}/forcefields/nep/nep.txt"},
-            MLFF.DeepMD: {"model": test_dir / "forcefields" / "deepmd_graph.pb"},
+            MLFF.DeepMD: {"model": get_deepmd_pretrained_model_path},
         }.get(mlff, {})
         static_maker = ForceFieldStaticMaker(
             name=f"{mlff} static",