Define needs_u_correction(comp: CompositionLike) -> set[str] utilit…

…y function (#3703)

* add needs_u_correction(comp: CompositionLike) -> set[str] utility function

* test needs_u_correction

* expected_u, actual_u ensure float

* make needs_u_correction accept a U-correction config dict, defaults to MP2020

* test_needs_u_correction pass different u_configs (MP and MP2020)

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ ci:
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.3.3
+    rev: v0.3.4
     hooks:
       - id: ruff
         args: [--fix, --unsafe-fixes]

pymatgen/entries/compatibility.py

@@ -19,7 +19,7 @@
 from uncertainties import ufloat
 
 from pymatgen.analysis.structure_analyzer import oxide_type, sulfide_type
-from pymatgen.core import SETTINGS, Element
+from pymatgen.core import SETTINGS, Composition, Element
 from pymatgen.entries.computed_entries import (
     CompositionEnergyAdjustment,
     ComputedEntry,
@@ -34,6 +34,8 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    from pymatgen.util.typing import CompositionLike
+
 __author__ = "Amanda Wang, Ryan Kingsbury, Shyue Ping Ong, Anubhav Jain, Stephen Dacek, Sai Jayaraman"
 __copyright__ = "Copyright 2012-2020, The Materials Project"
 __version__ = "1.0"
@@ -43,6 +45,13 @@
 
 MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
 MU_H2O = -2.4583  # Free energy of formation of water, eV/H2O, used by MaterialsProjectAqueousCompatibility
+MP2020_COMPAT_CONFIG = loadfn(f"{MODULE_DIR}/MP2020Compatibility.yaml")
+MP_COMPAT_CONFIG = loadfn(f"{MODULE_DIR}/MPCompatibility.yaml")
+
+assert (  # ping @janosh @rkingsbury on GitHub if this fails
+    MP2020_COMPAT_CONFIG["Corrections"]["GGAUMixingCorrections"]["O"]
+    == MP2020_COMPAT_CONFIG["Corrections"]["GGAUMixingCorrections"]["F"]
+), "MP2020Compatibility.yaml expected to have the same Hubbard U corrections for O and F"
 
 AnyComputedEntry = Union[ComputedEntry, ComputedStructureEntry]
 
@@ -790,13 +799,13 @@ def __init__(
         self.compat_type = compat_type
         self.correct_peroxide = correct_peroxide
         self.check_potcar_hash = check_potcar_hash
-        fp = f"{MODULE_DIR}/MPCompatibility.yaml"
+        file_path = f"{MODULE_DIR}/MPCompatibility.yaml"
         super().__init__(
             [
                 PotcarCorrection(MPRelaxSet, check_hash=check_potcar_hash),
-                GasCorrection(fp),
-                AnionCorrection(fp, correct_peroxide=correct_peroxide),
-                UCorrection(fp, MPRelaxSet, compat_type),
+                GasCorrection(file_path),
+                AnionCorrection(file_path, correct_peroxide=correct_peroxide),
+                UCorrection(file_path, MPRelaxSet, compat_type),
             ]
         )
 
@@ -891,21 +900,22 @@ def __init__(
         if config_file:
             if os.path.isfile(config_file):
                 self.config_file: str | None = config_file
-                c = loadfn(self.config_file)
+                config = loadfn(self.config_file)
             else:
                 raise ValueError(f"Custom MaterialsProject2020Compatibility {config_file=} does not exist.")
         else:
             self.config_file = None
-            c = loadfn(f"{MODULE_DIR}/MP2020Compatibility.yaml")
 
-        self.name = c["Name"]
-        self.comp_correction = c["Corrections"].get("CompositionCorrections", defaultdict(float))
-        self.comp_errors = c["Uncertainties"].get("CompositionCorrections", defaultdict(float))
+            config = MP2020_COMPAT_CONFIG
+
+        self.name = config["Name"]
+        self.comp_correction = config["Corrections"].get("CompositionCorrections", defaultdict(float))
+        self.comp_errors = config["Uncertainties"].get("CompositionCorrections", defaultdict(float))
 
         if self.compat_type == "Advanced":
             self.u_settings = MPRelaxSet.CONFIG["INCAR"]["LDAUU"]
-            self.u_corrections = c["Corrections"].get("GGAUMixingCorrections", defaultdict(float))
-            self.u_errors = c["Uncertainties"].get("GGAUMixingCorrections", defaultdict(float))
+            self.u_corrections = config["Corrections"].get("GGAUMixingCorrections", defaultdict(float))
+            self.u_errors = config["Uncertainties"].get("GGAUMixingCorrections", defaultdict(float))
         else:
             self.u_settings = {}
             self.u_corrections = {}
@@ -1072,10 +1082,10 @@ def get_adjustments(self, entry: AnyComputedEntry) -> list[EnergyAdjustment]:
         for el in comp.elements:
             symbol = el.symbol
             # Check for bad U values
-            expected_u = u_settings.get(symbol, 0)
-            actual_u = calc_u.get(symbol, 0)
+            expected_u = float(u_settings.get(symbol, 0))
+            actual_u = float(calc_u.get(symbol, 0))
             if actual_u != expected_u:
-                raise CompatibilityError(f"Invalid U value of {actual_u:.1f} on {symbol}, expected {expected_u:.1f}")
+                raise CompatibilityError(f"Invalid U value of {actual_u:.3} on {symbol}, expected {expected_u:.3}")
             if symbol in u_corrections:
                 adjustments.append(
                     CompositionEnergyAdjustment(
@@ -1457,3 +1467,29 @@ def process_entries(
             self.h2_energy = h2_entries[0].energy_per_atom  # type: ignore[assignment]
 
         return super().process_entries(entries, clean=clean, verbose=verbose, inplace=inplace, on_error=on_error)
+
+
+def needs_u_correction(
+    comp: CompositionLike,
+    u_config: dict[str, dict[str, float]] = MP2020_COMPAT_CONFIG["Corrections"]["GGAUMixingCorrections"],
+) -> set[str]:
+    """Check if a composition is Hubbard U-corrected in the Materials Project 2020
+    GGA/GGA+U mixing scheme.
+
+    Args:
+        comp (CompositionLike): The formula/composition to check.
+        u_config (dict): The U-correction configuration to use. Default is the
+            Materials Project 2020 configuration.
+
+    Returns:
+        set[str]: The subset of elements whose combination requires a U-correction. Pass
+            return value to bool(ret_val) if you just want True/False.
+    """
+    elements = set(map(str, Composition(comp).elements))
+    has_u_anion = set(u_config) & elements
+
+    u_corrected_cations = set(u_config["O"])
+    has_u_cation = u_corrected_cations & elements
+    if has_u_cation and has_u_anion:
+        return has_u_cation | has_u_anion
+    return set()

tests/entries/test_compatibility.py

@@ -7,6 +7,7 @@
 from collections import defaultdict
 from math import sqrt
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 from monty.json import MontyDecoder
@@ -18,6 +19,8 @@
 from pymatgen.core.lattice import Lattice
 from pymatgen.core.structure import Structure
 from pymatgen.entries.compatibility import (
+    MP2020_COMPAT_CONFIG,
+    MP_COMPAT_CONFIG,
     MU_H2O,
     AqueousCorrection,
     Compatibility,
@@ -27,10 +30,14 @@
     MaterialsProjectCompatibility,
     MITAqueousCompatibility,
     MITCompatibility,
+    needs_u_correction,
 )
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry, ConstantEnergyAdjustment
 from pymatgen.util.testing import TEST_FILES_DIR
 
+if TYPE_CHECKING:
+    from pymatgen.util.typing import CompositionLike
+
 
 class TestCorrectionSpecificity(unittest.TestCase):
     """Make sure corrections are only applied to GGA or GGA+U entries."""
@@ -2034,3 +2041,26 @@ def test_errors(self):
         ):
             corrected_entry = self.compat.process_entry(entry)
             assert corrected_entry.correction_uncertainty == approx(expected)
+
+
+@pytest.mark.parametrize(
+    "u_config",
+    [MP2020_COMPAT_CONFIG["Corrections"]["GGAUMixingCorrections"], MP_COMPAT_CONFIG["Advanced"]["UCorrections"]],
+)
+@pytest.mark.parametrize(
+    ("comp", "expected"),
+    [
+        ("Fe2O3", {"Fe", "O"}),
+        ("Fe3O4", {"Fe", "O"}),
+        ("FeS", set()),
+        ("FeF3", {"Fe", "F"}),
+        ("LiH", set()),
+        ("H", set()),
+        (Composition("MnO"), {"Mn", "O"}),
+        (Composition("MnO2"), {"Mn", "O"}),
+        (Composition("LiFePO4"), {"Fe", "O"}),
+        (Composition("LiFePS4"), set()),
+    ],
+)
+def test_needs_u_correction(comp: CompositionLike, expected: set[str], u_config: dict):
+    assert needs_u_correction(comp, u_config=u_config) == expected