port get_stability functionality from legacy api + test (#981)

Co-authored-by: esoteric-ephemera <[email protected]>

Author: tschaume

mp_api/client/mprester.py

@@ -12,6 +12,7 @@
 from emmet.core.mpid import MPID
 from emmet.core.settings import EmmetSettings
 from emmet.core.tasks import TaskDoc
+from emmet.core.thermo import ThermoType
 from emmet.core.vasp.calc_types import CalcType
 from monty.json import MontyDecoder
 from packaging import version
@@ -61,8 +62,10 @@
 from mp_api.client.routes.molecules import MoleculeRester
 
 if TYPE_CHECKING:
-    from typing import Literal
+    from typing import Any, Literal
 
+    from pymatgen.analysis.phase_diagram import PDEntry
+    from pymatgen.entries.computed_entries import ComputedEntry
 
 _EMMET_SETTINGS = EmmetSettings()
 _MAPI_SETTINGS = MAPIClientSettings()
@@ -1620,3 +1623,59 @@ def _get_cohesive_energy(
         elif normalization == "formula_unit":
             num_form_unit = comp.get_reduced_composition_and_factor()[1]
             return (energy_per_atom * natom - atomic_energy) / num_form_unit
+
+    def get_stability(
+        self,
+        entries: ComputedEntry | ComputedStructureEntry | PDEntry,
+        thermo_type: str | ThermoType = ThermoType.GGA_GGA_U,
+    ) -> list[dict[str, Any]] | None:
+        chemsys = set()
+        for entry in entries:
+            chemsys.update(entry.composition.elements)
+        chemsys_str = "-".join(sorted(str(ele) for ele in chemsys))
+
+        thermo_type = (
+            ThermoType(thermo_type) if isinstance(thermo_type, str) else thermo_type
+        )
+
+        corrector = None
+        if thermo_type == ThermoType.GGA_GGA_U:
+            from pymatgen.entries.compatibility import MaterialsProject2020Compatibility
+
+            corrector = MaterialsProject2020Compatibility()
+
+        elif thermo_type == ThermoType.GGA_GGA_U_R2SCAN:
+            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
+            )
+        except OSError:
+            pd = None
+
+        if not pd:
+            warnings.warn(
+                f"No phase diagram data available for chemical system {chemsys_str} "
+                f"and thermo type {thermo_type}."
+            )
+            return
+
+        if corrector:
+            corrected_entries = corrector.process_entries(entries + pd.all_entries)
+        else:
+            corrected_entries = [*entries, *pd.all_entries]
+
+        new_pd = PhaseDiagram(corrected_entries)
+
+        return [
+            {
+                "e_above_hull": new_pd.get_e_above_hull(entry),
+                "composition": entry.composition.as_dict(),
+                "energy": entry.energy,
+                "entry_id": getattr(entry, "entry_id", f"user-entry-{idx}"),
+            }
+            for idx, entry in enumerate(entries)
+        ]

tests/test_mprester.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pytest
 from emmet.core.tasks import TaskDoc
+from emmet.core.thermo import ThermoType
 from emmet.core.vasp.calc_types import CalcType
 from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.analysis.pourbaix_diagram import IonEntry, PourbaixDiagram, PourbaixEntry
@@ -18,7 +19,11 @@
     BandStructureSymmLine,
 )
 from pymatgen.electronic_structure.dos import CompleteDos
-from pymatgen.entries.compatibility import MaterialsProjectAqueousCompatibility
+from pymatgen.entries.compatibility import (
+    MaterialsProjectAqueousCompatibility,
+    MaterialsProject2020Compatibility,
+)
+from pymatgen.entries.mixing_scheme import MaterialsProjectDFTMixingScheme
 from pymatgen.entries.computed_entries import ComputedEntry, GibbsComputedStructureEntry
 from pymatgen.io.cif import CifParser
 from pymatgen.io.vasp import Chgcar
@@ -429,3 +434,78 @@ def test_get_cohesive_energy(self):
         assert all(
             v == pytest.approx(e_coh["noserial"][k]) for k, v in e_coh["serial"].items()
         )
+
+    @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"],
+        ],
+    )
+    def test_get_stability(self, chemsys, 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
+            )
+
+            modified_entries = [
+                ComputedEntry(
+                    entry.composition,
+                    entry.uncorrected_energy + 0.01,
+                    parameters=entry.parameters,
+                    entry_id=f"mod_{entry.entry_id}",
+                )
+                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",
+                    )
+                ]
+
+            if no_compound_entries:
+                with pytest.warns(UserWarning, 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"])