open_atmos_jupyter_utils.show_anim

 1"""inline animation function that allows rendering on github
 2 preview and showing click to download gif button"""
 3
 4#pylint: disable=consider-using-with
 5from typing import Iterable, Callable
 6import os
 7import tempfile
 8import base64
 9import imageio
10import matplotlib.pyplot as plt
11from IPython.display import HTML, display
12from open_atmos_jupyter_utils.temporary_file import TemporaryFile
13
14
15def show_anim(plot_func: Callable, frame_range: Iterable, duration=0.01, loop=0, gif_file=None):
16    """plot_func is called with one argument - the frame number from frame_range
17    and is expected to return a matplotlib figure instance (on which savefig()
18    and close() are subsequently called)"""
19    gif_file = TemporaryFile(suffix=".gif", filename=gif_file)
20    with tempfile.TemporaryDirectory() as tmpdirname:
21        for frame in frame_range:
22            fig = plot_func(frame)
23            fig.savefig(f"{tmpdirname}/{frame:05d}.png")
24            plt.close(fig)
25        __merge_images_into_gif(tmpdirname, gif_file, duration, loop)
26        show_gif(gif_file)
27
28
29def show_gif(gif_file):
30    """displays an animation storing the gif file data into the notebook using
31    base-64 encoding"""
32    b64 = base64.b64encode(open(gif_file.basename,'rb').read()).decode("ascii")
33    display(HTML(f'<img src="data:image/gif;base64,{b64}" />'))
34    display(gif_file.make_link_widget())
35
36
37def __merge_images_into_gif(image_folder, gif_name, duration, loop):
38    """creates a GIF file from a series of animation frames"""
39    with imageio.get_writer(
40        gif_name.basename, duration=duration, loop=loop, mode="I"
41    ) as writer:
42        for filename in sorted(os.listdir(image_folder)):
43            image = imageio.v3.imread(os.path.join(image_folder, filename))
44            writer.append_data(image)