will this fix?

Author: CompRhys

src/pymatgen/analysis/phase_diagram.py

@@ -375,15 +375,13 @@ def __init__(
             if not isinstance(computed_data, dict):
                 raise TypeError(f"computed_data should be dict, got {type(computed_data).__name__}")
 
-            # Update keys to be Element objects in case they are strings in pre-computed data
-            computed_data["el_refs"] = [(Element(el_str), entry) for el_str, entry in computed_data["el_refs"]]
         self.computed_data = computed_data
         self.facets = computed_data["facets"]
         self.simplexes = computed_data["simplexes"]
         self.all_entries = computed_data["all_entries"]
         self.qhull_data = computed_data["qhull_data"]
         self.dim = computed_data["dim"]
-        self.el_refs = dict(computed_data["el_refs"])
+        self.el_refs = {Element(el): entry for el, entry in computed_data["el_refs"]}
         self.qhull_entries = tuple(computed_data["qhull_entries"])
         self._qhull_spaces = tuple(frozenset(e.elements) for e in self.qhull_entries)
         self._stable_entries = tuple({self.qhull_entries[idx] for idx in set(itertools.chain(*self.facets))})
@@ -394,14 +392,29 @@ def as_dict(self):
 
         qhull_entry_indices = [self.all_entries.index(e) for e in self.qhull_entries]
 
+        # Create a copy of computed_data to avoid modifying the original
+        computed_data = self.computed_data.copy()
+        computed_data["elements"] = [el.symbol for el in self.elements]
+        computed_data["el_refs"] = [
+            (el.symbol if isinstance(el, Element) else el, entry.as_dict()) for el, entry in computed_data["el_refs"]
+        ]
+        computed_data["all_entries"] = [e.as_dict() for e in computed_data["all_entries"]]
+        computed_data["qhull_entries"] = qhull_entry_indices
+        computed_data["qhull_data"] = (
+            computed_data["qhull_data"].tolist()
+            if isinstance(computed_data["qhull_data"], np.ndarray)
+            else computed_data["qhull_data"]
+        )
+        computed_data["facets"] = [list(facet) for facet in computed_data["facets"]]
+        computed_data["simplexes"] = [
+            {**s.as_dict(), "coords": s.as_dict()["coords"].tolist()} for s in computed_data["simplexes"]
+        ]
+
         return {
             "@module": type(self).__module__,
             "@class": type(self).__name__,
-            "elements": [e.as_dict() for e in self.elements],
-            "computed_data": self.computed_data
-            | {
-                "qhull_entries": qhull_entry_indices,
-            },
+            "elements": [el.symbol for el in self.elements],
+            "computed_data": computed_data,
         }
 
     @classmethod
@@ -414,17 +427,22 @@ def from_dict(cls, dct: dict[str, Any]) -> Self:
             PhaseDiagram
         """
         computed_data = dct.get("computed_data")
-        elements = [Element.from_dict(elem) for elem in dct["elements"]]
+        elements = [Element(elem) for elem in dct["elements"]]
 
         # for backwards compatibility, check for old format
         if "all_entries" in dct:
             entries = [MontyDecoder().process_decoded(entry) for entry in dct["all_entries"]]
         else:
             entries = [MontyDecoder().process_decoded(entry) for entry in computed_data["all_entries"]]
 
-            complete_qhull_entries = [computed_data["all_entries"][i] for i in computed_data["qhull_entries"]]
-
-            computed_data = computed_data | {"qhull_entries": complete_qhull_entries}
+            # Reconstruct computed_data to match _compute() format: (str, Entry) tuples for el_refs
+            computed_data = computed_data.copy()
+            computed_data["qhull_entries"] = [entries[i] for i in computed_data["qhull_entries"]]
+            computed_data["elements"] = [Element(el) for el in computed_data["elements"]]
+            # Keep el_refs as (str, Entry) format to match _compute() output
+            computed_data["el_refs"] = [
+                (el_str, MontyDecoder().process_decoded(entry)) for el_str, entry in computed_data["el_refs"]
+            ]
 
         return cls(entries, elements, computed_data=computed_data)
 
@@ -497,8 +515,7 @@ def _compute(self) -> dict[str, Any]:
             "all_entries": all_entries,
             "qhull_data": qhull_data,
             "dim": dim,
-            # Dictionary with Element keys is not JSON-serializable
-            "el_refs": list(el_refs.items()),
+            "el_refs": [(el.symbol, entry) for el, entry in el_refs.items()],
             "qhull_entries": qhull_entries,
         }
 
@@ -1427,6 +1444,8 @@ def __init__(
     ):
         """Standard constructor for grand potential phase diagram.
 
+        TODO: update serialization here.
+
         Args:
             entries (Sequence[EntryLike]): A list of EntryLike objects having an
                 energy, energy_per_atom and composition.
@@ -1563,6 +1582,8 @@ def transform_entries(
         Li3PO4 is mapped into a Li2O:1.5, P2O5:0.5 composition. The terminal
         compositions are represented by DummySpecies.
 
+        TODO: update serialization here.
+
         Args:
             entries: Sequence of all input entries
             terminal_compositions: Terminal compositions of phase space.
@@ -1732,7 +1753,7 @@ def _compute(
             subspace_qhull_entry_indices = [pd.all_entries.index(entry) for entry in pd.qhull_entries]
 
             # Convert el_refs to indices into subspace's all_entries
-            subspace_el_refs = [(el.symbol, pd.all_entries.index(entry)) for el, entry in pd.computed_data["el_refs"]]
+            subspace_el_refs = [(el, pd.all_entries.index(entry)) for el, entry in pd.computed_data["el_refs"]]
 
             pds_computed_data[space] = {
                 "all_entries": subspace_all_entry_indices,
@@ -1921,7 +1942,7 @@ def as_dict(self) -> dict[str, Any]:
         # Update computed_data with deduplicated entries and remapped indices
         computed_data = self.computed_data.copy()
 
-        computed_data["elements"] = [e.as_dict() for e in self.elements]
+        computed_data["elements"] = [e.symbol for e in self.elements]
 
         # Convert qhull_entries to indices if they're entry objects
         qhull_entries_data = computed_data["qhull_entries"]
@@ -2043,7 +2064,7 @@ def as_dict(self) -> dict[str, Any]:
         return {
             "@module": type(self).__module__,
             "@class": type(self).__name__,
-            "elements": [e.as_dict() for e in self.elements],
+            "elements": [e.symbol for e in self.elements],
             "computed_data": computed_data | {"unique_entries": unique_entry_dicts},
         }
 
@@ -2058,7 +2079,7 @@ def from_dict(cls, dct: dict) -> Self:
             PatchedPhaseDiagram
         """
         computed_data = dct.get("computed_data")
-        elements = [Element.from_dict(elem) for elem in dct["elements"]]
+        elements = [Element(elem) for elem in dct["elements"]]
 
         # Handle new format with computed_data and unique_entries
         if computed_data and "unique_entries" in computed_data:
@@ -2069,15 +2090,15 @@ def from_dict(cls, dct: dict) -> Self:
             # Reconstruct computed_data with entry objects
             computed_data_reconstructed = computed_data.copy()
             computed_data_reconstructed["all_entries"] = all_entries
-            computed_data_reconstructed["elements"] = [Element.from_dict(e) for e in computed_data["elements"]]
+            computed_data_reconstructed["elements"] = [Element(elem) for elem in computed_data["elements"]]
             computed_data_reconstructed["spaces"] = computed_data["spaces"]
 
             # Reconstruct qhull_entries from indices
             computed_data_reconstructed["qhull_entries"] = [all_entries[idx] for idx in computed_data["qhull_entries"]]
 
             # Reconstruct el_refs from indices
             computed_data_reconstructed["el_refs"] = [
-                (Element(el_symbol), all_entries[idx]) for el_symbol, idx in computed_data["el_refs"]
+                (Element(elem), all_entries[idx]) for elem, idx in computed_data["el_refs"]
             ]
 
             # Reconstruct subspace PhaseDiagram computed_data

tests/analysis/test_phase_diagram.py

@@ -836,7 +836,7 @@ def test_as_from_dict(self):
         unique_entries = [MontyDecoder().process_decoded(entry) for entry in computed_data["unique_entries"]]
         reconstructed_entries = [unique_entries[idx] for idx in computed_data["all_entries"]]
         assert len(reconstructed_entries) == len(self.ppd.all_entries)
-        assert ppd_dict["elements"] == [elem.as_dict() for elem in self.ppd.elements]
+        assert ppd_dict["elements"] == [elem.symbol for elem in self.ppd.elements]
 
         reconstructed_ppd = PatchedPhaseDiagram.from_dict(ppd_dict)
         reconstructed_dict = reconstructed_ppd.as_dict()