fea: add verbose option to compatiblity

materialsproject · May 7, 2021 · 111f334 · 111f334
1 parent adee297
commit 111f334
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/pymatgen/entries/compatibility.py b/pymatgen/entries/compatibility.py
@@ -28,6 +28,7 @@
     TemperatureEnergyAdjustment,
 )
 from pymatgen.io.vasp.sets import MITRelaxSet, MPRelaxSet
+from pymatgen.util.sequence import PBar
 
 MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
 MU_H2O = -2.4583  # Free energy of formation of water, eV/H2O, used by MaterialsProjectAqueousCompatibility
@@ -532,7 +533,7 @@ def process_entry(self, entry):
             return self.process_entries(entry)[0]
         return None
 
-    def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = True):
+    def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = True, verbose: bool = False):
         """
         Process a sequence of entries with the chosen Compatibility scheme. Note
         that this method will change the data of the original entries.
@@ -542,6 +543,8 @@ def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = Tru
             clean: bool, whether to remove any previously-applied energy adjustments.
                 If True, all EnergyAdjustment are removed prior to processing the Entry.
                 Default is True.
+            verbose: bool, whether to display progress bar for processing multiple entries.
+                Default is False.
 
         Returns:
             A list of adjusted entries.  Entries in the original list which
@@ -553,7 +556,7 @@ def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = Tru
 
         processed_entry_list = []
 
-        for entry in entries:
+        for entry in PBar(entries, disable=(not verbose)):
             ignore_entry = False
             # if clean is True, remove all previous adjustments from the entry
             if clean:
@@ -1380,7 +1383,7 @@ def get_adjustments(self, entry: ComputedEntry):
 
         return adjustments
 
-    def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = False):
+    def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = False, verbose: bool = False):
         """
         Process a sequence of entries with the chosen Compatibility scheme.
 
@@ -1389,6 +1392,8 @@ def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = Fal
             clean: bool, whether to remove any previously-applied energy adjustments.
                 If True, all EnergyAdjustment are removed prior to processing the Entry.
                 Default is False.
+            verbose: bool, whether to display progress bar for processing multiple entries.
+                Default is False.
 
         Returns:
             A list of adjusted entries.  Entries in the original list which
@@ -1415,4 +1420,4 @@ def process_entries(self, entries: Union[ComputedEntry, list], clean: bool = Fal
                 self.h2o_energy = h2o_entries[0].energy_per_atom
                 self.h2o_adjustments = h2o_entries[0].correction / h2o_entries[0].composition.num_atoms
 
-        return super().process_entries(entries, clean=clean)
+        return super().process_entries(entries, clean=clean, verbose=verbose)