From 7dc5573bbd012707ba8ff19e5718d5b85adaa4e5 Mon Sep 17 00:00:00 2001 From: Yi Yao Date: Sat, 15 Feb 2025 22:05:44 -0600 Subject: [PATCH 1/3] 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 913482ed1c..4d1bf1bae4 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -233,27 +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): """Document containing information on molecule manipulation using ASE.""" diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 5ab0e71488..4f6aae8ea4 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__) @@ -48,7 +52,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: @@ -72,9 +77,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 @@ -118,7 +121,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" @@ -146,15 +149,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. @@ -170,7 +173,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, @@ -212,7 +215,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 +258,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 +294,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 +337,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 +375,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 +418,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 +454,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 +497,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 +532,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 +579,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 +619,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 +668,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 +710,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 +750,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 +786,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 5473933e5f..74c1ee1ac7 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -3,7 +3,7 @@ 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 @@ -23,7 +23,7 @@ NequipRelaxMaker, NequipStaticMaker, ) -from atomate2.forcefields.schemas import ForceFieldTaskDocument +from atomate2.forcefields.schemas import ForceFieldStructureTaskDocument def test_maker_initialization(): @@ -52,7 +52,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 @@ -113,7 +113,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 @@ -145,7 +145,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 @@ -172,7 +172,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 @@ -206,7 +206,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") @@ -292,7 +292,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) @@ -319,9 +319,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, @@ -333,12 +331,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") @@ -359,7 +372,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") @@ -393,7 +406,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) @@ -427,7 +440,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 @@ -469,7 +482,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 @@ -504,7 +517,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") @@ -543,7 +556,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 524f037f4c98b18fe59a78465a06dadeea6226d9 Mon Sep 17 00:00:00 2001 From: Yi Yao Date: Wed, 26 Feb 2025 11:25:29 -0600 Subject: [PATCH 2/3] 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 4d1bf1bae4..913482ed1c 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -233,6 +233,27 @@ 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): """Document containing information on molecule manipulation using ASE.""" From 4e9f06190e44463c0dc57469ba4872e7b8f99738 Mon Sep 17 00:00:00 2001 From: Yi Yao Date: Fri, 7 Mar 2025 11:18:21 -0600 Subject: [PATCH 3/3] 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 4f6aae8ea4..824f2d301e 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -53,7 +53,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: @@ -121,7 +122,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" @@ -215,7 +217,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" @@ -258,7 +261,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" @@ -294,7 +298,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" @@ -337,7 +342,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" @@ -375,7 +381,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" @@ -418,7 +425,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" @@ -454,7 +462,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" @@ -497,7 +506,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" @@ -532,7 +542,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" @@ -579,7 +590,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" @@ -619,7 +631,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" @@ -668,7 +681,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" @@ -710,7 +724,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" @@ -750,7 +765,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" @@ -786,7 +802,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"