Skip to content

Commit

Permalink
Fix el_refs dict keys in PhaseDiagram objects (#730)
Browse files Browse the repository at this point in the history
* Fix el_ref keys in client if strings

* Loop fix and linting

* Fix key ref
  • Loading branch information
Jason Munro authored Jan 19, 2023
1 parent 98eb360 commit ece1b3e
Showing 1 changed file with 15 additions and 2 deletions.
17 changes: 15 additions & 2 deletions mp_api/client/routes/thermo.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from mp_api.client.core.utils import validate_ids
from emmet.core.thermo import ThermoDoc, ThermoType
from pymatgen.analysis.phase_diagram import PhaseDiagram
from pymatgen.core import Element


class ThermoRester(BaseRester[ThermoDoc]):
Expand Down Expand Up @@ -173,6 +174,18 @@ def get_phase_diagram_from_chemsys(
use_document_model=False,
num_chunks=1,
chunk_size=1,
).get("data")
).get("data", [{}])

return response[0]["phase_diagram"] # type: ignore
pd = response[0].get("phase_diagram", None)

# Ensure el_ref keys are Element objects for PDPlotter.
# This should be fixed in pymatgen
if pd:
for key, entry in list(pd.el_refs.items()):
if not isinstance(key, str):
break

pd.el_refs[Element(str(key))] = entry
pd.el_refs.pop(key)

return pd # type: ignore

0 comments on commit ece1b3e

Please sign in to comment.