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,11 @@ def generate_plot(
     cols_in_legend_box: int = 3,
     legend_position: typing.Optional[str] = "best",
     show_plot_already: bool = True,
+    animate: bool = False,
+    animate_duration: int = 5,
+    animate_writer: typing.Tuple[str, typing.List[str]] = ("pillow", []),
     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 +129,19 @@ 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). Recommended to pass close_plt as True
+    :type animate: boolean
+    :param animate_duration: approx duration (seconds) of the animation. animate should be True (default: 5)
+    :type animate: int
+    :param animate_writer: specify different writer for saving animation (default: ("pillow", [])) :
+                See: https://matplotlib.org/stable/users/explain/animations/animations.html#saving-animations
+                format : writer=( "writer name", ["extra args list"] )
+                example : writer=( "imagemagick", ["-quality", "100"] )
+    :type animate_writer: tuple
     :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 +190,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 +199,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 +210,7 @@ def generate_plot(
                 label=label,
             )
         else:
-            plt.plot(
+            (artist,) = plt.plot(
                 xvalues[i],
                 yvalues[i],
                 marker=marker,
@@ -202,6 +219,7 @@ def generate_plot(
                 linewidth=linewidth,
                 label=label,
             )
+        artists.append(artist)
 
     if labels:
         if legend_position == "outer right":
@@ -240,6 +258,61 @@ def generate_plot(
         logger.info("Saved image to %s of plot: %s" % (save_figure_to, title))
 
     if show_plot_already:
+        if animate:
+            duration = animate_duration * 1000      # in ms
+            size = max(len(val) for val in xvalues) # maximum length
+            interval = 50                           # Delay between frames in milliseconds
+            pockets = duration // interval
+            skip = max(size // pockets, 1)
+            logger.info(
+                "Animation hyperparameters : duration=%sms, 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
+            )
+
+            if save_animation_to:
+                frame_length_threshold = 1500
+                if size > frame_length_threshold:
+                    minute = 1000
+                    logger.warning( # (approx 1 minute for 1000 frames)
+                        "Large Plot found!!! Saving animation will take approx %.2f minutes" % (size/minute))
+
+                logger.info("Saving animation to %s" %
+                            (save_animation_to))
+                writers = ["pillow", "html", "ffmpeg", "imagemagick"]
+                writer_name, writer_extra = animate_writer
+                if writer_name not in writers:
+                    writer_name = "pillow"
+                    writer_extra = []
+
+                try:
+                    ani.save(
+                        filename=save_animation_to,
+                        writer=writer_name,
+                        extra_args=writer_extra,  # pillow does not support extra_args
+                        progress_callback=lambda i, n: print(
+                            f'Saving frame {i+1}/{n}')
+                    )
+                except:
+                    ani.save(
+                        filename=save_animation_to,
+                        writer=writer_name,
+                        progress_callback=lambda i, n: print(
+                            f'Saving frame {i+1}/{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,128 @@
 
 
 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 animated",
+            xaxis="x",
+            yaxis="y",
+            grid=False,
+            show_plot_already=True,
+            animate=True,
+            legend_position="right",
+            save_animation_to=filename,
+            close_plot=True
+        )
+        self.assertIsFile(filename)
+        pl.Path(filename).unlink()
+
+    @pytest.mark.localonly
+    def test_generate_plot_animated_specify_writer(self):
+        """Test generate_plot function with specific writer."""
+        filename = "tests/plot/test_generate_plot_writer.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 animated writer",
+            xaxis="x",
+            yaxis="y",
+            grid=False,
+            show_plot_already=True,
+            animate=True,
+            legend_position="right",
+            save_animation_to=filename,
+            animate_writer=("ffmpeg", []),
+            close_plot=True
+        )
+        self.assertIsFile(filename)
+        pl.Path(filename).unlink()
+
+    @pytest.mark.localonly
+    def test_generate_plot_animated_should_default_pillow_when_writer_invalid(self):
+        """Test generate_plot function does not fail when writer is invalid."""
+        filename1 = "tests/plot/test_generate_plot_writer1.gif"
+        filename2 = "tests/plot/test_generate_plot_writer2.gif"
+
+        # remove the file first
+        try:
+            pl.Path(filename1).unlink()
+            pl.Path(filename2).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 animated writer",
+            xaxis="x",
+            yaxis="y",
+            grid=False,
+            show_plot_already=True,
+            animate=True,
+            legend_position="right",
+            save_animation_to=filename1,
+            # imagemagick is not a requirement in pyNeuroML
+            animate_writer=("imagemagick", ["-quality", "100"]),
+            close_plot=True  # without this, all plots not closed will also be plotted
+        )
+
+        generate_plot(
+            [xs, xs],
+            [ys, ys2],
+            "Test plot animated writer",
+            xaxis="x",
+            yaxis="y",
+            grid=False,
+            show_plot_already=True,
+            animate=True,
+            legend_position="right",
+            save_animation_to=filename2,
+            # tests will use pillow by default
+            animate_writer=("imaginary_writer", [""]),
+            close_plot=True  # without this, all plots not closed will also be plotted
+        )
+
+        self.assertIsFile(filename1)
+        self.assertIsFile(filename2)
+        pl.Path(filename1).unlink()
+        pl.Path(filename2).unlink()
+
     def test_generate_plot(self):
         """Test generate_plot function."""
         filename = "tests/plot/test_generate_plot.png"
@@ -50,6 +171,7 @@ def test_generate_plot(self):
             show_plot_already=False,
             save_figure_to=filename,
             legend_position="right",
+            close_plot=True
         )
         self.assertIsFile(filename)
         pl.Path(filename).unlink()