Faster Phonon Forcefield Workflow By Introducing Batch ASEMakers (especially relevant for remote execution)

src/atomate2/ase/jobs.py

@@ -8,11 +8,9 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
-from ase.io import Trajectory as AseTrajectory
 from emmet.core.vasp.calculation import StoreTrajectoryOption
 from jobflow import Maker, job
 from pymatgen.core import Molecule, Structure
-from pymatgen.core.trajectory import Trajectory as PmgTrajectory
 from pymatgen.io.ase import AseAtomsAdaptor
 
 from atomate2.ase.schemas import AseResult, AseTaskDoc
@@ -27,7 +25,9 @@
 
     from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc
 
-_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory]
+_ASE_DATA_OBJECTS = [
+    "output",  # will put everything in data store
+]
 
 
 @dataclass
@@ -119,7 +119,7 @@ def run_ase(
         self,
         mol_or_struct: Structure | Molecule,
         prev_dir: str | Path | None = None,
-    ) -> AseResult:
+    ) -> AseResult | list[AseResult]:
         """
         Run ASE, can be re-implemented in subclasses.
 
@@ -237,16 +237,16 @@ def make(
 
     def run_ase(
         self,
-        mol_or_struct: Structure | Molecule,
+        mol_or_struct: Structure | Molecule | list[Molecule] | list[Structure],
         prev_dir: str | Path | None = None,
-    ) -> AseResult:
+    ) -> AseResult | list[AseResult]:
         """
-        Relax a structure or molecule using ASE, not as a job.
+        Relax a structure, molecule or a batch of those using ASE, not as a job.
 
         Parameters
         ----------
-        mol_or_struct: .Molecule or .Structure
-            pymatgen molecule or structure
+        mol_or_struct: .Molecule or .Structure or list[.Molecule] or list[.Structure]
+            pymatgen molecule or structure or lists of those
         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.