Merge pull request #2894 from abillscmu/progress

add (optional) progress bar for cycling experiments

Author: valentinsulzer

pybamm/simulation.py

@@ -8,6 +8,7 @@
 import warnings
 import sys
 from functools import lru_cache
+import tqdm
 
 
 def is_notebook():
@@ -510,6 +511,7 @@ def solve(
         starting_solution=None,
         initial_soc=None,
         callbacks=None,
+        showprogress=False,
         **kwargs,
     ):
         """
@@ -557,6 +559,10 @@ def solve(
         callbacks : list of callbacks, optional
             A list of callbacks to be called at each time step. Each callback must
             implement all the methods defined in :class:`pybamm.callbacks.BaseCallback`.
+        showprogress : bool, optional
+            Whether to show a progress bar for cycling. If true, shows a progress bar
+            for cycles. Has no effect when not used with an experiment.
+            Default is False.
         **kwargs
             Additional key-word arguments passed to `solver.solve`.
             See :meth:`pybamm.BaseSolver.solve`.
@@ -698,7 +704,13 @@ def solve(
             num_cycles = len(self.experiment.cycle_lengths)
             feasible = True  # simulation will stop if experiment is infeasible
             for cycle_num, cycle_length in enumerate(
-                self.experiment.cycle_lengths, start=1
+                # tqdm is the progress bar.
+                tqdm.tqdm(
+                    self.experiment.cycle_lengths,
+                    disable=(not showprogress),
+                    desc="Cycling",
+                ),
+                start=1,
             ):
                 logs["cycle number"] = (
                     cycle_num + cycle_offset,

requirements.txt

@@ -9,6 +9,7 @@ imageio>=2.9.0
 pybtex>=0.24.0
 sympy >= 1.8
 bpx
+tqdm
 # Note: Matplotlib is loaded for debug plots but to ensure pybamm runs
 # on systems without an attached display it should never be imported
 # outside of plot() methods.

setup.py

@@ -215,6 +215,7 @@ def compile_KLU():
         "pybtex>=0.24.0",
         "sympy>=1.8",
         "bpx",
+        "tqdm",
         # Note: Matplotlib is loaded for debug plots, but to ensure pybamm runs
         # on systems without an attached display, it should never be imported
         # outside of plot() methods.

tests/unit/test_experiments/test_simulation_with_experiment.py

@@ -342,6 +342,21 @@ def test_run_experiment_termination_capacity(self):
         # all but the last value should be above the termination condition
         np.testing.assert_array_less(5.04, C[:-1])
 
+    def test_run_experiment_with_pbar(self):
+        # The only thing to test here is for errors.
+        experiment = pybamm.Experiment(
+            [
+                (
+                    "Discharge at 1C for 1 sec",
+                    "Charge at 1C for 1 sec",
+                ),
+            ]
+            * 10,
+        )
+        model = pybamm.lithium_ion.SPM()
+        sim = pybamm.Simulation(model, experiment=experiment)
+        sim.solve(showprogress=True)
+
     def test_run_experiment_termination_voltage(self):
         # with percent
         experiment = pybamm.Experiment(