Add custom as_dict/from_dict method for proper initialization of attributes of IcohpCollection (#4391)

* add custom from_dict method for proper initialization of attributes

* pre-commit auto-fixes

* add custom as_dict method for proper serialization

* adapt existing test to check for serialization

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: naik-aakash

src/pymatgen/electronic_structure/cohp.py

@@ -1387,6 +1387,75 @@ 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."""
+        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"]]  # 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]:
+                    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  # type: ignore[assignment]
+
+                    new_list_orb[bond_num][orb].update(sub_dict)
+
+        dct["list_icohp"] = list_icohp
+        dct["list_orb_icohp"] = new_list_orb
+
+        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(
     cohp: CompleteCohp,

tests/electronic_structure/test_cohp.py

@@ -235,29 +235,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]],
                         },
                     },
                 ],