Source code for plotnine.animation

from __future__ import annotations

import typing
from copy import copy

import pandas as pd
from matplotlib.animation import ArtistAnimation

from .exceptions import PlotnineError

if typing.TYPE_CHECKING:
    from typing import Iterable

    from plotnine.typing import (
        Artist,
        Axes,
        Figure,
        Ggplot,
        Scale,
    )

__all__ = ("PlotnineAnimation",)


[docs]class PlotnineAnimation(ArtistAnimation): """ Animation using ggplot objects Parameters ---------- plots : iterable ggplot objects that make up the the frames of the animation interval : number, optional Delay between frames in milliseconds. Defaults to 200. repeat_delay : number, optional If the animation in repeated, adds a delay in milliseconds before repeating the animation. Defaults to `None`. repeat : bool, optional Controls whether the animation should repeat when the sequence of frames is completed. Defaults to `True`. blit : bool, optional Controls whether blitting is used to optimize drawing. Defaults to `False`. Notes ----- 1. The plots should have the same `facet` and the facet should not have fixed x and y scales. 2. The scales of all the plots should have the same limits. It is a good idea to create a scale (with limits) for each aesthetic and add them to all the plots. """ def __init__( self, plots: Iterable[Ggplot], interval: int = 200, repeat_delay: int | None = None, repeat: bool = True, blit: bool = False, ): figure, artists = self._draw_plots(plots) ArtistAnimation.__init__( self, figure, artists, interval=interval, repeat_delay=repeat_delay, repeat=repeat, blit=blit, ) def _draw_plots( self, plots: Iterable[Ggplot] ) -> tuple[Figure, list[list[Artist]]]: with pd.option_context("mode.chained_assignment", None): return self.__draw_plots(plots) def __draw_plots( self, plots: Iterable[Ggplot] ) -> tuple[Figure, list[list[Artist]]]: """ Plot and return the figure and artists Parameters ---------- plots : iterable ggplot objects that make up the the frames of the animation Returns ------- figure : matplotlib.figure.Figure Matplotlib figure artists : list List of :class:`Matplotlib.artist.Artist` """ import matplotlib.pyplot as plt # For keeping track of artists for each frame artist_offsets: dict[str, list[int]] = { "collections": [], "patches": [], "lines": [], "texts": [], "artists": [], } scale_limits = {} def initialise_artist_offsets(n: int): """ Initilise artists_offsets arrays to zero Parameters ---------- n : int Number of axes to initialise artists for. The artists for each axes are tracked separately. """ for artist_type in artist_offsets: artist_offsets[artist_type] = [0] * n def get_frame_artists(axs: list[Axes]) -> list[Artist]: """ Artists shown in a given frame Parameters ---------- axs : list[Axes] Matplotlib axes that have had artists added to them. """ # The axes accumulate artists for all frames # For each frame we pickup the newly added artists # We use offsets to mark the end of the previous frame # e.g ax.collections[start:] frame_artists = [] for i, ax in enumerate(axs): for name in artist_offsets: start = artist_offsets[name][i] new_artists = getattr(ax, name)[start:] frame_artists.extend(new_artists) artist_offsets[name][i] += len(new_artists) return frame_artists def set_scale_limits(scales: list[Scale]): """ Set limits of all the scales in the animation Should be called before :func:`check_scale_limits`. Parameters ---------- scales : list[scales] List of scales the have been used in building a ggplot object. """ for sc in scales: ae = sc.aesthetics[0] scale_limits[ae] = sc.limits def check_scale_limits(scales: list[Scale], frame_no: int): """ Check limits of the scales of a plot in the animation Raises a PlotnineError if any of the scales has limits that do not match those of the first plot/frame. Should be called after :func:`set_scale_limits`. Parameters ---------- scales : list[scales] List of scales the have been used in building a ggplot object. frame_no : int Frame number """ if len(scale_limits) != len(scales): raise PlotnineError( "All plots must have the same number of scales " "as the first plot of the animation." ) for sc in scales: ae = sc.aesthetics[0] if ae not in scale_limits: raise PlotnineError( f"The plot for frame {frame_no} does not " f"have a scale for the {ae} aesthetic." ) if sc.limits != scale_limits[ae]: raise PlotnineError( f"The {ae} scale of plot for frame {frame_no} has " "different limits from those of the first frame." ) figure: Figure | None = None axs: list[Axes] = [] artists = [] scales = None # Will hold the scales of the first frame # The first ggplot creates the figure, axes and the initial # frame of the animation. The rest of the ggplots draw # onto the figure and axes created by the first ggplot and # they create the subsequent frames. for frame_no, p in enumerate(plots): if figure is None: figure = p.draw() axs = figure.get_axes() # pyright: ignore initialise_artist_offsets(len(axs)) scales = p._build_objs.scales set_scale_limits(scales) else: p = copy(p) plot = p._draw_using_figure(figure, axs) check_scale_limits(plot.scales, frame_no) artists.append(get_frame_artists(axs)) if figure is None: figure = plt.figure() # pyright: ignore assert figure is not None # Prevent Jupyter from plotting any static figure plt.close(figure) return figure, artists