jupyter.py 5.7 KB
Newer Older
1
import pystencils.plot as plt
Martin Bauer's avatar
Martin Bauer committed
2
3
4
5
import matplotlib.animation as animation
from IPython.display import HTML
from tempfile import NamedTemporaryFile
import base64
6
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
7

Martin Bauer's avatar
Martin Bauer committed
8
__all__ = ['log_progress', 'make_imshow_animation', 'display_animation', 'set_display_mode']
Martin Bauer's avatar
Martin Bauer committed
9

Martin Bauer's avatar
Martin Bauer committed
10

11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def log_progress(sequence, every=None, size=None, name='Items'):
    """Copied from https://github.com/alexanderkuk/log-progress"""
    from ipywidgets import IntProgress, HTML, VBox
    from IPython.display import display

    is_iterator = False
    if size is None:
        try:
            size = len(sequence)
        except TypeError:
            is_iterator = True
    if size is not None:
        if every is None:
            if size <= 200:
                every = 1
            else:
                every = int(size / 200)     # every 0.5%
    else:
        assert every is not None, 'sequence is iterator, set every'

    if is_iterator:
        progress = IntProgress(min=0, max=1, value=1)
        progress.bar_style = 'info'
    else:
        progress = IntProgress(min=0, max=size, value=0)
    label = HTML()
    box = VBox(children=[label, progress])
    display(box)

    index = 0
    try:
        for index, record in enumerate(sequence, 1):
            if index == 1 or index % every == 0:
                if is_iterator:
                    label.value = '{name}: {index} / ?'.format(
                        name=name,
                        index=index
                    )
                else:
                    progress.value = index
                    label.value = u'{name}: {index} / {size}'.format(
                        name=name,
                        index=index,
                        size=size
                    )
            yield record
    except:
        progress.bar_style = 'danger'
        raise
    else:
        progress.bar_style = 'success'
        progress.value = index
        label.value = "{name}: {index}".format(
            name=name,
            index=str(index or '?')
        )


Martin Bauer's avatar
Martin Bauer committed
69
70
71
72
73
74
VIDEO_TAG = """<video controls width="80%">
 <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
 Your browser does not support the video tag.
</video>"""


Martin Bauer's avatar
Martin Bauer committed
75
76
def __animation_to_html(animation, fps):
    if not hasattr(animation, 'encoded_video'):
Martin Bauer's avatar
Martin Bauer committed
77
        with NamedTemporaryFile(suffix='.mp4') as f:
Martin Bauer's avatar
Martin Bauer committed
78
79
            animation.save(f.name, fps=fps, extra_args=['-vcodec', 'libx264', '-pix_fmt',
                                                        'yuv420p', '-profile:v', 'baseline', '-level', '3.0'])
Martin Bauer's avatar
Martin Bauer committed
80
            video = open(f.name, "rb").read()
Martin Bauer's avatar
Martin Bauer committed
81
        animation.encoded_video = base64.b64encode(video).decode('ascii')
Martin Bauer's avatar
Martin Bauer committed
82

Martin Bauer's avatar
Martin Bauer committed
83
    return VIDEO_TAG.format(animation.encoded_video)
Martin Bauer's avatar
Martin Bauer committed
84
85


Martin Bauer's avatar
Martin Bauer committed
86
def make_imshow_animation(grid, grid_update_function, frames=90, **_):
Martin Bauer's avatar
Martin Bauer committed
87
88
89
90
    from functools import partial
    fig = plt.figure()
    im = plt.imshow(grid, interpolation='none')

Martin Bauer's avatar
Martin Bauer committed
91
    def update_figure(*_, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
92
        image = kwargs['image']
Martin Bauer's avatar
Martin Bauer committed
93
        image = grid_update_function(image)
Martin Bauer's avatar
Martin Bauer committed
94
95
96
        im.set_array(image)
        return im,

Martin Bauer's avatar
Martin Bauer committed
97
    return animation.FuncAnimation(fig, partial(update_figure, image=grid), frames=frames)
Martin Bauer's avatar
Martin Bauer committed
98
99
100
101


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

Martin Bauer's avatar
Martin Bauer committed
102
def display_as_html_video(animation, fps=30, show=True, **_):
Martin Bauer's avatar
Martin Bauer committed
103
    try:
Martin Bauer's avatar
Martin Bauer committed
104
105
        plt.close(animation._fig)
        res = __animation_to_html(animation, fps)
Martin Bauer's avatar
Martin Bauer committed
106
107
108
109
110
111
112
113
114
115
116
        if show:
            return HTML(res)
        else:
            return HTML("")
    except KeyboardInterrupt:
        pass


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


Martin Bauer's avatar
Martin Bauer committed
117
def display_in_extra_window(*_, **__):
Martin Bauer's avatar
Martin Bauer committed
118
119
    fig = plt.gcf()
    try:
Martin Bauer's avatar
Martin Bauer committed
120
        fig.canvas.manager.window.raise_()
Martin Bauer's avatar
Martin Bauer committed
121
    except Exception:
Martin Bauer's avatar
Martin Bauer committed
122
        pass
Martin Bauer's avatar
Martin Bauer committed
123
124
125
126
127
    plt.show()


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

128
def display_as_html_image(animation, show=True, *args, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
129
130
131
132
133
    from IPython import display

    try:
        if show:
            animation._init_draw()
134
        for _ in animation.frame_seq:
Martin Bauer's avatar
Martin Bauer committed
135
            if show:
Martin Bauer's avatar
Martin Bauer committed
136
                fig = plt.gcf()
Martin Bauer's avatar
Martin Bauer committed
137
138
139
140
141
142
143
144
145
146
                display.display(fig)
            animation._step()
            if show:
                display.clear_output(wait=True)
    except KeyboardInterrupt:
        display.clear_output(wait=False)


# Dispatcher

Martin Bauer's avatar
Martin Bauer committed
147
animation_display_mode = 'image_update'
Martin Bauer's avatar
Martin Bauer committed
148
149
150
display_animation_func = None


Martin Bauer's avatar
Martin Bauer committed
151
def display_animation(*args, **kwargs):
152
153
154
155
156
    from IPython import get_ipython
    ipython = get_ipython()
    if not ipython:
        return

Martin Bauer's avatar
Martin Bauer committed
157
    if not display_animation_func:
158
        raise Exception("Call set_display_mode first")
Martin Bauer's avatar
Martin Bauer committed
159
160
161
    return display_animation_func(*args, **kwargs)


Martin Bauer's avatar
Martin Bauer committed
162
def set_display_mode(mode):
Martin Bauer's avatar
Martin Bauer committed
163
164
    from IPython import get_ipython
    ipython = get_ipython()
Martin Bauer's avatar
Martin Bauer committed
165
166
    if not ipython:
        return
Martin Bauer's avatar
Martin Bauer committed
167
168
169
170
171
    global animation_display_mode
    global display_animation_func
    animation_display_mode = mode
    if animation_display_mode == 'video':
        ipython.magic("matplotlib inline")
Martin Bauer's avatar
Martin Bauer committed
172
        display_animation_func = display_as_html_video
Martin Bauer's avatar
Martin Bauer committed
173
174
    elif animation_display_mode == 'window':
        ipython.magic("matplotlib qt")
Martin Bauer's avatar
Martin Bauer committed
175
176
        display_animation_func = display_in_extra_window
    elif animation_display_mode == 'image_update':
Martin Bauer's avatar
Martin Bauer committed
177
        ipython.magic("matplotlib inline")
Martin Bauer's avatar
Martin Bauer committed
178
        display_animation_func = display_as_html_image
Martin Bauer's avatar
Martin Bauer committed
179
    else:
Martin Bauer's avatar
Martin Bauer committed
180
        raise Exception("Unknown mode. Available modes 'image_update', 'video' and 'window' ")
Martin Bauer's avatar
Martin Bauer committed
181
182


Martin Bauer's avatar
Martin Bauer committed
183
def activate_ipython():
184
185
186
    from IPython import get_ipython
    ipython = get_ipython()
    if ipython:
Martin Bauer's avatar
Martin Bauer committed
187
        set_display_mode('image_update')
188
189
190
191
192
        ipython.magic("config InlineBackend.rc = { }")
        ipython.magic("matplotlib inline")
        plt.rc('figure', figsize=(16, 6))
        sp.init_printing()

Martin Bauer's avatar
Martin Bauer committed
193
194

activate_ipython()