From 828a0479ddb97b56861ae9810b7df3edb7862ce4 Mon Sep 17 00:00:00 2001 From: Aakash Ashok Naik <91958822+naik-aakash@users.noreply.github.com> Date: Tue, 6 May 2025 22:03:13 +0200 Subject: [PATCH 1/4] add custom from_dict method for proper initialization of attributes --- src/pymatgen/electronic_structure/cohp.py | 35 +++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/src/pymatgen/electronic_structure/cohp.py b/src/pymatgen/electronic_structure/cohp.py index 20fdfcc9812..4af3d4c4b0a 100644 --- a/src/pymatgen/electronic_structure/cohp.py +++ b/src/pymatgen/electronic_structure/cohp.py @@ -1387,6 +1387,41 @@ def are_cobis(self) -> bool: """Whether this is COBI.""" return self._are_cobis + @classmethod + def from_dict(cls, dct: dict[str, Any]) -> Self: + """Generate IcohpCollection from a dict representation.""" + list_icohp = [] + for icohp_dict in dct["list_icohp"]: + new_icohp_dict = {} + for spin_key, value in icohp_dict.items(): + # Convert string/int to Spin enum + spin_enum = Spin(int(spin_key)) if isinstance(spin_key, int) or spin_key in ("1", "-1") else Spin[spin_key] + new_icohp_dict[spin_enum] = value + list_icohp.append(new_icohp_dict) + + new_list_orb = [{orb_label: {} for orb_label in bond} for bond in dct["list_orb_icohp"]] + for bond_num, lab_orb_icohp in enumerate(dct["list_orb_icohp"]): + for orb in lab_orb_icohp: + for key in lab_orb_icohp[orb]: + sub_dict = {} + if key == "icohp": + sub_dict[key] = {Spin.up : lab_orb_icohp[orb][key]["1"], + Spin.down : lab_orb_icohp[orb][key]["-1"]} + + if key == "orbitals": + orb_temp = [] + + for item in lab_orb_icohp[orb][key]: + item[1] = Orbital(item[1]) + orb_temp.append(item) + sub_dict[key] = orb_temp + + new_list_orb[bond_num][orb].update(sub_dict) + + dct["list_icohp"] = list_icohp + dct["list_orb_icohp"] = new_list_orb + + return cls(**dct) def get_integrated_cohp_in_energy_range( cohp: CompleteCohp, From 1e70091dc8bb90d12a3f29897bd52f9b05913fb8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 6 May 2025 20:04:25 +0000 Subject: [PATCH 2/4] pre-commit auto-fixes --- src/pymatgen/electronic_structure/cohp.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/pymatgen/electronic_structure/cohp.py b/src/pymatgen/electronic_structure/cohp.py index 4af3d4c4b0a..82fe871168d 100644 --- a/src/pymatgen/electronic_structure/cohp.py +++ b/src/pymatgen/electronic_structure/cohp.py @@ -1395,34 +1395,39 @@ def from_dict(cls, dct: dict[str, Any]) -> Self: new_icohp_dict = {} for spin_key, value in icohp_dict.items(): # Convert string/int to Spin enum - spin_enum = Spin(int(spin_key)) if isinstance(spin_key, int) or spin_key in ("1", "-1") else Spin[spin_key] + spin_enum = ( + Spin(int(spin_key)) if isinstance(spin_key, int) or spin_key in ("1", "-1") else Spin[spin_key] + ) new_icohp_dict[spin_enum] = value list_icohp.append(new_icohp_dict) - + new_list_orb = [{orb_label: {} for orb_label in bond} for bond in dct["list_orb_icohp"]] for bond_num, lab_orb_icohp in enumerate(dct["list_orb_icohp"]): for orb in lab_orb_icohp: for key in lab_orb_icohp[orb]: sub_dict = {} if key == "icohp": - sub_dict[key] = {Spin.up : lab_orb_icohp[orb][key]["1"], - Spin.down : lab_orb_icohp[orb][key]["-1"]} - + sub_dict[key] = { + Spin.up: lab_orb_icohp[orb][key]["1"], + Spin.down: lab_orb_icohp[orb][key]["-1"], + } + if key == "orbitals": orb_temp = [] - + for item in lab_orb_icohp[orb][key]: item[1] = Orbital(item[1]) orb_temp.append(item) sub_dict[key] = orb_temp - + new_list_orb[bond_num][orb].update(sub_dict) - + dct["list_icohp"] = list_icohp dct["list_orb_icohp"] = new_list_orb - + return cls(**dct) + def get_integrated_cohp_in_energy_range( cohp: CompleteCohp, label: str, From 4754806048526f497f2946a24254170c8484cc6c Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 12 May 2025 15:03:44 +0200 Subject: [PATCH 3/4] add custom as_dict method for proper serialization --- src/pymatgen/electronic_structure/cohp.py | 35 +++++++++++++++++++++-- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/src/pymatgen/electronic_structure/cohp.py b/src/pymatgen/electronic_structure/cohp.py index 82fe871168d..2782cda330d 100644 --- a/src/pymatgen/electronic_structure/cohp.py +++ b/src/pymatgen/electronic_structure/cohp.py @@ -1387,6 +1387,33 @@ def are_cobis(self) -> bool: """Whether this is COBI.""" return self._are_cobis + def as_dict(self) -> dict[str, Any]: + """JSON-serializable dict representation of COHP.""" + return { + "@module": type(self).__module__, + "@class": type(self).__name__, + "are_coops": self._are_coops, + "are_cobis": self._are_cobis, + "list_labels": self._list_labels, + "list_atom1": self._list_atom1, + "list_atom2": self._list_atom2, + "list_length": self._list_length, + "list_translation": self._list_translation, + "list_num": self._list_num, + "list_icohp": [{str(spin): value for spin, value in icohp.items()} for icohp in self._list_icohp], + "is_spin_polarized": self._is_spin_polarized, + "list_orb_icohp": [ + { + key: { + "icohp": {str(spin): value for spin, value in val["icohp"].items()}, + "orbitals": [[n, int(orb)] for n, orb in val["orbitals"]], + } + for key, val in entry.items() + } + for entry in self._list_orb_icohp + ], + } + @classmethod def from_dict(cls, dct: dict[str, Any]) -> Self: """Generate IcohpCollection from a dict representation.""" @@ -1401,7 +1428,7 @@ def from_dict(cls, dct: dict[str, Any]) -> Self: new_icohp_dict[spin_enum] = value list_icohp.append(new_icohp_dict) - new_list_orb = [{orb_label: {} for orb_label in bond} for bond in dct["list_orb_icohp"]] + new_list_orb = [{orb_label: {} for orb_label in bond} for bond in dct["list_orb_icohp"]] # type:ignore[var-annotated] for bond_num, lab_orb_icohp in enumerate(dct["list_orb_icohp"]): for orb in lab_orb_icohp: for key in lab_orb_icohp[orb]: @@ -1418,14 +1445,16 @@ def from_dict(cls, dct: dict[str, Any]) -> Self: for item in lab_orb_icohp[orb][key]: item[1] = Orbital(item[1]) orb_temp.append(item) - sub_dict[key] = orb_temp + sub_dict[key] = orb_temp # type: ignore[assignment] new_list_orb[bond_num][orb].update(sub_dict) dct["list_icohp"] = list_icohp dct["list_orb_icohp"] = new_list_orb - return cls(**dct) + init_dict = {k: v for k, v in dct.items() if "@" not in k} + + return cls(**init_dict) def get_integrated_cohp_in_energy_range( From fd13e102aa263bd1cb44f9d091fc9bc0399107cb Mon Sep 17 00:00:00 2001 From: naik-aakash Date: Mon, 12 May 2025 15:06:56 +0200 Subject: [PATCH 4/4] adapt existing test to check for serialization --- tests/electronic_structure/test_cohp.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/electronic_structure/test_cohp.py b/tests/electronic_structure/test_cohp.py index 17053dbfcb0..3d803dddc21 100644 --- a/tests/electronic_structure/test_cohp.py +++ b/tests/electronic_structure/test_cohp.py @@ -236,29 +236,29 @@ def setup_method(self): "list_translation": [[0, 0, -1], [0, 0, 0]], "list_num": [1, 1], "list_icohp": [ - {Spin.up: 0.29324, Spin.down: 0.29324}, - {Spin.up: 0.29324, Spin.down: 0.29324}, + {"1": 0.29324, "-1": 0.29324}, + {"1": 0.29324, "-1": 0.29324}, ], "is_spin_polarized": True, "list_orb_icohp": [ { "2s-6s": { - "icohp": {Spin.up: 0.0247, Spin.down: 0.0247}, - "orbitals": [[2, Orbital.s], [6, Orbital.s]], + "icohp": {"1": 0.0247, "-1": 0.0247}, + "orbitals": [[2, 0], [6, 0]], }, "2s-5py": { - "icohp": {Spin.up: 8e-05, Spin.down: 8e-05}, - "orbitals": [[2, Orbital.s], [5, Orbital.py]], + "icohp": {"1": 8e-05, "-1": 8e-05}, + "orbitals": [[2, 0], [5, 1]], }, }, { "2s-6s": { - "icohp": {Spin.up: 0.0247, Spin.down: 0.0247}, - "orbitals": [[2, Orbital.s], [6, Orbital.s]], + "icohp": {"1": 0.0247, "-1": 0.0247}, + "orbitals": [[2, 0], [6, 0]], }, "2s-5py": { - "icohp": {Spin.up: 0.5, Spin.down: 0}, - "orbitals": [[2, Orbital.s], [5, Orbital.py]], + "icohp": {"1": 0.5, "-1": 0}, + "orbitals": [[2, 0], [5, 1]], }, }, ],