retrieve prebuild phase diagrams from OpenData (#961)

* retrieve prebuild phase diagrams from OpenData * try fixing some tests * fix field typo in test
materialsproject · Jan 8, 2025 · c26cdd6 · c26cdd6
1 parent 0a67281
commit c26cdd6
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 21 deletions.
diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py
@@ -530,7 +530,7 @@ def get_material_id_references(self, material_id: str) -> list[str]:
         Returns:
             List of BibTeX references ([str])
         """
-        docs = self.provenance.search(material_ids=material_id)
+        docs = self.materials.provenance.search(material_ids=material_id)
 
         if not docs:
             return []

diff --git a/mp_api/client/routes/materials/thermo.py b/mp_api/client/routes/materials/thermo.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 from emmet.core.thermo import ThermoDoc, ThermoType
+from monty.json import MontyDecoder
 from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.core import Element
 
@@ -163,18 +164,14 @@ def get_phase_diagram_from_chemsys(
             )
 
         sorted_chemsys = "-".join(sorted(chemsys.split("-")))
-        phase_diagram_id = f"{sorted_chemsys}_{t_type}"
-
-        response = self._query_resource(
-            criteria={"phase_diagram_ids": phase_diagram_id},
-            fields=["phase_diagram"],
-            suburl="phase_diagram/",
-            use_document_model=False,
-            num_chunks=1,
-            chunk_size=1,
-        ).get("data", [{}])
-
-        pd = response[0].get("phase_diagram", None)
+        phdiag_id = f"thermo_type={t_type}/chemsys={sorted_chemsys}"
+        version = self.db_version.replace(".", "-")
+        obj_key = f"objects/{version}/phase-diagrams/{phdiag_id}.jsonl.gz"
+        pd = self._query_open_data(
+            bucket="materialsproject-build",
+            key=obj_key,
+            decoder=MontyDecoder().decode,
+        )[0][0].get("phase_diagram")
 
         # Ensure el_ref keys are Element objects for PDPlotter.
         # Ensure qhull_data is a numpy array

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -24,6 +24,7 @@
     "materials_electronic_structure_dos",
     "materials_substrates",
     "materials_synthesis",
+    "materials_similarity",
 ]
 
 special_resters = ["materials_charge_density", "doi"]

diff --git a/tests/test_mprester.py b/tests/test_mprester.py
@@ -139,10 +139,7 @@ def test_get_entries(self, mpr):
             assert e.data.get("energy_above_hull", None) is not None
 
         # Conventional structure
-        formula = "BiFeO3"
-        entry = mpr.get_entry_by_material_id(
-            "mp-22526", inc_structure=True, conventional_unit_cell=True
-        )[0]
+        entry = mpr.get_entry_by_material_id("mp-22526", conventional_unit_cell=True)[1]
 
         s = entry.structure
         assert pytest.approx(s.lattice.a) == s.lattice.b
@@ -152,9 +149,7 @@ def test_get_entries(self, mpr):
         assert pytest.approx(s.lattice.gamma) == 120
 
         # Ensure energy per atom is same
-        prim = mpr.get_entry_by_material_id(
-            "mp-22526", inc_structure=True, conventional_unit_cell=False
-        )[0]
+        prim = mpr.get_entry_by_material_id("mp-22526", conventional_unit_cell=False)[1]
 
         s = prim.structure
         assert pytest.approx(s.lattice.a) == s.lattice.b
@@ -327,7 +322,7 @@ def test_large_list(self, mpr):
                 chunk_size=1000, num_chunks=10, fields=["material_id"]
             )
         ]
-        docs = mpr.summary.search(material_ids=mpids, fields=["material_ids"])
+        docs = mpr.summary.search(material_ids=mpids, fields=["material_id"])
         assert len(docs) == 10000
 
     def test_get_api_key_endpoint_from_env_var(self, monkeypatch: pytest.MonkeyPatch):