1import base64 

2from tempfile import NamedTemporaryFile 

3 

4import matplotlib.animation as animation 

5import sympy as sp 

6from IPython.display import HTML 

7 

8import pystencils.plot as plt 

9 

10__all__ = ['make_imshow_animation', 'display_animation', 'set_display_mode'] 

11 

12 

13VIDEO_TAG = """<video controls width="80%"> 

14 <source src="data:video/x-m4v;base64,{0}" type="video/mp4"> 

15 Your browser does not support the video tag. 

16</video>""" 

17 

18 

19def __animation_to_html(animation, fps): 

20 if not hasattr(animation, 'encoded_video'): 

21 with NamedTemporaryFile(suffix='.mp4') as f: 

22 animation.save(f.name, fps=fps, extra_args=['-vcodec', 'libx264', '-pix_fmt', 

23 'yuv420p', '-profile:v', 'baseline', '-level', '3.0']) 

24 video = open(f.name, "rb").read() 

25 animation.encoded_video = base64.b64encode(video).decode('ascii') 

26 

27 return VIDEO_TAG.format(animation.encoded_video) 

28 

29 

30def make_imshow_animation(grid, grid_update_function, frames=90, **_): 

31 from functools import partial 

32 fig = plt.figure() 

33 im = plt.imshow(grid, interpolation='none') 

34 

35 def update_figure(*_, **kwargs): 

36 image = kwargs['image'] 

37 image = grid_update_function(image) 

38 im.set_array(image) 

39 return im, 

40 

41 return animation.FuncAnimation(fig, partial(update_figure, image=grid), frames=frames) 

42 

43 

44# ------- Version 1: Embed the animation as HTML5 video --------- ---------------------------------- 

45 

46def display_as_html_video(animation, fps=30, show=True, **_): 

47 try: 

48 plt.close(animation._fig) 

49 res = __animation_to_html(animation, fps) 

50 if show: 

51 return HTML(res) 

52 else: 

53 return HTML("") 

54 except KeyboardInterrupt: 

55 pass 

56 

57 

58# ------- Version 2: Animation is shown in extra matplotlib window ---------------------------------- 

59 

60 

61def display_in_extra_window(*_, **__): 

62 fig = plt.gcf() 

63 try: 

64 fig.canvas.manager.window.raise_() 

65 except Exception: 

66 pass 

67 plt.show() 

68 

69 

70# ------- Version 3: Animation is shown in images that are updated directly in website -------------- 

71 

72def display_as_html_image(animation, show=True, *args, **kwargs): 

73 from IPython import display 

74 

75 try: 

76 if show: 

77 animation._init_draw() 

78 for _ in animation.frame_seq: 

79 if show: 

80 fig = plt.gcf() 

81 display.display(fig) 

82 animation._step() 

83 if show: 

84 display.clear_output(wait=True) 

85 except KeyboardInterrupt: 

86 display.clear_output(wait=False) 

87 

88 

89# Dispatcher 

90 

91animation_display_mode = 'image_update' 

92display_animation_func = None 

93 

94 

95def display_animation(*args, **kwargs): 

96 from IPython import get_ipython 

97 ipython = get_ipython() 

98 if not ipython: 

99 return 

100 

101 if not display_animation_func: 

102 raise Exception("Call set_display_mode first") 

103 return display_animation_func(*args, **kwargs) 

104 

105 

106def set_display_mode(mode): 

107 from IPython import get_ipython 

108 ipython = get_ipython() 

109 if not ipython: 

110 return 

111 global animation_display_mode 

112 global display_animation_func 

113 animation_display_mode = mode 

114 if animation_display_mode == 'video': 

115 ipython.magic("matplotlib inline") 

116 display_animation_func = display_as_html_video 

117 elif animation_display_mode == 'window': 

118 ipython.magic("matplotlib qt") 

119 display_animation_func = display_in_extra_window 

120 elif animation_display_mode == 'image_update': 

121 ipython.magic("matplotlib inline") 

122 display_animation_func = display_as_html_image 

123 else: 

124 raise Exception("Unknown mode. Available modes 'image_update', 'video' and 'window' ") 

125 

126 

127def activate_ipython(): 

128 from IPython import get_ipython 

129 ipython = get_ipython() 

130 if ipython: 

131 set_display_mode('image_update') 

132 ipython.magic("config InlineBackend.rc = { }") 

133 ipython.magic("matplotlib inline") 

134 plt.rc('figure', figsize=(16, 6)) 

135 sp.init_printing() 

136 

137 

138activate_ipython()