FEAT: show and filepath support in Array.plot()

Author: gdementen

doc/source/changes/version_0_35.rst.inc

@@ -15,7 +15,10 @@ Syntax changes
 Backward incompatible changes
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-* other backward incompatible changes
+* Plots made with Array.plot() in a Python script will be shown by default,
+  unless either the filepath (see below) or ax arguments are used. Shown plots
+  will open a window and pause the running script until the window is closed by
+  the user. To revert to the previous behavior, use show=False.
 
 
 New features
@@ -30,6 +33,12 @@ New features
 
   will create an animated bar plot with one frame per year.
 
+* implemented Array.plot `filepath` argument to save plots to a file directly,
+  without having to use the matplotlib API.
+
+* implemented Array.plot `show` argument to display plots directly, without
+  having to use the matplotlib API. This is the new default behavior.
+
 * added a feature (see the :ref:`miscellaneous section <misc>` for details). It works on :ref:`api-axis` and
   :ref:`api-group` objects.
 

larray/core/array.py

@@ -7116,8 +7116,6 @@ def to_clipboard(self, *args, **kwargs) -> None:
     def plot(self) -> PlotObject:
         r"""Plot the data of the array into a graph (window pop-up).
 
-        The graph can be tweaked to achieve the desired formatting and can be saved to a .png file.
-
         Parameters
         ----------
         kind : str
@@ -7132,6 +7130,14 @@ def plot(self) -> PlotObject:
             - 'pie' : pie plot
             - 'scatter' : scatter plot (if array's dimensions >= 2)
             - 'hexbin' : hexbin plot (if array's dimensions >= 2)
+        filepath : str or Path, default None
+            Save plot as a file at `filepath`. Defaults to None (do not save).
+            When saving the plot to a file, the function returns None. In other
+            words, in that case, the plot is no longer available for further
+            tweaking or display.
+        show : bool, optional
+            Whether to display the plot directly.
+            Defaults to True if `filepath` is None and `ax` is None, False otherwise.
         ax : matplotlib axes object, default None
         subplots : boolean, Axis, int, str or tuple, default False
             Make several subplots.
@@ -7232,29 +7238,26 @@ def plot(self) -> PlotObject:
 
         Examples
         --------
-        >>> import matplotlib.pyplot as plt
-        >>> # let us define an array with some made up data
+        Let us first define an array with some made up data
+
         >>> import larray as la
-        >>> arr = la.Array([[5, 20, 5, 10],                 [6, 16, 8, 11]], 'gender=M,F;year=2018..2021')
+        >>> arr = la.Array([[5, 20, 5, 10],
+        ...              [6, 16, 8, 11]], 'gender=M,F;year=2018..2021')
 
         Simple line plot
 
         >>> arr.plot()
-        >>> # show figure (it also resets it after showing it! Do not call it before savefig)
-        >>> plt.show()
+        <Axes: xlabel='year'>
 
-        Line plot with grid and a title
+        Line plot with grid and a title, saved in a file
 
-        >>> arr.plot(grid=True, title='line plot')
-        >>> # save figure in a file (see matplotlib.pyplot.savefig documentation for more details)
-        >>> plt.savefig('my_file.png')
+        >>> arr.plot(grid=True, title='line plot', filepath='my_file.png')
 
         2 bar plots (one for each gender) sharing the same y axis, which makes sub plots
         easier to compare. By default sub plots are independant of each other and the axes
         ranges are computed to "fit" just the data for their individual plot.
 
-        >>> arr.plot.bar(subplots='gender', sharey=True)
-        >>> plt.show()
+        >>> arr.plot.bar(subplots='gender', sharey=True)                       # doctest: +SKIP
 
         A stacked bar plot (genders are stacked)
 
@@ -7263,10 +7266,11 @@ def plot(self) -> PlotObject:
         An animated bar chart (with two bars). We set explicit y bounds via ylim so that the
         same boundaries are used for the whole animation.
 
-        >>> arr.plot.bar(animate='year', ylim=(0, 22))                         # doctest: +SKIP
+        >>> arr.plot.bar(animate='year', ylim=(0, 22), filepath='myanim.avi')  # doctest: +SKIP
 
         Create a figure containing 2 x 2 graphs
 
+        >>> import matplotlib.pyplot as plt
         >>> # see matplotlib.pyplot.subplots documentation for more details
         >>> fig, ax = plt.subplots(2, 2, figsize=(10, 8), tight_layout=True)   # doctest: +SKIP
         >>> # line plot with 2 curves (Males and Females) in the top left corner (0, 0)

larray/core/plot.py

@@ -181,7 +181,7 @@ def _plot_array(array, *args, x=None, y=None, series=None, _x_axes_last=False, *
     @deprecate_kwarg('stacked', 'stack')
     def __call__(self, x=None, y=None, ax=None, subplots=False, layout=None, figsize=None,
                  sharex=None, sharey=False, tight_layout=None, constrained_layout=None, title=None, legend=None,
-                 animate=None, filepath=None, **kwargs):
+                 animate=None, filepath=None, show=None, **kwargs):
         from matplotlib import pyplot as plt
 
         array = self.array
@@ -200,14 +200,14 @@ def __call__(self, x=None, y=None, ax=None, subplots=False, layout=None, figsize
             kwargs['stacked'] = True
 
         animate_axes, subplot_axes, x, y, series_axes = PlotObject._handle_x_y_axes(array.axes, animate, subplots, x, y)
-
+        if show is None:
+            show = filepath is None and ax is None
         if constrained_layout is None and tight_layout is None:
             constrained_layout = True
 
         if ax is None:
             fig = plt.figure(figsize=figsize, tight_layout=tight_layout, constrained_layout=constrained_layout)
-
-            if subplots:
+            if subplots or layout is not None:
                 if layout is None:
                     subplots_shape = subplot_axes.shape
                     if len(subplots_shape) > 2:
@@ -217,37 +217,33 @@ def __call__(self, x=None, y=None, ax=None, subplots=False, layout=None, figsize
                         layout = subplot_axes.shape
                 if sharex is None:
                     sharex = True
-                ax = fig.subplots(*layout, sharex=sharex, sharey=sharey)
+                ax_to_return = fig.subplots(*layout, sharex=sharex, sharey=sharey)
+                ax = ax_to_return if subplots else ax_to_return.flat[0]
             else:
                 ax = fig.add_subplot()
+                ax_to_return = ax
+        else:
+            fig = ax.figure
+            ax_to_return = ax
 
+        anim_kwargs = kwargs.pop('anim_params', {})
         if animate:
-            import matplotlib.animation as animation
+            from matplotlib.animation import FuncAnimation
 
-            def run(t):
-                if subplots:
+            if subplots:
+                def run(t):
                     for subplot_ax in ax.flat:
                         subplot_ax.clear()
-                else:
+                    self._plot_many(array[t], ax, kwargs, series_axes, subplot_axes, title, x, y)
+            else:
+                def run(t):
                     ax.clear()
-                self._plot_many(array[t], ax, kwargs, series_axes, subplot_axes, title, x, y)
+                    self._plot_many(array[t], ax, kwargs, series_axes, subplot_axes, title, x, y)
             # TODO: add support for interpolation between frames/labels
             #  see https://github.com/julkaar9/pynimate for inspiration
-            ani = animation.FuncAnimation(fig, run, frames=animate_axes.iter_labels())
-            if not isinstance(filepath, Path):
-                filepath = Path(filepath)
-            print(f"Writing animation to {filepath} ...", end=' ', flush=True)
-            if '.htm' in filepath.suffix:
-                filepath.write_text(f'<html>{ani.to_html5_video()}</html>', encoding='utf8')
-            else:
-                # writer = self.writer
-                # if writer is None:
-                writer = 'pillow' if filepath.suffix == '.gif' else 'ffmpeg'
-                fps = 5
-                metadata = None
-                bitrate = None
-                ani.save(filepath, writer=writer, fps=fps, metadata=metadata, bitrate=bitrate)
+            ani = FuncAnimation(fig, run, frames=animate_axes.iter_labels())
         else:
+            ani = None
             self._plot_many(array, ax, kwargs, series_axes, subplot_axes, title, x, y)
 
         if legend or legend is None:
@@ -267,7 +263,56 @@ def run(t):
                 # use figure to place legend to add a single legend for all subplots
                 legend_parent = first_ax.figure if subplots else ax
                 legend_parent.legend(handles, labels, **legend_kwargs)
-        return ax
+
+        if filepath is not None:
+            if ani is None:
+                fig.savefig(filepath)
+            else:
+                if not isinstance(filepath, Path):
+                    filepath = Path(filepath)
+                if filepath.suffix in {'.htm', '.html'}:
+                    # TODO: we should offer the option to use to_jshtml instead of to_html5_video. Even if it makes the
+                    #       files (much) bigger (because they are stored as individual frames) it also adds some useful
+                    #       play/pause/next frame/... buttons.
+                    filepath.write_text(f'<html>{ani.to_html5_video()}</html>', encoding='utf8')
+                else:
+                    writer = anim_kwargs.pop('writer', None)
+                    fps = anim_kwargs.pop('fps', 5)
+                    metadata = anim_kwargs.pop('metadata', None)
+                    bitrate = anim_kwargs.pop('bitrate', None)
+                    if writer is None:
+                        # pillow only supports .gif, .png and .tiff ffmpeg supports .avi, .mov, .mp4 (but needs the
+                        # ffmpeg package installed)
+                        writer = 'pillow' if filepath.suffix in {'.gif', '.png', '.tiff'} else 'ffmpeg'
+                    from matplotlib.animation import writers
+                    if not writers.is_available(writer):
+                        raise Exception(f"Cannot write animation using '{filepath.suffix}' extension "
+                                        f"because '{writer}' writer is not available.\n"
+                                        "Installing an optional package is probably necessary.")
+
+                    ani.save(filepath, writer=writer, fps=fps, metadata=metadata, bitrate=bitrate)
+
+        if show:
+            # The following line displays the plot window. Note however that
+            # it is only blocking when no Qt loop is already running (i.e. we are
+            # not running inside the editor).
+            # When using the Qt backend this boils down to:
+            #     manager = fig.canvas.manager
+            #     manager.show()
+            #     if block:
+            #         manager.start_main_loop()
+            # the last line just gets the current Qt QApplication instance (created during
+            # the first canvas creation) and .exec() it
+            plt.show(block=True)
+            # It is important to return ani, because otherwise in the non-blocking case
+            # (i.e. when run in the editor), the animation is garbage-collected before
+            # it is drawn, and we get a blank animation.
+            return (ax_to_return, ani) if ani is not None else ax_to_return
+        elif filepath is not None:  # filepath and not show
+            plt.close(fig)
+            return None
+        else:                       # no filepath and not show
+            return (ax_to_return, ani) if ani is not None else ax_to_return
 
     def _plot_many(self, array, ax, kwargs, series_axes, subplot_axes, title, x, y):
         if len(subplot_axes):

larray/tests/test_array.py

@@ -3,6 +3,7 @@
 import pytest
 import numpy as np
 import pandas as pd
+import matplotlib.figure
 
 from io import StringIO
 
@@ -5377,11 +5378,22 @@ def test_broadcast_with():
 
 
 def test_plot():
-    pass
-    # small_h = small['c0']
-    # small_h.plot(kind='bar')
-    # small_h.plot()
-    # small_h.hist()
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots()  # doctest: +SKIP
+    assert isinstance(fig, matplotlib.figure.Figure)
+    arr = Array([[1, 5, 2],
+                 [3, 2, 4]], axes='a=a0,a1;b=b0,b1,b2')
+    arr.plot(kind='bar', ax=ax)
+    fig.savefig('bar.png')
+
+    fig, ax = plt.subplots()  # doctest: +SKIP
+    arr.plot(ax=ax)
+    fig.savefig('plot.png')
+
+    fig, ax = plt.subplots()  # doctest: +SKIP
+    arr.plot.hist(ax=ax)
+    fig.savefig('hist.png')
 
     # large_data = np.random.randn(1000)
     # tick_v = np.random.randint(ord('a'), ord('z'), size=1000)