diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 8513b3be45..0dae949244 100644 --- a/src/atomate2/ase/jobs.py +++ b/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. diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 35d5319f4e..3a5e0efbe6 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -11,7 +11,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any +from typing import Any, Optional from ase.stress import voigt_6_to_full_3x3_stress from ase.units import GPa @@ -93,18 +93,29 @@ class AseObject(ValueEnum): class AseBaseModel(BaseModel): """Base document class for ASE input and output.""" - mol_or_struct: Structure | Molecule | None = Field( - None, description="The molecule or structure at this step." + mol_or_struct: Structure | Molecule | list[Structure] | list[Molecule] | None = ( + Field(None, description="The molecule or structure at this step.") + ) + structure: Structure | list[Structure] | None = Field( + None, description="The structure at this step." + ) + molecule: Molecule | list[Molecule] | None = Field( + None, description="The molecule at this step." ) - structure: Structure | None = Field(None, description="The structure at this step.") - molecule: Molecule | None = Field(None, description="The molecule at this step.") def model_post_init(self, _context: Any) -> None: """Establish alias to structure and molecule fields.""" - if self.structure is None and isinstance(self.mol_or_struct, Structure): - self.structure = self.mol_or_struct - elif self.molecule is None and isinstance(self.mol_or_struct, Molecule): - self.molecule = self.mol_or_struct + if self.mol_or_struct is not None: + if self.structure is None and ( + isinstance(self.mol_or_struct, Structure) + or isinstance(self.mol_or_struct[0], Structure) + ): + self.structure = self.mol_or_struct + elif self.molecule is None and ( + isinstance(self.mol_or_struct, Molecule) + or isinstance(self.mol_or_struct[0], Molecule) + ): + self.molecule = self.mol_or_struct class IonicStep(AseBaseModel): @@ -121,15 +132,17 @@ class IonicStep(AseBaseModel): class OutputDoc(AseBaseModel): """The outputs of this job.""" - energy: float | None = Field(None, description="Total energy in units of eV.") + energy: float | list[float] | None = Field( + None, description="Total energy in units of eV." + ) - energy_per_atom: float | None = Field( + energy_per_atom: float | list[float] | None = Field( None, description="Energy per atom of the final molecule or structure " "in units of eV/atom.", ) - forces: list[Vector3D] | None = Field( + forces: list[Vector3D] | list[list[Vector3D]] | None = Field( None, description=( "The force on each atom in units of eV/A for the final molecule " @@ -139,24 +152,32 @@ class OutputDoc(AseBaseModel): # NOTE: units for stresses were converted to kbar (* -10 from standard output) # to comply with MP convention - stress: Matrix3D | None = Field( + stress: Matrix3D | list[Matrix3D] | None = Field( None, description="The stress on the cell in units of kbar (in Voigt notation)." ) # NOTE: the ionic_steps can also be a dict when these are in blob storage and # retrieved as objects. - ionic_steps: list[IonicStep] | dict | None = Field( - None, description="Step-by-step trajectory of the relaxation." + ionic_steps: list[IonicStep] | list[list[IonicStep]] | dict | list[dict] | None = ( + Field(None, description="Step-by-step trajectory of the relaxation.") ) - elapsed_time: float | None = Field( + elapsed_time: float | list[float] | None = Field( None, description="The time taken to run the ASE calculation in seconds." ) - n_steps: int | None = Field( + n_steps: int | list[int] | None = Field( None, description="total number of steps needed in the relaxation." ) + all_forces: list[list[Vector3D]] | None = Field( + None, + description=( + "The force on each atom in units of eV/A for the final molecules " + "or structures. Useful for batch calculations." + ), + ) + class InputDoc(AseBaseModel): """The inputs used to run this job.""" @@ -204,25 +225,26 @@ class AseStructureTaskDoc(StructureMetadata): description="name of the ASE calculator used in the calculation.", ) - dir_name: str | None = Field( + dir_name: str | list[str] | None = Field( None, description="Directory where the ASE calculations are performed." ) included_objects: list[AseObject] | None = Field( None, description="list of ASE objects included with this task document" ) - objects: dict[AseObject, Any] | None = Field( + # TODO: check if it needs to be a list + objects: dict[AseObject, Any] | list[dict[AseObject, Any]] | None = Field( None, description="ASE objects associated with this task" ) - is_force_converged: bool | None = Field( + is_force_converged: bool | list[bool] | None = Field( None, description=( "Whether the calculation is converged with respect to interatomic forces." ), ) - energy_downhill: bool | None = Field( + energy_downhill: bool | list[bool] | None = Field( None, description=( "Whether the final trajectory frame has lower total " @@ -257,7 +279,9 @@ def from_ase_task_doc( class AseMoleculeTaskDoc(MoleculeMetadata): """Document containing information on molecule manipulation using ASE.""" - molecule: Molecule = Field(None, description="Final output molecule from the task") + molecule: Molecule | list[Molecule] = Field( + None, description="Final output molecule from the task" + ) input: InputDoc = Field( None, description="The input information used to run this job." @@ -270,7 +294,7 @@ class AseMoleculeTaskDoc(MoleculeMetadata): description="name of the ASE calculator used in the calculation.", ) - dir_name: str | None = Field( + dir_name: str | list[str] | None = Field( None, description="Directory where the ASE calculations are performed." ) @@ -281,14 +305,14 @@ class AseMoleculeTaskDoc(MoleculeMetadata): None, description="ASE objects associated with this task" ) - is_force_converged: bool | None = Field( + is_force_converged: bool | list[bool] | None = Field( None, description=( "Whether the calculation is converged with respect to interatomic forces." ), ) - energy_downhill: bool | None = Field( + energy_downhill: bool | list[bool] | None = Field( None, description=( "Whether the total energy in the final frame " @@ -313,25 +337,26 @@ class AseTaskDoc(AseBaseModel): description="name of the ASE calculator used for this job.", ) - dir_name: str | None = Field( + dir_name: str | list[str] | None = Field( None, description="Directory where the ASE calculations are performed." ) included_objects: list[AseObject] | None = Field( None, description="list of ASE objects included with this task document" ) - objects: dict[AseObject, Any] | None = Field( + # maybe list + objects: dict[AseObject, Any] | list[dict[AseObject, Any]] | None = Field( None, description="ASE objects associated with this task" ) - is_force_converged: bool | None = Field( + is_force_converged: bool | list[bool] | None = Field( None, description=( "Whether the calculation is converged with respect to interatomic forces." ), ) - energy_downhill: bool | None = Field( + energy_downhill: bool | list[bool] | None = Field( None, description=( "Whether the total energy in the final frame " @@ -345,7 +370,7 @@ class AseTaskDoc(AseBaseModel): def from_ase_compatible_result( cls, ase_calculator_name: str, - result: AseResult, + result: AseResult | list[AseResult], steps: int, relax_kwargs: dict = None, optimizer_kwargs: dict = None, @@ -369,8 +394,8 @@ def from_ase_compatible_result( ---------- ase_calculator_name : str Name of the ASE calculator used. - result : AseResult - The output results from the task. + result : AseResult|list[AseResult] + The output results from the task. Can be a list for batch jobs. steps : int Maximum number of ionic steps allowed during relaxation. relax_cell : bool = True @@ -393,29 +418,164 @@ def from_ase_compatible_result( task_document_kwargs : dict Additional keyword args passed to :obj:`.AseTaskDoc()`. """ - trajectory = result.trajectory - - n_steps = None - input_mol_or_struct = None - if trajectory: - n_steps = len(trajectory) - - # NOTE: convert stress units from eV/A³ to kBar (* -1 from standard output) - # and to 3x3 matrix to comply with MP convention - if n_steps: - for idx in range(n_steps): - if trajectory.frame_properties[idx].get("stress") is not None: - trajectory.frame_properties[idx]["stress"] = ( - voigt_6_to_full_3x3_stress( - [ - val * -10 / GPa - for val in trajectory.frame_properties[idx]["stress"] - ] + is_list = isinstance(result, list) + results: list[AseResult] = result if isinstance(result, list) else [result] + + output_mol_or_struct: list[Molecule] | list[Structure] | None = [] + input_mol_or_struct: list[Molecule] | list[Structure] | None = [] + final_energy: list[float] = [] + final_forces: list[Vector3D] | list[list[Vector3D]] = [] + final_stress: list[Matrix3D] = [] + ionic_steps: list[Optional[list[IonicStep]]] = [] + n_steps = [] + objects = [] + + for result_item in results: + trajectory = result_item.trajectory + + # TODO: fix this + n_steps_here = None + input_mol_or_struct_here = None + if trajectory: + n_steps_here = len(trajectory) + + # NOTE: convert stress units from eV/A³ to kBar (* -1 from standard output) + # and to 3x3 matrix to comply with MP convention + if n_steps_here: + for idx in range(n_steps_here): + if trajectory.frame_properties[idx].get("stress") is not None: + trajectory.frame_properties[idx]["stress"] = ( + voigt_6_to_full_3x3_stress( + [ + val * -10 / GPa + for val in trajectory.frame_properties[idx][ + "stress" + ] + ] + ) ) - ) - input_mol_or_struct = trajectory[0] + input_mol_or_struct_here = trajectory[0] + + input_mol_or_struct.append(input_mol_or_struct_here) + + # Workaround for cases where the ASE optimizer does not correctly limit the + # number of steps for static calculations. + if (steps is not None) and steps <= 1: + steps = 1 + n_steps_here = 1 + + if isinstance(input_mol_or_struct_here, Structure): + traj_method = "from_structures" + elif isinstance(input_mol_or_struct_here, Molecule): + traj_method = "from_molecules" + trajectory = getattr(PmgTrajectory, traj_method)( + [input_mol_or_struct_here], + frame_properties=[trajectory.frame_properties[0]], + constant_lattice=False, + ) + output_mol_or_struct.append(input_mol_or_struct_here) + else: + output_mol_or_struct.append(result_item.final_mol_or_struct) + + n_steps.append(n_steps_here) + + if trajectory is None: + final_energy.append(result_item.final_energy) + final_forces.append(None) + final_stress.append(None) + ionic_steps.append(None) + + else: + final_energy.append(trajectory.frame_properties[-1]["energy"]) + final_forces.append(trajectory.frame_properties[-1]["forces"]) + final_stress.append(trajectory.frame_properties[-1].get("stress")) + + ionic_steps_individual: list[IonicStep] = [] + if ionic_step_data is not None and len(ionic_step_data) > 0: + for idx in range(n_steps_here): + _ionic_step_data = { + key: ( + trajectory.frame_properties[idx].get(key) + if key in ionic_step_data + else None + ) + for key in ("energy", "forces", "stress") + } + + current_mol_or_struct = ( + trajectory[idx] + if any( + v in ionic_step_data + for v in ("mol_or_struct", "structure", "molecule") + ) + else None + ) + + # include "magmoms" in `ionic_step` + # if the trajectory has "magmoms" + if "magmoms" in trajectory.frame_properties[idx]: + _ionic_step_data.update( + { + "magmoms": ( + trajectory.frame_properties[idx]["magmoms"] + if "magmoms" in ionic_step_data + else None + ) + } + ) + + ionic_step = IonicStep( + mol_or_struct=current_mol_or_struct, + **_ionic_step_data, + ) + + ionic_steps_individual.append(ionic_step) + ionic_steps.append(ionic_steps_individual) + + objects_structure: dict[AseObject, Any] = {} + if store_trajectory != StoreTrajectoryOption.NO: + # For VASP calculations, the PARTIAL trajectory option removes + # electronic step info. There is no equivalent for classical + # forcefields, so we just save the same info for FULL and + # PARTIAL options. + objects_structure[AseObject.TRAJECTORY] = trajectory # type: ignore[index] + objects.append(objects_structure) + if not is_list: + input_doc = InputDoc( + mol_or_struct=input_mol_or_struct[0], + relax_cell=relax_cell, + fix_symmetry=fix_symmetry, + symprec=symprec, + steps=steps, + relax_kwargs=relax_kwargs, + optimizer_kwargs=optimizer_kwargs, + ) + output_doc = OutputDoc( + mol_or_struct=output_mol_or_struct[0], + energy=final_energy[0], + energy_per_atom=final_energy[0] / len(output_mol_or_struct[0]), + forces=final_forces[0], + stress=final_stress[0], + ionic_steps=ionic_steps[0], + elapsed_time=results[0].elapsed_time, + n_steps=n_steps[0], + all_forces=final_forces if final_forces[0] is not None else None, + ) + return cls( + mol_or_struct=output_mol_or_struct[0], + input=input_doc, + output=output_doc, + ase_calculator_name=ase_calculator_name, + included_objects=list(objects[0].keys()), + objects=objects[0], + is_force_converged=results[0].is_force_converged, + energy_downhill=results[0].energy_downhill, + dir_name=results[0].dir_name, + tags=tags, + **task_document_kwargs, + ) input_doc = InputDoc( mol_or_struct=input_mol_or_struct, relax_cell=relax_cell, @@ -425,107 +585,33 @@ def from_ase_compatible_result( relax_kwargs=relax_kwargs, optimizer_kwargs=optimizer_kwargs, ) - - # Workaround for cases where the ASE optimizer does not correctly limit the - # number of steps for static calculations. - if (steps is not None) and steps <= 1: - steps = 1 - n_steps = 1 - - if isinstance(input_mol_or_struct, Structure): - traj_method = "from_structures" - elif isinstance(input_mol_or_struct, Molecule): - traj_method = "from_molecules" - - trajectory = getattr(PmgTrajectory, traj_method)( - [input_mol_or_struct], - frame_properties=[trajectory.frame_properties[0]], - constant_lattice=False, - ) - output_mol_or_struct = input_mol_or_struct - else: - output_mol_or_struct = result.final_mol_or_struct - - if trajectory is None: - final_energy = result.final_energy - final_forces = None - final_stress = None - ionic_steps = None - - else: - final_energy = trajectory.frame_properties[-1]["energy"] - final_forces = trajectory.frame_properties[-1]["forces"] - final_stress = trajectory.frame_properties[-1].get("stress") - - ionic_steps = [] - if ionic_step_data is not None and len(ionic_step_data) > 0: - for idx in range(n_steps): - _ionic_step_data = { - key: ( - trajectory.frame_properties[idx].get(key) - if key in ionic_step_data - else None - ) - for key in ("energy", "forces", "stress") - } - - current_mol_or_struct = ( - trajectory[idx] - if any( - v in ionic_step_data - for v in ("mol_or_struct", "structure", "molecule") - ) - else None - ) - - # include "magmoms" in `ionic_step` if the trajectory has "magmoms" - if "magmoms" in trajectory.frame_properties[idx]: - _ionic_step_data.update( - { - "magmoms": ( - trajectory.frame_properties[idx]["magmoms"] - if "magmoms" in ionic_step_data - else None - ) - } - ) - - ionic_step = IonicStep( - mol_or_struct=current_mol_or_struct, - **_ionic_step_data, - ) - - ionic_steps.append(ionic_step) - - objects: dict[AseObject, Any] = {} - if store_trajectory != StoreTrajectoryOption.NO: - # For VASP calculations, the PARTIAL trajectory option removes - # electronic step info. There is no equivalent for classical - # forcefields, so we just save the same info for FULL and - # PARTIAL options. - objects[AseObject.TRAJECTORY] = trajectory # type: ignore[index] - output_doc = OutputDoc( mol_or_struct=output_mol_or_struct, energy=final_energy, - energy_per_atom=final_energy / len(output_mol_or_struct), + energy_per_atom=[ + final_energy_item / len(output_mol_or_struct_item) + for final_energy_item, output_mol_or_struct_item in zip( + final_energy, output_mol_or_struct, strict=False + ) + ], forces=final_forces, stress=final_stress, ionic_steps=ionic_steps, - elapsed_time=result.elapsed_time, + elapsed_time=[result.elapsed_time for result in results], n_steps=n_steps, + all_forces=final_forces, ) return cls( - mol_or_struct=output_mol_or_struct, + mol_or_struct=output_mol_or_struct[-1], # last structure by default input=input_doc, output=output_doc, ase_calculator_name=ase_calculator_name, - included_objects=list(objects.keys()), + included_objects=list(objects[0].keys()), objects=objects, - is_force_converged=result.is_force_converged, - energy_downhill=result.energy_downhill, - dir_name=result.dir_name, + is_force_converged=[result.is_force_converged for result in results], + energy_downhill=[result.energy_downhill for result in results], + dir_name=[result.dir_name for result in results], tags=tags, **task_document_kwargs, ) @@ -534,7 +620,7 @@ def from_ase_compatible_result( def to_mol_or_struct_metadata_doc( cls, ase_calculator_name: str, - result: AseResult, + result: AseResult | list[AseResult], steps: int | None = None, **task_document_kwargs, ) -> AseStructureTaskDoc | AseMoleculeTaskDoc: @@ -560,14 +646,20 @@ def to_mol_or_struct_metadata_doc( ase_calculator_name, result, steps, **task_document_kwargs ) kwargs = {k: getattr(task_doc, k, None) for k in _task_doc_translation_keys} - if isinstance(task_doc.mol_or_struct, Structure): + + if isinstance(task_doc.mol_or_struct, Structure) or isinstance( + task_doc.mol_or_struct[0], Structure + ): meta_class = AseStructureTaskDoc k = "structure" if relax_cell := getattr(task_doc, "relax_cell", None): kwargs.update({"relax_cell": relax_cell}) - elif isinstance(task_doc.mol_or_struct, Molecule): + elif isinstance(task_doc.mol_or_struct, Molecule) or isinstance( + task_doc.mol_or_struct[0], Molecule + ): meta_class = AseMoleculeTaskDoc k = "molecule" + kwargs.update({k: task_doc.mol_or_struct, f"meta_{k}": task_doc.mol_or_struct}) return getattr(meta_class, f"from_{k}")(**kwargs) diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index d9106ef8ee..6afc949ea6 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -325,7 +325,12 @@ def __init__( def relax( self, - atoms: Atoms | Structure | Molecule, + atoms: Atoms + | Structure + | Molecule + | list[Atoms] + | list[Molecule] + | list[Structure], fmax: float = 0.1, steps: int = 500, traj_file: str = None, @@ -334,13 +339,14 @@ def relax( verbose: bool = False, cell_filter: Filter = FrechetCellFilter, **kwargs, - ) -> AseResult: + ) -> AseResult | list[AseResult]: """ - Relax the molecule or structure. + Relax the molecule or structure or a list of those. Parameters ---------- - atoms : ASE Atoms, pymatgen Structure, or pymatgen Molecule + atoms : ASE Atoms, pymatgen .Structure, or pymatgen .Molecule + or lists of those. The atoms for relaxation. fmax : float Total force tolerance for relaxation convergence. @@ -360,56 +366,69 @@ def relax( Returns ------- dict including optimized structure and the trajectory + or a list of those dicts """ - is_mol = isinstance(atoms, Molecule) or ( - isinstance(atoms, Atoms) and all(not pbc for pbc in atoms.pbc) - ) + is_list = isinstance(atoms[0], (Atoms | Molecule | Structure)) - if isinstance(atoms, Structure | Molecule): - atoms = self.ase_adaptor.get_atoms(atoms) - - input_atoms = atoms.copy() - if self.fix_symmetry: - atoms.set_constraint(FixSymmetry(atoms, symprec=self.symprec)) - atoms.calc = self.calculator - with contextlib.redirect_stdout(sys.stdout if verbose else io.StringIO()): - obs = TrajectoryObserver(atoms) - if self.relax_cell and (not is_mol): - atoms = cell_filter(atoms) - optimizer = self.opt_class(atoms, **kwargs) - optimizer.attach(obs, interval=interval) - t_i = time.perf_counter() - optimizer.run(fmax=fmax, steps=steps) - t_f = time.perf_counter() - obs() - if traj_file is not None: - obs.save(traj_file) - if isinstance(atoms, cell_filter): - atoms = atoms.atoms - - struct = self.ase_adaptor.get_structure( - atoms, cls=Molecule if is_mol else Structure - ) - traj = obs.to_pymatgen_trajectory(None) - is_force_conv = all( - np.linalg.norm(traj.frame_properties[-1]["forces"][idx]) < abs(fmax) - for idx in range(len(struct)) - ) + list_atoms: list[Atoms] = [atoms] if not is_list else atoms - if final_atoms_object_file is not None: - if steps <= 1: - write_atoms = input_atoms - write_atoms.calc = self.calculator - else: - write_atoms = atoms - write(final_atoms_object_file, write_atoms, format="extxyz", append=True) - - return AseResult( - final_mol_or_struct=struct, - trajectory=traj, - is_force_converged=is_force_conv, - energy_downhill=traj.frame_properties[-1]["energy"] - < traj.frame_properties[0]["energy"], - dir_name=os.getcwd(), - elapsed_time=t_f - t_i, - ) + list_ase_results = [] + for atoms_item_start in list_atoms: + atoms_item = atoms_item_start.copy() + is_mol = isinstance(atoms_item, Molecule) or ( + isinstance(atoms_item, Atoms) and all(not pbc for pbc in atoms_item.pbc) + ) + if isinstance(atoms_item, Structure | Molecule): + atoms_item = self.ase_adaptor.get_atoms(atoms_item) + atoms_item_start_0 = atoms_item.copy() + if self.fix_symmetry: + atoms_item.set_constraint(FixSymmetry(atoms_item, symprec=self.symprec)) + atoms_item.calc = self.calculator + with contextlib.redirect_stdout(sys.stdout if verbose else io.StringIO()): + obs = TrajectoryObserver(atoms_item) + if self.relax_cell and (not is_mol): + atoms_item = cell_filter(atoms_item) + optimizer = self.opt_class(atoms_item, **kwargs) + optimizer.attach(obs, interval=interval) + t_i = time.perf_counter() + optimizer.run(fmax=fmax, steps=steps) + t_f = time.perf_counter() + obs() + if traj_file is not None: + obs.save(traj_file) + if final_atoms_object_file is not None: + if steps <= 1: + write_atoms = atoms_item_start_0 + write_atoms.calc = self.calculator + else: + write_atoms = atoms_item + if isinstance(write_atoms, cell_filter): + write_atoms = write_atoms.atoms + + write( + final_atoms_object_file, write_atoms, format="extxyz", append=True + ) + + if isinstance(atoms_item, cell_filter): + atoms_item = atoms_item.atoms + + struct = self.ase_adaptor.get_structure( + atoms_item, cls=Molecule if is_mol else Structure + ) + traj = obs.to_pymatgen_trajectory(None) + is_force_conv = all( + np.linalg.norm(traj.frame_properties[-1]["forces"][idx]) < abs(fmax) + for idx in range(len(struct)) + ) + list_ase_results.append( + AseResult( + final_mol_or_struct=struct, + trajectory=traj, + is_force_converged=is_force_conv, + energy_downhill=traj.frame_properties[-1]["energy"] + < traj.frame_properties[0]["energy"], + dir_name=os.getcwd(), + elapsed_time=t_f - t_i, + ) + ) + return list_ase_results[0] if not is_list else list_ase_results diff --git a/src/atomate2/common/flows/phonons.py b/src/atomate2/common/flows/phonons.py index 9e6a21c90b..9a19eb6ee2 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 mode for the calculation """ name: str = "phonon" @@ -166,6 +166,7 @@ class BasePhononMaker(Maker, ABC): code: str = None store_force_constants: bool = True socket: bool = False + chunk_size = 2 def make( self, @@ -305,7 +306,6 @@ def make( jobs.append(compute_total_energy_job) total_dft_energy = compute_total_energy_job.output - # get a phonon object from phonopy displacements = generate_phonon_displacements( structure=structure, supercell_matrix=supercell_matrix, @@ -318,7 +318,6 @@ def make( ) jobs.append(displacements) - # perform the phonon displacement calculations displacement_calcs = run_phonon_displacements( displacements=displacements.output, structure=structure, @@ -329,7 +328,6 @@ def make( prev_dir=prev_dir, ) jobs.append(displacement_calcs) - # Computation of BORN charges born_run_job_dir = None born_run_uuid = None diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index 41b7eafb10..900a362968 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -15,6 +15,7 @@ from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine from pymatgen.phonon.dos import PhononDos +from atomate2.ase.jobs import AseMaker from atomate2.common.schemas.phonons import ForceConstants, PhononBSDOSDoc, get_factor from atomate2.common.utils import get_supercell_matrix @@ -174,7 +175,14 @@ def generate_phonon_displacements( @job( output_schema=PhononBSDOSDoc, - data=[PhononDos, PhononBandStructureSymmLine, ForceConstants], + data=[ + PhononDos, + PhononBandStructureSymmLine, + ForceConstants, + Structure, + "jobdirs", + "uuids", + ], ) def generate_frequencies_eigenvectors( structure: Structure, @@ -242,7 +250,7 @@ def generate_frequencies_eigenvectors( ) -@job(data=["forces", "displaced_structures"]) +@job(data=["forces", "displaced_structures", "uuids", "dirs", "displacement_number"]) def run_phonon_displacements( displacements: list[Structure], structure: Structure, @@ -293,9 +301,11 @@ def run_phonon_displacements( "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 issubclass(phonon_maker.__class__, AseMaker): + 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) diff --git a/src/atomate2/common/schemas/phonons.py b/src/atomate2/common/schemas/phonons.py index 99e95c4c10..ee977d7b41 100644 --- a/src/atomate2/common/schemas/phonons.py +++ b/src/atomate2/common/schemas/phonons.py @@ -527,6 +527,13 @@ def from_forces_born( volume_per_formula_unit = structure.volume / formula_units + if displacement_data["dirs"] and isinstance(displacement_data["dirs"][0], list): + displacement_data["dirs"] = [ + dir_item + for dir_list in displacement_data["dirs"] + for dir_item in dir_list + ] + doc = cls.from_structure( structure=structure, meta_structure=structure, diff --git a/src/atomate2/forcefields/flows/phonons.py b/src/atomate2/forcefields/flows/phonons.py index 75a79c6f71..ae53319ff3 100644 --- a/src/atomate2/forcefields/flows/phonons.py +++ b/src/atomate2/forcefields/flows/phonons.py @@ -111,7 +111,7 @@ class PhononMaker(BasePhononMaker): 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 mode for the calculation """ name: str = "phonon" @@ -140,6 +140,7 @@ class PhononMaker(BasePhononMaker): store_force_constants: bool = True code: str = "forcefields" born_maker: ForceFieldStaticMaker | None = None + socket = True @property def prev_calc_dir_argname(self) -> None: diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 63bdf23b45..2abf132af6 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -7,11 +7,10 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING -from ase.io import Trajectory as AseTrajectory from ase.units import GPa as _GPa_to_eV_per_A3 from jobflow import job from monty.dev import deprecated -from pymatgen.core.trajectory import Trajectory as PmgTrajectory +from pymatgen.core.structure import Structure from atomate2.ase.jobs import AseRelaxMaker from atomate2.forcefields import MLFF, _get_formatted_ff_name @@ -27,7 +26,10 @@ logger = logging.getLogger(__name__) -_FORCEFIELD_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, "ionic_steps"] +_FORCEFIELD_DATA_OBJECTS = [ + "output", # will put everything in the data store +] + _DEFAULT_CALCULATOR_KWARGS = { MLFF.CHGNet: {"stress_weight": _GPa_to_eV_per_A3}, diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 14f737ec09..1e9b359915 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -49,18 +49,18 @@ class ForceFieldTaskDocument(AseStructureTaskDoc): description="version of the interatomic potential used for relaxation.", ) - dir_name: Optional[str] = Field( + dir_name: Optional[str] | list[Optional[str]] = Field( None, description="Directory where the force field calculations are performed." ) included_objects: Optional[list[AseObject]] = Field( None, description="list of forcefield objects included with this task document" ) - objects: Optional[dict[AseObject, Any]] = Field( - None, description="Forcefield objects associated with this task" + objects: Optional[dict[AseObject, Any]] | list[Optional[dict[AseObject, Any]]] = ( + Field(None, description="Forcefield objects associated with this task") ) - is_force_converged: Optional[bool] = Field( + is_force_converged: Optional[bool] | list[Optional[bool]] = Field( None, description=( "Whether the calculation is converged with respect to interatomic forces." @@ -71,7 +71,7 @@ class ForceFieldTaskDocument(AseStructureTaskDoc): def from_ase_compatible_result( cls, ase_calculator_name: str, - result: AseResult, + result: AseResult | list[AseResult], steps: int, relax_kwargs: dict = None, optimizer_kwargs: dict = None, @@ -156,6 +156,8 @@ def from_ase_compatible_result( return cls.from_ase_task_doc(ase_task_doc, **ff_kwargs) @property - def forcefield_objects(self) -> Optional[dict[AseObject, Any]]: + def forcefield_objects( + self, + ) -> Optional[dict[AseObject, Any]] | list[Optional[dict[AseObject, Any]]]: """Alias `objects` attr for backwards compatibility.""" return self.objects diff --git a/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index 10c41e0ef1..d624d95f20 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -47,6 +47,27 @@ def test_base_maker(test_dir): assert isinstance(output, AseStructureTaskDoc) +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 output.output.structure[0].volume == pytest.approx(22.304245) + assert output.output.structure[1].volume == pytest.approx(22.304245) + assert output.output.energy[0] == pytest.approx(-0.018494767) + assert output.output.energy[1] == pytest.approx(-0.018494767) + assert isinstance(output, AseStructureTaskDoc) + assert fcc_ne_structure.matches(output.output.structure[0]), ( + f"{output.output.structure[0]} != {fcc_ne_structure}" + ) + + def test_lennard_jones_relax_maker(lj_fcc_ne_pars, fcc_ne_structure): job = LennardJonesRelaxMaker( calculator_kwargs=lj_fcc_ne_pars, relax_kwargs={"fmax": 0.001} diff --git a/tests/forcefields/flows/test_phonon.py b/tests/forcefields/flows/test_phonon.py index 340e70164a..42e96dd06a 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -17,13 +17,16 @@ from atomate2.forcefields.flows.phonons import PhononMaker -@pytest.mark.parametrize("from_name", [False, True]) +@pytest.mark.parametrize( + "from_name, socket", [(False, True), (False, False), (True, False), (True, True)] +) 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 phonon_kwargs = dict( + socket=socket, use_symmetrized_structure="conventional", create_thermal_displacements=False, store_force_constants=False, @@ -102,3 +105,34 @@ def test_phonon_wf_force_field( # check phonon plots exist assert os.path.isfile(filename_bs) assert os.path.isfile(filename_dos) + + +@pytest.mark.parametrize("socket", [True, False]) +def test_phonon_wf_force_field_diamond( + clean_dir, si_diamond: Structure, tmp_path: Path, socket: bool +): + # TODO brittle due to inability to adjust dtypes in CHGNetRelaxMaker + + phonon_kwargs = dict( + socket=socket, + use_symmetrized_structure="conventional", + create_thermal_displacements=False, + store_force_constants=False, + prefer_90_degrees=False, + generate_frequencies_eigenvectors_kwargs={ + "tstep": 100, + "filename_bs": f"{tmp_path}/phonon_bs_test.png", + "filename_dos": f"{tmp_path}/phonon_dos_test.pdf", + }, + ) + + phonon_maker = PhononMaker(**phonon_kwargs) + + flow = phonon_maker.make(si_diamond) + + # run the flow or job and ensure that it finished running successfully + responses = run_locally(flow, create_folders=True, ensure_success=True) + + # validate the outputs + ph_bs_dos_doc = responses[flow[-1].uuid][1].output + assert isinstance(ph_bs_dos_doc, PhononBSDOSDoc) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index e280f6c705..a1eb5cb82c 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -133,6 +133,32 @@ def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): CHGNetRelaxMaker() +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 + output1 = responses[job.uuid][1].output + assert isinstance(output1, ForceFieldTaskDocument) + + assert len(output1.structure) == 2 + assert output1.output.energy[0] == approx(-9.96250, rel=1e-2) + assert output1.output.energy[1] == approx(-9.4781, rel=1e-2) + + # check the force_field_task_doc attributes + assert Path(responses[job.uuid][1].output.dir_name[0]).exists() + assert Path(responses[job.uuid][1].output.dir_name[1]).exists() + + @pytest.mark.skip(reason="M3GNet requires DGL which is PyTorch 2.4 incompatible") def test_m3gnet_static_maker(si_structure): # generate job