Alternative ASE/MLFF batch mode

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

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

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
 

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"

src/atomate2/common/jobs/phonons.py

@@ -21,18 +21,18 @@
 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
 
     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)