from-python-to-numpy/code/smoke-interactive.py at master · B3QL/from-python-to-numpy · 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
# -----------------------------------------------------------------------------
# From Numpy to Python
# Copyright (2017) Nicolas P. Rougier - BSD license
# More information at https://github.com/rougier/numpy-book
# -----------------------------------------------------------------------------
import numpy as np
from smoke_solver import vel_step, dens_step

N = 64
size = N + 2
dt = 0.1
diff = 0.0
visc = 0.0
force = 5.0
source = 100.0
dvel = False

mouse = {"ox": 0.0, "oy": 0.0,
         "x": 0.0,  "y": 0.0,
         "button": None}

u = np.zeros((size, size), np.float32)  # velocity
u_prev = np.zeros((size, size), np.float32)

v = np.zeros((size, size), np.float32)  # velocity
v_prev = np.zeros((size, size), np.float32)

dens = np.zeros((size, size), np.float32)  # density
dens_prev = np.zeros((size, size), np.float32)


def initialization():
    global u, v, u_prev, v_prev, dens, dens_prev, size

    u[:, :] = 0.0
    v[:, :] = 0.0
    u_prev[:, :] = 0.0
    v_prev[:, :] = 0.0
    dens[:, :] = 0.0
    dens_prev[:, :] = 0.0


def user_step(d, u, v):
    global mouse

    d[:, :] = 0.0
    u[:, :] = 0.0
    v[:, :] = 0.0

    if mouse["button"] not in [1, 3]:
        return
    if mouse["x"] is None or mouse["y"] is None:
        return

    i = int(mouse["y"]*N) + 1
    j = int(mouse["x"]*N) + 1
    if not 0 < i < N+1 and not 0 < j < N+1:
        return

    if mouse["button"] == 3:
        d[i, j] = source
    elif mouse["button"] == 1:
        u[i, j] = force * (mouse["y"] - mouse["oy"])*200
        v[i, j] = force * (mouse["x"] - mouse["ox"])*200

    mouse["ox"] = mouse["x"]
    mouse["oy"] = mouse["y"]


def update(*args):
    global im, dens, dens_prev, u, u_prev, v, v_prev, N, visc, dt, diff

    user_step(dens_prev, u_prev, v_prev)
    vel_step(N, u, v, u_prev, v_prev, visc, dt)
    dens_step(N, dens, dens_prev, u, v, diff, dt)
    im.set_data(dens)
    # im.set_clim(vmin=dens.min(), vmax=dens.max())


def on_button_press(event):
    global mouse
    mouse["ox"] = mouse["x"] = event.xdata
    mouse["oy"] = mouse["y"] = event.ydata
    mouse["button"] = event.button


def on_button_release(event):
    global mouse
    mouse["ox"] = mouse["x"] = event.xdata
    mouse["oy"] = mouse["y"] = event.ydata
    mouse["button"] = None


def on_motion(event):
    global mouse
    mouse["x"] = event.xdata
    mouse["y"] = event.ydata


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation

    fig = plt.figure(figsize=(5, 5))
    ax = fig.add_axes([0, 0, 1, 1], frameon=False)

    cid = fig.canvas.mpl_connect('button_press_event', on_button_press)
    cid = fig.canvas.mpl_connect('button_release_event', on_button_release)
    cid = fig.canvas.mpl_connect('motion_notify_event', on_motion)

    ax.set_xlim(0, 1)
    ax.set_xticks([])
    ax.set_ylim(0, 1)
    ax.set_yticks([])

    initialization()
    im = ax.imshow(dens[1:-1, 1:-1],
                   interpolation='bicubic', extent=[0, 1, 0, 1],
                   cmap=plt.cm.gray, origin="lower", vmin=0, vmax=1)
    animation = FuncAnimation(fig, update, interval=10, frames=800)
    plt.show()