diff --git a/src/tikzplotlib/_save.py b/src/tikzplotlib/_save.py index d89cadfa..de4eab35 100644 --- a/src/tikzplotlib/_save.py +++ b/src/tikzplotlib/_save.py @@ -3,10 +3,13 @@ import enum import tempfile import warnings +from typing import Literal from pathlib import Path import matplotlib as mpl import matplotlib.pyplot as plt +import matplotlib.animation as animation +import matplotlib.figure as figure from . import _axes from . import _image as img @@ -17,7 +20,7 @@ def get_tikz_code( - figure="gcf", + figure: Literal["gcf"] | figure.Figure | animation.TimedAnimation = "gcf", filepath: str | Path | None = None, axis_width: str | None = None, axis_height: str | None = None, @@ -32,6 +35,11 @@ def get_tikz_code( extra_axis_parameters: list | set | None = None, extra_groupstyle_parameters: dict = {}, extra_tikzpicture_parameters: list | set | None = None, + extra_animation_parameters: list | set | None = [ + "autoplay", + "autoresume", + "controls", + ], extra_lines_start: list | set | None = None, dpi: int | None = None, show_info: bool = False, @@ -44,7 +52,7 @@ def get_tikz_code( """Main function. Here, the recursion into the image starts and the contents are picked up. The actual file gets written in this routine. - :param figure: either a Figure object or 'gcf' (default). + :param figure: either a Figure object or an animation or 'gcf' (default). :param axis_width: If not ``None``, this will be used as figure width within the TikZ/PGFPlots output. If ``axis_height`` is not given, @@ -109,6 +117,10 @@ def get_tikz_code( (as a set) to pgfplots. :type extra_tikzpicture_parameters: a set of strings for the pfgplots tikzpicture. + :param extra_animation_parameters: Extra animation options to be passed + (as a set) to animateinline when an animation is passed. + :type extra_animation_parameters: a set of strings for the animateinline animation. + :param dpi: The resolution in dots per inch of the rendered image in case of QuadMesh plots. If ``None`` it will default to the value ``savefig.dpi`` from matplotlib.rcParams. Default is ``None``. @@ -150,6 +162,7 @@ def get_tikz_code( if figure == "gcf": figure = plt.gcf() data = {} + data["animation"] = isinstance(figure, animation.Animation) data["axis width"] = axis_width data["axis height"] = axis_height data["rel data path"] = ( @@ -209,39 +222,89 @@ def get_tikz_code( if show_info: _print_pgfplot_libs_message(data) - # gather the file content - data, content = _recurse(data, figure) - - # Check if there is still an open groupplot environment. This occurs if not - # all of the group plot slots are used. - if "is_in_groupplot_env" in data and data["is_in_groupplot_env"]: - content.extend(data["flavor"].end("groupplot") + "\n\n") + def get_figure_tikz_code( + data, + figure, + wrap: bool = wrap, + include_disclaimer: bool = include_disclaimer, + ): + # gather the file content + data, content = _recurse(data, figure) + + # Check if there is still an open groupplot environment. This occurs if not + # all of the group plot slots are used. + if "is_in_groupplot_env" in data and data["is_in_groupplot_env"]: + content.extend(data["flavor"].end("groupplot") + "\n\n") + data["is_in_groupplot_env"] = False + + code = """""" + + if include_disclaimer: + disclaimer = f"This file was created with tikzplotlib v{__version__}." + code += _tex_comment(disclaimer) + + # write the contents + if wrap and add_axis_environment: + code += data["flavor"].start("tikzpicture") + if extra_tikzpicture_parameters: + code += "[\n" + ",\n".join(extra_tikzpicture_parameters) + "\n]" + code += "\n" + if extra_lines_start: + code += "\n".join(extra_lines_start) + "\n" + code += "\n" + + coldefs = _get_color_definitions(data) + if coldefs: + code += "\n".join(coldefs) + "\n\n" + + code += "".join(content) + + if wrap and add_axis_environment: + code += data["flavor"].end("tikzpicture") + "\n" + + return data, content, code + + if isinstance(figure, animation.TimedAnimation): + extra_animation_parameters = list(extra_animation_parameters or []) + if figure._repeat and "loop" not in extra_animation_parameters: + extra_animation_parameters.append("loop") + + data["framerate"] = 1000 / figure._interval + + frames = [] + + for frame in figure.new_frame_seq(): + figure._draw_frame(frame) + data, content, code = get_figure_tikz_code( + data, + figure._fig, + wrap=True, + include_disclaimer=False, + ) + frames.append(f"% Frame {frame + 1}\n{code}\n") - # write disclaimer to the file header - code = """""" + code = """""" - if include_disclaimer: - disclaimer = f"This file was created with tikzplotlib v{__version__}." - code += _tex_comment(disclaimer) + if include_disclaimer: + disclaimer = f"This file was created with tikzplotlib v{__version__}." + code += _tex_comment(disclaimer) - # write the contents - if wrap and add_axis_environment: - code += data["flavor"].start("tikzpicture") - if extra_tikzpicture_parameters: - code += "[\n" + ",\n".join(extra_tikzpicture_parameters) + "\n]" - code += "\n" - if extra_lines_start: - code += "\n".join(extra_lines_start) + "\n" - code += "\n" + # write the contents + if wrap: + code += data["flavor"].start("animateinline") + if extra_animation_parameters: + code += "[\n" + ",\n".join(extra_animation_parameters) + "\n]" + code += f"{{{data['framerate']}}}" + code += "\n" + code += "\n" - coldefs = _get_color_definitions(data) - if coldefs: - code += "\n".join(coldefs) + "\n\n" + code += "\n\\newframe\n".join(frames) - code += "".join(content) + if wrap: + code += data["flavor"].end("animateinline") + "\n" - if wrap and add_axis_environment: - code += data["flavor"].end("tikzpicture") + "\n" + else: + data, content, code = get_figure_tikz_code(data, figure) if standalone: # When using pdflatex, \\DeclareUnicodeCharacter is necessary. @@ -406,6 +469,7 @@ class Flavors(enum.Enum): \\usetikzlibrary{{{tikzlibs}}} \\pgfplotsset{{compat=newest}} """, + r"\usepackage{{{}}}", ) context = ( r"\start{}", @@ -422,6 +486,7 @@ class Flavors(enum.Enum): \\unexpanded\\def\\startgroupplot{{\\groupplot}} \\unexpanded\\def\\stopgroupplot{{\\endgroupplot}} """, + r"\usemodule[{}]", ) def start(self, what): @@ -430,6 +495,9 @@ def start(self, what): def end(self, what): return self.value[1].format(what) + def usepackage(self, *what): + return self.value[4] + def preamble(self, data=None): if data is None: data = { @@ -438,7 +506,13 @@ def preamble(self, data=None): } pgfplotslibs = ",".join(data["pgfplots libs"]) tikzlibs = ",".join(data["tikz libs"]) - return self.value[3].format(pgfplotslibs=pgfplotslibs, tikzlibs=tikzlibs) + extra_imports = ( + self.usepackage("animate") + "\n" if data.get("animation", False) else "" + ) + return ( + self.value[3].format(pgfplotslibs=pgfplotslibs, tikzlibs=tikzlibs) + + extra_imports + ) def standalone(self, code): docenv = self.value[2] diff --git a/tests/test_animation.py b/tests/test_animation.py new file mode 100644 index 00000000..c136568a --- /dev/null +++ b/tests/test_animation.py @@ -0,0 +1,28 @@ +#%% +def plot(): + import numpy as np + from matplotlib import pyplot as plt + import matplotlib.animation as animation + + fig, ax = plt.subplots() + scat = ax.scatter(range(10), [0] * 10) + ax.set_xlim(0, 9) + ax.set_ylim(0, 10) + + def update(frame): + y_data = [yi + frame * 0.2 for yi in range(10)] + scat.set_offsets(list(zip(range(10), y_data))) + ax.set_title(f"Frame {frame+1}/20") + + ani = animation.FuncAnimation(fig, update, frames=20, repeat=False) + return ani + + +def test(): + from .helpers import assert_equality + + assert_equality(plot, __file__[:-3] + "_reference.tex") + +# %% +ani = plot() +# %% diff --git a/tests/test_animation_reference.tex b/tests/test_animation_reference.tex new file mode 100644 index 00000000..b1858066 --- /dev/null +++ b/tests/test_animation_reference.tex @@ -0,0 +1,688 @@ +\begin{animateinline}[ +autoplay, +autoresume, +controls +]{5.0} + +\definecolor{darkgray176}{RGB}{176,176,176} +\definecolor{steelblue31119180}{RGB}{31,119,180} + +% Frame 1 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 1/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 0 +1 1 +2 2 +3 3 +4 4 +5 5 +6 6 +7 7 +8 8 +9 9 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 2 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 2/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 0.2 +1 1.2 +2 2.2 +3 3.2 +4 4.2 +5 5.2 +6 6.2 +7 7.2 +8 8.2 +9 9.2 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 3 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 3/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 0.4 +1 1.4 +2 2.4 +3 3.4 +4 4.4 +5 5.4 +6 6.4 +7 7.4 +8 8.4 +9 9.4 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 4 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 4/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 0.6 +1 1.6 +2 2.6 +3 3.6 +4 4.6 +5 5.6 +6 6.6 +7 7.6 +8 8.6 +9 9.6 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 5 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 5/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 0.8 +1 1.8 +2 2.8 +3 3.8 +4 4.8 +5 5.8 +6 6.8 +7 7.8 +8 8.8 +9 9.8 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 6 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 6/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 1 +1 2 +2 3 +3 4 +4 5 +5 6 +6 7 +7 8 +8 9 +9 10 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 7 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 7/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 1.2 +1 2.2 +2 3.2 +3 4.2 +4 5.2 +5 6.2 +6 7.2 +7 8.2 +8 9.2 +9 10.2 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 8 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 8/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 1.4 +1 2.4 +2 3.4 +3 4.4 +4 5.4 +5 6.4 +6 7.4 +7 8.4 +8 9.4 +9 10.4 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 9 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 9/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 1.6 +1 2.6 +2 3.6 +3 4.6 +4 5.6 +5 6.6 +6 7.6 +7 8.6 +8 9.6 +9 10.6 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 10 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 10/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 1.8 +1 2.8 +2 3.8 +3 4.8 +4 5.8 +5 6.8 +6 7.8 +7 8.8 +8 9.8 +9 10.8 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 11 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 11/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 2 +1 3 +2 4 +3 5 +4 6 +5 7 +6 8 +7 9 +8 10 +9 11 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 12 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 12/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 2.2 +1 3.2 +2 4.2 +3 5.2 +4 6.2 +5 7.2 +6 8.2 +7 9.2 +8 10.2 +9 11.2 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 13 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 13/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 2.4 +1 3.4 +2 4.4 +3 5.4 +4 6.4 +5 7.4 +6 8.4 +7 9.4 +8 10.4 +9 11.4 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 14 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 14/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 2.6 +1 3.6 +2 4.6 +3 5.6 +4 6.6 +5 7.6 +6 8.6 +7 9.6 +8 10.6 +9 11.6 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 15 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 15/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 2.8 +1 3.8 +2 4.8 +3 5.8 +4 6.8 +5 7.8 +6 8.8 +7 9.8 +8 10.8 +9 11.8 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 16 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 16/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 3 +1 4 +2 5 +3 6 +4 7 +5 8 +6 9 +7 10 +8 11 +9 12 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 17 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 17/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 3.2 +1 4.2 +2 5.2 +3 6.2 +4 7.2 +5 8.2 +6 9.2 +7 10.2 +8 11.2 +9 12.2 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 18 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 18/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 3.4 +1 4.4 +2 5.4 +3 6.4 +4 7.4 +5 8.4 +6 9.4 +7 10.4 +8 11.4 +9 12.4 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 19 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 19/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 3.6 +1 4.6 +2 5.6 +3 6.6 +4 7.6 +5 8.6 +6 9.6 +7 10.6 +8 11.6 +9 12.6 +}; +\end{axis} + +\end{tikzpicture} + + +\newframe +% Frame 20 +\begin{tikzpicture} + +\begin{axis}[ +tick align=outside, +tick pos=left, +title={Frame 20/20}, +x grid style={darkgray176}, +xmin=0, xmax=9, +xtick style={color=black}, +y grid style={darkgray176}, +ymin=0, ymax=10, +ytick style={color=black} +] +\addplot [draw=steelblue31119180, fill=steelblue31119180, mark=*, only marks] +table{% +x y +0 3.8 +1 4.8 +2 5.8 +3 6.8 +4 7.8 +5 8.8 +6 9.8 +7 10.8 +8 11.8 +9 12.8 +}; +\end{axis} + +\end{tikzpicture} + +\end{animateinline}