diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 362e566869..9301af96e7 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -54,8 +54,7 @@ class AseMaker(Maker, ABC): class EMTStaticMaker(AseMaker): name: str = "EMT static maker" - @property - def calculator(self): + def _get_calculator(self): return EMT() ``` @@ -95,27 +94,44 @@ def calculator(self): store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO tags: list[str] | None = None + def __post_init__(self) -> None: + """Enable caching of the ASE calculator via private attribute.""" + self._calculator: Calculator | None = None + @job(data=_ASE_DATA_OBJECTS) def make( self, - mol_or_struct: Molecule | Structure, + mol_or_struct: Molecule | Structure | list[Molecule | Structure], prev_dir: str | Path | None = None, - ) -> AseStructureTaskDoc | AseMoleculeTaskDoc: + ) -> ( + AseStructureTaskDoc + | AseMoleculeTaskDoc + | list[AseStructureTaskDoc | AseMoleculeTaskDoc] + ): """ Run ASE as job, can be re-implemented in subclasses. Parameters ---------- - mol_or_struct: .Molecule or .Structure - pymatgen molecule or structure + mol_or_struct: .Molecule, .Structure, or a list thereof + pymatgen molecule(s) or structure(s) 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. + + Returns + ------- + AseStructureTaskDoc, AseMoleculeTaskDoc, or list thereof. """ - return AseTaskDoc.to_mol_or_struct_metadata_doc( - getattr(self.calculator, "name", type(self.calculator).__name__), - self.run_ase(mol_or_struct, prev_dir=prev_dir), - ) + batch_mode = isinstance(mol_or_struct, list) + results = [ + AseTaskDoc.to_mol_or_struct_metadata_doc( + getattr(self.calculator, "name", type(self.calculator).__name__), + self.run_ase(atoms, prev_dir=prev_dir), + ) + for atoms in (mol_or_struct if batch_mode else [mol_or_struct]) + ] + return results if batch_mode else results[0] def run_ase( self, @@ -148,11 +164,18 @@ def run_ase( elapsed_time=t_f - t_i, ) - @property @abstractmethod + def _get_calculator(self) -> Calculator: + """Load ASE calculator, to be implemented by the user.""" + + @property def calculator(self) -> Calculator: - """ASE calculator, method to be implemented in subclasses.""" - raise NotImplementedError + """Retrieve cached ASE calculator.""" + if getattr(self, "_calculator", None) is None: + self._calculator = self._get_calculator() + if self._calculator is None: + raise ValueError("ASE calculator not properly initialized.") + return self._calculator @dataclass @@ -208,8 +231,7 @@ class AseRelaxMaker(AseMaker): def __post_init__(self) -> None: """Ensure that physical relaxation settings are used.""" - if hasattr(super(), "__post_init__"): - super().__post_init__() # type: ignore[misc] + super().__post_init__() if self.relax_cell and self.relax_shape: raise ValueError( "You have set both `relax_cell` (relaxing the cell shape and volume) " @@ -220,38 +242,48 @@ def __post_init__(self) -> None: @job(data=_ASE_DATA_OBJECTS) def make( self, - mol_or_struct: Molecule | Structure, + mol_or_struct: Molecule | Structure | list[Molecule | Structure], prev_dir: str | Path | None = None, - ) -> AseStructureTaskDoc | AseMoleculeTaskDoc: + ) -> ( + AseStructureTaskDoc + | AseMoleculeTaskDoc + | list[AseStructureTaskDoc | AseMoleculeTaskDoc] + ): """ Relax a structure or molecule using ASE as a job. Parameters ---------- - mol_or_struct: .Molecule or .Structure - pymatgen molecule or structure + mol_or_struct: .Molecule or .Structure, or list thereof + pymatgen molecule(s) or structure(s) 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. Returns ------- - AseStructureTaskDoc or AseMoleculeTaskDoc + AseStructureTaskDoc or AseMoleculeTaskDoc, or list thereof """ - return AseTaskDoc.to_mol_or_struct_metadata_doc( - getattr(self.calculator, "name", type(self.calculator).__name__), - self.run_ase(mol_or_struct, prev_dir=prev_dir), - self.steps, - relax_kwargs=self.relax_kwargs, - optimizer_kwargs=self.optimizer_kwargs, - relax_cell=self.relax_cell, - relax_shape=self.relax_shape, - fix_symmetry=self.fix_symmetry, - symprec=self.symprec if self.fix_symmetry else None, - ionic_step_data=self.ionic_step_data, - store_trajectory=self.store_trajectory, - tags=self.tags, - ) + batch_mode = isinstance(mol_or_struct, list) + + results = [ + AseTaskDoc.to_mol_or_struct_metadata_doc( + getattr(self.calculator, "name", type(self.calculator).__name__), + self.run_ase(atoms, prev_dir=prev_dir), + self.steps, + relax_kwargs=self.relax_kwargs, + optimizer_kwargs=self.optimizer_kwargs, + relax_cell=self.relax_cell, + relax_shape=self.relax_shape, + fix_symmetry=self.fix_symmetry, + symprec=self.symprec if self.fix_symmetry else None, + ionic_step_data=self.ionic_step_data, + store_trajectory=self.store_trajectory, + tags=self.tags, + ) + for atoms in (mol_or_struct if batch_mode else [mol_or_struct]) + ] + return results if batch_mode else results[0] def run_ase( self, @@ -299,8 +331,7 @@ class EmtRelaxMaker(AseRelaxMaker): name: str = "EMT relaxation" - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> Calculator: """EMT calculator.""" from ase.calculators.emt import EMT @@ -320,8 +351,7 @@ class LennardJonesRelaxMaker(AseRelaxMaker): name: str = "Lennard-Jones 6-12 relaxation" - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> None: """Lennard-Jones calculator.""" from ase.calculators.lj import LennardJones @@ -378,8 +408,7 @@ class GFNxTBRelaxMaker(AseRelaxMaker): } ) - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> None: """GFN-xTB / TBLite calculator.""" try: from tblite.ase import TBLite diff --git a/src/atomate2/ase/md.py b/src/atomate2/ase/md.py index 670b5322e4..4a52c29e2d 100644 --- a/src/atomate2/ase/md.py +++ b/src/atomate2/ase/md.py @@ -8,7 +8,7 @@ import os import sys import time -from abc import ABC, abstractmethod +from abc import ABC from collections.abc import Sequence from dataclasses import dataclass, field from enum import Enum @@ -189,6 +189,7 @@ class AseMDMaker(AseMaker, ABC): def __post_init__(self) -> None: """Ensure that ensemble is an enum.""" + super().__post_init__() if isinstance(self.ensemble, str): self.ensemble = MDEnsemble(self.ensemble.split("MDEnsemble.")[-1]) @@ -444,12 +445,6 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: elapsed_time=t_f - t_i, ) - @property - @abstractmethod - def calculator(self) -> Calculator: - """ASE calculator, to be overwritten by user.""" - raise NotImplementedError - @dataclass class LennardJonesMDMaker(AseMDMaker): @@ -461,8 +456,7 @@ class LennardJonesMDMaker(AseMDMaker): name: str = "Lennard-Jones 6-12 MD" - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> Calculator: """Lennard-Jones calculator.""" from ase.calculators.lj import LennardJones @@ -495,8 +489,7 @@ class GFNxTBMDMaker(AseMDMaker): } ) - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> Calculator: """GFN-xTB / TBLite calculator.""" try: from tblite.ase import TBLite diff --git a/src/atomate2/ase/neb.py b/src/atomate2/ase/neb.py index 3cfd8304f6..7977f5bf77 100644 --- a/src/atomate2/ase/neb.py +++ b/src/atomate2/ase/neb.py @@ -257,8 +257,7 @@ class EmtNebFromImagesMaker(AseNebFromImagesMaker): name: str = "EMT NEB from images maker" - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> Calculator: """EMT calculator.""" from ase.calculators.emt import EMT diff --git a/src/atomate2/common/flows/phonons.py b/src/atomate2/common/flows/phonons.py index 3766c83edd..ae69cb6104 100644 --- a/src/atomate2/common/flows/phonons.py +++ b/src/atomate2/common/flows/phonons.py @@ -132,7 +132,7 @@ class BasePhononMaker(Maker, ABC): store_force_constants: bool if True, force constants will be stored socket: bool - If True, use the socket for the calculation + If True, use the socket/batch for the calculation """ name: str = "phonon" diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index 243381bff9..a83578e970 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -21,8 +21,11 @@ from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine from pymatgen.phonon.dos import PhononDos +from atomate2.ase.jobs import AseRelaxMaker from atomate2.common.schemas.phonons import ForceConstants, PhononBSDOSDoc, get_factor from atomate2.common.utils import get_supercell_matrix +from atomate2.forcefields.jobs import ForceFieldRelaxMaker +from atomate2.vasp.jobs.base import BaseVaspMaker if TYPE_CHECKING: from pathlib import Path @@ -30,9 +33,6 @@ from emmet.core.math import Matrix3D from atomate2.aims.jobs.base import BaseAimsMaker - from atomate2.forcefields.jobs import ForceFieldStaticMaker - from atomate2.vasp.jobs.base import BaseVaspMaker - logger = logging.getLogger(__name__) @@ -253,7 +253,10 @@ def run_phonon_displacements( displacements: list[Structure], structure: Structure, supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker | BaseAimsMaker = None, + phonon_maker: BaseVaspMaker + | AseRelaxMaker + | ForceFieldRelaxMaker + | BaseAimsMaker = None, prev_dir: str | Path = None, prev_dir_argname: str = None, socket: bool = False, @@ -272,14 +275,16 @@ def run_phonon_displacements( Fully optimized structure used for phonon computations. supercell_matrix: Matrix3D supercell matrix for meta data - phonon_maker : .BaseVaspMaker or .ForceFieldStaticMaker or .BaseAimsMaker - A maker to use to generate dispacement calculations + phonon_maker : .BaseVaspMaker, .AseRelaxMaker, + .ForceFieldRelaxMaker, or .BaseAimsMaker + A maker to use to generate dispacement calculations. + NB: this should be a static maker. prev_dir: str or Path The previous working directory prev_dir_argname: str argument name for the prev_dir variable socket: bool - If True use the socket-io interface to increase performance + If True use the socket-io (batch-mode) interface to increase performance """ phonon_jobs = [] outputs: dict[str, list] = { @@ -292,28 +297,39 @@ def run_phonon_displacements( if prev_dir is not None and prev_dir_argname is not None: phonon_job_kwargs[prev_dir_argname] = prev_dir + num_disp = len(displacements) if socket: + if isinstance(phonon_maker, BaseVaspMaker): + raise ValueError("VASP makers do not currently support socket/batch mode.") + phonon_job = phonon_maker.make(displacements, **phonon_job_kwargs) info = { "original_structure": structure, "supercell_matrix": supercell_matrix, "displaced_structures": displacements, } - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, dict_mod=True - ) + if not isinstance(phonon_maker, AseRelaxMaker | ForceFieldRelaxMaker): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) + phonon_jobs.append(phonon_job) - outputs["displacement_number"] = list(range(len(displacements))) - outputs["uuids"] = [phonon_job.output.uuid] * len(displacements) - outputs["dirs"] = [phonon_job.output.dir_name] * len(displacements) - outputs["forces"] = phonon_job.output.output.all_forces + outputs["displacement_number"] = list(range(num_disp)) + if isinstance(phonon_maker, AseRelaxMaker | ForceFieldRelaxMaker): + outputs["uuids"] = [phonon_job.output[0].uuid] * num_disp + outputs["dirs"] = [phonon_job.output[0].dir_name] * num_disp + outputs["forces"] = [ + phonon_job.output[idx].output.forces for idx in range(num_disp) + ] + else: + outputs["uuids"] = [phonon_job.output.uuid] * num_disp + outputs["dirs"] = [phonon_job.output.dir_name] * num_disp + outputs["forces"] = phonon_job.output.output.all_forces else: for idx, displacement in enumerate(displacements): - if prev_dir is not None: - phonon_job = phonon_maker.make(displacement, prev_dir=prev_dir) - else: - phonon_job = phonon_maker.make(displacement) - phonon_job.append_name(f" {idx + 1}/{len(displacements)}") + phonon_job = phonon_maker.make(displacement, prev_dir=prev_dir) + phonon_job.append_name(f" {idx + 1}/{num_disp}") # we will add some meta data info = { @@ -323,10 +339,11 @@ def run_phonon_displacements( "displaced_structure": displacement, } with contextlib.suppress(Exception): - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, - dict_mod=True, - ) + if not isinstance(phonon_maker, AseRelaxMaker | ForceFieldRelaxMaker): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) phonon_jobs.append(phonon_job) outputs["displacement_number"].append(idx) outputs["uuids"].append(phonon_job.output.uuid) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index ba181b4d19..6ac7295b39 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -122,21 +122,29 @@ class ForceFieldRelaxMaker(ForceFieldMixin, AseRelaxMaker): @forcefield_job def make( - self, structure: Molecule | Structure, prev_dir: str | Path | None = None - ) -> ForceFieldTaskDocument | ForceFieldMoleculeTaskDocument: + self, + structure: Molecule | Structure | list[Molecule | Structure], + prev_dir: str | Path | None = None, + ) -> ( + ForceFieldTaskDocument + | ForceFieldMoleculeTaskDocument + | list[ForceFieldTaskDocument | ForceFieldMoleculeTaskDocument] + ): """ Perform a relaxation of a structure using a force field. Parameters ---------- - structure: .Structure or Molecule - pymatgen structure or molecule. + structure: .Molecule or .Structure, or a list thereof + pymatgen molecule(s) or structure(s) 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. - """ - ase_result = self._run_ase_safe(structure, prev_dir=prev_dir) + Returns + ------- + ForceFieldTaskDocument, ForceFieldMoleculeTaskDocument, or a list thereof + """ if len(self.task_document_kwargs) > 0: warnings.warn( "`task_document_kwargs` is now deprecated, please use the top-level " @@ -145,22 +153,33 @@ def make( stacklevel=1, ) - return ForceFieldTaskDocument.from_ase_compatible_result( - self.ase_calculator_name, - ase_result, - self.steps, - calculator_meta=self.calculator_meta, - relax_kwargs=self.relax_kwargs, - optimizer_kwargs=self.optimizer_kwargs, - relax_cell=self.relax_cell, - relax_shape=self.relax_shape, - fix_symmetry=self.fix_symmetry, - symprec=self.symprec if self.fix_symmetry else None, - ionic_step_data=self.ionic_step_data, - store_trajectory=self.store_trajectory, - tags=self.tags, - **self.task_document_kwargs, - ) + batch_mode = isinstance(structure, list) + + ase_results = [ + self._run_ase_safe(atoms, prev_dir=prev_dir) + for atoms in (structure if batch_mode else [structure]) + ] + + task_docs = [ + ForceFieldTaskDocument.from_ase_compatible_result( + self.ase_calculator_name, + ase_result, + self.steps, + calculator_meta=self.calculator_meta, + relax_kwargs=self.relax_kwargs, + optimizer_kwargs=self.optimizer_kwargs, + relax_cell=self.relax_cell, + relax_shape=self.relax_shape, + fix_symmetry=self.fix_symmetry, + symprec=self.symprec if self.fix_symmetry else None, + ionic_step_data=self.ionic_step_data, + store_trajectory=self.store_trajectory, + tags=self.tags, + **self.task_document_kwargs, + ) + for ase_result in ase_results + ] + return task_docs if batch_mode else task_docs[0] @dataclass diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index fc1f0d6250..8d57091ccd 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -210,8 +210,7 @@ def _run_ase_safe(self, *args, **kwargs) -> AseResult: with revert_default_dtype(): return self.run_ase(*args, **kwargs) - @property - def calculator(self) -> Calculator: + def _get_calculator(self) -> Calculator: """ASE calculator, can be overwritten by user.""" return ase_calculator( self.calculator_meta, diff --git a/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index 0a8870aff4..5f7fe893d1 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -30,8 +30,7 @@ class EMTStaticMaker(AseMaker): name: str = "EMT static maker" - @property - def calculator(self): + def _get_calculator(self): return EMT() @@ -39,8 +38,7 @@ def calculator(self): class EMTRelaxMaker(AseRelaxMaker): name: str = "EMT relax maker" - @property - def calculator(self): + def _get_calculator(self): return EMT() @@ -91,6 +89,27 @@ def test_lennard_jones_relax_maker(lj_fcc_ne_pars, fcc_ne_structure): ) +def test_lennard_jones_batch_relax_maker( + lj_fcc_ne_pars, fcc_ne_structure, memory_jobstore +): + job = LennardJonesRelaxMaker( + calculator_kwargs=lj_fcc_ne_pars, relax_kwargs={"fmax": 0.001} + ).make([fcc_ne_structure, fcc_ne_structure]) + + response = run_locally(job, store=memory_jobstore) + + output = response[job.uuid][1].output + + assert [calc.output.structure.volume for calc in output] == pytest.approx( + [22.304245, 22.304245] + ) + assert [calc.output.energy for calc in output] == pytest.approx( + [-0.018494767, -0.018494767] + ) + assert all(isinstance(calc, AseStructureTaskDoc) for calc in output) + assert fcc_ne_structure.matches(output[0].output.structure) + + def test_lennard_jones_static_maker(lj_fcc_ne_pars, fcc_ne_structure): job = LennardJonesStaticMaker(calculator_kwargs=lj_fcc_ne_pars).make( fcc_ne_structure diff --git a/tests/ase/test_neb.py b/tests/ase/test_neb.py index ab0806d40e..bec2d43376 100644 --- a/tests/ase/test_neb.py +++ b/tests/ase/test_neb.py @@ -47,8 +47,7 @@ class EmtNebFromEndpointsMaker(AseNebFromEndpointsMaker): default_factory=EmtRelaxMaker, ) - @property - def calculator(self): + def _get_calculator(self): return EMT(**self.calculator_kwargs) diff --git a/tests/forcefields/flows/test_phonon.py b/tests/forcefields/flows/test_phonon.py index f6b85dc495..3909068b77 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -1,5 +1,6 @@ import os from importlib.util import find_spec +from itertools import product from pathlib import Path from tempfile import TemporaryDirectory @@ -109,10 +110,12 @@ def test_phonon_maker_initialization_with_all_mlff( ) from exc -@pytest.mark.skipif(not mlff_is_installed("CHGNet"), reason="matgl is not installed") -@pytest.mark.parametrize("from_name", [False, True]) +@pytest.mark.skipif( + not mlff_is_installed("CHGNet"), reason="matgl/chgnet is not installed" +) +@pytest.mark.parametrize("from_name, socket", list(product(*[[True, False]] * 2))) def test_phonon_wf_force_field( - clean_dir, si_structure: Structure, tmp_path: Path, from_name: bool + clean_dir, si_structure: Structure, tmp_path: Path, from_name: bool, socket: bool ): # TODO brittle due to inability to adjust dtypes in CHGNetRelaxMaker @@ -148,6 +151,7 @@ def test_phonon_wf_force_field( "filename_bs": (filename_bs := f"{tmp_path}/phonon_bs_test.png"), "filename_dos": (filename_dos := f"{tmp_path}/phonon_dos_test.pdf"), }, + socket=socket, ) if from_name: diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 7cbe644e60..9efb4abb57 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -171,6 +171,36 @@ def test_chgnet_relax_maker( assert Path(responses[job.uuid][1].output.dir_name).exists() +@pytest.mark.skipif( + not mlff_is_installed("CHGNet"), + reason="Required packages (chgnet or matgl/dgl) are not installed", +) +def test_chgnet_batch_static_maker(si_structure: Structure, memory_jobstore): + # translate one atom to ensure a small number of relaxation steps are taken + si_structure2 = si_structure.copy() + si_structure.translate_sites(0, [0, 0, 0.1]) + si_structure2.translate_sites(0, [0.1, 0, 0.1]) + + # generate job + job = ForceFieldStaticMaker( + force_field_name="CHGNet", + ).make([si_structure, si_structure2]) + + # run the flow or job and ensure that it finished running successfully + responses = run_locally(job, ensure_success=True, store=memory_jobstore) + # validate job outputs + output = responses[job.uuid][1].output + assert all(isinstance(calc, ForceFieldTaskDocument) for calc in output) + + assert len(output) == 2 + assert [calc.output.energy for calc in output] == approx( + [-9.96250, -9.4781], rel=1e-2 + ) + + # check the force_field_task_doc attributes + assert all(Path(calc.dir_name).exists() for calc in output) + + @pytest.mark.xfail( reason="M3GNet tests not working consistently in CI vs local", strict=False,