Merge remote-tracking branch 'origin/main'

Author: TomGeorge1234

README.md

@@ -1,15 +1,4 @@
-⚠️⚠️⚠️WARNING: THIS BRANCH IS DEPRECATED⚠️⚠️⚠️
-
-The new default branch is `origin/main`, please make all PRs etc. here. To reassign your local branches upstream to be `origin/main` do 
-
-```
-git branch --set-upstream-to=origin/main <name-of-your-local-branch.
-```
-
-
-
-
-# RatInABox ![Tests](https://github.com/RatInABox-Lab/RatInABox/actions/workflows/test.yml/badge.svg)   [![PyPI version](https://badge.fury.io/py/ratinabox.svg)](https://badge.fury.io/py/ratinabox) [![Downloads](https://static.pepy.tech/badge/ratinabox)](https://pepy.tech/project/ratinabox)<img align="right" src=".images/readme/logo.png" width=150> 
+# RatInABox ![Tests](https://github.com/RatInABox-Lab/RatInABox/actions/workflows/test.yml/badge.svg)   [![PyPI version](https://badge.fury.io/py/ratinabox.svg)](https://badge.fury.io/py/ratinabox) [![Downloads](https://static.pepy.tech/badge/ratinabox)](https://pepy.tech/project/ratinabox) [![PyPI Downloads](https://static.pepy.tech/badge/ratinabox/week)](https://pepy.tech/projects/ratinabox)<img align="right" src=".images/readme/logo.png" width=150> 
 
 `RatInABox` (see [paper](https://elifesciences.org/articles/85274)) is a toolkit for generating synthetic behaviour and neural data for spatially and/or velocity selective cell types in complex continuous environments. 
 

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),
         )