diff --git a/animatplot/animation.py b/animatplot/animation.py index ce19522..7bf99fe 100644 --- a/animatplot/animation.py +++ b/animatplot/animation.py @@ -4,6 +4,7 @@ import numpy as np from animatplot import Timeline +from animatplot.blocks.base import Block class Animation: @@ -26,6 +27,11 @@ class Animation: a matplotlib animation returned from FuncAnimation """ def __init__(self, blocks, timeline=None, fig=None): + self.fig = plt.gcf() if fig is None else fig + + self.animation = self._animate(blocks, timeline) + + def _animate(self, blocks, timeline): if timeline is None: self.timeline = Timeline(range(len(blocks[0]))) elif not isinstance(timeline, Timeline): @@ -36,14 +42,14 @@ def __init__(self, blocks, timeline=None, fig=None): _len_time = len(self.timeline) for block in blocks: if len(block) != _len_time: - raise ValueError("All blocks must animate for the same amount of time") + raise ValueError("All blocks must animate for the same amount " + "of time") self.blocks = blocks - self.fig = plt.gcf() if fig is None else fig self._has_slider = False self._pause = False - def animate(i): + def update_all(i): updates = [] for block in self.blocks: updates.append(block._update(self.timeline.index)) @@ -52,11 +58,8 @@ def animate(i): self.timeline._update() return updates - self.animation = FuncAnimation( - self.fig, animate, - frames=self.timeline._len, - interval=1000/self.timeline.fps - ) + return FuncAnimation(self.fig, update_all, frames=self.timeline._len, + interval=1000 / self.timeline.fps) def toggle(self, ax=None): """Creates a play/pause button to start/stop the animation @@ -176,7 +179,8 @@ def save_gif(self, filename): the name of the file to be created without the file extension """ self.timeline.index -= 1 # required for proper starting point for save - self.animation.save(filename+'.gif', writer=PillowWriter(fps=self.timeline.fps)) + self.animation.save(filename+'.gif', + writer=PillowWriter(fps=self.timeline.fps)) def save(self, *args, **kwargs): """Saves an animation @@ -185,3 +189,39 @@ def save(self, *args, **kwargs): """ self.timeline.index -= 1 # required for proper starting point for save self.animation.save(*args, **kwargs) + + def add(self, new): + """ + Updates the animation object by adding additional blocks. + + The new blocks can be passed as a list, or as part of a second animaion. + If passed as part of a new animation, the timeline of this new + animation object will replace the old one. + + Parameters + ---------- + new : amp.animation.Animation, or list of amp.block.Block objects + Either blocks to add to animation instance, or another animation + instance whose blocks should be combined with this animation. + """ + + if isinstance(new, Animation): + new_blocks = new.blocks + new_timeline = new.timeline + + else: + if not isinstance(new, list): + new_blocks = [new] + else: + new_blocks = new + new_timeline = self.timeline + + for i, block in enumerate(new_blocks): + if not isinstance(block, Block): + raise TypeError(f"Block number {i} passed is of type " + f"{type(block)}, not of type " + f"animatplot.blocks.Block (or a subclass)") + + self.blocks.append(block) + + self.animation = self._animate(self.blocks, new_timeline) diff --git a/tests/test_animation.py b/tests/test_animation.py index 203bf74..da3475a 100644 --- a/tests/test_animation.py +++ b/tests/test_animation.py @@ -1,14 +1,21 @@ -from matplotlib.testing import setup -setup() import os + import pytest -import numpy as np + +import matplotlib.testing import matplotlib.pyplot as plt from matplotlib.animation import PillowWriter + +import numpy as np +import numpy.testing as npt + import animatplot as amp from tests.tools import animation_compare +matplotlib.testing.setup() + + @pytest.mark.xfail @animation_compare(baseline_images='Animation/controls', nframes=5, tol=.5) def test_controls(): @@ -41,3 +48,90 @@ def test_save(): anim.save_gif(base+'save') plt.close('all') assert os.path.exists(base+'save.gif') + + +@pytest.fixture() +def line_block(): + def make_line_block(t_length=5): + x = np.linspace(0, 1, 10) + t = np.linspace(0, 1, t_length) + x_grid, t_grid = np.meshgrid(x, t) + y_data = np.sin(2 * np.pi * (x_grid + t_grid)) + + return amp.blocks.Line(x, y_data) + return make_line_block + + +@pytest.fixture() +def line_anim(): + def make_line_anim(t_length=5, timeline=False): + x = np.linspace(0, 1, 10) + t = np.linspace(0, 1, t_length) + x_grid, t_grid = np.meshgrid(x, t) + y_data = np.sin(2 * np.pi * (x_grid + t_grid)) + + block = amp.blocks.Line(x, y_data) + + if timeline: + return amp.Animation([block], timeline=amp.Timeline(t)) + else: + return amp.Animation([block]) + + return make_line_anim + + +class TestAddBlocks: + def test_add_blocks(self, line_block): + anim = amp.Animation([line_block()]) + anim.add(line_block()) + + assert isinstance(anim, amp.Animation) + assert len(anim.blocks) == 2 + for actual in anim.blocks: + assert len(actual) == 5 + npt.assert_equal(actual.line.get_xdata(), + np.linspace(0, 1, 10)) + + def test_wrong_length_block(self, line_block): + anim = amp.Animation([line_block()]) + + with pytest.raises(ValueError): + anim.add(line_block(t_length=6)) + + def test_wrong_type(self, line_block): + anim = amp.Animation([line_block()]) + + with pytest.raises(TypeError): + anim.add('not a block') + + +class TestAddAnimations: + def test_add_animations(self, line_anim): + anim = line_anim() + anim.add(line_anim()) + + assert isinstance(anim, amp.Animation) + + def test_add_animation_with_timeline(self, line_anim): + anim = line_anim() + anim.add(line_anim(timeline=True)) + + assert isinstance(anim, amp.Animation) + assert len(anim.timeline) == 5 + + def test_add_animations_both_with_timelines(self, line_anim): + anim = line_anim(timeline=True) + anim2 = line_anim() + t = 10*np.arange(5) + anim2.timeline = amp.Timeline(t) + + anim.add(anim2) + + assert isinstance(anim, amp.Animation) + assert len(anim.timeline) == 5 + npt.assert_equal(anim.timeline.t, t) + + def test_add_animations_different_lengths(self, line_anim): + anim = line_anim() + with pytest.raises(ValueError): + anim.add(line_anim(t_length=6))