Suppress pymatgen element data warnings when imputing

hackingmaterials · Apr 10, 2024 · be1531c · be1531c
1 parent d75d28a
commit be1531c
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/matminer/utils/data.py b/matminer/utils/data.py
@@ -357,7 +357,12 @@ def get_elemental_property(self, elem, property_name):
             block_key = {"s": 1.0, "p": 2.0, "d": 3.0, "f": 3.0}
             return block_key[getattr(elem, property_name)]
         else:
-            value = getattr(elem, property_name)
+            # Suppress pymatgen warnings for missing properties if we are going to impute anyway
+            with warnings.catch_warnings():
+                if self.impute_nan:
+                    warnings.simplefilter("ignore", category=UserWarning)
+                value = getattr(elem, property_name)
+
             if self.impute_nan:
                 if value and not pd.isnull(value):
                     return value