From e58d54eaab480f78b04e56cb4f7f9ea38721a61d Mon Sep 17 00:00:00 2001 From: anaik Date: Tue, 21 Nov 2023 17:17:56 +0100 Subject: [PATCH] add POSCAR.lobster support in featurizer --- lobsterpy/featurize/batch.py | 129 +++++++++++++----- lobsterpy/featurize/core.py | 34 +++-- .../mp-2176/{POSCAR.gz => POSCAR.lobster.gz} | Bin 3 files changed, 116 insertions(+), 47 deletions(-) rename tests/test_data/Featurizer_test_data/Lobster_calcs/mp-2176/{POSCAR.gz => POSCAR.lobster.gz} (100%) diff --git a/lobsterpy/featurize/batch.py b/lobsterpy/featurize/batch.py index aea3faad..8e1a4daf 100644 --- a/lobsterpy/featurize/batch.py +++ b/lobsterpy/featurize/batch.py @@ -134,12 +134,23 @@ def _featurizecoxx(self, path_to_lobster_calc) -> pd.DataFrame: "icoxxlist_path": "ICOHPLIST.lobster", } for file, default_value in req_files.items(): - file_path = dir_name / default_value - req_files[file] = file_path # type: ignore - if not file_path.exists(): - gz_file_path = Path(zpath(file_path)) - if gz_file_path.exists(): - req_files[file] = gz_file_path # type: ignore + # Check if "POSCAR" exists, and if not, check for "POSCAR.lobster" + if file == "structure_path": + for filename in [default_value, "POSCAR.lobster"]: + poscar_path = dir_name / filename + req_files[file] = poscar_path # type: ignore + if not poscar_path.exists(): + gz_file_path = Path(zpath(poscar_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore + break + else: + file_path = dir_name / default_value + req_files[file] = file_path # type: ignore + if not file_path.exists(): + gz_file_path = Path(zpath(file_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore coxxcar_path = req_files.get("coxxcar_path") structure_path = req_files.get("structure_path") @@ -262,12 +273,22 @@ def _featurizecharges(self, path_to_lobster_calc) -> pd.DataFrame: "structure_path": "POSCAR", } for file, default_value in req_files.items(): - file_path = dir_name / default_value - req_files[file] = file_path # type: ignore - if not file_path.exists(): - gz_file_path = Path(zpath(file_path)) - if gz_file_path.exists(): - req_files[file] = gz_file_path # type: ignore + if file == "structure_path": + for filename in [default_value, "POSCAR.lobster"]: + poscar_path = dir_name / filename + req_files[file] = poscar_path # type: ignore + if not poscar_path.exists(): + gz_file_path = Path(zpath(poscar_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore + break + else: + file_path = dir_name / default_value + req_files[file] = file_path # type: ignore + if not file_path.exists(): + gz_file_path = Path(zpath(file_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore charge_path = req_files.get("charge_path") structure_path = req_files.get("structure_path") @@ -633,11 +654,19 @@ def _fingerprint_df(self, path_to_lobster_calc) -> pd.DataFrame: are_cobis = False are_coops = False - structure_path = dir_name / "POSCAR" - if not structure_path.exists(): - gz_file_path = Path(zpath(structure_path)) - if gz_file_path.exists(): - structure_path = gz_file_path + for filename in ["POSCAR", "POSCAR.lobster"]: + structure_path = dir_name / filename + if not structure_path.exists(): + gz_file_path = Path(zpath(structure_path)) + if gz_file_path.exists(): + structure_path = gz_file_path # type: ignore + break + + # structure_path = dir_name / "POSCAR" + # if not structure_path.exists(): + # gz_file_path = Path(zpath(structure_path)) + # if gz_file_path.exists(): + # structure_path = gz_file_path coxx = FeaturizeCOXX( path_to_coxxcar=str(coxxcar_path), @@ -750,12 +779,22 @@ def _get_sg_df(self, path_to_lobster_calc) -> pd.DataFrame: } for file, default_value in req_files.items(): - file_path = dir_name / default_value - req_files[file] = file_path # type: ignore - if not file_path.exists(): - gz_file_path = Path(zpath(file_path)) - if gz_file_path.exists(): - req_files[file] = gz_file_path # type: ignore + if file == "structure_path": + for filename in [default_value, "POSCAR.lobster"]: + poscar_path = dir_name / filename + req_files[file] = poscar_path # type: ignore + if not poscar_path.exists(): + gz_file_path = Path(zpath(poscar_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore + break + else: + file_path = dir_name / default_value + req_files[file] = file_path # type: ignore + if not file_path.exists(): + gz_file_path = Path(zpath(file_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore charge_path = str(req_files.get("charge_path")) cohpcar_path = str(req_files.get("cohpcar_path")) @@ -887,12 +926,22 @@ def _get_dos_moments_df(self, path_to_lobster_calc) -> pd.DataFrame: "structure_path": "POSCAR", } for file, default_value in req_files.items(): - file_path = dir_name / default_value - req_files[file] = file_path # type: ignore - if not file_path.exists(): - gz_file_path = Path(zpath(file_path)) - if gz_file_path.exists(): - req_files[file] = gz_file_path # type: ignore + if file == "structure_path": + for filename in [default_value, "POSCAR.lobster"]: + poscar_path = dir_name / filename + req_files[file] = poscar_path # type: ignore + if not poscar_path.exists(): + gz_file_path = Path(zpath(poscar_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore + break + else: + file_path = dir_name / default_value + req_files[file] = file_path # type: ignore + if not file_path.exists(): + gz_file_path = Path(zpath(file_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore doscar_path = req_files.get("doscar_path") structure_path = req_files.get("structure_path") @@ -928,12 +977,22 @@ def _get_dos_fingerprints_df(self, path_to_lobster_calc) -> pd.DataFrame: "structure_path": "POSCAR", } for file, default_value in req_files.items(): - file_path = dir_name / default_value - req_files[file] = file_path # type: ignore - if not file_path.exists(): - gz_file_path = Path(zpath(file_path)) - if gz_file_path.exists(): - req_files[file] = gz_file_path # type: ignore + if file == "structure_path": + for filename in [default_value, "POSCAR.lobster"]: + poscar_path = dir_name / filename + req_files[file] = poscar_path # type: ignore + if not poscar_path.exists(): + gz_file_path = Path(zpath(poscar_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore + break + else: + file_path = dir_name / default_value + req_files[file] = file_path # type: ignore + if not file_path.exists(): + gz_file_path = Path(zpath(file_path)) + if gz_file_path.exists(): + req_files[file] = gz_file_path # type: ignore doscar_path = req_files.get("doscar_path") structure_path = req_files.get("structure_path") diff --git a/lobsterpy/featurize/core.py b/lobsterpy/featurize/core.py index a440e013..d0042f66 100644 --- a/lobsterpy/featurize/core.py +++ b/lobsterpy/featurize/core.py @@ -384,18 +384,28 @@ def get_lobsterpy_cba_dict( } for file, default_value in req_files_lobsterpy.items(): - file_path = dir_name / default_value - req_files_lobsterpy[file] = file_path # type: ignore - if not file_path.exists(): - gz_file_path = Path(zpath(file_path)) - if gz_file_path.exists(): - req_files_lobsterpy[file] = gz_file_path # type: ignore - else: - raise Exception( - "Path provided for Lobster calc directory seems incorrect." - "It does not contain COHPCAR.lobster, ICOHPLIST.lobster, POSCAR and " - "CHARGE.lobster files needed for automatic analysis using LobsterPy" - ) + if file == "structure_path": + for filename in [default_value, "POSCAR.lobster"]: + poscar_path = dir_name / filename + req_files_lobsterpy[file] = poscar_path # type: ignore + if not poscar_path.exists(): + gz_file_path = Path(zpath(poscar_path)) + if gz_file_path.exists(): + req_files_lobsterpy[file] = gz_file_path # type: ignore + break + else: + file_path = dir_name / default_value + req_files_lobsterpy[file] = file_path # type: ignore + if not file_path.exists(): + gz_file_path = Path(zpath(file_path)) + if gz_file_path.exists(): + req_files_lobsterpy[file] = gz_file_path # type: ignore + else: + raise Exception( + "Path provided for Lobster calc directory seems incorrect." + "It does not contain COHPCAR.lobster, ICOHPLIST.lobster, POSCAR and " + "CHARGE.lobster files needed for automatic analysis using LobsterPy" + ) cohpcar_path = req_files_lobsterpy.get("cohpcar_path") charge_path = req_files_lobsterpy.get("charge_path") diff --git a/tests/test_data/Featurizer_test_data/Lobster_calcs/mp-2176/POSCAR.gz b/tests/test_data/Featurizer_test_data/Lobster_calcs/mp-2176/POSCAR.lobster.gz similarity index 100% rename from tests/test_data/Featurizer_test_data/Lobster_calcs/mp-2176/POSCAR.gz rename to tests/test_data/Featurizer_test_data/Lobster_calcs/mp-2176/POSCAR.lobster.gz