Serve the pre-built PD for property_data / conventional_unit_cell queries

mp_api/client/mprester.py

@@ -6,6 +6,7 @@
 import re
 import warnings
 from collections import defaultdict
+from copy import deepcopy
 from functools import cache, lru_cache
 from typing import TYPE_CHECKING
 from urllib.parse import urlencode
@@ -675,29 +676,13 @@ def get_entries(
                 ):  # merge property_data, retaining entry data (e.g. `oxidation_states`)
                     entry_dict["data"] |= {prop: doc[prop] for prop in property_data}
 
+                entry = TypeAdapter(ComputedStructureEntryType).validate_python(
+                    entry_dict
+                )
                 if conventional_unit_cell:
-                    entry_struct = Structure.from_dict(entry_dict["structure"])
-                    s = SpacegroupAnalyzer(
-                        entry_struct
-                    ).get_conventional_standard_structure()
-                    site_ratio = len(s) / len(entry_struct)
-                    new_energy = entry_dict["energy"] * site_ratio
-
-                    entry_dict["energy"] = new_energy
-                    entry_dict["structure"] = s.as_dict()
-                    entry_dict["correction"] = 0.0
-
-                    for element in entry_dict["composition"]:
-                        entry_dict["composition"][element] *= site_ratio
+                    entry = self._get_conventional_cell_entry(entry)
 
-                    for correction in entry_dict["energy_adjustments"]:
-                        if "n_atoms" in correction:
-                            correction["n_atoms"] *= site_ratio
-
-                # Need to store object to permit de-duplication
-                entries.add(
-                    TypeAdapter(ComputedStructureEntryType).validate_python(entry_dict)
-                )
+                entries.add(entry)  # object permits de-duplication
 
         return list(entries)
 
@@ -1089,6 +1074,31 @@ def _get_unmixed_entries(
             )
         ]
 
+    @staticmethod
+    def _get_conventional_cell_entry(
+        entry: ComputedStructureEntry,
+    ) -> ComputedStructureEntry:
+        """Rebuild ``entry`` on the standard conventional unit cell, scaling the energy
+        and energy adjustments accordingly.
+        """
+        conventional_structure = SpacegroupAnalyzer(
+            entry.structure
+        ).get_conventional_standard_structure()
+        site_ratio = len(conventional_structure) / len(entry.structure)
+
+        energy_adjustments = deepcopy(entry.energy_adjustments)
+        for adjustment in energy_adjustments:  # adjustment values are extensive
+            adjustment.normalize(1 / site_ratio)
+
+        return ComputedStructureEntry(
+            conventional_structure,
+            entry.uncorrected_energy * site_ratio,
+            energy_adjustments=energy_adjustments,
+            parameters=entry.parameters,
+            data=entry.data,
+            entry_id=entry.entry_id,
+        )
+
     def get_entries_in_chemsys(
         self,
         elements: str | list[str],
@@ -1110,11 +1120,15 @@ def get_entries_in_chemsys(
 
         Mixed entries are taken from the MP-built phase diagram for the whole chemical
         system, so they share one energy scale and reproduce the hull shown on
-        https://materialsproject.org. Narrowing the query with `additional_criteria`,
-        or passing ``compatible_only = False``, cannot be served that way and returns
-        entries that are *not* immediately suitable for constructing a phase diagram;
-        ``property_data`` and ``conventional_unit_cell`` re-apply the mixing scheme here
-        instead, which can differ slightly from MP. Warnings are thrown for these cases.
+        https://materialsproject.org; ``property_data`` fields are attached to the
+        served entries after the fact, and any further ``additional_criteria`` narrow
+        the served entries to the materials matching them (with MP's own semantics,
+        i.e. ``is_stable`` / ``energy_above_hull`` refer to each material's `own` thermo
+        doc; entries with no mixed-type thermo doc, which can occur when the mixing
+        scheme is re-applied locally when MP has no pre-built diagram, are then dropped).
+        Passing ``compatible_only = False`` cannot be served that way and returns
+        entries that are *not* immediately suitable for constructing a phase diagram,
+        with a warning.
 
         Args:
             elements (str or [str]): Parent chemical system string comprising element
@@ -1181,25 +1195,38 @@ def get_entries_in_chemsys(
         # (issue #1104). Thus we serve MP's phase diagram; built with mixing applied across the
         # full system, thus self-consistent by construction and identical to the MP website:
         mixed = set(additional_criteria["thermo_types"]) == {"GGA_GGA+U_R2SCAN"}
-        consistent = (
-            mixed and compatible_only and set(additional_criteria) == {"thermo_types"}
-        )
+        consistent = mixed and compatible_only
+        extra_criteria = {
+            k: v for k, v in additional_criteria.items() if k != "thermo_types"
+        }
 
         entries: list[ComputedStructureEntry] | None = None
         if consistent:
-            if not (property_data or conventional_unit_cell):
-                phase_diagram = self.materials.thermo.get_phase_diagram_from_chemsys(
-                    "-".join(sorted(elements_set)),
-                    thermo_type=additional_criteria["thermo_types"][0],
-                )  # default, mixed thermotype; takes a single type, not a list
-                if phase_diagram is not None:
-                    entries = list(phase_diagram.all_entries)
+            phase_diagram = self.materials.thermo.get_phase_diagram_from_chemsys(
+                "-".join(sorted(elements_set)),
+                thermo_type=additional_criteria["thermo_types"][0],
+            )  # default, mixed thermotype; takes a single type, not a list
+            if phase_diagram is not None:
+                entries = list(phase_diagram.all_entries)
+
+                if property_data:  # decorate the served entries post-hoc
+                    docs = self.materials.thermo.search(
+                        chemsys=all_chemsyses,
+                        thermo_types=additional_criteria["thermo_types"],
+                        all_fields=False,
+                        fields=["material_id", *property_data],
+                    )
+                    props = {
+                        str(doc["material_id"]): {p: doc[p] for p in property_data}
+                        for doc in docs
+                    }
+                    for entry in entries:  # served entries carry `material_id`
+                        entry.data |= props[str(entry.data["material_id"])]
 
             if entries is None:
-                # MP has no pre-built diagram for this system, or the entries need reshaping
-                # first, so redo the mixing here as MP does when building PDs. Mixing scheme
-                # is chemical-system dependent, so this can anchor on a different hull than
-                # MP did/would, and it drops entries it cannot place:
+                # MP has no pre-built diagram for this system, so redo the mixing here as MP does when
+                # building PDs. Mixing scheme is chemical-system dependent, so this can anchor on a
+                # different hull than MP did/would, and it drops entries it cannot place:
                 from pymatgen.entries.mixing_scheme import (
                     MaterialsProjectDFTMixingScheme,
                 )
@@ -1217,18 +1244,36 @@ def get_entries_in_chemsys(
                     self._get_unmixed_entries(
                         all_chemsyses,
                         property_data=property_data,
-                        conventional_unit_cell=conventional_unit_cell,
                         **kwargs,
                     )
                 )
+
+            if extra_criteria:
+                # narrow the common-scale entries post-hoc, with MP's own criteria semantics
+                # (querying with the criteria directly returns per-material corrected entries
+                # which do not share a common energy scale):
+                matching_ids = {
+                    str(doc["material_id"])
+                    for doc in self.materials.thermo.search(
+                        chemsys=all_chemsyses,
+                        thermo_types=additional_criteria["thermo_types"],
+                        all_fields=False,
+                        fields=["material_id"],
+                        **extra_criteria,
+                    )
+                }
+                entries = [
+                    entry
+                    for entry in entries
+                    if str(entry.data["material_id"]) in matching_ids
+                ]
         else:  # non-consistent
             if mixed:
                 warnings.warn(
                     "Mixed GGA(+U)/r2SCAN entries can only be placed on a common energy scale "
-                    "when the whole chemical system is retrieved with `compatible_only = True`, "
-                    "so these entries are not suitable for constructing a phase diagram. Either "
-                    "drop the extra `additional_criteria` (and filter the returned entries "
-                    "instead), or request a single functional with "
+                    "with `compatible_only = True`, so these uncorrected entries are not "
+                    "suitable for constructing a phase diagram. Either use "
+                    "`compatible_only = True`, or request a single functional with "
                     '`additional_criteria = {"thermo_types": ["GGA_GGA+U"]}`.',
                     category=MPRestWarning,
                     stacklevel=2,
@@ -1238,11 +1283,15 @@ def get_entries_in_chemsys(
                 all_chemsyses,
                 compatible_only=compatible_only,
                 property_data=property_data,
-                conventional_unit_cell=conventional_unit_cell,
                 additional_criteria=additional_criteria,
                 **kwargs,
             )
 
+        if conventional_unit_cell:
+            # reshaped here rather than in the queries above, so that structure matching in the mixing
+            # scheme sees the original cells, and so that every energy adjustment is scaled appropriately:
+            entries = [self._get_conventional_cell_entry(entry) for entry in entries]
+
         if use_gibbs:
             # replace the entries with GibbsComputedStructureEntry
             from pymatgen.entries.computed_entries import GibbsComputedStructureEntry

tests/client/test_mprester.py

@@ -2,6 +2,7 @@
 import itertools
 import os
 import random
+import warnings
 from collections import defaultdict
 from tempfile import NamedTemporaryFile
 
@@ -29,7 +30,11 @@
     MaterialsProject2020Compatibility,
     MaterialsProjectAqueousCompatibility,
 )
-from pymatgen.entries.computed_entries import ComputedEntry, GibbsComputedStructureEntry
+from pymatgen.entries.computed_entries import (
+    ComputedEntry,
+    ConstantEnergyAdjustment,
+    GibbsComputedStructureEntry,
+)
 from pymatgen.entries.mixing_scheme import MaterialsProjectDFTMixingScheme
 from pymatgen.io.cif import CifParser
 from pymatgen.io.vasp import Chgcar
@@ -236,15 +241,14 @@ def test_get_entries(self, mpr):
         non_standardized = mpr.get_entry_by_material_id(
             thermo_docs[3].material_id, conventional_unit_cell=False
         )
-        assert all(
-            e.uncorrected_energy_per_atom
-            == pytest.approx(
-                next(
-                    f for f in non_standardized if f.entry_id == e.entry_id
-                ).uncorrected_energy_per_atom
+        for e in as_conv:
+            ref = next(f for f in non_standardized if f.entry_id == e.entry_id)
+            assert e.uncorrected_energy_per_atom == pytest.approx(
+                ref.uncorrected_energy_per_atom
             )
-            for e in as_conv
-        )
+            # corrected too: every adjustment must scale with the cell, including
+            # extensive ones with no ``n_atoms``, e.g. the r2SCAN mixing correction
+            assert e.energy_per_atom == pytest.approx(ref.energy_per_atom)
 
         # Additional criteria
         entry = mpr.get_entries(
@@ -309,6 +313,19 @@ def test_get_entries_in_chemsys_mixed_hull(self, mpr):
         host = next(e for e in entries if e.composition.reduced_formula == "Cs2TiI6")
         assert phase_diagram.get_e_above_hull(host) == pytest.approx(0.0, abs=1e-6)
 
+        # extra criteria narrow the served (common-scale) entries, with MP's own semantics
+        with warnings.catch_warnings(record=True) as record:
+            stable = mpr.get_entries_in_chemsys(
+                "Cs-Ti-I", additional_criteria={"is_stable": True}
+            )
+        assert not [w for w in record if issubclass(w.category, MPRestWarning)]
+        assert 0 < len(stable) < len(entries)
+        assert {str(e.data["material_id"]) for e in stable} == {
+            str(e.data["material_id"])
+            for e in entries
+            if phase_diagram.get_e_above_hull(e) == pytest.approx(0.0, abs=1e-6)
+        }
+
         # hull distances must match the ones MP serves, and no material may go missing --
         # both fail if this silently falls through to re-applying the mixing scheme here
         docs = mpr.materials.thermo.search(
@@ -333,11 +350,47 @@ def test_get_entries_in_chemsys_mixed_hull(self, mpr):
                 for e in hull_entries
             )
 
-        # a narrowed query cannot be placed on a common scale, so it must say so
+        # uncorrected mixed entries cannot be placed on a common scale, so a warning is thrown:
         with pytest.warns(MPRestWarning, match="common energy scale"):
-            mpr.get_entries_in_chemsys(
-                "Cs-Ti-I", additional_criteria={"is_stable": True}
+            mpr.get_entries_in_chemsys("Cs-Ti-I", compatible_only=False)
+
+    def test_get_entries_in_chemsys_decorated_served_pd(self, mpr):
+        """
+        ``property_data`` / ``conventional_unit_cell`` requests are also served from
+        the pre-built phase diagram (and decorated post-hoc), rather than falling back
+        to re-applying the mixing scheme locally (which can differ from MP's hull).
+        """
+        entries = mpr.get_entries_in_chemsys("H-O")
+        decorated = mpr.get_entries_in_chemsys(
+            "H-O", property_data=["energy_above_hull"], conventional_unit_cell=True
+        )
+        served = {
+            str(e.entry_id): (
+                e.energy_per_atom,
+                len(e.structure),
+                e.composition.reduced_formula,
             )
+            for e in entries
+        }
+
+        # same served entry set, same energy scale, and the reshaping is not a no-op
+        assert {str(e.entry_id) for e in decorated} == set(served)
+        assert any(len(e.structure) != served[str(e.entry_id)][1] for e in decorated)
+        # not vacuous: served r2SCAN entries carry their mixing correction as an
+        # extensive ``ConstantEnergyAdjustment``, which has no ``n_atoms`` to scale
+        assert any(
+            isinstance(adj, ConstantEnergyAdjustment)
+            for e in entries
+            for adj in e.energy_adjustments
+        )
+
+        for entry in decorated:
+            energy, _, formula = served[str(entry.entry_id)]
+            assert entry.data["energy_above_hull"] >= 0
+            # conventional reshaping must preserve per-atom corrected energies (including
+            # the extensive mixing-scheme adjustments on r2SCAN entries) and stoichiometry:
+            assert entry.energy_per_atom == pytest.approx(energy, abs=1e-8)
+            assert entry.composition.reduced_formula == formula
 
     @pytest.mark.skipif(
         contribs_client is None,