refactor get stability test to not rely on golden test data

Author: esoteric-ephemera

mp_api/client/mprester.py

@@ -1601,32 +1601,34 @@ def get_stability(
         }
         chemsys_str = "-".join(sorted(str(ele) for ele in chemsys))
 
-        thermo_type = (
-            ThermoType(thermo_type) if isinstance(thermo_type, str) else thermo_type
+        thermo_type_valid_str: str = (
+            ThermoType(thermo_type).value
+            if (isinstance(thermo_type, str) and thermo_type != "r2SCAN")
+            else str(thermo_type)
         )
 
         corrector: Compatibility | None = None
-        if thermo_type == ThermoType.GGA_GGA_U:
+        if thermo_type_valid_str == ThermoType.GGA_GGA_U.value:
             from pymatgen.entries.compatibility import MaterialsProject2020Compatibility
 
             corrector = MaterialsProject2020Compatibility()
 
-        elif thermo_type == ThermoType.GGA_GGA_U_R2SCAN:
+        elif thermo_type_valid_str == ThermoType.GGA_GGA_U_R2SCAN.value:
             from pymatgen.entries.mixing_scheme import MaterialsProjectDFTMixingScheme
 
             corrector = MaterialsProjectDFTMixingScheme(run_type_2="r2SCAN")
 
         try:
             pd = self.materials.thermo.get_phase_diagram_from_chemsys(
-                chemsys_str, thermo_type=thermo_type
+                chemsys_str, thermo_type=thermo_type_valid_str
             )
         except MPRestError:
             pd = None
 
         if not pd:
             warnings.warn(
                 f"No phase diagram data available for chemical system {chemsys_str} "
-                f"and thermo type {thermo_type}.",
+                f"and thermo type {thermo_type_valid_str}.",
                 category=MPRestWarning,
                 stacklevel=2,
             )

mp_api/client/routes/materials/thermo.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from typing import TYPE_CHECKING
 
 import numpy as np
 from emmet.core.thermo import ThermoDoc, validate_thermo_id
@@ -14,12 +15,45 @@
 from mp_api.client.core.exceptions import MPRestError
 from mp_api.client.core.utils import validate_ids
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from enums import Enum
+
 
 class ThermoRester(BaseRester):
     suffix = "materials/thermo"
     document_model = ThermoDoc  # type: ignore
     primary_key = "material_id"
 
+    @staticmethod
+    def _check_thermo_types(thermo_types: Sequence[str | Enum]) -> set[str]:
+        """Check if a user has input any invalid thermo types.
+
+        Args:
+        thermo_types (Sequence of str or Enum) : list of thermo types
+            the user has queried for
+
+        phase-diagram tbl has "r2SCAN", not "R2SCAN"
+        mixing of ThermoType/RunType in emmet -_-
+        TODO: coerce upstream? allow case-insensitivity in emmet?
+
+        Returns:
+        set of str: validated thermo types
+
+        Raises:
+        ValueError if any invalid thermo types are input
+        """
+        t_types: set[str] = {t if isinstance(t, str) else t.value for t in thermo_types}
+        t_types = {"r2SCAN" if t == "R2SCAN" else t for t in t_types}
+        valid_types = {"r2SCAN", *map(str, ThermoType.__members__.values())}
+
+        if invalid_types := t_types - valid_types:
+            raise ValueError(
+                f"Invalid thermo type(s) passed: {invalid_types}, valid types are: {valid_types}"
+            )
+        return t_types
+
     def search(
         self,
         material_ids: str | list[str] | None = None,
@@ -57,7 +91,7 @@ def search(
             material_ids (List[str]): List of Materials Project IDs to return data for.
             thermo_ids (List[str]): List of thermo IDs to return data for. This is a combination of the Materials
                 Project ID and thermo type (e.g. mp-149_GGA_GGA+U).
-            thermo_types (List[ThermoType]): List of thermo types to return data for (e.g. ThermoType.GGA_GGA_U).
+            thermo_types (List[ThermoType or str]): List of thermo/run types to return data for (e.g. ThermoType.GGA_GGA_U).
             num_elements (Tuple[int,int]): Minimum and maximum number of elements in the material to consider.
             total_energy (Tuple[float,float]): Minimum and maximum corrected total energy in eV/atom to consider.
             uncorrected_energy (Tuple[float,float]): Minimum and maximum uncorrected total
@@ -105,13 +139,9 @@ def search(
                 )
 
         if thermo_types:
-            t_types = {t if isinstance(t, str) else t.value for t in thermo_types}
-            valid_types = {*map(str, ThermoType.__members__.values())}
-            if invalid_types := t_types - valid_types:
-                raise ValueError(
-                    f"Invalid thermo type(s) passed: {invalid_types}, valid types are: {valid_types}"
-                )
-            query_params.update({"thermo_types": ",".join(t_types)})
+            query_params.update(
+                {"thermo_types": ",".join(self._check_thermo_types(thermo_types))}
+            )
 
         if num_elements:
             if isinstance(num_elements, int):
@@ -168,12 +198,7 @@ def get_phase_diagram_from_chemsys(
         Returns:
             (PhaseDiagram): Pymatgen phase diagram object.
         """
-        t_type = thermo_type if isinstance(thermo_type, str) else thermo_type.value
-        valid_types = {*map(str, ThermoType.__members__.values())}
-        if invalid_types := {t_type} - valid_types:
-            raise ValueError(
-                f"Invalid thermo type(s) passed: {invalid_types}, valid types are: {valid_types}"
-            )
+        validated_thermo_type = self._check_thermo_types([thermo_type]).pop()
 
         sorted_chemsys = "-".join(sorted(chemsys.split("-")))
         version = self.db_version.replace(".", "-")
@@ -182,18 +207,12 @@ def get_phase_diagram_from_chemsys(
             "materialsproject-build", "objects/phase-diagrams", label="phase_diagrams"
         )
 
-        # phase-diagram tbl has r2SCAN, not R2SCAN
-        # mixing of ThermoType/RunType in emmet -_-
-        # TODO: coerce upstream? allow case-insensitivity in emmet?
-        if thermo_type == ThermoType.R2SCAN:
-            thermo_type = "r2SCAN"
-
         query = f"""
             SELECT phase_diagram
             FROM   {pd_lbl}
             WHERE  chemsys='{sorted_chemsys}'
               AND  version='{version}'
-              AND  thermo_type='{thermo_type}'
+              AND  thermo_type='{validated_thermo_type}'
         """
         table = self._query_delta_single(query)
         as_py = table["phase_diagram"].to_pylist(maps_as_pydicts="strict")

tests/client/test_mprester.py

@@ -564,81 +564,83 @@ def test_get_cohesive_energy(self):
                 mpr.get_cohesive_energy("mp-1")
 
     @pytest.mark.parametrize(
-        "chemsys, thermo_type",
-        [
-            [("Fe", "P"), "GGA_GGA+U"],
-            [("Li", "S"), ThermoType.GGA_GGA_U_R2SCAN],
-            [("Ni", "Se"), ThermoType.R2SCAN],
-            [("Ni", "Kr"), "R2SCAN"],
-        ],
+        "thermo_type", ["GGA_GGA+U", ThermoType.GGA_GGA_U_R2SCAN, "r2SCAN"]
     )
-    def test_get_stability(self, chemsys, thermo_type):
+    def test_get_stability(self, thermo_type):
         """
         This test is adapted from the pymatgen one - the scope is broadened
         to include more diverse chemical environments and thermo types which
         reflect the scope of the current MP database.
         """
         with MPRester() as mpr:
-            entries = mpr.get_entries_in_chemsys(
-                chemsys, additional_criteria={"thermo_types": [thermo_type]}
-            )
 
-            no_compound_entries = all(
-                len(entry.composition.elements) == 1 for entry in entries
-            )
+            # No golden test data. Always test on fetched thermo data
+            chemsys_to_test: set[str] = {
+                doc.chemsys
+                for doc in mpr.materials.thermo.search(
+                    thermo_types=[thermo_type],
+                    num_elements=2,
+                    num_chunks=1,
+                    chunk_size=4,
+                    fields=["chemsys"],
+                )
+            }
+
+            for chemsys in chemsys_to_test:
 
-            modified_entries = [
-                ComputedEntry(
-                    entry.composition,
-                    entry.uncorrected_energy + 0.01,
-                    parameters=entry.parameters,
-                    entry_id=f"mod_{entry.entry_id}",
+                entries = mpr.get_entries_in_chemsys(
+                    chemsys, additional_criteria={"thermo_types": [thermo_type]}
                 )
-                for entry in entries
-                if entry.composition.reduced_formula in ["Fe2P", "".join(chemsys)]
-            ]
 
-            if len(modified_entries) == 0:
-                # create fake entry to get PD retrieval to fail
                 modified_entries = [
                     ComputedEntry(
-                        "".join(chemsys),
-                        np.average([entry.energy for entry in entries]),
-                        entry_id=f"hypothetical",
+                        entry.composition,
+                        entry.uncorrected_energy + 0.01,
+                        parameters=entry.parameters,
+                        entry_id=f"mod_{entry.entry_id}",
                     )
+                    for entry in entries
+                    if entry.entry_id == entries[0].entry_id
                 ]
 
-            if no_compound_entries:
-                with pytest.warns(
-                    MPRestWarning, match="No phase diagram data available"
+                if (
+                    all(len(entry.composition.elements) == 1 for entry in entries)
+                    and chemsys.count("-") > 0
                 ):
-                    mpr.get_stability(modified_entries, thermo_type=thermo_type)
-                return
-
-            else:
-                rester_ehulls = mpr.get_stability(
-                    modified_entries, thermo_type=thermo_type
-                )
+                    # For a multi-element chemsys with no multinaries, only elementals,
+                    # there should be no phase diagram data available.
+                    with pytest.warns(
+                        MPRestWarning, match="No phase diagram data available"
+                    ):
+                        mpr.get_stability(modified_entries, thermo_type=thermo_type)
+                    return
+
+                else:
+                    rester_ehulls = mpr.get_stability(
+                        modified_entries, thermo_type=thermo_type
+                    )
 
-        all_entries = entries + modified_entries
-
-        compat = None
-        if thermo_type == "GGA_GGA+U":
-            compat = MaterialsProject2020Compatibility()
-        elif thermo_type == "GGA_GGA+U_R2SCAN":
-            compat = MaterialsProjectDFTMixingScheme(run_type_2="r2SCAN")
-
-        if compat:
-            all_entries = compat.process_entries(all_entries)
-
-        pd = PhaseDiagram(all_entries)
-        for entry in all_entries:
-            if str(entry.entry_id).startswith("mod"):
-                for dct in rester_ehulls:
-                    if dct["entry_id"] == entry.entry_id:
-                        data = dct
-                        break
-                assert pd.get_e_above_hull(entry) == pytest.approx(data["e_above_hull"])
+            all_entries = entries + modified_entries
+
+            compat = None
+            if thermo_type == "GGA_GGA+U":
+                compat = MaterialsProject2020Compatibility()
+            elif thermo_type == "GGA_GGA+U_R2SCAN":
+                compat = MaterialsProjectDFTMixingScheme(run_type_2="r2SCAN")
+
+            if compat:
+                all_entries = compat.process_entries(all_entries)
+
+            pd = PhaseDiagram(all_entries)
+            for entry in all_entries:
+                if str(entry.entry_id).startswith("mod"):
+                    for dct in rester_ehulls:
+                        if dct["entry_id"] == entry.entry_id:
+                            data = dct
+                            break
+                    assert pd.get_e_above_hull(entry) == pytest.approx(
+                        data["e_above_hull"]
+                    )
 
     @pytest.mark.parametrize(
         "mpid, working_ion, thermo_type",