Skip to content

Molecule for forcefield #1127

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Molecule for forcefield #1127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7af3143
output ForceFieldStructureTaskDocument or ForceFieldMoleculeTaskDocum…
yaoyi92 Feb 16, 2025
18534e0
add from_ase_task_doc func back in ase/schemas.py
yaoyi92 Feb 26, 2025
07ffd13
add or ForceFieldMoleculeTaskDocument in the function descriptions
yaoyi92 Mar 7, 2025
c040c8e
update ForceFieldTaskDocument to ForceFieldStructureTaskDocument
yaoyi92 Apr 18, 2025
95fd127
pick the correct mdareporter version
yaoyi92 Apr 21, 2025
fcea4a6
Merge branch 'main' into mol_for_forcefield
JaGeo Jun 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,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==2024.11.13", # TODO: open ff is extremely sensitive to pymatgen version
]
Expand Down
74 changes: 47 additions & 27 deletions src/atomate2/forcefields/jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,19 @@

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:
from collections.abc import Callable
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__)

Expand Down Expand Up @@ -58,7 +62,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
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:
Expand All @@ -82,9 +88,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
Expand Down Expand Up @@ -128,7 +132,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = "Force field relax"
Expand Down Expand Up @@ -156,15 +161,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.
Expand All @@ -180,7 +185,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,
Expand Down Expand Up @@ -222,7 +227,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:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = "Force field static"
Expand Down Expand Up @@ -265,7 +271,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.CHGNet} relax"
Expand Down Expand Up @@ -301,7 +308,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.CHGNet} static"
Expand Down Expand Up @@ -344,7 +352,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.M3GNet} relax"
Expand Down Expand Up @@ -382,7 +391,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.M3GNet} static"
Expand Down Expand Up @@ -425,7 +435,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.NEP} relax"
Expand Down Expand Up @@ -461,7 +472,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.NEP} static"
Expand Down Expand Up @@ -504,7 +516,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.Nequip} relax"
Expand Down Expand Up @@ -539,7 +552,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.Nequip} static"
Expand Down Expand Up @@ -586,7 +600,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.MACE_MP_0} relax"
Expand Down Expand Up @@ -626,7 +641,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.MACE_MP_0} static"
Expand Down Expand Up @@ -675,7 +691,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.SevenNet} relax"
Expand Down Expand Up @@ -717,7 +734,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.SevenNet} static"
Expand Down Expand Up @@ -757,7 +775,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.GAP} relax"
Expand Down Expand Up @@ -793,7 +812,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:`.ForceFieldTaskDocument()`.
Additional keyword args passed to :obj:`.ForceFieldStructureTaskDocument()` or
:obj: `ForceFieldMoleculeTaskDocument`.
"""

name: str = f"{MLFF.GAP} static"
Expand Down
17 changes: 10 additions & 7 deletions src/atomate2/forcefields/md.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure why this isn't done for all other ForceFieldMakers? There, I only see a replacement with ForceFieldStructureTaskDocument.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Which makers are you referring to? I only found two makers that have an explicit make function, ForceFieldRelaxMaker and ForceFieldMDMaker. It seems, the other makers inherent from these two.

"""
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
Expand All @@ -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),
Expand Down
Loading
Loading