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