From 3c44175c8b2b49859f314a7309dc96c5a8277714 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Wed, 6 Nov 2024 14:49:21 -0800 Subject: [PATCH] lint --- emmet-core/emmet/core/neb.py | 92 +++++++++++++++++++++--------------- emmet-core/tests/test_neb.py | 29 ++++++++---- 2 files changed, 74 insertions(+), 47 deletions(-) diff --git a/emmet-core/emmet/core/neb.py b/emmet-core/emmet/core/neb.py index 1a14ff9435..fcd8546096 100644 --- a/emmet-core/emmet/core/neb.py +++ b/emmet-core/emmet/core/neb.py @@ -43,18 +43,16 @@ class NebTaskDoc(BaseModel, extra="allow"): None, description="The initial and final configurations (reactants and products) of the barrier.", ) - endpoint_energies : Optional[Sequence[float]] = Field( - None, - description="Energies of the endpoint structures." + endpoint_energies: Optional[Sequence[float]] = Field( + None, description="Energies of the endpoint structures." ) - endpoint_calculations : Optional[list[Calculation]] = Field( - None, - description = "Calculation information for the endpoint structures" + endpoint_calculations: Optional[list[Calculation]] = Field( + None, description="Calculation information for the endpoint structures" ) - endpoint_objects : Optional[list[dict]] = Field( + endpoint_objects: Optional[list[dict]] = Field( None, description="VASP objects for each endpoint calculation." ) - endpoint_directories : Optional[list[str]] = Field( + endpoint_directories: Optional[list[str]] = Field( None, description="List of the directories for the endpoint calculations." ) @@ -140,25 +138,33 @@ def set_barriers(self) -> Self: def num_images(self) -> int: """Return the number of VASP calculations / number of images performed.""" return len(self.image_directories) - + @property def energies(self) -> list[float]: """Return the endpoint (optional) and image energies.""" if self.endpoint_energies is not None: - return [self.endpoint_energies[0], *self.image_energies, self.endpoint_energies[1]] + return [ + self.endpoint_energies[0], + *self.image_energies, + self.endpoint_energies[1], + ] return self.image_energies @property def structures(self) -> list[Structure]: """Return the endpoint and image structures.""" - return [self.endpoint_structures[0], *self.image_structures, self.endpoint_structures[1]] - + return [ + self.endpoint_structures[0], + *self.image_structures, + self.endpoint_structures[1], + ] + @classmethod def from_directory( cls, dir_name: Union[Path, str], volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES, - store_calculations : bool = True, + store_calculations: bool = True, **neb_task_doc_kwargs, ) -> Self: """ @@ -172,7 +178,7 @@ def from_directory( neb_directories = sorted(dir_name.glob("[0-9][0-9]")) - if (ep_calcs := neb_task_doc_kwargs.pop("endpoint_calculations", None) ) is None: + if (ep_calcs := neb_task_doc_kwargs.pop("endpoint_calculations", None)) is None: endpoint_directories = [neb_directories[0], neb_directories[-1]] endpoint_structures = [ Structure.from_file(zpath(f"{endpoint_dir}/POSCAR")) @@ -181,12 +187,8 @@ def from_directory( endpoint_energies = None else: endpoint_directories = neb_task_doc_kwargs.pop("endpoint_directories") - endpoint_structures = [ - ep_calc.output.structure for ep_calc in ep_calcs - ] - endpoint_energies = [ - ep_calc.output.energy for ep_calc in ep_calcs - ] + endpoint_structures = [ep_calc.output.structure for ep_calc in ep_calcs] + endpoint_energies = [ep_calc.output.energy for ep_calc in ep_calcs] image_directories = neb_directories[1:-1] @@ -216,8 +218,7 @@ def from_directory( task_state = ( TaskState.SUCCESS if all( - calc.has_vasp_completed == TaskState.SUCCESS - for calc in calcs_to_check + calc.has_vasp_completed == TaskState.SUCCESS for calc in calcs_to_check ) else TaskState.FAILED ) @@ -247,11 +248,11 @@ def from_directory( return cls( endpoint_structures=endpoint_structures, - endpoint_energies = endpoint_energies, - endpoint_directories = [str(ep_dir) for ep_dir in endpoint_directories], - endpoint_calculations = ep_calcs if store_calculations else None, + endpoint_energies=endpoint_energies, + endpoint_directories=[str(ep_dir) for ep_dir in endpoint_directories], + endpoint_calculations=ep_calcs if store_calculations else None, image_calculations=image_calculations if store_calculations else None, - image_structures = image_structures, + image_structures=image_structures, dir_name=str(dir_name), image_directories=[str(img_dir) for img_dir in image_directories], orig_inputs=inputs["orig_inputs"], @@ -271,7 +272,7 @@ def from_directories( endpoint_directories: list[str | Path], neb_directory: str | Path, volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES, - **neb_task_doc_kwargs + **neb_task_doc_kwargs, ) -> Self: """ Return an NebTaskDoc from endpoint and NEB calculation directories. @@ -282,12 +283,26 @@ def from_directories( endpoint_calculations = [None for _ in range(2)] endpoint_objects = [None for _ in range(2)] for idx, endpoint_dir in enumerate(endpoint_directories): - vasp_files = _find_vasp_files(endpoint_dir, volumetric_files=volumetric_files) - ep_key = "standard" if vasp_files.get("standard") else "relax" + str(max( - int(k.split("relax")[-1]) for k in vasp_files if k.startswith("relax") - )) + vasp_files = _find_vasp_files( + endpoint_dir, volumetric_files=volumetric_files + ) + ep_key = ( + "standard" + if vasp_files.get("standard") + else "relax" + + str( + max( + int(k.split("relax")[-1]) + for k in vasp_files + if k.startswith("relax") + ) + ) + ) - endpoint_calculations[idx], endpoint_objects[idx] = Calculation.from_vasp_files( + ( + endpoint_calculations[idx], + endpoint_objects[idx], + ) = Calculation.from_vasp_files( dir_name=endpoint_dir, task_name=f"NEB endpoint {idx + 1}", vasprun_file=vasp_files[ep_key]["vasprun_file"], @@ -299,16 +314,17 @@ def from_directories( "parse_potcar_file": False, }, ) - + return cls.from_directory( neb_directory, volumetric_files=volumetric_files, - endpoint_calculations = endpoint_calculations, - endpoint_objects = endpoint_objects, - endpoint_directories = endpoint_directories, - **neb_task_doc_kwargs + endpoint_calculations=endpoint_calculations, + endpoint_objects=endpoint_objects, + endpoint_directories=endpoint_directories, + **neb_task_doc_kwargs, ) - + + def neb_barrier_spline_fit( energies: Sequence[float], spline_kwargs: dict | None = None, diff --git a/emmet-core/tests/test_neb.py b/emmet-core/tests/test_neb.py index 85b854de32..72c808b545 100644 --- a/emmet-core/tests/test_neb.py +++ b/emmet-core/tests/test_neb.py @@ -37,7 +37,9 @@ def test_neb_doc(test_dir, from_dir: bool): assert neb_doc.num_images == 3 assert len(neb_doc.image_structures) == neb_doc.num_images assert len(neb_doc.energies) == neb_doc.num_images - assert len(neb_doc.structures) == neb_doc.num_images + 2 # always includes endpoints + assert ( + len(neb_doc.structures) == neb_doc.num_images + 2 + ) # always includes endpoints assert isinstance(neb_doc.orig_inputs, OrigInputs) # test that NEB image calculations are all VASP Calculation objects @@ -77,8 +79,8 @@ def test_neb_doc(test_dir, from_dir: bool): ) assert len(neb_doc.image_energies) == neb_doc.num_images -def test_from_directories(test_dir): +def test_from_directories(test_dir): with TemporaryDirectory() as tmpdir: tmpdir = Path(tmpdir) shutil.unpack_archive(test_dir / "neb_sample_calc.zip", tmpdir, "zip") @@ -87,23 +89,32 @@ def test_from_directories(test_dir): tmpdir / "neb", ) - assert all(isinstance(ep_calc,Calculation) for ep_calc in neb_doc.endpoint_calculations) - assert all( - "relax_endpoint_" in ep_dir for ep_dir in neb_doc.endpoint_directories + isinstance(ep_calc, Calculation) for ep_calc in neb_doc.endpoint_calculations ) + assert all("relax_endpoint_" in ep_dir for ep_dir in neb_doc.endpoint_directories) + assert len(neb_doc.energies) == neb_doc.num_images + 2 assert len(neb_doc.structures) == neb_doc.num_images + 2 - assert isinstance(neb_doc.barrier_analysis,dict) + assert isinstance(neb_doc.barrier_analysis, dict) assert all( neb_doc.barrier_analysis.get(k) is not None - for k in ("energies","frame_index","cubic_spline_pars","ts_frame_index","ts_energy","ts_in_frames","forward_barrier","reverse_barrier") + for k in ( + "energies", + "frame_index", + "cubic_spline_pars", + "ts_frame_index", + "ts_energy", + "ts_in_frames", + "forward_barrier", + "reverse_barrier", + ) ) assert all( - getattr(neb_doc,f"{direction}_barrier") == neb_doc.barrier_analysis[f"{direction}_barrier"] + getattr(neb_doc, f"{direction}_barrier") + == neb_doc.barrier_analysis[f"{direction}_barrier"] for direction in ("forward", "reverse") ) - \ No newline at end of file