From 4a10bb39418393f9fb41c4136e21cdeb8b869e4b Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 4 Mar 2024 14:37:54 +0100 Subject: [PATCH 01/39] temp mods commit --- src/atomate2/common/flows/phonons.py | 85 +++- src/atomate2/common/jobs/combined_job.py | 457 +++++++++++++++++++++ src/atomate2/common/jobs/phonons.py | 492 +++++++++++++++++++++++ src/atomate2/forcefields/flows/test.py | 25 ++ 4 files changed, 1037 insertions(+), 22 deletions(-) create mode 100644 src/atomate2/common/jobs/combined_job.py create mode 100644 src/atomate2/forcefields/flows/test.py diff --git a/src/atomate2/common/flows/phonons.py b/src/atomate2/common/flows/phonons.py index 46da4837ad..064f7d180c 100644 --- a/src/atomate2/common/flows/phonons.py +++ b/src/atomate2/common/flows/phonons.py @@ -13,7 +13,12 @@ generate_phonon_displacements, get_supercell_size, get_total_energy_per_cell, + all_jobs, run_phonon_displacements, + run_phonon_displacements_mod, + chunk_and_aggregate_recur, + chunk_and_aggregate, + chunk_and_aggregate2 ) from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive @@ -147,6 +152,7 @@ class BasePhononMaker(Maker, ABC): code: str = None store_force_constants: bool = True socket: bool = False + chunk_size = 2 def make( self, @@ -285,30 +291,65 @@ def make( jobs.append(compute_total_energy_job) total_dft_energy = compute_total_energy_job.output + # gen_run_disp_calc = chunk_and_aggregate2(displacement=self.displacement, + # sym_reduce=self.sym_reduce, + # structure=structure, + # supercell_matrix=supercell_matrix, + # symprec=self.symprec, + # use_symmetrized_structure=self.use_symmetrized_structure, + # kpath_scheme=self.kpath_scheme, + # code=self.code, + # phonon_maker=self.phonon_displacement_maker, + # chunk_size=self.chunk_size) + + gen_run_disp_calc = chunk_and_aggregate_recur(displacement=self.displacement, + sym_reduce=self.sym_reduce, + structure=structure, + supercell_matrix=supercell_matrix, + symprec=self.symprec, + use_symmetrized_structure=self.use_symmetrized_structure, + kpath_scheme=self.kpath_scheme, + code=self.code, + phonon_maker=self.phonon_displacement_maker, + chunk_size=self.chunk_size) + + jobs.append(gen_run_disp_calc) + + # get a phonon object from phonopy - displacements = generate_phonon_displacements( - structure=structure, - supercell_matrix=supercell_matrix, - displacement=self.displacement, - sym_reduce=self.sym_reduce, - symprec=self.symprec, - use_symmetrized_structure=self.use_symmetrized_structure, - kpath_scheme=self.kpath_scheme, - code=self.code, - ) - jobs.append(displacements) + # displacements = generate_phonon_displacements( + # structure=structure, + # supercell_matrix=supercell_matrix, + # displacement=self.displacement, + # sym_reduce=self.sym_reduce, + # symprec=self.symprec, + # use_symmetrized_structure=self.use_symmetrized_structure, + # kpath_scheme=self.kpath_scheme, + # code=self.code, + # ) + # jobs.append(displacements) # perform the phonon displacement calculations - displacement_calcs = run_phonon_displacements( - displacements=displacements.output, - structure=structure, - supercell_matrix=supercell_matrix, - phonon_maker=self.phonon_displacement_maker, - socket=self.socket, - prev_dir_argname=self.prev_calc_dir_argname, - prev_dir=prev_dir, - ) - jobs.append(displacement_calcs) + # displacement_calcs = run_phonon_displacements( + # displacements=displacements.output, + # structure=structure, + # supercell_matrix=supercell_matrix, + # phonon_maker=self.phonon_displacement_maker, + # socket=self.socket, + # prev_dir_argname=self.prev_calc_dir_argname, + # prev_dir=prev_dir, + # ) + + # displacement_calcs = chunk_and_aggregate( + # displacements=displacements.output, + # structure=structure, + # supercell_matrix=supercell_matrix, + # phonon_maker=self.phonon_displacement_maker, + # # socket=self.socket, + # # prev_dir_argname=self.prev_calc_dir_argname, + # # prev_dir=prev_dir, + # ) + # jobs.append(displacement_calcs) # Computation of BORN charges born_run_job_dir = None @@ -338,7 +379,7 @@ def make( kpath_scheme=self.kpath_scheme, code=self.code, structure=structure, - displacement_data=displacement_calcs.output, + displacement_data=gen_run_disp_calc.output, epsilon_static=epsilon_static, born=born, total_dft_energy=total_dft_energy, diff --git a/src/atomate2/common/jobs/combined_job.py b/src/atomate2/common/jobs/combined_job.py new file mode 100644 index 0000000000..b8c0b74a76 --- /dev/null +++ b/src/atomate2/common/jobs/combined_job.py @@ -0,0 +1,457 @@ + +from __future__ import annotations + +import contextlib +from phonopy import Phonopy +from pymatgen.core import Structure +from typing import TYPE_CHECKING +import warnings +from monty.json import jsanitize +from monty.serialization import dumpfn +from jobflow import job +from pymatgen.symmetry.analyzer import SpacegroupAnalyzer +from pymatgen.transformations.advanced_transformations import ( + CubicSupercellTransformation, +) +import numpy as np +from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive +from atomate2.forcefields.utils import Relaxer +from atomate2.forcefields.schemas import ForceFieldTaskDocument +from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine +from pymatgen.phonon.dos import PhononDos +from atomate2.common.schemas.phonons import ForceConstants, PhononBSDOSDoc, get_factor +from pymatgen.io.phonopy import get_phonopy_structure, get_pmg_structure + +if TYPE_CHECKING: + from pathlib import Path + + from emmet.core.math import Matrix3D + from pymatgen.core.structure import Structure + + from atomate2.aims.jobs.base import BaseAimsMaker + from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker + from atomate2.vasp.jobs.base import BaseVaspMaker + +SUPPORTED_CODES = ["vasp", "aims", "forcefields"] + + +def get_supercell_size( + structure: Structure, min_length: float, prefer_90_degrees: bool, **kwargs +) -> list[list[float]]: + """ + Determine supercell size with given min_length. + + Parameters + ---------- + structure: Structure Object + Input structure that will be used to determine supercell + min_length: float + minimum length of cell in Angstrom + prefer_90_degrees: bool + if True, the algorithm will try to find a cell with 90 degree angles first + **kwargs: + Additional parameters that can be set. + """ + kwargs.setdefault("min_atoms", None) + kwargs.setdefault("force_diagonal", False) + + if not prefer_90_degrees: + kwargs.setdefault("max_atoms", None) + transformation = CubicSupercellTransformation( + min_length=min_length, + min_atoms=kwargs["min_atoms"], + max_atoms=kwargs["max_atoms"], + force_diagonal=kwargs["force_diagonal"], + force_90_degrees=False, + ) + transformation.apply_transformation(structure=structure) + else: + max_atoms = kwargs.get("max_atoms", 1000) + kwargs.setdefault("angle_tolerance", 1e-2) + try: + transformation = CubicSupercellTransformation( + min_length=min_length, + min_atoms=kwargs["min_atoms"], + max_atoms=max_atoms, + force_diagonal=kwargs["force_diagonal"], + force_90_degrees=True, + angle_tolerance=kwargs["angle_tolerance"], + ) + transformation.apply_transformation(structure=structure) + + except AttributeError: + kwargs.setdefault("max_atoms", None) + + transformation = CubicSupercellTransformation( + min_length=min_length, + min_atoms=kwargs["min_atoms"], + max_atoms=kwargs["max_atoms"], + force_diagonal=kwargs["force_diagonal"], + force_90_degrees=False, + ) + transformation.apply_transformation(structure=structure) + + return transformation.transformation_matrix.tolist() + +def structure_to_conventional( + structure: Structure, symprec: float = 1e-4 +) -> Structure: + """ + Job hat creates a standard conventional structure. + + Parameters + ---------- + structure: Structure object + input structure that will be transformed + symprec: float + precision to determine symmetry + + Returns + ------- + .Structure + """ + sga = SpacegroupAnalyzer(structure, symprec=symprec) + return sga.get_conventional_standard_structure() + +def structure_to_primitive( + structure: Structure, symprec: float = 1e-4 +) -> Structure: + """ + Job that creates a standard primitive structure. + + Parameters + ---------- + structure: Structure object + input structure that will be transformed + symprec: float + precision to determine symmetry + + Returns + ------- + .Structure + """ + sga = SpacegroupAnalyzer(structure, symprec=symprec) + return sga.get_primitive_standard_structure() + +def gap_relax_calculator(potential_param_file_name,potential_kwargs={}, optimizer_kwargs={},relax_cell=False): + from quippy.potential import Potential + + calculator = Potential( + args_str="IP GAP", + param_filename=str(potential_param_file_name), + **potential_kwargs, + ) + relaxer = Relaxer( + calculator, **optimizer_kwargs, relax_cell=relax_cell + ) + return relaxer + +def gap_static_calculator(potential_param_file_name, potential_kwargs={}): + from quippy.potential import Potential + + calculator = Potential( + args_str="IP GAP", + param_filename=str(potential_param_file_name), + **potential_kwargs, + ) + relaxer = Relaxer(calculator, relax_cell=False) + return relaxer + + +def generate_frequencies_eigenvectors( + structure: Structure, + supercell_matrix: np.array, + displacement: float, + sym_reduce: bool, + symprec: float, + use_symmetrized_structure: str | None, + kpath_scheme: str, + code: str, + displacement_data: dict[str, list], + total_dft_energy: float, + epsilon_static: Matrix3D = None, + born: Matrix3D = None, + **kwargs, +) -> PhononBSDOSDoc: + """ + Analyze the phonon runs and summarize the results. + + Parameters + ---------- + structure: Structure object + Fully optimized structure used for phonon runs + supercell_matrix: np.array + array to describe supercell + displacement: float + displacement in Angstrom used for supercell computation + sym_reduce: bool + if True, symmetry will be used in phonopy + symprec: float + precision to determine symmetry + use_symmetrized_structure: str + primitive, conventional, None are allowed + kpath_scheme: str + kpath scheme for phonon band structure computation + code: str + code to run computations + displacement_data: dict + outputs from displacements + total_dft_energy: float + total DFT energy in eV per cell + epsilon_static: Matrix3D + The high-frequency dielectric constant + born: Matrix3D + Born charges + kwargs: dict + Additional parameters that are passed to PhononBSDOSDoc.from_forces_born + """ + return PhononBSDOSDoc.from_forces_born( + structure=structure.remove_site_property(property_name="magmom") + if "magmom" in structure.site_properties + else structure, + supercell_matrix=supercell_matrix, + displacement=displacement, + sym_reduce=sym_reduce, + symprec=symprec, + use_symmetrized_structure=use_symmetrized_structure, + kpath_scheme=kpath_scheme, + code=code, + displacement_data=displacement_data, + total_dft_energy=total_dft_energy, + epsilon_static=epsilon_static, + born=born, + **kwargs, + ) + + +@job( + output_schema=PhononBSDOSDoc, + data=[PhononDos, PhononBandStructureSymmLine, ForceConstants], +) +def phonon_job( + structure: Structure, + potential_parameter_filename: str, + supercell_size_kwargs: dict={}, + static_kwargs: dict={}, + optimizer_kwargs: dict={}, + bulk_potential_kwargs:dict={}, + relax_steps: int = 5000, + relax_kwargs: dict = {"interval": 5000, "fmax": 0.00001}, + generate_frequencies_eigenvectors_kwargs: dict = {"npoints_band": 50}, + sym_reduce: bool = True, + symprec: float = 1e-4, + displacement: float = 0.01, + min_length: float | None = 20.0, + prefer_90_degrees: bool = True, + use_symmetrized_structure: str | None = None, + store_force_constants=True, + bulk_relax: bool = True, + calculate_static_energy: bool = True, + create_thermal_displacements= True, + prev_dir: str | Path | None = None, + born: list[Matrix3D] | None = None, + bulk_relax_cell=True, + static_relax_cell=False, + epsilon_static: Matrix3D | None = None, + total_dft_energy_per_formula_unit: float | None = None, + supercell_matrix: Matrix3D | None = None, + kpath_scheme: str = "seekpath", + code: str = None + ) -> PhononBSDOSDoc: + use_symmetrized_structure = use_symmetrized_structure + kpath_scheme = kpath_scheme + valid_structs = (None, "primitive", "conventional") + if use_symmetrized_structure not in valid_structs: + raise ValueError( + f"Invalid {use_symmetrized_structure=}, use one of {valid_structs}" + ) + + if use_symmetrized_structure != "primitive" and kpath_scheme != "seekpath": + raise ValueError( + f"You can't use {kpath_scheme=} with the primitive standard " + "structure, please use seekpath" + ) + + valid_schemes = ("seekpath", "hinuma", "setyawan_curtarolo", "latimer_munro") + if kpath_scheme not in valid_schemes: + raise ValueError( + f"{kpath_scheme=} is not implemented, use one of {valid_schemes}" + ) + + if code is None or code not in SUPPORTED_CODES: + raise ValueError( + "The code variable must be passed and it must be a supported code." + f" Supported codes are: {SUPPORTED_CODES}" + ) + + if use_symmetrized_structure == "primitive": + # These structures are compatible with many + # of the kpath algorithms that are used for Materials Project + prim_job = structure_to_primitive(structure, symprec) + structure = prim_job + + elif use_symmetrized_structure == "conventional": + # it could be beneficial to use conventional standard structures to arrive + # faster at supercells with right angles + conv_job = structure_to_conventional(structure, symprec) + structure = conv_job + + optimization_run_job_dir = None + optimization_run_uuid = None + + if bulk_relax: + # optionally relax the structure + # load potential + bulk_relax_calculator = gap_relax_calculator(potential_param_file_name=potential_parameter_filename, + optimizer_kwargs=optimizer_kwargs, potential_kwargs=bulk_potential_kwargs, + relax_cell=bulk_relax_cell) + + bulk = bulk_relax_calculator.relax(structure, steps=relax_steps, **relax_kwargs) + bulk_task_doc = ForceFieldTaskDocument.from_ase_compatible_result(result=bulk, forcefield_name='GAP', relax_cell=bulk_relax_cell, steps=relax_steps, + relax_kwargs=relax_kwargs, optimizer_kwargs=optimizer_kwargs) + structure = bulk_task_doc.output.structure + + if supercell_matrix is None: + supercell_matrix = get_supercell_size( + structure, + min_length, + prefer_90_degrees, + **supercell_size_kwargs, + ) + + # Initialize static calculator + static_energy_calculator = gap_static_calculator( + potential_param_file_name=potential_parameter_filename, + ) + + # Computation of static energy + total_dft_energy = None + static_run_job_dir = None + static_run_uuid = None + if calculate_static_energy and ( + total_dft_energy_per_formula_unit is None + ): + static_job_kwargs = {} + # static_energy_calculator = gap_static_calculator( + # potential_param_file_name=potential_parameter_filename, + # ) + static_run = static_energy_calculator.relax(structure, steps=1, **static_kwargs) + static_task_doc = ForceFieldTaskDocument.from_ase_compatible_result(result=static_run, forcefield_name='GAP', + relax_cell=static_relax_cell, + steps=1, + relax_kwargs=static_kwargs, + optimizer_kwargs=optimizer_kwargs) + total_dft_energy = static_task_doc.output.energy + + warnings.warn( + "Initial magnetic moments will not be considered for the determination " + "of the symmetry of the structure and thus will be removed now.", + stacklevel=1, + ) + cell = get_phonopy_structure( + structure.remove_site_property(property_name="magmom") + if "magmom" in structure.site_properties + else structure + ) + factor = get_factor(code) + + if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": + primitive_matrix: np.ndarray | str = np.eye(3) + else: + primitive_matrix = "auto" + + # TARP: THIS IS BAD! Including for discussions sake + if cell.magnetic_moments is not None and primitive_matrix == "auto": + if np.any(cell.magnetic_moments != 0.0): + raise ValueError( + "For materials with magnetic moments specified " + "use_symmetrized_structure must be 'primitive'" + ) + cell.magnetic_moments = None + + phonon = Phonopy( + cell, + supercell_matrix, + primitive_matrix=primitive_matrix, + factor=factor, + symprec=symprec, + is_symmetry=sym_reduce, + ) + phonon.generate_displacements(distance=displacement) + + supercells = phonon.supercells_with_displacements + + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [], + } + + for idx, sc in enumerate(supercells): + if prev_dir is not None: + phonon_job = static_energy_calculator.relax(get_pmg_structure(sc), steps=1) + else: + phonon_job = static_energy_calculator.relax(get_pmg_structure(sc), steps=1) + + phonon_job_taskdoc = ForceFieldTaskDocument.from_ase_compatible_result(result=phonon_job, forcefield_name='GAP', + relax_cell=static_relax_cell, + steps=1, + relax_kwargs=static_kwargs, + optimizer_kwargs=optimizer_kwargs) + print(f'Running supercell {idx+1}/{len(supercells)}') + # we will add some metadata + info = { + "displacement_number": idx, + "original_structure": structure, + "supercell_matrix": supercell_matrix, + "displaced_structure": sc, + } + with contextlib.suppress(Exception): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) + outputs["displacement_number"].append(idx) + # outputs["uuids"].append(phonon_job_taskdoc.output.uuid) + # outputs["dirs"].append(phonon_job_taskdoc.output.dir_name) + outputs["forces"].append(phonon_job_taskdoc.output.forces) + + born_run_job_dir = None + born_run_uuid = None + + phonon_collect = generate_frequencies_eigenvectors( + supercell_matrix=supercell_matrix, + displacement=displacement, + sym_reduce=sym_reduce, + symprec=symprec, + use_symmetrized_structure=use_symmetrized_structure, + kpath_scheme=kpath_scheme, + code=code, + structure=structure, + displacement_data=outputs, + epsilon_static=epsilon_static, + born=born, + total_dft_energy=total_dft_energy, + **{'static_run_job_dir': static_run_job_dir, + 'static_run_uuid': static_run_uuid, + 'born_run_job_dir': born_run_job_dir, + 'born_run_uuid': born_run_uuid, + 'optimization_run_job_dir': optimization_run_job_dir, + 'optimization_run_uuid': optimization_run_uuid, + 'create_thermal_displacements':create_thermal_displacements, + 'store_force_constants':store_force_constants}, + **generate_frequencies_eigenvectors_kwargs, + ) + + # save forces for later use ?? + forces_doc = jsanitize(outputs, strict=False, allow_bson=False, enum_values=False, recursive_msonable=False) + dumpfn(forces_doc, fn='forces_data.json.gz') + + # jsanited_doc = jsanitize(phonon_collect, strict=False, allow_bson=False, + # enum_values=False, recursive_msonable=False) + + # dumpfn(jsanited_doc, fn='phononbsdostaskdoc.json.gz') + + return phonon_collect + diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index d273f1546a..aa742163fd 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -348,3 +348,495 @@ def run_phonon_displacements( displacement_flow = Flow(phonon_jobs, outputs) return Response(replace=displacement_flow) + + +@job(data=["forces", "displaced_structures"]) +def run_phonon_displacements_mod( + displacements: list[Structure], + structure: Structure, + supercell_matrix: Matrix3D, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + start_inx=0, + batch_size=2, + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [] + }, + stop_inx=None, +): + """ + Run phonon displacements. + + Note, this job will replace itself with N displacement calculations. + + Parameters + ---------- + displacements + structure: Structure object + Fully optimized structure used for phonon computations. + supercell_matrix: Matrix3D + supercell matrix for metadata + phonon_maker : .BaseVaspMaker + A VaspMaker to use to generate the elastic relaxation jobs. + prev_dir : str or Path or None + A previous vasp calculation directory to use for copying outputs. + """ + stop_inx = stop_inx if stop_inx is not None else len(displacements) + + if start_inx < stop_inx: + new_inx = start_inx + batch_size + jobs = [] + for idx, displacement in enumerate(displacements[start_inx:new_inx]): + 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 + start_inx}/{len(displacements)}") + # print(idx+start_inx) + # we will add some meta data + info = { + "displacement_number": idx + start_inx, + "original_structure": structure, + "supercell_matrix": supercell_matrix, + "displaced_structure": displacement, + } + with contextlib.suppress(Exception): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) + # outputs.append(idx+start_inx) + outputs["displacement_number"].append(idx + start_inx) + outputs["uuids"].append(phonon_job.output.uuid) + outputs["dirs"].append(phonon_job.output.dir_name) + outputs["forces"].append(phonon_job.output.output.forces) + print(outputs) + # print(outputs['uuids']) + jobs.append(phonon_job) + + new_job = run_phonon_displacements_mod(structure=structure, supercell_matrix=supercell_matrix, + displacements=displacements, start_inx=new_inx, stop_inx=stop_inx, + phonon_maker=phonon_maker, prev_dir=prev_dir) + + return Response(addition=[new_job, *jobs], output=outputs) + +@job(data=["forces", "displaced_structures"]) +def run_phonon_displacements_mod2( + total_displacements:int, + displacements: list[Structure], + structure: Structure, + supercell_matrix: Matrix3D, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + indexs=list[int], + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [] + }, +): + """ + Run phonon displacements. + + Note, this job will replace itself with N displacement calculations. + + Parameters + ---------- + displacements + structure: Structure object + Fully optimized structure used for phonon computations. + supercell_matrix: Matrix3D + supercell matrix for metadata + phonon_maker : .BaseVaspMaker + A VaspMaker to use to generate the elastic relaxation jobs. + prev_dir : str or Path or None + A previous vasp calculation directory to use for copying outputs. + """ + jobs = [] + 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" {indexs[idx]}/{total_displacements}") + # print(idx+start_inx) + # we will add some meta data + info = { + "displacement_number": indexs[idx]-1, + "original_structure": structure, + "supercell_matrix": supercell_matrix, + "displaced_structures": displacement, + } + with contextlib.suppress(Exception): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) + # outputs.append(idx+start_inx) + outputs["displacement_number"].append(indexs[idx]-1) + outputs["uuids"].append(phonon_job.output.uuid) + outputs["dirs"].append(phonon_job.output.dir_name) + outputs["forces"].append(phonon_job.output.output.forces) + # print(outputs['uuids']) + jobs.append(phonon_job) + + displacement_flow = Flow(jobs) + return Response(replace=displacement_flow, output=outputs) + +@job(data=["forces", "displaced_structures"]) +def run_phonon_displacements_recur( + total_displacements:int, + index_list, + chunks_list, + structure: Structure, + supercell_matrix: Matrix3D, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + start:int=0, + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [] + }, +): + """ + Run phonon displacements. + + Note, this job will replace itself with N displacement calculations. + + Parameters + ---------- + displacements + structure: Structure object + Fully optimized structure used for phonon computations. + supercell_matrix: Matrix3D + supercell matrix for metadata + phonon_maker : .BaseVaspMaker + A VaspMaker to use to generate the elastic relaxation jobs. + prev_dir : str or Path or None + A previous vasp calculation directory to use for copying outputs. + """ + index_list_here=index_list[start] + disp_here = chunks_list[start] + jobs = [] + for idx, displacement in enumerate(disp_here): + 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" {index_list_here[idx]}/{total_displacements} : batch {start+1}") + info = { + "displacement_number": index_list_here[idx] - 1, + "original_structure": structure, + "supercell_matrix": supercell_matrix, + "displaced_structures": displacement, + } + with contextlib.suppress(Exception): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) + jobs.append(phonon_job) + outputs["displacement_number"].append(index_list_here[idx] - 1) + outputs["uuids"].append(phonon_job.output.uuid) + outputs["dirs"].append(phonon_job.output.dir_name) + outputs["forces"].append(phonon_job.output.output.forces) + + + if start+1 != len(index_list): + start = start+1 + new_job = run_phonon_displacements_recur(structure=structure, supercell_matrix=supercell_matrix, + chunks_list=chunks_list, index_list=index_list, + total_displacements=total_displacements, + phonon_maker=phonon_maker, prev_dir=prev_dir,start=start) + + displacement_flow = Flow([*jobs, new_job]) + return Response(addition=displacement_flow, output=outputs) + else: + displacement_flow = Flow(jobs) + return Response(addition=displacement_flow, output=outputs) + + +def chunks(lst, n): + """Yield successive n-sized chunks from lst along with their start and end indices.""" + for start_index, i in enumerate(range(0, len(lst), n)): + end_index = min(i + n, len(lst)) + yield lst[i:end_index], i, end_index + + +@job(data=["forces"]) +def chunk_and_aggregate2(displacement: float, + sym_reduce: bool, + symprec: float, + use_symmetrized_structure: str | None, + kpath_scheme: str, + code: str, + structure, + supercell_matrix, + phonon_maker, + chunk_size=3, + ): + warnings.warn( + "Initial magnetic moments will not be considered for the determination " + "of the symmetry of the structure and thus will be removed now.", + stacklevel=1, + ) + cell = get_phonopy_structure( + structure.remove_site_property(property_name="magmom") + if "magmom" in structure.site_properties + else structure + ) + factor = get_factor(code) + + # a bit of code repetition here as I currently + # do not see how to pass the phonopy object? + if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": + primitive_matrix: np.ndarray | str = np.eye(3) + else: + primitive_matrix = "auto" + + # TARP: THIS IS BAD! Including for discussions sake + if cell.magnetic_moments is not None and primitive_matrix == "auto": + if np.any(cell.magnetic_moments != 0.0): + raise ValueError( + "For materials with magnetic moments specified " + "use_symmetrized_structure must be 'primitive'" + ) + cell.magnetic_moments = None + + phonon = Phonopy( + cell, + supercell_matrix, + primitive_matrix=primitive_matrix, + factor=factor, + symprec=symprec, + is_symmetry=sym_reduce, + ) + phonon.generate_displacements(distance=displacement) + + supercells = phonon.supercells_with_displacements + + displacements = [get_pmg_structure(cell) for cell in supercells] + + jobs = [] + outputs: dict[str, list] = { + "displacement_number": None, + "forces": None, + "uuids": None, + "dirs": None + } + for chunk, start_index, end_index in chunks(displacements, chunk_size): + indexs = list(range(start_index + 1, end_index + 1, 1)) + + job = run_phonon_displacements_mod2(total_displacements=len(displacements), displacements=chunk, structure=structure, + supercell_matrix=supercell_matrix,phonon_maker=phonon_maker, indexs=indexs) + jobs.append(job) + outputs["displacement_number"] = job.output["displacement_number"] + outputs["uuids"] = job.output["uuids"] + outputs["dirs"] = job.output["dirs"] + outputs["forces"] = job.output["forces"] + + displacement_flow = Flow(jobs) + return Response(replace=displacement_flow, output=outputs) # Response(addition=jobs, output=outputs) + + +@job +def chunk_and_aggregate(displacements, + structure, + supercell_matrix, + phonon_maker, + chunk_size=3, + ): + + jobs = [] + outputs: dict[str, list] = { + "displacement_number": None, + "forces": None, + "uuids": None, + "dirs": None + } + for chunk, start_index, end_index in chunks(displacements, chunk_size): + indexs = list(range(start_index + 1, end_index + 1, 1)) + + job = run_phonon_displacements_mod2(total_displacements=len(displacements), displacements=chunk, + structure=structure, + supercell_matrix=supercell_matrix, phonon_maker=phonon_maker, indexs=indexs) + jobs.append(job) + outputs["displacement_number"] = job.output["displacement_number"] + outputs["uuids"] = job.output["uuids"] + outputs["dirs"] = job.output["dirs"] + outputs["forces"] = job.output["forces"] + + return Response(replace=jobs, output=outputs) + + +@job(data=["forces", Structure, 'chunks_list', 'index_list']) +def chunk_and_aggregate_recur(displacement: float, + sym_reduce: bool, + symprec: float, + use_symmetrized_structure: str | None, + kpath_scheme: str, + code: str, + structure, + supercell_matrix, + phonon_maker, + chunk_size=3, + ): + warnings.warn( + "Initial magnetic moments will not be considered for the determination " + "of the symmetry of the structure and thus will be removed now.", + stacklevel=1, + ) + cell = get_phonopy_structure( + structure.remove_site_property(property_name="magmom") + if "magmom" in structure.site_properties + else structure + ) + factor = get_factor(code) + + # a bit of code repetition here as I currently + # do not see how to pass the phonopy object? + if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": + primitive_matrix: np.ndarray | str = np.eye(3) + else: + primitive_matrix = "auto" + + # TARP: THIS IS BAD! Including for discussions sake + if cell.magnetic_moments is not None and primitive_matrix == "auto": + if np.any(cell.magnetic_moments != 0.0): + raise ValueError( + "For materials with magnetic moments specified " + "use_symmetrized_structure must be 'primitive'" + ) + cell.magnetic_moments = None + + phonon = Phonopy( + cell, + supercell_matrix, + primitive_matrix=primitive_matrix, + factor=factor, + symprec=symprec, + is_symmetry=sym_reduce, + ) + phonon.generate_displacements(distance=displacement) + + supercells = phonon.supercells_with_displacements + + displacements = [get_pmg_structure(cell) for cell in supercells] + + total_displacements = len(displacements) + recur_function_inputs = {'chunks_list': [], + 'index_list': []} + # chunks_list = [] + # index_list = [] + for chunk, start_index, end_index in chunks(displacements, chunk_size): + indexs = list(range(start_index + 1, end_index + 1, 1)) + recur_function_inputs['chunks_list'].append(chunk) + recur_function_inputs['index_list'].append(indexs) + # chunks_list.append(chunk) + # index_list.append(indexs) + + del supercells, displacements + + displ_job = run_phonon_displacements_recur(total_displacements=total_displacements, structure=structure, + supercell_matrix=supercell_matrix,phonon_maker=phonon_maker, + chunks_list=recur_function_inputs['chunks_list'], + index_list=recur_function_inputs['index_list'], start=0) + + displacement_flow = Flow(displ_job, output=displ_job.output) + return Response(replace=displacement_flow) # Response(addition=jobs, output=outputs) + + +@job(data=["forces", Structure, 'chunks_list', 'index_list']) +def all_jobs(displacement: float, + sym_reduce: bool, + symprec: float, + kpath_scheme: str, + code: str, + structure, + supercell_matrix, + use_symmetrized_structure: str | None, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + chunk_size=3 + ): + warnings.warn( + "Initial magnetic moments will not be considered for the determination " + "of the symmetry of the structure and thus will be removed now.", + stacklevel=1, + ) + cell = get_phonopy_structure( + structure.remove_site_property(property_name="magmom") + if "magmom" in structure.site_properties + else structure + ) + factor = get_factor(code) + + # a bit of code repetition here as I currently + # do not see how to pass the phonopy object? + if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": + primitive_matrix: np.ndarray | str = np.eye(3) + else: + primitive_matrix = "auto" + + # TARP: THIS IS BAD! Including for discussions sake + if cell.magnetic_moments is not None and primitive_matrix == "auto": + if np.any(cell.magnetic_moments != 0.0): + raise ValueError( + "For materials with magnetic moments specified " + "use_symmetrized_structure must be 'primitive'" + ) + cell.magnetic_moments = None + + phonon = Phonopy( + cell, + supercell_matrix, + primitive_matrix=primitive_matrix, + factor=factor, + symprec=symprec, + is_symmetry=sym_reduce, + ) + phonon.generate_displacements(distance=displacement) + + supercells = phonon.supercells_with_displacements + + displacements = [get_pmg_structure(cell) for cell in supercells] + + del supercells + + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [] + } + + 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)}") + info = { + "displacement_number": idx, + "original_structure": structure, + "supercell_matrix": supercell_matrix, + "displaced_structures": displacement, + } + with contextlib.suppress(Exception): + phonon_job.update_maker_kwargs( + {"_set": {"write_additional_data->phonon_info:json": info}}, + dict_mod=True, + ) + outputs["displacement_number"].append(idx - 1) + outputs["uuids"].append(phonon_job.output.uuid) + outputs["dirs"].append(phonon_job.output.dir_name) + outputs["forces"].append(phonon_job.output.output.forces) + + displacement_flow = Flow(phonon_job) + return Response(addition=displacement_flow) diff --git a/src/atomate2/forcefields/flows/test.py b/src/atomate2/forcefields/flows/test.py new file mode 100644 index 0000000000..ca2075accb --- /dev/null +++ b/src/atomate2/forcefields/flows/test.py @@ -0,0 +1,25 @@ +import json +from pymatgen.io.ase import AseAtomsAdaptor +from ase.io import read + +from atomate2.forcefields.flows.phonons import PhononMaker +from atomate2.forcefields.jobs import GAPRelaxMaker +from atomate2.forcefields.jobs import GAPStaticMaker +from atomate2.forcefields.jobs import CHGNetRelaxMaker, CHGNetStaticMaker +from pymatgen.core import Structure +from jobflow.managers.fireworks import flow_to_workflow +from fireworks import LaunchPad +from jobflow import run_locally +import os +from jobflow import SETTINGS + +from mp_api.client import MPRester + +mpr = MPRester(api_key='ajGziP3VMy57gFn9yzlJSuWQMNdjDa8q') + +st = mpr.get_structure_by_material_id('mp-10635') + +phonons = PhononMaker(generate_frequencies_eigenvectors_kwargs={"npoints_band": 50}, bulk_relax_maker=None).make(structure=st) + + +resp = run_locally(phonons, create_folders=False, allow_external_references=True, store=SETTINGS.JOB_STORE) \ No newline at end of file From 0ea96a5e4c34308d802972ec3c2dbf510ce2bebd Mon Sep 17 00:00:00 2001 From: Aakash Ashok Naik <91958822+naik-aakash@users.noreply.github.com> Date: Tue, 6 May 2025 08:53:44 +0200 Subject: [PATCH 02/39] Delete debug file src/atomate2/forcefields/flows/test.py --- src/atomate2/forcefields/flows/test.py | 25 ------------------------- 1 file changed, 25 deletions(-) delete mode 100644 src/atomate2/forcefields/flows/test.py diff --git a/src/atomate2/forcefields/flows/test.py b/src/atomate2/forcefields/flows/test.py deleted file mode 100644 index ca2075accb..0000000000 --- a/src/atomate2/forcefields/flows/test.py +++ /dev/null @@ -1,25 +0,0 @@ -import json -from pymatgen.io.ase import AseAtomsAdaptor -from ase.io import read - -from atomate2.forcefields.flows.phonons import PhononMaker -from atomate2.forcefields.jobs import GAPRelaxMaker -from atomate2.forcefields.jobs import GAPStaticMaker -from atomate2.forcefields.jobs import CHGNetRelaxMaker, CHGNetStaticMaker -from pymatgen.core import Structure -from jobflow.managers.fireworks import flow_to_workflow -from fireworks import LaunchPad -from jobflow import run_locally -import os -from jobflow import SETTINGS - -from mp_api.client import MPRester - -mpr = MPRester(api_key='ajGziP3VMy57gFn9yzlJSuWQMNdjDa8q') - -st = mpr.get_structure_by_material_id('mp-10635') - -phonons = PhononMaker(generate_frequencies_eigenvectors_kwargs={"npoints_band": 50}, bulk_relax_maker=None).make(structure=st) - - -resp = run_locally(phonons, create_folders=False, allow_external_references=True, store=SETTINGS.JOB_STORE) \ No newline at end of file From 8fca55e2f22e9cf17d09f069844f01b95e4a9987 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Tue, 6 May 2025 21:01:02 +0200 Subject: [PATCH 03/39] add a batch option to the AseMaker --- src/atomate2/ase/jobs.py | 54 ++++++++++++++-------- src/atomate2/ase/utils.py | 97 +++++++++++++++++++++------------------ tests/ase/test_jobs.py | 20 ++++++++ 3 files changed, 108 insertions(+), 63 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 8513b3be45..9252c000c3 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -221,32 +221,50 @@ def make( ------- AseStructureTaskDoc or AseMoleculeTaskDoc """ - 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, - 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, - ) + result=self.run_ase(mol_or_struct, prev_dir=prev_dir) + is_list = isinstance(result, list) + + if not is_list: + return AseTaskDoc.to_mol_or_struct_metadata_doc( + getattr(self.calculator, "name", type(self.calculator).__name__), + result, + self.steps, + relax_kwargs=self.relax_kwargs, + optimizer_kwargs=self.optimizer_kwargs, + relax_cell=self.relax_cell, + 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, + ) + else: + return [AseTaskDoc.to_mol_or_struct_metadata_doc( + getattr(self.calculator, "name", type(self.calculator).__name__), + resul, + self.steps, + relax_kwargs=self.relax_kwargs, + optimizer_kwargs=self.optimizer_kwargs, + relax_cell=self.relax_cell, + 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 resul in result] 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/utils.py b/src/atomate2/ase/utils.py index 2ffadc4ab5..0bddab683d 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -324,7 +324,7 @@ 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, @@ -332,13 +332,13 @@ 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 corresponding lists The atoms for relaxation. fmax : float Total force tolerance for relaxation convergence. @@ -355,46 +355,53 @@ def relax( Returns ------- - dict including optimized structure and the trajectory + 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) or isinstance(atoms[0], Molecule) or isinstance(atoms[0], Structure) - if isinstance(atoms, Structure | Molecule): - atoms = self.ase_adaptor.get_atoms(atoms) - 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)) - ) - 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, - ) + if not is_list: + atoms = [atoms] + + list_ase_results = [] + for atom in atoms: + is_mol = isinstance(atoms, Molecule) or ( + isinstance(atoms, Atoms) and all(not pbc for pbc in atoms.pbc) + ) + if isinstance(atoms, Structure | Molecule): + atoms = self.ase_adaptor.get_atoms(atoms) + 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_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/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index 10c41e0ef1..8ac8fdf150 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -47,6 +47,26 @@ def test_base_maker(test_dir): assert isinstance(output, AseStructureTaskDoc) +def test_lennard_jones_batch_relax_maker(lj_fcc_ne_pars, fcc_ne_structure): + job = LennardJonesRelaxMaker( + calculator_kwargs=lj_fcc_ne_pars, relax_kwargs={"fmax": 0.001} + ).make([fcc_ne_structure,fcc_ne_structure]) + + response = run_locally(job) + output = response[job.uuid][1].output + assert len(output) ==2 + + assert output[0].structure.volume == pytest.approx(22.304245) + assert output[0].output.energy == pytest.approx(-0.018494767) + assert isinstance(output[0], AseStructureTaskDoc) + assert isinstance(output[1], AseStructureTaskDoc) + assert fcc_ne_structure.matches(output[0].structure), ( + f"{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} From 5fd612f0f2a4b4d976a12c580cc552e15d433975 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 16:46:10 +0200 Subject: [PATCH 04/39] fix first test and adapt aseoutput documents --- src/atomate2/ase/jobs.py | 22 +- src/atomate2/ase/schemas.py | 393 +++++++++++-------- src/atomate2/ase/utils.py | 7 +- src/atomate2/common/flows/phonons.py | 90 ++--- src/atomate2/common/jobs/combined_job.py | 457 ----------------------- src/atomate2/common/jobs/phonons.py | 7 +- tests/forcefields/flows/test_phonon.py | 5 +- 7 files changed, 270 insertions(+), 711 deletions(-) delete mode 100644 src/atomate2/common/jobs/combined_job.py diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 9252c000c3..5aadbcc393 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -221,13 +221,9 @@ def make( ------- AseStructureTaskDoc or AseMoleculeTaskDoc """ - result=self.run_ase(mol_or_struct, prev_dir=prev_dir) - is_list = isinstance(result, list) - - if not is_list: - return AseTaskDoc.to_mol_or_struct_metadata_doc( + return AseTaskDoc.to_mol_or_struct_metadata_doc( getattr(self.calculator, "name", type(self.calculator).__name__), - result, + self.run_ase(mol_or_struct, prev_dir=prev_dir), self.steps, relax_kwargs=self.relax_kwargs, optimizer_kwargs=self.optimizer_kwargs, @@ -238,20 +234,6 @@ def make( store_trajectory=self.store_trajectory, tags=self.tags, ) - else: - return [AseTaskDoc.to_mol_or_struct_metadata_doc( - getattr(self.calculator, "name", type(self.calculator).__name__), - resul, - self.steps, - relax_kwargs=self.relax_kwargs, - optimizer_kwargs=self.optimizer_kwargs, - relax_cell=self.relax_cell, - 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 resul in result] def run_ase( self, diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 35d5319f4e..8abf0b273d 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -93,17 +93,17 @@ class AseObject(ValueEnum): class AseBaseModel(BaseModel): """Base document class for ASE input and output.""" - mol_or_struct: Structure | Molecule | None = Field( + mol_or_struct: Structure | Molecule|list[Structure]|list[Molecule] | None = Field( None, description="The molecule or structure 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.") + 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.") 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): + 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): + 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 @@ -121,15 +121,15 @@ 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 +139,31 @@ 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( + 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. Only present for batch calculations." + ), + ) class InputDoc(AseBaseModel): """The inputs used to run this job.""" @@ -189,7 +196,7 @@ class InputDoc(AseBaseModel): class AseStructureTaskDoc(StructureMetadata): """Document containing information on structure manipulation using ASE.""" - structure: Structure = Field( + structure: Structure|list[Structure] = Field( None, description="Final output structure from the task" ) @@ -204,25 +211,26 @@ class AseStructureTaskDoc(StructureMetadata): description="name of the ASE calculator used in the calculation.", ) - dir_name: str | None = Field( + dir_name: str |list[dirname] |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" ) + # TODO: check if it needs to be a list objects: 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 +265,7 @@ 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." @@ -265,12 +273,12 @@ class AseMoleculeTaskDoc(MoleculeMetadata): output: OutputDoc = Field(None, description="The output information from this job.") - ase_calculator_name: str = Field( + ase_calculator_name: str = Field( None, 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 +289,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 +321,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 +354,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 +378,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,148 +402,202 @@ 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 = not isinstance(result, AseResult) + + results= result if is_list else [result] + + output_mol_or_struct=[] + input_mol_or_struct=[] + final_energy=[] + final_forces=[] + final_stress=[] + ionic_steps =[] + n_steps = [] + objects =[] + + for result in results: + trajectory = result.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_doc = InputDoc( - mol_or_struct=input_mol_or_struct, - relax_cell=relax_cell, - fix_symmetry=fix_symmetry, - symprec=symprec, - steps=steps, - 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 + input_mol_or_struct_here=trajectory[0] + + input_mol_or_struct.append(input_mol_or_struct_here) + n_steps.append(n_steps_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.final_mol_or_struct) + + if trajectory is None: + final_energy.append(result.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_structure = [] + 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 ) - 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 - ) - } + + # 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_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), - forces=final_forces, - stress=final_stress, - ionic_steps=ionic_steps, - elapsed_time=result.elapsed_time, - n_steps=n_steps, - ) + ionic_steps_structure.append(ionic_step) + ionic_steps.append(ionic_steps_structure) + + 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], + ) - return cls( - mol_or_struct=output_mol_or_struct, - input=input_doc, - output=output_doc, - ase_calculator_name=ase_calculator_name, - included_objects=list(objects.keys()), - objects=objects, - is_force_converged=result.is_force_converged, - energy_downhill=result.energy_downhill, - dir_name=result.dir_name, - tags=tags, - **task_document_kwargs, - ) + 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, + ) + else: + input_doc = InputDoc( + mol_or_struct=input_mol_or_struct, + 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, + energy=final_energy, + energy_per_atom=[final_energy_here/len(output_mol_or_struct_here) for final_energy_here, output_mol_or_struct_here in zip(final_energy, output_mol_or_struct)], + forces=final_forces, + stress=final_stress, + ionic_steps=ionic_steps, + 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[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, + ) @classmethod 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 +623,24 @@ 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 0bddab683d..a0156639af 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -360,10 +360,11 @@ def relax( is_list=isinstance(atoms[0], Atoms) or isinstance(atoms[0], Molecule) or isinstance(atoms[0], Structure) if not is_list: - atoms = [atoms] - + list_atoms = [atoms] + else: + list_atoms=atoms list_ase_results = [] - for atom in atoms: + for atoms in list_atoms: is_mol = isinstance(atoms, Molecule) or ( isinstance(atoms, Atoms) and all(not pbc for pbc in atoms.pbc) ) diff --git a/src/atomate2/common/flows/phonons.py b/src/atomate2/common/flows/phonons.py index 167ac9f606..193edfc9c5 100644 --- a/src/atomate2/common/flows/phonons.py +++ b/src/atomate2/common/flows/phonons.py @@ -14,13 +14,9 @@ generate_phonon_displacements, get_supercell_size, get_total_energy_per_cell, - all_jobs, run_phonon_displacements, - run_phonon_displacements_mod, chunk_and_aggregate_recur, - chunk_and_aggregate, - chunk_and_aggregate2 -) + ) from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive if TYPE_CHECKING: @@ -137,7 +133,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" @@ -311,66 +307,28 @@ def make( jobs.append(compute_total_energy_job) total_dft_energy = compute_total_energy_job.output - # gen_run_disp_calc = chunk_and_aggregate2(displacement=self.displacement, - # sym_reduce=self.sym_reduce, - # structure=structure, - # supercell_matrix=supercell_matrix, - # symprec=self.symprec, - # use_symmetrized_structure=self.use_symmetrized_structure, - # kpath_scheme=self.kpath_scheme, - # code=self.code, - # phonon_maker=self.phonon_displacement_maker, - # chunk_size=self.chunk_size) - - gen_run_disp_calc = chunk_and_aggregate_recur(displacement=self.displacement, - sym_reduce=self.sym_reduce, - structure=structure, - supercell_matrix=supercell_matrix, - symprec=self.symprec, - use_symmetrized_structure=self.use_symmetrized_structure, - kpath_scheme=self.kpath_scheme, - code=self.code, - phonon_maker=self.phonon_displacement_maker, - chunk_size=self.chunk_size) - - jobs.append(gen_run_disp_calc) - - - # get a phonon object from phonopy - # displacements = generate_phonon_displacements( - # structure=structure, - # supercell_matrix=supercell_matrix, - # displacement=self.displacement, - # sym_reduce=self.sym_reduce, - # symprec=self.symprec, - # use_symmetrized_structure=self.use_symmetrized_structure, - # kpath_scheme=self.kpath_scheme, - # code=self.code, - # ) - # jobs.append(displacements) - - # perform the phonon displacement calculations - # displacement_calcs = run_phonon_displacements( - # displacements=displacements.output, - # structure=structure, - # supercell_matrix=supercell_matrix, - # phonon_maker=self.phonon_displacement_maker, - # socket=self.socket, - # prev_dir_argname=self.prev_calc_dir_argname, - # prev_dir=prev_dir, - # ) - - # displacement_calcs = chunk_and_aggregate( - # displacements=displacements.output, - # structure=structure, - # supercell_matrix=supercell_matrix, - # phonon_maker=self.phonon_displacement_maker, - # # socket=self.socket, - # # prev_dir_argname=self.prev_calc_dir_argname, - # # prev_dir=prev_dir, - # ) - # jobs.append(displacement_calcs) + displacements = generate_phonon_displacements( + structure=structure, + supercell_matrix=supercell_matrix, + displacement=self.displacement, + sym_reduce=self.sym_reduce, + symprec=self.symprec, + use_symmetrized_structure=self.use_symmetrized_structure, + kpath_scheme=self.kpath_scheme, + code=self.code, + ) + jobs.append(displacements) + displacement_calcs = run_phonon_displacements( + displacements=displacements.output, + structure=structure, + supercell_matrix=supercell_matrix, + phonon_maker=self.phonon_displacement_maker, + socket=self.socket, + prev_dir_argname=self.prev_calc_dir_argname, + prev_dir=prev_dir, + ) + jobs.append(displacement_calcs) # Computation of BORN charges born_run_job_dir = None born_run_uuid = None @@ -399,7 +357,7 @@ def make( kpath_scheme=self.kpath_scheme, code=self.code, structure=structure, - displacement_data=gen_run_disp_calc.output, + displacement_data=displacement_calcs.output, epsilon_static=epsilon_static, born=born, total_dft_energy=total_dft_energy, diff --git a/src/atomate2/common/jobs/combined_job.py b/src/atomate2/common/jobs/combined_job.py deleted file mode 100644 index b8c0b74a76..0000000000 --- a/src/atomate2/common/jobs/combined_job.py +++ /dev/null @@ -1,457 +0,0 @@ - -from __future__ import annotations - -import contextlib -from phonopy import Phonopy -from pymatgen.core import Structure -from typing import TYPE_CHECKING -import warnings -from monty.json import jsanitize -from monty.serialization import dumpfn -from jobflow import job -from pymatgen.symmetry.analyzer import SpacegroupAnalyzer -from pymatgen.transformations.advanced_transformations import ( - CubicSupercellTransformation, -) -import numpy as np -from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive -from atomate2.forcefields.utils import Relaxer -from atomate2.forcefields.schemas import ForceFieldTaskDocument -from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine -from pymatgen.phonon.dos import PhononDos -from atomate2.common.schemas.phonons import ForceConstants, PhononBSDOSDoc, get_factor -from pymatgen.io.phonopy import get_phonopy_structure, get_pmg_structure - -if TYPE_CHECKING: - from pathlib import Path - - from emmet.core.math import Matrix3D - from pymatgen.core.structure import Structure - - from atomate2.aims.jobs.base import BaseAimsMaker - from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker - from atomate2.vasp.jobs.base import BaseVaspMaker - -SUPPORTED_CODES = ["vasp", "aims", "forcefields"] - - -def get_supercell_size( - structure: Structure, min_length: float, prefer_90_degrees: bool, **kwargs -) -> list[list[float]]: - """ - Determine supercell size with given min_length. - - Parameters - ---------- - structure: Structure Object - Input structure that will be used to determine supercell - min_length: float - minimum length of cell in Angstrom - prefer_90_degrees: bool - if True, the algorithm will try to find a cell with 90 degree angles first - **kwargs: - Additional parameters that can be set. - """ - kwargs.setdefault("min_atoms", None) - kwargs.setdefault("force_diagonal", False) - - if not prefer_90_degrees: - kwargs.setdefault("max_atoms", None) - transformation = CubicSupercellTransformation( - min_length=min_length, - min_atoms=kwargs["min_atoms"], - max_atoms=kwargs["max_atoms"], - force_diagonal=kwargs["force_diagonal"], - force_90_degrees=False, - ) - transformation.apply_transformation(structure=structure) - else: - max_atoms = kwargs.get("max_atoms", 1000) - kwargs.setdefault("angle_tolerance", 1e-2) - try: - transformation = CubicSupercellTransformation( - min_length=min_length, - min_atoms=kwargs["min_atoms"], - max_atoms=max_atoms, - force_diagonal=kwargs["force_diagonal"], - force_90_degrees=True, - angle_tolerance=kwargs["angle_tolerance"], - ) - transformation.apply_transformation(structure=structure) - - except AttributeError: - kwargs.setdefault("max_atoms", None) - - transformation = CubicSupercellTransformation( - min_length=min_length, - min_atoms=kwargs["min_atoms"], - max_atoms=kwargs["max_atoms"], - force_diagonal=kwargs["force_diagonal"], - force_90_degrees=False, - ) - transformation.apply_transformation(structure=structure) - - return transformation.transformation_matrix.tolist() - -def structure_to_conventional( - structure: Structure, symprec: float = 1e-4 -) -> Structure: - """ - Job hat creates a standard conventional structure. - - Parameters - ---------- - structure: Structure object - input structure that will be transformed - symprec: float - precision to determine symmetry - - Returns - ------- - .Structure - """ - sga = SpacegroupAnalyzer(structure, symprec=symprec) - return sga.get_conventional_standard_structure() - -def structure_to_primitive( - structure: Structure, symprec: float = 1e-4 -) -> Structure: - """ - Job that creates a standard primitive structure. - - Parameters - ---------- - structure: Structure object - input structure that will be transformed - symprec: float - precision to determine symmetry - - Returns - ------- - .Structure - """ - sga = SpacegroupAnalyzer(structure, symprec=symprec) - return sga.get_primitive_standard_structure() - -def gap_relax_calculator(potential_param_file_name,potential_kwargs={}, optimizer_kwargs={},relax_cell=False): - from quippy.potential import Potential - - calculator = Potential( - args_str="IP GAP", - param_filename=str(potential_param_file_name), - **potential_kwargs, - ) - relaxer = Relaxer( - calculator, **optimizer_kwargs, relax_cell=relax_cell - ) - return relaxer - -def gap_static_calculator(potential_param_file_name, potential_kwargs={}): - from quippy.potential import Potential - - calculator = Potential( - args_str="IP GAP", - param_filename=str(potential_param_file_name), - **potential_kwargs, - ) - relaxer = Relaxer(calculator, relax_cell=False) - return relaxer - - -def generate_frequencies_eigenvectors( - structure: Structure, - supercell_matrix: np.array, - displacement: float, - sym_reduce: bool, - symprec: float, - use_symmetrized_structure: str | None, - kpath_scheme: str, - code: str, - displacement_data: dict[str, list], - total_dft_energy: float, - epsilon_static: Matrix3D = None, - born: Matrix3D = None, - **kwargs, -) -> PhononBSDOSDoc: - """ - Analyze the phonon runs and summarize the results. - - Parameters - ---------- - structure: Structure object - Fully optimized structure used for phonon runs - supercell_matrix: np.array - array to describe supercell - displacement: float - displacement in Angstrom used for supercell computation - sym_reduce: bool - if True, symmetry will be used in phonopy - symprec: float - precision to determine symmetry - use_symmetrized_structure: str - primitive, conventional, None are allowed - kpath_scheme: str - kpath scheme for phonon band structure computation - code: str - code to run computations - displacement_data: dict - outputs from displacements - total_dft_energy: float - total DFT energy in eV per cell - epsilon_static: Matrix3D - The high-frequency dielectric constant - born: Matrix3D - Born charges - kwargs: dict - Additional parameters that are passed to PhononBSDOSDoc.from_forces_born - """ - return PhononBSDOSDoc.from_forces_born( - structure=structure.remove_site_property(property_name="magmom") - if "magmom" in structure.site_properties - else structure, - supercell_matrix=supercell_matrix, - displacement=displacement, - sym_reduce=sym_reduce, - symprec=symprec, - use_symmetrized_structure=use_symmetrized_structure, - kpath_scheme=kpath_scheme, - code=code, - displacement_data=displacement_data, - total_dft_energy=total_dft_energy, - epsilon_static=epsilon_static, - born=born, - **kwargs, - ) - - -@job( - output_schema=PhononBSDOSDoc, - data=[PhononDos, PhononBandStructureSymmLine, ForceConstants], -) -def phonon_job( - structure: Structure, - potential_parameter_filename: str, - supercell_size_kwargs: dict={}, - static_kwargs: dict={}, - optimizer_kwargs: dict={}, - bulk_potential_kwargs:dict={}, - relax_steps: int = 5000, - relax_kwargs: dict = {"interval": 5000, "fmax": 0.00001}, - generate_frequencies_eigenvectors_kwargs: dict = {"npoints_band": 50}, - sym_reduce: bool = True, - symprec: float = 1e-4, - displacement: float = 0.01, - min_length: float | None = 20.0, - prefer_90_degrees: bool = True, - use_symmetrized_structure: str | None = None, - store_force_constants=True, - bulk_relax: bool = True, - calculate_static_energy: bool = True, - create_thermal_displacements= True, - prev_dir: str | Path | None = None, - born: list[Matrix3D] | None = None, - bulk_relax_cell=True, - static_relax_cell=False, - epsilon_static: Matrix3D | None = None, - total_dft_energy_per_formula_unit: float | None = None, - supercell_matrix: Matrix3D | None = None, - kpath_scheme: str = "seekpath", - code: str = None - ) -> PhononBSDOSDoc: - use_symmetrized_structure = use_symmetrized_structure - kpath_scheme = kpath_scheme - valid_structs = (None, "primitive", "conventional") - if use_symmetrized_structure not in valid_structs: - raise ValueError( - f"Invalid {use_symmetrized_structure=}, use one of {valid_structs}" - ) - - if use_symmetrized_structure != "primitive" and kpath_scheme != "seekpath": - raise ValueError( - f"You can't use {kpath_scheme=} with the primitive standard " - "structure, please use seekpath" - ) - - valid_schemes = ("seekpath", "hinuma", "setyawan_curtarolo", "latimer_munro") - if kpath_scheme not in valid_schemes: - raise ValueError( - f"{kpath_scheme=} is not implemented, use one of {valid_schemes}" - ) - - if code is None or code not in SUPPORTED_CODES: - raise ValueError( - "The code variable must be passed and it must be a supported code." - f" Supported codes are: {SUPPORTED_CODES}" - ) - - if use_symmetrized_structure == "primitive": - # These structures are compatible with many - # of the kpath algorithms that are used for Materials Project - prim_job = structure_to_primitive(structure, symprec) - structure = prim_job - - elif use_symmetrized_structure == "conventional": - # it could be beneficial to use conventional standard structures to arrive - # faster at supercells with right angles - conv_job = structure_to_conventional(structure, symprec) - structure = conv_job - - optimization_run_job_dir = None - optimization_run_uuid = None - - if bulk_relax: - # optionally relax the structure - # load potential - bulk_relax_calculator = gap_relax_calculator(potential_param_file_name=potential_parameter_filename, - optimizer_kwargs=optimizer_kwargs, potential_kwargs=bulk_potential_kwargs, - relax_cell=bulk_relax_cell) - - bulk = bulk_relax_calculator.relax(structure, steps=relax_steps, **relax_kwargs) - bulk_task_doc = ForceFieldTaskDocument.from_ase_compatible_result(result=bulk, forcefield_name='GAP', relax_cell=bulk_relax_cell, steps=relax_steps, - relax_kwargs=relax_kwargs, optimizer_kwargs=optimizer_kwargs) - structure = bulk_task_doc.output.structure - - if supercell_matrix is None: - supercell_matrix = get_supercell_size( - structure, - min_length, - prefer_90_degrees, - **supercell_size_kwargs, - ) - - # Initialize static calculator - static_energy_calculator = gap_static_calculator( - potential_param_file_name=potential_parameter_filename, - ) - - # Computation of static energy - total_dft_energy = None - static_run_job_dir = None - static_run_uuid = None - if calculate_static_energy and ( - total_dft_energy_per_formula_unit is None - ): - static_job_kwargs = {} - # static_energy_calculator = gap_static_calculator( - # potential_param_file_name=potential_parameter_filename, - # ) - static_run = static_energy_calculator.relax(structure, steps=1, **static_kwargs) - static_task_doc = ForceFieldTaskDocument.from_ase_compatible_result(result=static_run, forcefield_name='GAP', - relax_cell=static_relax_cell, - steps=1, - relax_kwargs=static_kwargs, - optimizer_kwargs=optimizer_kwargs) - total_dft_energy = static_task_doc.output.energy - - warnings.warn( - "Initial magnetic moments will not be considered for the determination " - "of the symmetry of the structure and thus will be removed now.", - stacklevel=1, - ) - cell = get_phonopy_structure( - structure.remove_site_property(property_name="magmom") - if "magmom" in structure.site_properties - else structure - ) - factor = get_factor(code) - - if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": - primitive_matrix: np.ndarray | str = np.eye(3) - else: - primitive_matrix = "auto" - - # TARP: THIS IS BAD! Including for discussions sake - if cell.magnetic_moments is not None and primitive_matrix == "auto": - if np.any(cell.magnetic_moments != 0.0): - raise ValueError( - "For materials with magnetic moments specified " - "use_symmetrized_structure must be 'primitive'" - ) - cell.magnetic_moments = None - - phonon = Phonopy( - cell, - supercell_matrix, - primitive_matrix=primitive_matrix, - factor=factor, - symprec=symprec, - is_symmetry=sym_reduce, - ) - phonon.generate_displacements(distance=displacement) - - supercells = phonon.supercells_with_displacements - - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [], - } - - for idx, sc in enumerate(supercells): - if prev_dir is not None: - phonon_job = static_energy_calculator.relax(get_pmg_structure(sc), steps=1) - else: - phonon_job = static_energy_calculator.relax(get_pmg_structure(sc), steps=1) - - phonon_job_taskdoc = ForceFieldTaskDocument.from_ase_compatible_result(result=phonon_job, forcefield_name='GAP', - relax_cell=static_relax_cell, - steps=1, - relax_kwargs=static_kwargs, - optimizer_kwargs=optimizer_kwargs) - print(f'Running supercell {idx+1}/{len(supercells)}') - # we will add some metadata - info = { - "displacement_number": idx, - "original_structure": structure, - "supercell_matrix": supercell_matrix, - "displaced_structure": sc, - } - with contextlib.suppress(Exception): - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, - dict_mod=True, - ) - outputs["displacement_number"].append(idx) - # outputs["uuids"].append(phonon_job_taskdoc.output.uuid) - # outputs["dirs"].append(phonon_job_taskdoc.output.dir_name) - outputs["forces"].append(phonon_job_taskdoc.output.forces) - - born_run_job_dir = None - born_run_uuid = None - - phonon_collect = generate_frequencies_eigenvectors( - supercell_matrix=supercell_matrix, - displacement=displacement, - sym_reduce=sym_reduce, - symprec=symprec, - use_symmetrized_structure=use_symmetrized_structure, - kpath_scheme=kpath_scheme, - code=code, - structure=structure, - displacement_data=outputs, - epsilon_static=epsilon_static, - born=born, - total_dft_energy=total_dft_energy, - **{'static_run_job_dir': static_run_job_dir, - 'static_run_uuid': static_run_uuid, - 'born_run_job_dir': born_run_job_dir, - 'born_run_uuid': born_run_uuid, - 'optimization_run_job_dir': optimization_run_job_dir, - 'optimization_run_uuid': optimization_run_uuid, - 'create_thermal_displacements':create_thermal_displacements, - 'store_force_constants':store_force_constants}, - **generate_frequencies_eigenvectors_kwargs, - ) - - # save forces for later use ?? - forces_doc = jsanitize(outputs, strict=False, allow_bson=False, enum_values=False, recursive_msonable=False) - dumpfn(forces_doc, fn='forces_data.json.gz') - - # jsanited_doc = jsanitize(phonon_collect, strict=False, allow_bson=False, - # enum_values=False, recursive_msonable=False) - - # dumpfn(jsanited_doc, fn='phononbsdostaskdoc.json.gz') - - return phonon_collect - diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index b60d18bf72..e5c0422cdf 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -293,9 +293,10 @@ 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 - ) + + #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/tests/forcefields/flows/test_phonon.py b/tests/forcefields/flows/test_phonon.py index 340e70164a..9afd87fda6 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -17,13 +17,14 @@ from atomate2.forcefields.flows.phonons import PhononMaker -@pytest.mark.parametrize("from_name", [False, True]) +@pytest.mark.parametrize("from_name, socket", [(False, True), (True,False)]) 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, From 39b8485ad1dc79892b70a4ab6ba151a98d9cc3bd Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 17:30:52 +0200 Subject: [PATCH 05/39] fix some more --- src/atomate2/ase/schemas.py | 18 +++++++++--------- src/atomate2/common/schemas/phonons.py | 5 +++++ src/atomate2/forcefields/schemas.py | 6 +++--- tests/ase/test_jobs.py | 17 +++++++++-------- tests/forcefields/flows/test_phonon.py | 2 +- 5 files changed, 27 insertions(+), 21 deletions(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 8abf0b273d..dc99f185af 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -66,7 +66,7 @@ class AseResult(BaseModel): ), ) - dir_name: str | Path | None = Field( + dir_name: str | Path | list[str]|list[Path] |None = Field( None, description="The directory where the calculation was run" ) @@ -196,7 +196,7 @@ class InputDoc(AseBaseModel): class AseStructureTaskDoc(StructureMetadata): """Document containing information on structure manipulation using ASE.""" - structure: Structure|list[Structure] = Field( + structure: Structure = Field( None, description="Final output structure from the task" ) @@ -211,7 +211,7 @@ class AseStructureTaskDoc(StructureMetadata): description="name of the ASE calculator used in the calculation.", ) - dir_name: str |list[dirname] |None = Field( + dir_name: str |list[str] |None = Field( None, description="Directory where the ASE calculations are performed." ) @@ -219,7 +219,7 @@ class AseStructureTaskDoc(StructureMetadata): None, description="list of ASE objects included with this task document" ) # TODO: check if it needs to be a list - objects: dict[AseObject, Any] | None = Field( + objects: dict[AseObject, Any] | list[dict[AseObject, Any]]|None = Field( None, description="ASE objects associated with this task" ) @@ -580,15 +580,15 @@ def from_ase_compatible_result( ) return cls( - mol_or_struct=output_mol_or_struct[0], + 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[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, + objects=objects, + 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, ) diff --git a/src/atomate2/common/schemas/phonons.py b/src/atomate2/common/schemas/phonons.py index 99e95c4c10..d7fe5f5890 100644 --- a/src/atomate2/common/schemas/phonons.py +++ b/src/atomate2/common/schemas/phonons.py @@ -527,6 +527,11 @@ def from_forces_born( volume_per_formula_unit = structure.volume / formula_units + if displacement_data["dirs"] is not None: + if isinstance(displacement_data["dirs"][0], list): + displacement_data["dirs"] = [dir for dir_list in displacement_data["dirs"] for dir in dir_list] + + doc = cls.from_structure( structure=structure, meta_structure=structure, diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 14f737ec09..ef212d36e5 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( + 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." diff --git a/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index 8ac8fdf150..b9994b9cb0 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -53,15 +53,16 @@ def test_lennard_jones_batch_relax_maker(lj_fcc_ne_pars, fcc_ne_structure): ).make([fcc_ne_structure,fcc_ne_structure]) response = run_locally(job) + output = response[job.uuid][1].output - assert len(output) ==2 - - assert output[0].structure.volume == pytest.approx(22.304245) - assert output[0].output.energy == pytest.approx(-0.018494767) - assert isinstance(output[0], AseStructureTaskDoc) - assert isinstance(output[1], AseStructureTaskDoc) - assert fcc_ne_structure.matches(output[0].structure), ( - f"{output.structure[0]} != {fcc_ne_structure}" + + 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}" ) diff --git a/tests/forcefields/flows/test_phonon.py b/tests/forcefields/flows/test_phonon.py index 9afd87fda6..094d168dc0 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -17,7 +17,7 @@ from atomate2.forcefields.flows.phonons import PhononMaker -@pytest.mark.parametrize("from_name, socket", [(False, True), (True,False)]) +@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, socket: bool ): From 383a2c3ccb52c4567e87f81ab9ecd387c91d38e9 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 17:33:03 +0200 Subject: [PATCH 06/39] fix documentation --- src/atomate2/forcefields/flows/phonons.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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: From 29da6e413cb75ea09622c075f2989015fcff9543 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 17:43:21 +0200 Subject: [PATCH 07/39] fix a few linter problems --- src/atomate2/ase/jobs.py | 28 +-- src/atomate2/ase/schemas.py | 200 +++++++++-------- src/atomate2/ase/utils.py | 37 ++-- src/atomate2/common/flows/phonons.py | 3 +- src/atomate2/common/jobs/phonons.py | 286 ++++++++++++++----------- src/atomate2/common/schemas/phonons.py | 5 +- src/atomate2/forcefields/schemas.py | 8 +- tests/ase/test_jobs.py | 4 +- tests/forcefields/flows/test_phonon.py | 4 +- 9 files changed, 319 insertions(+), 256 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 5aadbcc393..9839004a0b 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -222,24 +222,24 @@ def make( AseStructureTaskDoc or AseMoleculeTaskDoc """ 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, - 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, - ) + 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, + 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, + ) def run_ase( self, - mol_or_struct: Structure | Molecule|list[Molecule]|list[Structure], + mol_or_struct: Structure | Molecule | list[Molecule] | list[Structure], prev_dir: str | Path | None = None, - ) -> AseResult|list[AseResult]: + ) -> AseResult | list[AseResult]: """ Relax a structure, molecule or a batch of those using ASE, not as a job. diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index dc99f185af..060953f2fc 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -43,22 +43,22 @@ class AseResult(BaseModel): None, description="The molecule or structure in the final trajectory frame." ) - final_energy: float | None = Field( + final_energy: float | list[float] | None = Field( None, description="The final total energy from the calculation." ) - trajectory: PmgTrajectory | None = Field( + trajectory: PmgTrajectory | list[PmgTrajectory] | None = Field( None, description="The relaxation or molecular dynamics trajectory." ) - 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 " @@ -66,11 +66,11 @@ class AseResult(BaseModel): ), ) - dir_name: str | Path | list[str]|list[Path] |None = Field( + dir_name: str | Path | list[str] | list[Path] | None = Field( None, description="The directory where the calculation was run" ) - elapsed_time: float | None = Field( + elapsed_time: float | list[float] | None = Field( None, description="The time taken to run the ASE calculation in seconds." ) @@ -93,17 +93,27 @@ class AseObject(ValueEnum): class AseBaseModel(BaseModel): """Base document class for ASE input and output.""" - mol_or_struct: Structure | Molecule|list[Structure]|list[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|list[Structure] | None = Field(None, description="The structure at this step.") - molecule: Molecule|list[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) or isinstance(self.mol_or_struct[0], Structure)): + 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)): + 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 @@ -121,15 +131,17 @@ class IonicStep(AseBaseModel): class OutputDoc(AseBaseModel): """The outputs of this job.""" - energy: float|list[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|list[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]|list[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,21 +151,21 @@ class OutputDoc(AseBaseModel): # NOTE: units for stresses were converted to kbar (* -10 from standard output) # to comply with MP convention - stress: Matrix3D | list[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]|list[list[IonicStep]] | dict|list[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 | list[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 | list[int]|None = Field( + n_steps: int | list[int] | None = Field( None, description="total number of steps needed in the relaxation." ) @@ -165,6 +177,7 @@ class OutputDoc(AseBaseModel): ), ) + class InputDoc(AseBaseModel): """The inputs used to run this job.""" @@ -211,7 +224,7 @@ class AseStructureTaskDoc(StructureMetadata): description="name of the ASE calculator used in the calculation.", ) - dir_name: str |list[str] |None = Field( + dir_name: str | list[str] | None = Field( None, description="Directory where the ASE calculations are performed." ) @@ -219,18 +232,18 @@ class AseStructureTaskDoc(StructureMetadata): None, description="list of ASE objects included with this task document" ) # TODO: check if it needs to be a list - objects: dict[AseObject, Any] | list[dict[AseObject, Any]]|None = Field( + objects: dict[AseObject, Any] | list[dict[AseObject, Any]] | None = Field( None, description="ASE objects associated with this task" ) - is_force_converged: bool | list[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|list[bool] | None = Field( + energy_downhill: bool | list[bool] | None = Field( None, description=( "Whether the final trajectory frame has lower total " @@ -265,7 +278,9 @@ def from_ase_task_doc( class AseMoleculeTaskDoc(MoleculeMetadata): """Document containing information on molecule manipulation using ASE.""" - molecule: Molecule|list[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." @@ -273,12 +288,12 @@ class AseMoleculeTaskDoc(MoleculeMetadata): output: OutputDoc = Field(None, description="The output information from this job.") - ase_calculator_name: str = Field( + ase_calculator_name: str = Field( None, description="name of the ASE calculator used in the calculation.", ) - dir_name: str | list[str]|None = Field( + dir_name: str | list[str] | None = Field( None, description="Directory where the ASE calculations are performed." ) @@ -289,14 +304,14 @@ class AseMoleculeTaskDoc(MoleculeMetadata): None, description="ASE objects associated with this task" ) - is_force_converged: bool|list[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|list[bool] | None = Field( + energy_downhill: bool | list[bool] | None = Field( None, description=( "Whether the total energy in the final frame " @@ -321,7 +336,7 @@ class AseTaskDoc(AseBaseModel): description="name of the ASE calculator used for this job.", ) - dir_name: str |list[str]| None = Field( + dir_name: str | list[str] | None = Field( None, description="Directory where the ASE calculations are performed." ) @@ -329,18 +344,18 @@ class AseTaskDoc(AseBaseModel): None, description="list of ASE objects included with this task document" ) # maybe list - objects: dict[AseObject, Any] |list[dict[AseObject, Any]] | None = Field( + objects: dict[AseObject, Any] | list[dict[AseObject, Any]] | None = Field( None, description="ASE objects associated with this task" ) - is_force_converged: bool | list[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|list[bool] | None = Field( + energy_downhill: bool | list[bool] | None = Field( None, description=( "Whether the total energy in the final frame " @@ -354,7 +369,7 @@ class AseTaskDoc(AseBaseModel): def from_ase_compatible_result( cls, ase_calculator_name: str, - result: AseResult|list[AseResult], + result: AseResult | list[AseResult], steps: int, relax_kwargs: dict = None, optimizer_kwargs: dict = None, @@ -404,16 +419,16 @@ def from_ase_compatible_result( """ is_list = not isinstance(result, AseResult) - results= result if is_list else [result] + results = result if is_list else [result] - output_mol_or_struct=[] - input_mol_or_struct=[] - final_energy=[] - final_forces=[] - final_stress=[] - ionic_steps =[] + output_mol_or_struct = [] + input_mol_or_struct = [] + final_energy = [] + final_forces: list[Vector3D] | list[list[Vector3D]] = [] + final_stress = [] + ionic_steps = [] n_steps = [] - objects =[] + objects = [] for result in results: trajectory = result.trajectory @@ -422,7 +437,7 @@ def from_ase_compatible_result( n_steps_here = None input_mol_or_struct_here = None if trajectory: - n_steps_here=len(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 @@ -433,12 +448,14 @@ def from_ase_compatible_result( voigt_6_to_full_3x3_stress( [ val * -10 / GPa - for val in trajectory.frame_properties[idx]["stress"] + for val in trajectory.frame_properties[idx][ + "stress" + ] ] ) ) - input_mol_or_struct_here=trajectory[0] + input_mol_or_struct_here = trajectory[0] input_mol_or_struct.append(input_mol_or_struct_here) n_steps.append(n_steps_here) @@ -453,7 +470,7 @@ def from_ase_compatible_result( traj_method = "from_structures" elif isinstance(input_mol_or_struct_here, Molecule): traj_method = "from_molecules" - trajectory=getattr(PmgTrajectory, traj_method)( + trajectory = getattr(PmgTrajectory, traj_method)( [input_mol_or_struct_here], frame_properties=[trajectory.frame_properties[0]], constant_lattice=False, @@ -523,7 +540,6 @@ def from_ase_compatible_result( 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, @@ -557,47 +573,51 @@ def from_ase_compatible_result( tags=tags, **task_document_kwargs, ) - else: - input_doc = InputDoc( - mol_or_struct=input_mol_or_struct, - 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, - energy=final_energy, - energy_per_atom=[final_energy_here/len(output_mol_or_struct_here) for final_energy_here, output_mol_or_struct_here in zip(final_energy, output_mol_or_struct)], - forces=final_forces, - stress=final_stress, - ionic_steps=ionic_steps, - elapsed_time=[result.elapsed_time for result in results], - n_steps=n_steps, - all_forces=final_forces, - ) + input_doc = InputDoc( + mol_or_struct=input_mol_or_struct, + 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, + energy=final_energy, + energy_per_atom=[ + final_energy_here / len(output_mol_or_struct_here) + for final_energy_here, output_mol_or_struct_here 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 for result in results], + n_steps=n_steps, + all_forces=final_forces, + ) - return cls( - 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[0].keys()), - objects=objects, - 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, - ) + return cls( + 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[0].keys()), + objects=objects, + 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, + ) @classmethod def to_mol_or_struct_metadata_doc( cls, ase_calculator_name: str, - result: AseResult|list[AseResult], + result: AseResult | list[AseResult], steps: int | None = None, **task_document_kwargs, ) -> AseStructureTaskDoc | AseMoleculeTaskDoc: @@ -624,23 +644,19 @@ def to_mol_or_struct_metadata_doc( ) kwargs = {k: getattr(task_doc, k, None) for k in _task_doc_translation_keys} - - - if isinstance(task_doc.mol_or_struct, Structure) or isinstance(task_doc.mol_or_struct[0], 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) or isinstance(task_doc.mol_or_struct[0], 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 a0156639af..9d0de17b72 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -324,7 +324,12 @@ def __init__( def relax( self, - atoms: Atoms | Structure | Molecule | list[Atoms] |list[Molecule] |list[Structure], + atoms: Atoms + | Structure + | Molecule + | list[Atoms] + | list[Molecule] + | list[Structure], fmax: float = 0.1, steps: int = 500, traj_file: str = None, @@ -332,7 +337,7 @@ def relax( verbose: bool = False, cell_filter: Filter = FrechetCellFilter, **kwargs, - ) -> AseResult|list[AseResult]: + ) -> AseResult | list[AseResult]: """ Relax the molecule or structure or a list of those. @@ -357,12 +362,16 @@ def relax( ------- dict including optimized structure and the trajectory or a list of those dicts """ - is_list=isinstance(atoms[0], Atoms) or isinstance(atoms[0], Molecule) or isinstance(atoms[0], Structure) + is_list = ( + isinstance(atoms[0], Atoms) + or isinstance(atoms[0], Molecule) + or isinstance(atoms[0], Structure) + ) if not is_list: list_atoms = [atoms] else: - list_atoms=atoms + list_atoms = atoms list_ase_results = [] for atoms in list_atoms: is_mol = isinstance(atoms, Molecule) or ( @@ -396,13 +405,15 @@ def relax( 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, - )) + 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 193edfc9c5..9a19eb6ee2 100644 --- a/src/atomate2/common/flows/phonons.py +++ b/src/atomate2/common/flows/phonons.py @@ -15,8 +15,7 @@ get_supercell_size, get_total_energy_per_cell, run_phonon_displacements, - chunk_and_aggregate_recur, - ) +) from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive if TYPE_CHECKING: diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index e5c0422cdf..95fb123f2b 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -293,10 +293,10 @@ def run_phonon_displacements( "supercell_matrix": supercell_matrix, "displaced_structures": displacements, } - - #phonon_job.update_maker_kwargs( + + # 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) @@ -334,20 +334,20 @@ def run_phonon_displacements( @job(data=["forces", "displaced_structures"]) def run_phonon_displacements_mod( - displacements: list[Structure], - structure: Structure, - supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - start_inx=0, - batch_size=2, - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [] - }, - stop_inx=None, + displacements: list[Structure], + structure: Structure, + supercell_matrix: Matrix3D, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + start_inx=0, + batch_size=2, + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [], + }, + stop_inx=None, ): """ Run phonon displacements. @@ -399,27 +399,34 @@ def run_phonon_displacements_mod( # print(outputs['uuids']) jobs.append(phonon_job) - new_job = run_phonon_displacements_mod(structure=structure, supercell_matrix=supercell_matrix, - displacements=displacements, start_inx=new_inx, stop_inx=stop_inx, - phonon_maker=phonon_maker, prev_dir=prev_dir) + new_job = run_phonon_displacements_mod( + structure=structure, + supercell_matrix=supercell_matrix, + displacements=displacements, + start_inx=new_inx, + stop_inx=stop_inx, + phonon_maker=phonon_maker, + prev_dir=prev_dir, + ) return Response(addition=[new_job, *jobs], output=outputs) + @job(data=["forces", "displaced_structures"]) def run_phonon_displacements_mod2( - total_displacements:int, - displacements: list[Structure], - structure: Structure, - supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - indexs=list[int], - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [] - }, + total_displacements: int, + displacements: list[Structure], + structure: Structure, + supercell_matrix: Matrix3D, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + indexs=list[int], + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [], + }, ): """ Run phonon displacements. @@ -448,7 +455,7 @@ def run_phonon_displacements_mod2( # print(idx+start_inx) # we will add some meta data info = { - "displacement_number": indexs[idx]-1, + "displacement_number": indexs[idx] - 1, "original_structure": structure, "supercell_matrix": supercell_matrix, "displaced_structures": displacement, @@ -459,7 +466,7 @@ def run_phonon_displacements_mod2( dict_mod=True, ) # outputs.append(idx+start_inx) - outputs["displacement_number"].append(indexs[idx]-1) + outputs["displacement_number"].append(indexs[idx] - 1) outputs["uuids"].append(phonon_job.output.uuid) outputs["dirs"].append(phonon_job.output.dir_name) outputs["forces"].append(phonon_job.output.output.forces) @@ -469,22 +476,23 @@ def run_phonon_displacements_mod2( displacement_flow = Flow(jobs) return Response(replace=displacement_flow, output=outputs) + @job(data=["forces", "displaced_structures"]) def run_phonon_displacements_recur( - total_displacements:int, - index_list, - chunks_list, - structure: Structure, - supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - start:int=0, - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [] - }, + total_displacements: int, + index_list, + chunks_list, + structure: Structure, + supercell_matrix: Matrix3D, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + start: int = 0, + outputs: dict[str, list] = { + "displacement_number": [], + "forces": [], + "uuids": [], + "dirs": [], + }, ): """ Run phonon displacements. @@ -503,7 +511,7 @@ def run_phonon_displacements_recur( prev_dir : str or Path or None A previous vasp calculation directory to use for copying outputs. """ - index_list_here=index_list[start] + index_list_here = index_list[start] disp_here = chunks_list[start] jobs = [] for idx, displacement in enumerate(disp_here): @@ -511,7 +519,9 @@ def run_phonon_displacements_recur( phonon_job = phonon_maker.make(displacement, prev_dir=prev_dir) else: phonon_job = phonon_maker.make(displacement) - phonon_job.append_name(f" {index_list_here[idx]}/{total_displacements} : batch {start+1}") + phonon_job.append_name( + f" {index_list_here[idx]}/{total_displacements} : batch {start + 1}" + ) info = { "displacement_number": index_list_here[idx] - 1, "original_structure": structure, @@ -529,19 +539,23 @@ def run_phonon_displacements_recur( outputs["dirs"].append(phonon_job.output.dir_name) outputs["forces"].append(phonon_job.output.output.forces) - - if start+1 != len(index_list): - start = start+1 - new_job = run_phonon_displacements_recur(structure=structure, supercell_matrix=supercell_matrix, - chunks_list=chunks_list, index_list=index_list, - total_displacements=total_displacements, - phonon_maker=phonon_maker, prev_dir=prev_dir,start=start) + if start + 1 != len(index_list): + start = start + 1 + new_job = run_phonon_displacements_recur( + structure=structure, + supercell_matrix=supercell_matrix, + chunks_list=chunks_list, + index_list=index_list, + total_displacements=total_displacements, + phonon_maker=phonon_maker, + prev_dir=prev_dir, + start=start, + ) displacement_flow = Flow([*jobs, new_job]) return Response(addition=displacement_flow, output=outputs) - else: - displacement_flow = Flow(jobs) - return Response(addition=displacement_flow, output=outputs) + displacement_flow = Flow(jobs) + return Response(addition=displacement_flow, output=outputs) def chunks(lst, n): @@ -552,17 +566,18 @@ def chunks(lst, n): @job(data=["forces"]) -def chunk_and_aggregate2(displacement: float, - sym_reduce: bool, - symprec: float, - use_symmetrized_structure: str | None, - kpath_scheme: str, - code: str, - structure, - supercell_matrix, - phonon_maker, - chunk_size=3, - ): +def chunk_and_aggregate2( + displacement: float, + sym_reduce: bool, + symprec: float, + use_symmetrized_structure: str | None, + kpath_scheme: str, + code: str, + structure, + supercell_matrix, + phonon_maker, + chunk_size=3, +): warnings.warn( "Initial magnetic moments will not be considered for the determination " "of the symmetry of the structure and thus will be removed now.", @@ -610,13 +625,19 @@ def chunk_and_aggregate2(displacement: float, "displacement_number": None, "forces": None, "uuids": None, - "dirs": None + "dirs": None, } for chunk, start_index, end_index in chunks(displacements, chunk_size): indexs = list(range(start_index + 1, end_index + 1, 1)) - job = run_phonon_displacements_mod2(total_displacements=len(displacements), displacements=chunk, structure=structure, - supercell_matrix=supercell_matrix,phonon_maker=phonon_maker, indexs=indexs) + job = run_phonon_displacements_mod2( + total_displacements=len(displacements), + displacements=chunk, + structure=structure, + supercell_matrix=supercell_matrix, + phonon_maker=phonon_maker, + indexs=indexs, + ) jobs.append(job) outputs["displacement_number"] = job.output["displacement_number"] outputs["uuids"] = job.output["uuids"] @@ -624,30 +645,37 @@ def chunk_and_aggregate2(displacement: float, outputs["forces"] = job.output["forces"] displacement_flow = Flow(jobs) - return Response(replace=displacement_flow, output=outputs) # Response(addition=jobs, output=outputs) + return Response( + replace=displacement_flow, output=outputs + ) # Response(addition=jobs, output=outputs) @job -def chunk_and_aggregate(displacements, - structure, - supercell_matrix, - phonon_maker, - chunk_size=3, - ): - +def chunk_and_aggregate( + displacements, + structure, + supercell_matrix, + phonon_maker, + chunk_size=3, +): jobs = [] outputs: dict[str, list] = { "displacement_number": None, "forces": None, "uuids": None, - "dirs": None + "dirs": None, } for chunk, start_index, end_index in chunks(displacements, chunk_size): indexs = list(range(start_index + 1, end_index + 1, 1)) - job = run_phonon_displacements_mod2(total_displacements=len(displacements), displacements=chunk, - structure=structure, - supercell_matrix=supercell_matrix, phonon_maker=phonon_maker, indexs=indexs) + job = run_phonon_displacements_mod2( + total_displacements=len(displacements), + displacements=chunk, + structure=structure, + supercell_matrix=supercell_matrix, + phonon_maker=phonon_maker, + indexs=indexs, + ) jobs.append(job) outputs["displacement_number"] = job.output["displacement_number"] outputs["uuids"] = job.output["uuids"] @@ -657,18 +685,19 @@ def chunk_and_aggregate(displacements, return Response(replace=jobs, output=outputs) -@job(data=["forces", Structure, 'chunks_list', 'index_list']) -def chunk_and_aggregate_recur(displacement: float, - sym_reduce: bool, - symprec: float, - use_symmetrized_structure: str | None, - kpath_scheme: str, - code: str, - structure, - supercell_matrix, - phonon_maker, - chunk_size=3, - ): +@job(data=["forces", Structure, "chunks_list", "index_list"]) +def chunk_and_aggregate_recur( + displacement: float, + sym_reduce: bool, + symprec: float, + use_symmetrized_structure: str | None, + kpath_scheme: str, + code: str, + structure, + supercell_matrix, + phonon_maker, + chunk_size=3, +): warnings.warn( "Initial magnetic moments will not be considered for the determination " "of the symmetry of the structure and thus will be removed now.", @@ -712,41 +741,48 @@ def chunk_and_aggregate_recur(displacement: float, displacements = [get_pmg_structure(cell) for cell in supercells] total_displacements = len(displacements) - recur_function_inputs = {'chunks_list': [], - 'index_list': []} + recur_function_inputs = {"chunks_list": [], "index_list": []} # chunks_list = [] # index_list = [] for chunk, start_index, end_index in chunks(displacements, chunk_size): indexs = list(range(start_index + 1, end_index + 1, 1)) - recur_function_inputs['chunks_list'].append(chunk) - recur_function_inputs['index_list'].append(indexs) + recur_function_inputs["chunks_list"].append(chunk) + recur_function_inputs["index_list"].append(indexs) # chunks_list.append(chunk) # index_list.append(indexs) del supercells, displacements - displ_job = run_phonon_displacements_recur(total_displacements=total_displacements, structure=structure, - supercell_matrix=supercell_matrix,phonon_maker=phonon_maker, - chunks_list=recur_function_inputs['chunks_list'], - index_list=recur_function_inputs['index_list'], start=0) + displ_job = run_phonon_displacements_recur( + total_displacements=total_displacements, + structure=structure, + supercell_matrix=supercell_matrix, + phonon_maker=phonon_maker, + chunks_list=recur_function_inputs["chunks_list"], + index_list=recur_function_inputs["index_list"], + start=0, + ) displacement_flow = Flow(displ_job, output=displ_job.output) - return Response(replace=displacement_flow) # Response(addition=jobs, output=outputs) - - -@job(data=["forces", Structure, 'chunks_list', 'index_list']) -def all_jobs(displacement: float, - sym_reduce: bool, - symprec: float, - kpath_scheme: str, - code: str, - structure, - supercell_matrix, - use_symmetrized_structure: str | None, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - chunk_size=3 - ): + return Response( + replace=displacement_flow + ) # Response(addition=jobs, output=outputs) + + +@job(data=["forces", Structure, "chunks_list", "index_list"]) +def all_jobs( + displacement: float, + sym_reduce: bool, + symprec: float, + kpath_scheme: str, + code: str, + structure, + supercell_matrix, + use_symmetrized_structure: str | None, + phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, + prev_dir: str | Path = None, + chunk_size=3, +): warnings.warn( "Initial magnetic moments will not be considered for the determination " "of the symmetry of the structure and thus will be removed now.", @@ -795,7 +831,7 @@ def all_jobs(displacement: float, "displacement_number": [], "forces": [], "uuids": [], - "dirs": [] + "dirs": [], } for idx, displacement in enumerate(displacements): @@ -803,7 +839,7 @@ def all_jobs(displacement: float, 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.append_name(f" {[idx + 1]}/{len(displacements)}") info = { "displacement_number": idx, "original_structure": structure, diff --git a/src/atomate2/common/schemas/phonons.py b/src/atomate2/common/schemas/phonons.py index d7fe5f5890..3fbd70d3e0 100644 --- a/src/atomate2/common/schemas/phonons.py +++ b/src/atomate2/common/schemas/phonons.py @@ -529,8 +529,9 @@ def from_forces_born( if displacement_data["dirs"] is not None: if isinstance(displacement_data["dirs"][0], list): - displacement_data["dirs"] = [dir for dir_list in displacement_data["dirs"] for dir in dir_list] - + displacement_data["dirs"] = [ + dir for dir_list in displacement_data["dirs"] for dir in dir_list + ] doc = cls.from_structure( structure=structure, diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index ef212d36e5..bf9233ece0 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]|list[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]] |list[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]|list[Optional[bool]] = Field( + is_force_converged: Optional[bool] | list[Optional[bool]] = Field( None, description=( "Whether the calculation is converged with respect to interatomic forces." diff --git a/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index b9994b9cb0..43e04f15b8 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -50,7 +50,7 @@ def test_base_maker(test_dir): def test_lennard_jones_batch_relax_maker(lj_fcc_ne_pars, fcc_ne_structure): job = LennardJonesRelaxMaker( calculator_kwargs=lj_fcc_ne_pars, relax_kwargs={"fmax": 0.001} - ).make([fcc_ne_structure,fcc_ne_structure]) + ).make([fcc_ne_structure, fcc_ne_structure]) response = run_locally(job) @@ -66,8 +66,6 @@ def test_lennard_jones_batch_relax_maker(lj_fcc_ne_pars, 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 094d168dc0..9295c4c066 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -17,7 +17,9 @@ from atomate2.forcefields.flows.phonons import PhononMaker -@pytest.mark.parametrize("from_name, socket", [(False, True), (False, False), (True,False), (True,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, socket: bool ): From 7d5a19658c8c95aa051d9536a77990fe636ab3b0 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 17:52:28 +0200 Subject: [PATCH 08/39] fix code --- src/atomate2/common/jobs/phonons.py | 528 ---------------------------- 1 file changed, 528 deletions(-) diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index 95fb123f2b..b60c204b10 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -330,531 +330,3 @@ def run_phonon_displacements( displacement_flow = Flow(phonon_jobs, outputs) return Response(replace=displacement_flow) - - -@job(data=["forces", "displaced_structures"]) -def run_phonon_displacements_mod( - displacements: list[Structure], - structure: Structure, - supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - start_inx=0, - batch_size=2, - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [], - }, - stop_inx=None, -): - """ - Run phonon displacements. - - Note, this job will replace itself with N displacement calculations. - - Parameters - ---------- - displacements - structure: Structure object - Fully optimized structure used for phonon computations. - supercell_matrix: Matrix3D - supercell matrix for metadata - phonon_maker : .BaseVaspMaker - A VaspMaker to use to generate the elastic relaxation jobs. - prev_dir : str or Path or None - A previous vasp calculation directory to use for copying outputs. - """ - stop_inx = stop_inx if stop_inx is not None else len(displacements) - - if start_inx < stop_inx: - new_inx = start_inx + batch_size - jobs = [] - for idx, displacement in enumerate(displacements[start_inx:new_inx]): - 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 + start_inx}/{len(displacements)}") - # print(idx+start_inx) - # we will add some meta data - info = { - "displacement_number": idx + start_inx, - "original_structure": structure, - "supercell_matrix": supercell_matrix, - "displaced_structure": displacement, - } - with contextlib.suppress(Exception): - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, - dict_mod=True, - ) - # outputs.append(idx+start_inx) - outputs["displacement_number"].append(idx + start_inx) - outputs["uuids"].append(phonon_job.output.uuid) - outputs["dirs"].append(phonon_job.output.dir_name) - outputs["forces"].append(phonon_job.output.output.forces) - print(outputs) - # print(outputs['uuids']) - jobs.append(phonon_job) - - new_job = run_phonon_displacements_mod( - structure=structure, - supercell_matrix=supercell_matrix, - displacements=displacements, - start_inx=new_inx, - stop_inx=stop_inx, - phonon_maker=phonon_maker, - prev_dir=prev_dir, - ) - - return Response(addition=[new_job, *jobs], output=outputs) - - -@job(data=["forces", "displaced_structures"]) -def run_phonon_displacements_mod2( - total_displacements: int, - displacements: list[Structure], - structure: Structure, - supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - indexs=list[int], - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [], - }, -): - """ - Run phonon displacements. - - Note, this job will replace itself with N displacement calculations. - - Parameters - ---------- - displacements - structure: Structure object - Fully optimized structure used for phonon computations. - supercell_matrix: Matrix3D - supercell matrix for metadata - phonon_maker : .BaseVaspMaker - A VaspMaker to use to generate the elastic relaxation jobs. - prev_dir : str or Path or None - A previous vasp calculation directory to use for copying outputs. - """ - jobs = [] - 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" {indexs[idx]}/{total_displacements}") - # print(idx+start_inx) - # we will add some meta data - info = { - "displacement_number": indexs[idx] - 1, - "original_structure": structure, - "supercell_matrix": supercell_matrix, - "displaced_structures": displacement, - } - with contextlib.suppress(Exception): - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, - dict_mod=True, - ) - # outputs.append(idx+start_inx) - outputs["displacement_number"].append(indexs[idx] - 1) - outputs["uuids"].append(phonon_job.output.uuid) - outputs["dirs"].append(phonon_job.output.dir_name) - outputs["forces"].append(phonon_job.output.output.forces) - # print(outputs['uuids']) - jobs.append(phonon_job) - - displacement_flow = Flow(jobs) - return Response(replace=displacement_flow, output=outputs) - - -@job(data=["forces", "displaced_structures"]) -def run_phonon_displacements_recur( - total_displacements: int, - index_list, - chunks_list, - structure: Structure, - supercell_matrix: Matrix3D, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - start: int = 0, - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [], - }, -): - """ - Run phonon displacements. - - Note, this job will replace itself with N displacement calculations. - - Parameters - ---------- - displacements - structure: Structure object - Fully optimized structure used for phonon computations. - supercell_matrix: Matrix3D - supercell matrix for metadata - phonon_maker : .BaseVaspMaker - A VaspMaker to use to generate the elastic relaxation jobs. - prev_dir : str or Path or None - A previous vasp calculation directory to use for copying outputs. - """ - index_list_here = index_list[start] - disp_here = chunks_list[start] - jobs = [] - for idx, displacement in enumerate(disp_here): - 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" {index_list_here[idx]}/{total_displacements} : batch {start + 1}" - ) - info = { - "displacement_number": index_list_here[idx] - 1, - "original_structure": structure, - "supercell_matrix": supercell_matrix, - "displaced_structures": displacement, - } - with contextlib.suppress(Exception): - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, - dict_mod=True, - ) - jobs.append(phonon_job) - outputs["displacement_number"].append(index_list_here[idx] - 1) - outputs["uuids"].append(phonon_job.output.uuid) - outputs["dirs"].append(phonon_job.output.dir_name) - outputs["forces"].append(phonon_job.output.output.forces) - - if start + 1 != len(index_list): - start = start + 1 - new_job = run_phonon_displacements_recur( - structure=structure, - supercell_matrix=supercell_matrix, - chunks_list=chunks_list, - index_list=index_list, - total_displacements=total_displacements, - phonon_maker=phonon_maker, - prev_dir=prev_dir, - start=start, - ) - - displacement_flow = Flow([*jobs, new_job]) - return Response(addition=displacement_flow, output=outputs) - displacement_flow = Flow(jobs) - return Response(addition=displacement_flow, output=outputs) - - -def chunks(lst, n): - """Yield successive n-sized chunks from lst along with their start and end indices.""" - for start_index, i in enumerate(range(0, len(lst), n)): - end_index = min(i + n, len(lst)) - yield lst[i:end_index], i, end_index - - -@job(data=["forces"]) -def chunk_and_aggregate2( - displacement: float, - sym_reduce: bool, - symprec: float, - use_symmetrized_structure: str | None, - kpath_scheme: str, - code: str, - structure, - supercell_matrix, - phonon_maker, - chunk_size=3, -): - warnings.warn( - "Initial magnetic moments will not be considered for the determination " - "of the symmetry of the structure and thus will be removed now.", - stacklevel=1, - ) - cell = get_phonopy_structure( - structure.remove_site_property(property_name="magmom") - if "magmom" in structure.site_properties - else structure - ) - factor = get_factor(code) - - # a bit of code repetition here as I currently - # do not see how to pass the phonopy object? - if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": - primitive_matrix: np.ndarray | str = np.eye(3) - else: - primitive_matrix = "auto" - - # TARP: THIS IS BAD! Including for discussions sake - if cell.magnetic_moments is not None and primitive_matrix == "auto": - if np.any(cell.magnetic_moments != 0.0): - raise ValueError( - "For materials with magnetic moments specified " - "use_symmetrized_structure must be 'primitive'" - ) - cell.magnetic_moments = None - - phonon = Phonopy( - cell, - supercell_matrix, - primitive_matrix=primitive_matrix, - factor=factor, - symprec=symprec, - is_symmetry=sym_reduce, - ) - phonon.generate_displacements(distance=displacement) - - supercells = phonon.supercells_with_displacements - - displacements = [get_pmg_structure(cell) for cell in supercells] - - jobs = [] - outputs: dict[str, list] = { - "displacement_number": None, - "forces": None, - "uuids": None, - "dirs": None, - } - for chunk, start_index, end_index in chunks(displacements, chunk_size): - indexs = list(range(start_index + 1, end_index + 1, 1)) - - job = run_phonon_displacements_mod2( - total_displacements=len(displacements), - displacements=chunk, - structure=structure, - supercell_matrix=supercell_matrix, - phonon_maker=phonon_maker, - indexs=indexs, - ) - jobs.append(job) - outputs["displacement_number"] = job.output["displacement_number"] - outputs["uuids"] = job.output["uuids"] - outputs["dirs"] = job.output["dirs"] - outputs["forces"] = job.output["forces"] - - displacement_flow = Flow(jobs) - return Response( - replace=displacement_flow, output=outputs - ) # Response(addition=jobs, output=outputs) - - -@job -def chunk_and_aggregate( - displacements, - structure, - supercell_matrix, - phonon_maker, - chunk_size=3, -): - jobs = [] - outputs: dict[str, list] = { - "displacement_number": None, - "forces": None, - "uuids": None, - "dirs": None, - } - for chunk, start_index, end_index in chunks(displacements, chunk_size): - indexs = list(range(start_index + 1, end_index + 1, 1)) - - job = run_phonon_displacements_mod2( - total_displacements=len(displacements), - displacements=chunk, - structure=structure, - supercell_matrix=supercell_matrix, - phonon_maker=phonon_maker, - indexs=indexs, - ) - jobs.append(job) - outputs["displacement_number"] = job.output["displacement_number"] - outputs["uuids"] = job.output["uuids"] - outputs["dirs"] = job.output["dirs"] - outputs["forces"] = job.output["forces"] - - return Response(replace=jobs, output=outputs) - - -@job(data=["forces", Structure, "chunks_list", "index_list"]) -def chunk_and_aggregate_recur( - displacement: float, - sym_reduce: bool, - symprec: float, - use_symmetrized_structure: str | None, - kpath_scheme: str, - code: str, - structure, - supercell_matrix, - phonon_maker, - chunk_size=3, -): - warnings.warn( - "Initial magnetic moments will not be considered for the determination " - "of the symmetry of the structure and thus will be removed now.", - stacklevel=1, - ) - cell = get_phonopy_structure( - structure.remove_site_property(property_name="magmom") - if "magmom" in structure.site_properties - else structure - ) - factor = get_factor(code) - - # a bit of code repetition here as I currently - # do not see how to pass the phonopy object? - if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": - primitive_matrix: np.ndarray | str = np.eye(3) - else: - primitive_matrix = "auto" - - # TARP: THIS IS BAD! Including for discussions sake - if cell.magnetic_moments is not None and primitive_matrix == "auto": - if np.any(cell.magnetic_moments != 0.0): - raise ValueError( - "For materials with magnetic moments specified " - "use_symmetrized_structure must be 'primitive'" - ) - cell.magnetic_moments = None - - phonon = Phonopy( - cell, - supercell_matrix, - primitive_matrix=primitive_matrix, - factor=factor, - symprec=symprec, - is_symmetry=sym_reduce, - ) - phonon.generate_displacements(distance=displacement) - - supercells = phonon.supercells_with_displacements - - displacements = [get_pmg_structure(cell) for cell in supercells] - - total_displacements = len(displacements) - recur_function_inputs = {"chunks_list": [], "index_list": []} - # chunks_list = [] - # index_list = [] - for chunk, start_index, end_index in chunks(displacements, chunk_size): - indexs = list(range(start_index + 1, end_index + 1, 1)) - recur_function_inputs["chunks_list"].append(chunk) - recur_function_inputs["index_list"].append(indexs) - # chunks_list.append(chunk) - # index_list.append(indexs) - - del supercells, displacements - - displ_job = run_phonon_displacements_recur( - total_displacements=total_displacements, - structure=structure, - supercell_matrix=supercell_matrix, - phonon_maker=phonon_maker, - chunks_list=recur_function_inputs["chunks_list"], - index_list=recur_function_inputs["index_list"], - start=0, - ) - - displacement_flow = Flow(displ_job, output=displ_job.output) - return Response( - replace=displacement_flow - ) # Response(addition=jobs, output=outputs) - - -@job(data=["forces", Structure, "chunks_list", "index_list"]) -def all_jobs( - displacement: float, - sym_reduce: bool, - symprec: float, - kpath_scheme: str, - code: str, - structure, - supercell_matrix, - use_symmetrized_structure: str | None, - phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None, - prev_dir: str | Path = None, - chunk_size=3, -): - warnings.warn( - "Initial magnetic moments will not be considered for the determination " - "of the symmetry of the structure and thus will be removed now.", - stacklevel=1, - ) - cell = get_phonopy_structure( - structure.remove_site_property(property_name="magmom") - if "magmom" in structure.site_properties - else structure - ) - factor = get_factor(code) - - # a bit of code repetition here as I currently - # do not see how to pass the phonopy object? - if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath": - primitive_matrix: np.ndarray | str = np.eye(3) - else: - primitive_matrix = "auto" - - # TARP: THIS IS BAD! Including for discussions sake - if cell.magnetic_moments is not None and primitive_matrix == "auto": - if np.any(cell.magnetic_moments != 0.0): - raise ValueError( - "For materials with magnetic moments specified " - "use_symmetrized_structure must be 'primitive'" - ) - cell.magnetic_moments = None - - phonon = Phonopy( - cell, - supercell_matrix, - primitive_matrix=primitive_matrix, - factor=factor, - symprec=symprec, - is_symmetry=sym_reduce, - ) - phonon.generate_displacements(distance=displacement) - - supercells = phonon.supercells_with_displacements - - displacements = [get_pmg_structure(cell) for cell in supercells] - - del supercells - - outputs: dict[str, list] = { - "displacement_number": [], - "forces": [], - "uuids": [], - "dirs": [], - } - - 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)}") - info = { - "displacement_number": idx, - "original_structure": structure, - "supercell_matrix": supercell_matrix, - "displaced_structures": displacement, - } - with contextlib.suppress(Exception): - phonon_job.update_maker_kwargs( - {"_set": {"write_additional_data->phonon_info:json": info}}, - dict_mod=True, - ) - outputs["displacement_number"].append(idx - 1) - outputs["uuids"].append(phonon_job.output.uuid) - outputs["dirs"].append(phonon_job.output.dir_name) - outputs["forces"].append(phonon_job.output.output.forces) - - displacement_flow = Flow(phonon_job) - return Response(addition=displacement_flow) From 3be6b52a4301a371acf98799bc3c09e363f3e73b Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 18:39:54 +0200 Subject: [PATCH 09/39] fix ase jobs tests --- src/atomate2/ase/schemas.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 060953f2fc..a1c6329cc8 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -105,16 +105,17 @@ class AseBaseModel(BaseModel): 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) - 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 + + 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): @@ -169,11 +170,11 @@ class OutputDoc(AseBaseModel): None, description="total number of steps needed in the relaxation." ) - all_forces: list[list[Vector3D]] | None = Field( + 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. Only present for batch calculations." + "or structures. Useful for batch calculations." ), ) @@ -417,8 +418,7 @@ def from_ase_compatible_result( task_document_kwargs : dict Additional keyword args passed to :obj:`.AseTaskDoc()`. """ - is_list = not isinstance(result, AseResult) - + is_list = isinstance(result, list) results = result if is_list else [result] output_mol_or_struct = [] @@ -558,6 +558,7 @@ def from_ase_compatible_result( 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( From fa03cf2c5bd0c4e9b52b7ff2c9013e2356b5ca66 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Wed, 7 May 2025 19:39:18 +0200 Subject: [PATCH 10/39] fi nearly all mypy errors --- src/atomate2/ase/jobs.py | 2 +- src/atomate2/ase/schemas.py | 44 ++++++++++++++--------------- src/atomate2/forcefields/schemas.py | 6 ++-- 3 files changed, 27 insertions(+), 25 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 9839004a0b..705eb8ea96 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -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. diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index a1c6329cc8..0033f72800 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -43,22 +43,22 @@ class AseResult(BaseModel): None, description="The molecule or structure in the final trajectory frame." ) - final_energy: float | list[float] | None = Field( + final_energy: float | None = Field( None, description="The final total energy from the calculation." ) - trajectory: PmgTrajectory | list[PmgTrajectory] | None = Field( + trajectory: PmgTrajectory | None = Field( None, description="The relaxation or molecular dynamics trajectory." ) - is_force_converged: bool | list[bool] | None = Field( + is_force_converged: bool | None = Field( None, description=( "Whether the calculation is converged with respect to interatomic forces." ), ) - energy_downhill: bool | list[bool] | None = Field( + energy_downhill: bool | None = Field( None, description=( "Whether the final trajectory frame has lower total " @@ -66,11 +66,11 @@ class AseResult(BaseModel): ), ) - dir_name: str | Path | list[str] | list[Path] | None = Field( + dir_name: str | Path | None = Field( None, description="The directory where the calculation was run" ) - elapsed_time: float | list[float] | None = Field( + elapsed_time: float | None = Field( None, description="The time taken to run the ASE calculation in seconds." ) @@ -105,19 +105,19 @@ class AseBaseModel(BaseModel): def model_post_init(self, _context: Any) -> None: """Establish alias to structure and molecule fields.""" - 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) + 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) + isinstance(self.mol_or_struct, Molecule) + or isinstance(self.mol_or_struct[0], Molecule) ): self.molecule = self.mol_or_struct - class IonicStep(AseBaseModel): """Document defining the information at each ionic step.""" @@ -170,7 +170,7 @@ class OutputDoc(AseBaseModel): None, description="total number of steps needed in the relaxation." ) - all_forces: list[list[Vector3D]] | None = Field( + all_forces: list[list[Vector3D]] | None = Field( None, description=( "The force on each atom in units of eV/A for the final molecules " @@ -419,14 +419,14 @@ def from_ase_compatible_result( Additional keyword args passed to :obj:`.AseTaskDoc()`. """ is_list = isinstance(result, list) - results = result if is_list else [result] + results: list[AseResult] = result if isinstance(result, list) else [result] - output_mol_or_struct = [] - input_mol_or_struct = [] - final_energy = [] + 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 = [] - ionic_steps = [] + final_stress: list[Matrix3D] = [] + ionic_steps: list[list[IonicStep]] | list[dict] = [] n_steps = [] objects = [] @@ -490,7 +490,7 @@ def from_ase_compatible_result( final_forces.append(trajectory.frame_properties[-1]["forces"]) final_stress.append(trajectory.frame_properties[-1].get("stress")) - ionic_steps_structure = [] + 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 = { @@ -528,8 +528,8 @@ def from_ase_compatible_result( **_ionic_step_data, ) - ionic_steps_structure.append(ionic_step) - ionic_steps.append(ionic_steps_structure) + ionic_steps_individual.append(ionic_step) + ionic_steps.append(ionic_steps_individual) objects_structure: dict[AseObject, Any] = {} if store_trajectory != StoreTrajectoryOption.NO: @@ -587,8 +587,8 @@ def from_ase_compatible_result( mol_or_struct=output_mol_or_struct, energy=final_energy, energy_per_atom=[ - final_energy_here / len(output_mol_or_struct_here) - for final_energy_here, output_mol_or_struct_here in zip( + 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 ) ], diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index bf9233ece0..1e9b359915 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -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 From e3805666dfaa8f68ceec172b8fa16748e74db6cd Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 07:30:43 +0200 Subject: [PATCH 11/39] fix some more style problems and linting errors --- src/atomate2/ase/schemas.py | 11 ++++--- src/atomate2/ase/utils.py | 44 ++++++++++++-------------- src/atomate2/common/schemas/phonons.py | 11 ++++--- 3 files changed, 32 insertions(+), 34 deletions(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 0033f72800..f75a730ec7 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -430,8 +430,8 @@ def from_ase_compatible_result( n_steps = [] objects = [] - for result in results: - trajectory = result.trajectory + for result_item in results: + trajectory = result_item.trajectory # TODO: fix this n_steps_here = None @@ -477,10 +477,10 @@ def from_ase_compatible_result( ) output_mol_or_struct.append(input_mol_or_struct_here) else: - output_mol_or_struct.append(result.final_mol_or_struct) + output_mol_or_struct.append(result_item.final_mol_or_struct) if trajectory is None: - final_energy.append(result.final_energy) + final_energy.append(result_item.final_energy) final_forces.append(None) final_stress.append(None) ionic_steps.append(None) @@ -511,7 +511,8 @@ def from_ase_compatible_result( else None ) - # include "magmoms" in `ionic_step` if the trajectory has "magmoms" + # include "magmoms" in `ionic_step` + # if the trajectory has "magmoms" if "magmoms" in trajectory.frame_properties[idx]: _ionic_step_data.update( { diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index 9d0de17b72..36bbfe8bed 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -343,7 +343,8 @@ def relax( Parameters ---------- - atoms : ASE Atoms, pymatgen .Structure, or pymatgen .Molecule or corresponding lists + 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,33 +361,28 @@ def relax( Returns ------- - dict including optimized structure and the trajectory or a list of those dicts + dict including optimized structure and the trajectory + or a list of those dicts """ - is_list = ( - isinstance(atoms[0], Atoms) - or isinstance(atoms[0], Molecule) - or isinstance(atoms[0], Structure) - ) + is_list = isinstance(atoms[0], (Atoms, Molecule, Structure)) + + list_atoms: list[Atoms] = [atoms] if not is_list else atoms - if not is_list: - list_atoms = [atoms] - else: - list_atoms = atoms list_ase_results = [] - for atoms in list_atoms: - is_mol = isinstance(atoms, Molecule) or ( - isinstance(atoms, Atoms) and all(not pbc for pbc in atoms.pbc) + for atoms_item in list_atoms: + is_mol = isinstance(atoms_item, Molecule) or ( + isinstance(atoms_item, Atoms) and all(not pbc for pbc in atoms_item.pbc) ) - if isinstance(atoms, Structure | Molecule): - atoms = self.ase_adaptor.get_atoms(atoms) + if isinstance(atoms_item, Structure | Molecule): + atoms_item = self.ase_adaptor.get_atoms(atoms_item) if self.fix_symmetry: - atoms.set_constraint(FixSymmetry(atoms, symprec=self.symprec)) - atoms.calc = self.calculator + 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) + obs = TrajectoryObserver(atoms_item) if self.relax_cell and (not is_mol): - atoms = cell_filter(atoms) - optimizer = self.opt_class(atoms, **kwargs) + 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) @@ -394,11 +390,11 @@ def relax( obs() if traj_file is not None: obs.save(traj_file) - if isinstance(atoms, cell_filter): - atoms = atoms.atoms + if isinstance(atoms_item, cell_filter): + atoms_item = atoms_item.atoms struct = self.ase_adaptor.get_structure( - atoms, cls=Molecule if is_mol else Structure + atoms_item, cls=Molecule if is_mol else Structure ) traj = obs.to_pymatgen_trajectory(None) is_force_conv = all( diff --git a/src/atomate2/common/schemas/phonons.py b/src/atomate2/common/schemas/phonons.py index 3fbd70d3e0..ee977d7b41 100644 --- a/src/atomate2/common/schemas/phonons.py +++ b/src/atomate2/common/schemas/phonons.py @@ -527,11 +527,12 @@ def from_forces_born( volume_per_formula_unit = structure.volume / formula_units - if displacement_data["dirs"] is not None: - if isinstance(displacement_data["dirs"][0], list): - displacement_data["dirs"] = [ - dir for dir_list in displacement_data["dirs"] for dir in dir_list - ] + 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, From b3de466de17382d4eccc1440b5de5b985f4db471 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 07:36:15 +0200 Subject: [PATCH 12/39] avoid overwriting of a for loop variable --- src/atomate2/ase/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index 36bbfe8bed..0ad20852b3 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -364,12 +364,13 @@ def relax( dict including optimized structure and the trajectory or a list of those dicts """ - is_list = isinstance(atoms[0], (Atoms, Molecule, Structure)) + is_list = isinstance(atoms[0], (Atoms | Molecule | Structure)) list_atoms: list[Atoms] = [atoms] if not is_list else atoms list_ase_results = [] - for atoms_item in list_atoms: + for atoms_item_start in list_atoms: + atoms_item = atoms_item_start is_mol = isinstance(atoms_item, Molecule) or ( isinstance(atoms_item, Atoms) and all(not pbc for pbc in atoms_item.pbc) ) From 5a814b2885a7c6a50537e72d39716f7bef0a651d Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 07:45:28 +0200 Subject: [PATCH 13/39] fix hopefully last linting error --- src/atomate2/ase/schemas.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index f75a730ec7..e18e12e231 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 @@ -426,7 +426,7 @@ def from_ase_compatible_result( final_energy: list[float] = [] final_forces: list[Vector3D] | list[list[Vector3D]] = [] final_stress: list[Matrix3D] = [] - ionic_steps: list[list[IonicStep]] | list[dict] = [] + ionic_steps: list[Optional[list[IonicStep]]] = [] n_steps = [] objects = [] From 8399a66439f6ef31afba31dd4eb7713d2d78ef94 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 08:02:31 +0200 Subject: [PATCH 14/39] add another test for the forcefields --- tests/forcefields/test_jobs.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index dc68c160ee..feb2ef16c9 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -133,6 +133,33 @@ def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): CHGNetRelaxMaker() +def test_chgnet_batch_static_maker(si_structure: Structure): + # 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) + + # 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 From 86b7d1b1f9760fc6bb10dd75f78bceebe9713174 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 08:14:42 +0200 Subject: [PATCH 15/39] add uncommented update back and add another test --- src/atomate2/common/jobs/phonons.py | 10 +++++---- tests/forcefields/flows/test_phonon.py | 31 ++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index b60c204b10..5a7f6d350c 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 @@ -293,10 +294,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/tests/forcefields/flows/test_phonon.py b/tests/forcefields/flows/test_phonon.py index 9295c4c066..53caa214bd 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -105,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( + 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": (filename_bs := f"{tmp_path}/phonon_bs_test.png"), + "filename_dos": (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) From 3e896ef75fcec7f8c16469ce06ade8a165d571a6 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 08:28:35 +0200 Subject: [PATCH 16/39] new linting fixes --- tests/forcefields/flows/test_phonon.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/forcefields/flows/test_phonon.py b/tests/forcefields/flows/test_phonon.py index 53caa214bd..42e96dd06a 100644 --- a/tests/forcefields/flows/test_phonon.py +++ b/tests/forcefields/flows/test_phonon.py @@ -108,7 +108,7 @@ def test_phonon_wf_force_field( @pytest.mark.parametrize("socket", [True, False]) -def test_phonon_wf_force_field( +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 @@ -121,8 +121,8 @@ def test_phonon_wf_force_field( prefer_90_degrees=False, generate_frequencies_eigenvectors_kwargs={ "tstep": 100, - "filename_bs": (filename_bs := f"{tmp_path}/phonon_bs_test.png"), - "filename_dos": (filename_dos := f"{tmp_path}/phonon_dos_test.pdf"), + "filename_bs": f"{tmp_path}/phonon_bs_test.png", + "filename_dos": f"{tmp_path}/phonon_dos_test.pdf", }, ) From eea7efac5cfc051dedf4fb81ab6d44dd9f72c94e Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 09:21:06 +0200 Subject: [PATCH 17/39] fix n_steps error --- src/atomate2/ase/schemas.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index e18e12e231..3a5e0efbe6 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -458,7 +458,6 @@ def from_ase_compatible_result( input_mol_or_struct_here = trajectory[0] input_mol_or_struct.append(input_mol_or_struct_here) - n_steps.append(n_steps_here) # Workaround for cases where the ASE optimizer does not correctly limit the # number of steps for static calculations. @@ -479,6 +478,8 @@ def from_ase_compatible_result( 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) From 0f2586a0d7fef9c12a22ae2f4aaffed17fa78415 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 13:15:07 +0200 Subject: [PATCH 18/39] move OutputDoc into datastore --- src/atomate2/ase/jobs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 705eb8ea96..000d4dd2f4 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -15,7 +15,7 @@ from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor -from atomate2.ase.schemas import AseResult, AseTaskDoc +from atomate2.ase.schemas import AseResult, AseTaskDoc, OutputDoc from atomate2.ase.utils import AseRelaxer logger = logging.getLogger(__name__) @@ -27,7 +27,7 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc -_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory] +_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, OutputDoc] @dataclass From 697a07bbe6fa7be826c6d6535c4ad52ad5354f53 Mon Sep 17 00:00:00 2001 From: "J. George" Date: Thu, 8 May 2025 13:42:08 +0200 Subject: [PATCH 19/39] Update pyproject.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 5dd8717c9a..eb303ccfb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,7 @@ docs = [ "sphinx==8.1.3", "sphinx_design==0.6.1", "jupyterlab==4.4.1", + "snowballstemmer==2.2.0", # release from 2021 ] dev = ["pre-commit>=2.12.1"] tests = [ From 3ecb448ffb23e285edd54bf72abdffc0e52a2bca Mon Sep 17 00:00:00 2001 From: "J. George" Date: Thu, 8 May 2025 14:25:31 +0200 Subject: [PATCH 20/39] Update pyproject.toml --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index eb303ccfb9..5dd8717c9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,7 +82,6 @@ docs = [ "sphinx==8.1.3", "sphinx_design==0.6.1", "jupyterlab==4.4.1", - "snowballstemmer==2.2.0", # release from 2021 ] dev = ["pre-commit>=2.12.1"] tests = [ From a158db73a4a3ebcc57c3dc9fc177fe8fa92d66c8 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 8 May 2025 21:50:11 +0200 Subject: [PATCH 21/39] move InputDoc to Datastore --- src/atomate2/ase/jobs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 000d4dd2f4..ae21189788 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -15,7 +15,7 @@ from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor -from atomate2.ase.schemas import AseResult, AseTaskDoc, OutputDoc +from atomate2.ase.schemas import AseResult, AseTaskDoc, OutputDoc, InputDoc from atomate2.ase.utils import AseRelaxer logger = logging.getLogger(__name__) @@ -27,7 +27,7 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc -_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, OutputDoc] +_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, OutputDoc, InputDoc] @dataclass From 9965ba39502e8e45d3369c4c46cb034a5464d7b9 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Fri, 9 May 2025 07:25:18 +0200 Subject: [PATCH 22/39] fix order issue that does not show up locally --- src/atomate2/ase/jobs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index ae21189788..30b174a93b 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -15,7 +15,7 @@ from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor -from atomate2.ase.schemas import AseResult, AseTaskDoc, OutputDoc, InputDoc +from atomate2.ase.schemas import AseResult, AseTaskDoc, InputDoc, OutputDoc from atomate2.ase.utils import AseRelaxer logger = logging.getLogger(__name__) @@ -27,7 +27,7 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc -_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, OutputDoc, InputDoc] +_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, InputDoc, OutputDoc] @dataclass From fe384939b962fc52de42a8ce6ca41a3a9408f51e Mon Sep 17 00:00:00 2001 From: JaGeo Date: Fri, 9 May 2025 14:15:09 +0200 Subject: [PATCH 23/39] fix data storage --- src/atomate2/ase/jobs.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 30b174a93b..f16df49cb8 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -27,7 +27,14 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc -_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, InputDoc, OutputDoc] +_ASE_DATA_OBJECTS = [ + PmgTrajectory, + AseTrajectory, + InputDoc, + OutputDoc, + "is_force_converged", + "energy_downhill", +] @dataclass From b2061f39b83ed006d6cd22070b2b313fd9612535 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 12 May 2025 07:09:24 +0200 Subject: [PATCH 24/39] update data --- src/atomate2/forcefields/jobs.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 63bdf23b45..da27f95bba 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -14,6 +14,7 @@ from pymatgen.core.trajectory import Trajectory as PmgTrajectory from atomate2.ase.jobs import AseRelaxMaker +from atomate2.ase.schemas import InputDoc, OutputDoc from atomate2.forcefields import MLFF, _get_formatted_ff_name from atomate2.forcefields.schemas import ForceFieldTaskDocument from atomate2.forcefields.utils import ase_calculator, revert_default_dtype @@ -27,7 +28,15 @@ logger = logging.getLogger(__name__) -_FORCEFIELD_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, "ionic_steps"] +_FORCEFIELD_DATA_OBJECTS = [ + PmgTrajectory, + AseTrajectory, + "ionic_steps", + InputDoc, + OutputDoc, + "is_force_converged", + "energy_downhill", +] _DEFAULT_CALCULATOR_KWARGS = { MLFF.CHGNet: {"stress_weight": _GPa_to_eV_per_A3}, From 452f74709b4b86ea65be2340871266db7deea2e0 Mon Sep 17 00:00:00 2001 From: jgeorge Date: Thu, 15 May 2025 12:55:34 +0200 Subject: [PATCH 25/39] fix output doc --- src/atomate2/forcefields/jobs.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index da27f95bba..3d16223189 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -36,7 +36,15 @@ OutputDoc, "is_force_converged", "energy_downhill", -] + "forces", + "input", + "energy", + "stress", + "energy_per_atom", + "elapsed_time", + "n_steps", + "all_forces" + ] _DEFAULT_CALCULATOR_KWARGS = { MLFF.CHGNet: {"stress_weight": _GPa_to_eV_per_A3}, From 73adc85cf5ce73c306cd83ed5983541e55a60033 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Fri, 16 May 2025 06:51:09 +0200 Subject: [PATCH 26/39] change data store --- src/atomate2/forcefields/jobs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 3d16223189..6bd6ef3864 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -40,11 +40,11 @@ "input", "energy", "stress", - "energy_per_atom", + "energy_per_atom", "elapsed_time", "n_steps", - "all_forces" - ] + "all_forces", +] _DEFAULT_CALCULATOR_KWARGS = { MLFF.CHGNet: {"stress_weight": _GPa_to_eV_per_A3}, From 2af321b0df0941179a74847eaf4329c920051979 Mon Sep 17 00:00:00 2001 From: jgeorge Date: Wed, 21 May 2025 18:16:47 +0200 Subject: [PATCH 27/39] test out more with regard to the datastore --- src/atomate2/forcefields/jobs.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 3d16223189..552ade5247 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -18,6 +18,8 @@ from atomate2.forcefields import MLFF, _get_formatted_ff_name from atomate2.forcefields.schemas import ForceFieldTaskDocument from atomate2.forcefields.utils import ase_calculator, revert_default_dtype +from pymatgen.core.structure import Structure + if TYPE_CHECKING: from collections.abc import Callable @@ -43,7 +45,12 @@ "energy_per_atom", "elapsed_time", "n_steps", - "all_forces" + "all_forces", + "dir_name", + "objects", + "mol_or_struct", + "structure", + "output", ] _DEFAULT_CALCULATOR_KWARGS = { From 86b4c7eeb694b63aa5683b960acd8ba949708094 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 22 May 2025 08:49:57 +0200 Subject: [PATCH 28/39] fix some other points --- src/atomate2/ase/schemas.py | 24 +++++++++++++++++++++++- src/atomate2/common/jobs/phonons.py | 2 +- src/atomate2/forcefields/jobs.py | 1 + tests/forcefields/test_jobs.py | 13 ++++++++++--- 4 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 3a5e0efbe6..20f1ddef8e 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -12,7 +12,7 @@ from pathlib import Path from typing import Any, Optional - +from ase.io import Trajectory as AseTrajectory from ase.stress import voigt_6_to_full_3x3_stress from ase.units import GPa from emmet.core.math import Matrix3D, Vector3D @@ -22,6 +22,7 @@ from pydantic import BaseModel, Field from pymatgen.core import Molecule, Structure from pymatgen.core.trajectory import Trajectory as PmgTrajectory +from atomate2.ase.schemas import InputDoc, OutputDoc _task_doc_translation_keys = { "input", @@ -33,6 +34,27 @@ "is_force_converged", "energy_downhill", "tags", + PmgTrajectory, + AseTrajectory, + "ionic_steps", + InputDoc, + OutputDoc, + "is_force_converged", + "energy_downhill", + "forces", + "input", + "energy", + "stress", + "energy_per_atom", + "elapsed_time", + "n_steps", + "all_forces", + "dir_name", + "objects", + "mol_or_struct", + "structure", + "output", + Structure, } diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index 5a7f6d350c..68bb97d2d0 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -243,7 +243,7 @@ def generate_frequencies_eigenvectors( ) -@job(data=["forces", "displaced_structures"]) +@job(data=["forces", "displaced_structures", "uuids", "dirs"]) def run_phonon_displacements( displacements: list[Structure], structure: Structure, diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index ca1f6794a6..1361abeae8 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -51,6 +51,7 @@ "mol_or_struct", "structure", "output", + Structure ] diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index feb2ef16c9..2c49f30154 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -133,7 +133,7 @@ def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): CHGNetRelaxMaker() -def test_chgnet_batch_static_maker(si_structure: Structure): +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]) @@ -145,8 +145,8 @@ def test_chgnet_batch_static_maker(si_structure: Structure): ).make([si_structure, si_structure2]) # run the flow or job and ensure that it finished running successfully - responses = run_locally(job, ensure_success=True) - + responses = run_locally(job, ensure_success=True, store=memory_jobstore) + print(responses) # validate job outputs output1 = responses[job.uuid][1].output assert isinstance(output1, ForceFieldTaskDocument) @@ -158,6 +158,13 @@ def test_chgnet_batch_static_maker(si_structure: Structure): # 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() + result = memory_jobstore.query_one( + {"name": "Force field static"}, + load=False, + sort={"completed_at": -1} # to get the latest computation + ) + + print(result) @pytest.mark.skip(reason="M3GNet requires DGL which is PyTorch 2.4 incompatible") From 38f94ecd61fc1df2ba09dcf63eb251d968f0ea63 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 22 May 2025 09:24:06 +0200 Subject: [PATCH 29/39] fix some other points --- src/atomate2/ase/schemas.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 20f1ddef8e..9459afa1c1 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -22,7 +22,6 @@ from pydantic import BaseModel, Field from pymatgen.core import Molecule, Structure from pymatgen.core.trajectory import Trajectory as PmgTrajectory -from atomate2.ase.schemas import InputDoc, OutputDoc _task_doc_translation_keys = { "input", From d672fe7908212854dbfd9b252b71ff4bbaa7a856 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 22 May 2025 09:26:17 +0200 Subject: [PATCH 30/39] fix some other points --- src/atomate2/ase/schemas.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 9459afa1c1..6aaaa5424d 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -23,6 +23,7 @@ from pymatgen.core import Molecule, Structure from pymatgen.core.trajectory import Trajectory as PmgTrajectory + _task_doc_translation_keys = { "input", "output", @@ -36,8 +37,6 @@ PmgTrajectory, AseTrajectory, "ionic_steps", - InputDoc, - OutputDoc, "is_force_converged", "energy_downhill", "forces", From 600857fa497eab2d6ab46023a1a4abb14ce7f60f Mon Sep 17 00:00:00 2001 From: JaGeo Date: Thu, 22 May 2025 09:52:36 +0200 Subject: [PATCH 31/39] fix mix up between datastore and other ase objects --- src/atomate2/ase/jobs.py | 15 +++++++++++++++ src/atomate2/ase/schemas.py | 19 ------------------- tests/ase/test_jobs.py | 4 ++-- 3 files changed, 17 insertions(+), 21 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index f16df49cb8..f96774cb00 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -30,10 +30,25 @@ _ASE_DATA_OBJECTS = [ PmgTrajectory, AseTrajectory, + "ionic_steps", InputDoc, OutputDoc, "is_force_converged", "energy_downhill", + "forces", + "input", + "energy", + "stress", + "energy_per_atom", + "elapsed_time", + "n_steps", + "all_forces", + "dir_name", + "objects", + "mol_or_struct", + "structure", + "output", + Structure ] diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 6aaaa5424d..0a3d26dc32 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -34,25 +34,6 @@ "is_force_converged", "energy_downhill", "tags", - PmgTrajectory, - AseTrajectory, - "ionic_steps", - "is_force_converged", - "energy_downhill", - "forces", - "input", - "energy", - "stress", - "energy_per_atom", - "elapsed_time", - "n_steps", - "all_forces", - "dir_name", - "objects", - "mol_or_struct", - "structure", - "output", - Structure, } diff --git a/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index 43e04f15b8..3f21cc5d7d 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -47,12 +47,12 @@ def test_base_maker(test_dir): assert isinstance(output, AseStructureTaskDoc) -def test_lennard_jones_batch_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) + response = run_locally(job, store=memory_jobstore) output = response[job.uuid][1].output From 3d5d2d2efaae0a2bf847a13350714fc85e373cee Mon Sep 17 00:00:00 2001 From: JaGeo Date: Fri, 23 May 2025 09:17:27 +0200 Subject: [PATCH 32/39] add more to datastore --- src/atomate2/common/jobs/phonons.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index 68bb97d2d0..a2877934e4 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -175,7 +175,7 @@ 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, @@ -243,7 +243,7 @@ def generate_frequencies_eigenvectors( ) -@job(data=["forces", "displaced_structures", "uuids", "dirs"]) +@job(data=["forces", "displaced_structures", "uuids", "dirs", "displacement_number"]) def run_phonon_displacements( displacements: list[Structure], structure: Structure, From b10fb0a44ff411246840dabb5740d58e45910db0 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 21:54:52 +0200 Subject: [PATCH 33/39] more fixes --- src/atomate2/ase/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index dfa0759665..298a398dfb 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -381,6 +381,7 @@ def relax( ) 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 @@ -398,7 +399,7 @@ def relax( obs.save(traj_file) if final_atoms_object_file is not None: if steps <= 1: - write_atoms = atoms_item_start + write_atoms = atoms_item_start_0 write_atoms.calc = self.calculator else: write_atoms = atoms_item From 38ad652b22e2c872b6da766215e6d9d85209d415 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 22:04:12 +0200 Subject: [PATCH 34/39] fix conflicts --- src/atomate2/ase/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index 298a398dfb..e0756106a2 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -403,8 +403,12 @@ def relax( 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 From 2ee706708a327d003b796e2e40b5e5c53695aa77 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 22:06:55 +0200 Subject: [PATCH 35/39] simplify --- src/atomate2/forcefields/jobs.py | 22 +--------------------- 1 file changed, 1 insertion(+), 21 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 1361abeae8..fe4fb81740 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -31,28 +31,8 @@ logger = logging.getLogger(__name__) _FORCEFIELD_DATA_OBJECTS = [ - PmgTrajectory, - AseTrajectory, - "ionic_steps", - InputDoc, - OutputDoc, - "is_force_converged", - "energy_downhill", - "forces", - "input", - "energy", - "stress", - "energy_per_atom", - "elapsed_time", - "n_steps", - "all_forces", - "dir_name", - "objects", - "mol_or_struct", - "structure", "output", - Structure -] + ] _DEFAULT_CALCULATOR_KWARGS = { From e3c813035fba83cded7dfbdcc9500a38de36e5ee Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 22:08:39 +0200 Subject: [PATCH 36/39] simplify --- src/atomate2/ase/jobs.py | 23 ++--------------------- src/atomate2/forcefields/jobs.py | 2 +- 2 files changed, 3 insertions(+), 22 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index f96774cb00..600bea5447 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -28,27 +28,8 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc _ASE_DATA_OBJECTS = [ - PmgTrajectory, - AseTrajectory, - "ionic_steps", - InputDoc, - OutputDoc, - "is_force_converged", - "energy_downhill", - "forces", - "input", - "energy", - "stress", - "energy_per_atom", - "elapsed_time", - "n_steps", - "all_forces", - "dir_name", - "objects", - "mol_or_struct", - "structure", - "output", - Structure + "output", # will put everything in data store + ] diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index fe4fb81740..0db13ce8c5 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -31,7 +31,7 @@ logger = logging.getLogger(__name__) _FORCEFIELD_DATA_OBJECTS = [ - "output", + "output", # will put everything in the data store ] From db4449bda01f17c316cad039ee881a0123c135ea Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 22:17:29 +0200 Subject: [PATCH 37/39] fix linting --- src/atomate2/ase/jobs.py | 7 ++----- src/atomate2/ase/schemas.py | 3 +-- src/atomate2/ase/utils.py | 8 ++++---- src/atomate2/common/jobs/phonons.py | 9 ++++++++- src/atomate2/forcefields/jobs.py | 10 +++------- tests/ase/test_jobs.py | 4 +++- tests/forcefields/test_jobs.py | 2 +- 7 files changed, 22 insertions(+), 21 deletions(-) diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 600bea5447..0dae949244 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -8,14 +8,12 @@ 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, InputDoc, OutputDoc +from atomate2.ase.schemas import AseResult, AseTaskDoc from atomate2.ase.utils import AseRelaxer logger = logging.getLogger(__name__) @@ -28,8 +26,7 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc _ASE_DATA_OBJECTS = [ - "output", # will put everything in data store - + "output", # will put everything in data store ] diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 0a3d26dc32..3a5e0efbe6 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -12,7 +12,7 @@ from pathlib import Path from typing import Any, Optional -from ase.io import Trajectory as AseTrajectory + from ase.stress import voigt_6_to_full_3x3_stress from ase.units import GPa from emmet.core.math import Matrix3D, Vector3D @@ -23,7 +23,6 @@ from pymatgen.core import Molecule, Structure from pymatgen.core.trajectory import Trajectory as PmgTrajectory - _task_doc_translation_keys = { "input", "output", diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index e0756106a2..6afc949ea6 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -372,7 +372,6 @@ def relax( list_atoms: list[Atoms] = [atoms] if not is_list else atoms - list_ase_results = [] for atoms_item_start in list_atoms: atoms_item = atoms_item_start.copy() @@ -381,7 +380,7 @@ def relax( ) if isinstance(atoms_item, Structure | Molecule): atoms_item = self.ase_adaptor.get_atoms(atoms_item) - atoms_item_start_0=atoms_item.copy() + 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 @@ -406,8 +405,9 @@ def relax( if isinstance(write_atoms, cell_filter): write_atoms = write_atoms.atoms - write(final_atoms_object_file, write_atoms, format="extxyz", append=True) - + write( + final_atoms_object_file, write_atoms, format="extxyz", append=True + ) if isinstance(atoms_item, cell_filter): atoms_item = atoms_item.atoms diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index a2877934e4..900a362968 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -175,7 +175,14 @@ def generate_phonon_displacements( @job( output_schema=PhononBSDOSDoc, - data=[PhononDos, PhononBandStructureSymmLine, ForceConstants, Structure, "jobdirs", "uuids"], + data=[ + PhononDos, + PhononBandStructureSymmLine, + ForceConstants, + Structure, + "jobdirs", + "uuids", + ], ) def generate_frequencies_eigenvectors( structure: Structure, diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 0db13ce8c5..2abf132af6 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -7,19 +7,15 @@ 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.ase.schemas import InputDoc, OutputDoc from atomate2.forcefields import MLFF, _get_formatted_ff_name from atomate2.forcefields.schemas import ForceFieldTaskDocument from atomate2.forcefields.utils import ase_calculator, revert_default_dtype -from pymatgen.core.structure import Structure - if TYPE_CHECKING: from collections.abc import Callable @@ -31,8 +27,8 @@ logger = logging.getLogger(__name__) _FORCEFIELD_DATA_OBJECTS = [ - "output", # will put everything in the data store - ] + "output", # will put everything in the data store +] _DEFAULT_CALCULATOR_KWARGS = { diff --git a/tests/ase/test_jobs.py b/tests/ase/test_jobs.py index 3f21cc5d7d..d624d95f20 100644 --- a/tests/ase/test_jobs.py +++ b/tests/ase/test_jobs.py @@ -47,7 +47,9 @@ 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): +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]) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 97e8c2cd37..2783bd21ed 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -161,7 +161,7 @@ def test_chgnet_batch_static_maker(si_structure: Structure, memory_jobstore): result = memory_jobstore.query_one( {"name": "Force field static"}, load=False, - sort={"completed_at": -1} # to get the latest computation + sort={"completed_at": -1}, # to get the latest computation ) print(result) From c657a0d0fe5cd848957187b01e3288cc91154f30 Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 22:31:38 +0200 Subject: [PATCH 38/39] remove print --- tests/forcefields/test_jobs.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 2783bd21ed..928cd2e8f7 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -146,7 +146,6 @@ def test_chgnet_batch_static_maker(si_structure: Structure, memory_jobstore): # run the flow or job and ensure that it finished running successfully responses = run_locally(job, ensure_success=True, store=memory_jobstore) - print(responses) # validate job outputs output1 = responses[job.uuid][1].output assert isinstance(output1, ForceFieldTaskDocument) @@ -164,8 +163,6 @@ def test_chgnet_batch_static_maker(si_structure: Structure, memory_jobstore): sort={"completed_at": -1}, # to get the latest computation ) - print(result) - @pytest.mark.skip(reason="M3GNet requires DGL which is PyTorch 2.4 incompatible") def test_m3gnet_static_maker(si_structure): From bfde3b69bb758d727c4952e617d8d53ead71a8bd Mon Sep 17 00:00:00 2001 From: JaGeo Date: Mon, 9 Jun 2025 22:42:26 +0200 Subject: [PATCH 39/39] remove print --- tests/forcefields/test_jobs.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 928cd2e8f7..a1eb5cb82c 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -157,11 +157,6 @@ def test_chgnet_batch_static_maker(si_structure: Structure, memory_jobstore): # 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() - result = memory_jobstore.query_one( - {"name": "Force field static"}, - load=False, - sort={"completed_at": -1}, # to get the latest computation - ) @pytest.mark.skip(reason="M3GNet requires DGL which is PyTorch 2.4 incompatible")