Add additional_criteria option for get_entries_in_chemsys.

materialsproject · May 19, 2022 · ea688a4 · ea688a4
1 parent ec94645
commit ea688a4
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/pymatgen/ext/matproj.py b/pymatgen/ext/matproj.py
@@ -854,13 +854,12 @@ def get_entries_in_chemsys(
         inc_structure=None,
         property_data=None,
         conventional_unit_cell=False,
+        additional_criteria=None,
     ):
         """
-        Helper method to get a list of ComputedEntries in a chemical system.
-        For example, elements = ["Li", "Fe", "O"] will return a list of all
-        entries in the Li-Fe-O chemical system, i.e., all LixOy,
-        FexOy, LixFey, LixFeyOz, Li, Fe and O phases. Extremely useful for
-        creating phase diagrams of entire chemical systems.
+        Helper method to get a list of ComputedEntries in a chemical system. For example, elements = ["Li", "Fe", "O"]
+        will return a list of all entries in the Li-Fe-O chemical system, i.e., all LixOy, FexOy, LixFey, LixFeyOz,
+        Li, Fe and O phases. Extremely useful for creating phase diagrams of entire chemical systems.
 
         Args:
             elements (str or [str]): Chemical system string comprising element
@@ -882,6 +881,8 @@ def get_entries_in_chemsys(
                 supported_properties.
             conventional_unit_cell (bool): Whether to get the standard
                 conventional unit cell
+            additional_criteria (dict): Any additional criteria to pass. For instance, if you are only interested in
+                stable entries, you can pass {"e_above_hull": {"$lte": 0.001}}.
 
         Returns:
             List of ComputedEntries.
@@ -895,8 +896,12 @@ def get_entries_in_chemsys(
             for els in itertools.combinations(elements, i + 1):
                 all_chemsyses.append("-".join(sorted(els)))
 
+        criteria = {"chemsys": {"$in": all_chemsyses}}
+        if additional_criteria:
+            criteria.update(additional_criteria)
+
         entries = self.get_entries(
-            {"chemsys": {"$in": all_chemsyses}},
+            criteria,
             compatible_only=compatible_only,
             inc_structure=inc_structure,
             property_data=property_data,

diff --git a/pymatgen/ext/tests/test_matproj.py b/pymatgen/ext/tests/test_matproj.py
@@ -157,6 +157,9 @@ def test_get_entries_in_chemsys(self):
         e2 = {i.entry_id for i in entries2}
         self.assertTrue(e1 == e2)
 
+        stable_entries = self.rester.get_entries_in_chemsys(syms, additional_criteria={"e_above_hull": {"$lte": 0.001}})
+        self.assertTrue(len(stable_entries) < len(entries))
+
     def test_get_structure_by_material_id(self):
         s1 = self.rester.get_structure_by_material_id("mp-1")
         self.assertEqual(s1.formula, "Cs1")