batch mode extension in phonons + tests adapted from #1196

Author: esoteric-ephemera

src/atomate2/ase/jobs.py

@@ -101,24 +101,37 @@ 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]
+    ):
         """
         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,
@@ -229,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,

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)

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

tests/ase/test_jobs.py

@@ -89,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

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: