animate matplotlib plots

pyneuroml/plot/Plot.py

@@ -12,6 +12,7 @@
 import typing
 import matplotlib
 import matplotlib.axes
+import matplotlib.animation as animation
 from typing import Optional
 
 logger = logging.getLogger(__name__)
@@ -42,7 +43,9 @@ def generate_plot(
     cols_in_legend_box: int = 3,
     legend_position: typing.Optional[str] = "best",
     show_plot_already: bool = True,
+    animate: bool = False,
     save_figure_to: typing.Optional[str] = None,
+    save_animation_to: typing.Optional[str] = None,
     title_above_plot: bool = False,
     verbose: bool = False,
     close_plot: bool = False,
@@ -124,8 +127,12 @@ def generate_plot(
     :type legend_position: str
     :param show_plot_already: if plot should be shown when created (default: True)
     :type show_plot_already: boolean
+    :param animate: if shown plot should be animated. show_splot_already should be True (default: False)
+    :type animate: boolean
     :param save_figure_to: location to save generated figure to (default: None)
     :type save_figure_to: str
+    :param save_animation_to: location to save generated animation to (default: None)
+    :type save_animation_to: str
     :param title_above_plot: enable/disable title above the plot (default: False)
     :type title_above_plot: boolean
     :param verbose: enable/disable verbose logging (default: False)
@@ -174,6 +181,7 @@ def generate_plot(
     if not show_yticklabels:
         ax.set_yticklabels([])
 
+    artists = []
     for i in range(len(xvalues)):
         linestyle = rcParams["lines.linestyle"] if not linestyles else linestyles[i]
         label = "" if not labels else labels[i]
@@ -182,7 +190,7 @@ def generate_plot(
         markersize = rcParams["lines.markersize"] if not markersizes else markersizes[i]
 
         if colors:
-            plt.plot(
+            (artist,) = plt.plot(
                 xvalues[i],
                 yvalues[i],
                 marker=marker,
@@ -193,7 +201,7 @@ def generate_plot(
                 label=label,
             )
         else:
-            plt.plot(
+            (artist,) = plt.plot(
                 xvalues[i],
                 yvalues[i],
                 marker=marker,
@@ -202,6 +210,7 @@ def generate_plot(
                 linewidth=linewidth,
                 label=label,
             )
+        artists.append(artist)
 
     if labels:
         if legend_position == "outer right":
@@ -240,6 +249,41 @@ def generate_plot(
         logger.info("Saved image to %s of plot: %s" % (save_figure_to, title))
 
     if show_plot_already:
+        if animate:
+            d = 5
+            duration = d * 1000 # in ms
+            size = len(xvalues[0])
+            interval = 50  # Delay between frames in milliseconds
+            pockets = duration // interval
+            skip = max(size // pockets, 1)
+            logger.info(
+                "Animation hyperparameters : duration=%s, size=%s, interval=%s, pockets=%s, skip=%s" % (duration, size, interval, pockets,  skip))
+
+            def update(frame):
+                for i, artist in enumerate(artists):
+                    artist.set_xdata(xvalues[i][:frame*skip])
+                    artist.set_ydata(yvalues[i][:frame*skip])
+                return artists
+
+            ani = animation.FuncAnimation(
+                fig=fig,
+                frames=size-1,
+                func=update,
+                interval=interval,
+                blit=True,
+                cache_frame_data=False
+            )
+
+            frame_length_threshold = 5000
+            if len(xvalues[0]) < frame_length_threshold and save_animation_to:
+                logger.info("Saving animation to %s" %
+                            (save_animation_to))
+                ani.save(
+                    filename=save_animation_to,
+                    writer="pillow",
+                    progress_callback=lambda i, n: print(
+                        f'Saving frame {i}/{n}')
+                )
         plt.show()
 
     if close_plot:

tests/plot/test_plot.py

@@ -7,6 +7,8 @@
 Copyright 2023 NeuroML contributors
 """
 
+import random
+import pytest
 import unittest
 import logging
 import pathlib as pl
@@ -19,9 +21,39 @@
 
 
 class TestPlot(BaseTestCase):
-
     """Test Plot module"""
 
+    @pytest.mark.localonly
+    def test_generate_plot_animated(self):
+        """Test generate_plot function."""
+        filename = "tests/plot/test_generate_plot.gif"
+
+        # remove the file first
+        try:
+            pl.Path(filename).unlink()
+        except FileNotFoundError:
+            pass
+
+        numpoints = 100
+        xs = list(range(0, numpoints))
+        ys = random.choices(list(range(0, 1000)), k=numpoints)
+        ys2 = random.choices(list(range(0, 1500)), k=numpoints)
+
+        generate_plot(
+            [xs, xs],
+            [ys, ys2],
+            "Test plot",
+            xaxis="x",
+            yaxis="y",
+            grid=False,
+            show_plot_already=True,
+            animate=True,
+            legend_position="right",
+            save_animation_to=filename
+        )
+        self.assertIsFile(filename)
+        pl.Path(filename).unlink()
+
     def test_generate_plot(self):
         """Test generate_plot function."""
         filename = "tests/plot/test_generate_plot.png"