test_needs_u_correction pass different u_configs (MP and MP2020)

materialsproject · Mar 26, 2024 · 20e60b2 · 20e60b2
1 parent 3f592fd
commit 20e60b2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/pymatgen/entries/compatibility.py b/pymatgen/entries/compatibility.py
@@ -46,6 +46,7 @@
 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"]
@@ -798,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),
             ]
         )
 

diff --git a/tests/entries/test_compatibility.py b/tests/entries/test_compatibility.py
@@ -19,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,
@@ -2041,6 +2043,10 @@ def test_errors(self):
             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"),
     [
@@ -2056,5 +2062,5 @@ def test_errors(self):
         (Composition("LiFePS4"), set()),
     ],
 )
-def test_needs_u_correction(comp: CompositionLike, expected: set[str]):
-    assert needs_u_correction(comp) == expected
+def test_needs_u_correction(comp: CompositionLike, expected: set[str], u_config: dict):
+    assert needs_u_correction(comp, u_config=u_config) == expected