Merge pull request #128 from colleenjg/cjg-dev

Progress bar for creating animations

Author: TomGeorge1234

ratinabox/Agent.py

@@ -843,6 +843,7 @@ def animate_trajectory(
         t_end=None, 
         fps=15, 
         speed_up=5, #by default the animation is 5x faster than real time 
+        progress_bar=False,
         autosave=None, 
         **kwargs
     ):
@@ -852,6 +853,7 @@ def animate_trajectory(
             t_end (_type_, optional): _description_. Defaults to None.
             fps: frames per second of end video
             speed_up: #times real speed animation should come out at
+            progress_bar (bool): if True, a progress bar will be shown as the animation is created. Defaults to False.
             autosave (bool): whether to automatical try and save this. Defaults to None in which case looks for global constant ratinabox.autosave_plots
             kwargs: passed to trajectory plotting function (chuck anything you wish in here). A particularly useful kwarg is 'additional_plot_func': any function which takes a fig, ax and t as input. The animation wll be passed through this each time after plotting the trajectory, use it to modify your animations however you like
 
@@ -895,13 +897,18 @@ def animate_(i, fig, ax, t_start, t_max, speed_up, dt, kwargs):
             t_start=0, t_end=10 * self.dt, xlim=t_end / 60, autosave=False, **kwargs
         )
 
+        frames = int((t_end - t_start) / (dt * speed_up))
+        if progress_bar:
+            from tqdm import tqdm
+            frames = tqdm(range(frames), position=0, leave=True)
+
         from matplotlib import animation
 
         anim = matplotlib.animation.FuncAnimation(
             fig,
             animate_,
             interval=1000 * dt,
-            frames=int((t_end - t_start) / (dt * speed_up)),
+            frames=frames,
             blit=False,
             fargs=(fig, ax, t_start, t_end, speed_up, dt, kwargs),
         )

ratinabox/Neurons.py

@@ -698,6 +698,7 @@ def animate_rate_timeseries(
         chosen_neurons="all",
         fps=15,
         speed_up=1,
+        progress_bar=False,
         autosave=None,
         **kwargs,
     ):
@@ -710,8 +711,10 @@ def animate_rate_timeseries(
         Args:
             • t_end (_type_, optional): _description_. Defaults to None.
             • chosen_neurons: Which neurons to plot. string "10" or 10 will plot ten of them, "all" will plot all of them, "12rand" will plot 12 random ones. A list like [1,4,5] will plot cells indexed 1, 4 and 5. Defaults to "all".
+            • fps: frames per second of end video. Defaults to 15.
+            • speed_up: #times real speed animation should come out at. Defaults to 1.
+            • progress_bar: if True, a progress bar will be shown as the animation is created. Default to False.
 
-            • speed_up: #times real speed animation should come out at.
 
         Returns:
             animation
@@ -753,13 +756,18 @@ def animate_(i, fig, ax, chosen_neurons, t_start, t_max, dt, speed_up):
             **kwargs,
         )
 
+        frames = int((t_end - t_start) / (dt * speed_up))
+        if progress_bar:
+            from tqdm import tqdm
+            frames = tqdm(range(frames), position=0, leave=True)
+
         from matplotlib import animation
 
         anim = matplotlib.animation.FuncAnimation(
             fig,
             animate_,
             interval=1000 * dt,
-            frames=int((t_end - t_start) / (dt * speed_up)),
+            frames=frames,
             blit=False,
             fargs=(fig, ax, chosen_neurons, t_start, t_end, dt, speed_up),
         )