WCONViewer/Player.py at main · openworm/WCONViewer · GitHub

1
2
3
4
5
6
7
8
9
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
# Based on example at https://stackoverflow.com/questions/46325447/animated-interactive-plot-using-matplotlib

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import mpl_toolkits.axes_grid1
import matplotlib.widgets


class Player(FuncAnimation):
    def __init__(
        self,
        fig,
        func,
        frames=None,
        init_func=None,
        fargs=None,
        save_count=None,
        mini=0,
        maxi=100,
        pos=(0.125, 0.92),
        times=None,
        t_units="",
        **kwargs,
    ):
        self.i = 0
        self.min = mini
        self.max = maxi
        self.runs = True
        self.forwards = True
        self.fig = fig
        self.func = func
        self.setup(pos)
        FuncAnimation.__init__(
            self,
            self.fig,
            self.update,
            frames=self.play(),
            init_func=init_func,
            fargs=fargs,
            save_count=save_count,
            **kwargs,
        )
        self.times = times
        self.t_units = t_units

    def play(self):
        while self.runs:
            self.i = self.i + self.forwards - (not self.forwards)
            if self.i > self.min and self.i < self.max:
                yield self.i
            else:
                self.stop()
                yield self.i

    def start(self):
        self.runs = True
        self.event_source.start()

    def stop(self, event=None):
        self.runs = False
        self.event_source.stop()

    def forward(self, event=None):
        self.forwards = True
        self.start()

    def backward(self, event=None):
        self.forwards = False
        self.start()

    def oneforward(self, event=None):
        self.forwards = True
        self.onestep()

    def onebackward(self, event=None):
        self.forwards = False
        self.onestep()

    def onestep(self):
        if self.i > self.min and self.i < self.max:
            self.i = self.i + self.forwards - (not self.forwards)
        elif self.i == self.min and self.forwards:
            self.i += 1
        elif self.i == self.max and not self.forwards:
            self.i -= 1
        self.func(self.i)
        self.slider.set_val(self.i)

        self.fig.canvas.draw_idle()

    def setup(self, pos):
        playerax = self.fig.add_axes([pos[0], pos[1], 0.64, 0.04])
        divider = mpl_toolkits.axes_grid1.make_axes_locatable(playerax)
        bax = divider.append_axes("right", size="80%", pad=0.05)
        sax = divider.append_axes("right", size="80%", pad=0.05)
        fax = divider.append_axes("right", size="80%", pad=0.05)
        ofax = divider.append_axes("right", size="100%", pad=0.05)
        sliderax = divider.append_axes("right", size="500%", pad=0.07)
        self.button_oneback = matplotlib.widgets.Button(playerax, label="$\u29cf$")
        self.button_back = matplotlib.widgets.Button(bax, label="$\u25c0$")
        self.button_stop = matplotlib.widgets.Button(sax, label="$\u25a0$")
        self.button_forward = matplotlib.widgets.Button(fax, label="$\u25b6$")
        self.button_oneforward = matplotlib.widgets.Button(ofax, label="$\u29d0$")
        self.button_oneback.on_clicked(self.onebackward)
        self.button_back.on_clicked(self.backward)
        self.button_stop.on_clicked(self.stop)
        self.button_forward.on_clicked(self.forward)
        self.button_oneforward.on_clicked(self.oneforward)
        self.slider = matplotlib.widgets.Slider(
            sliderax, "", self.min, self.max, valinit=self.i
        )
        self.slider.valtext.set_visible(False)
        self.slider.on_changed(self.set_pos)
        self.slider.label.set_horizontalalignment("left")
        label_position = self.slider.label.get_position()  # Get the current position
        self.slider.label.set_position(
            (label_position[0] + 1.1, label_position[1])
        )  # Adjust x-coordinate

        # self.text_box = matplotlib.widgets.TextBox(textax, label="", initial="0 ms")

    def set_pos(self, i):
        self.i = int(self.slider.val)
        self.func(self.i)
        self.set_time_text(self.i)

    def set_time_text(self, i):
        time_text = str(i)
        if self.times is not None:
            time_text = "%s%s" % (self.times[i], self.t_units)

        # self.text_box.set_val(time_text)
        self.slider.label.set_text(time_text)

    def update(self, i):
        self.slider.set_val(i)
        self.set_time_text(i)


if __name__ == "__main__":
    ### using this class is as easy as using FuncAnimation:

    fig, ax = plt.subplots()
    x = np.linspace(0, 6 * np.pi, num=100)
    y = np.sin(x)

    ax.plot(x, y)
    (point,) = ax.plot([], [], marker="o", color="crimson", ms=15)

    def update(i):
        point.set_data(x[i], y[i])

    ani = Player(fig, update, maxi=len(y) - 1)

    plt.show()