/
example_diffusion.py
204 lines (172 loc) · 8.13 KB
/
example_diffusion.py
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
"""
NOTE: This is a legacy example. For more up-to-date examples
please see the tutorials under `examples/cfd/`.
This example encodes multiple ways to solve the 2D diffusion equations
using an explicit finite difference scheme with fixed boundary values
and a given initial value for the density. The example also includes
demo implementations using more traditional methods, including pure
Python, vectorized NumPy and a vectorized symbolic implementation
using SymPy's `lambdify()` functionality. For a visual demonstration
and comparison between pure Python and NumPy, for example, run:
python example_diffusion.py run -m python numpy --show
"""
import time
from argparse import ArgumentParser
import numpy as np
import sympy
from devito import Grid, Eq, Operator, TimeFunction, solve
from devito.logger import log
def ring_initial(spacing=0.01):
"""Initialise grid with initial condition ("ring")"""
nx, ny = int(1 / spacing), int(1 / spacing)
xx, yy = np.meshgrid(np.linspace(0., 1., nx, dtype=np.float32),
np.linspace(0., 1., ny, dtype=np.float32))
ui = np.zeros((nx, ny), dtype=np.float32)
r = (xx - .5)**2. + (yy - .5)**2.
ui[np.logical_and(.05 <= r, r <= .1)] = 1.
return ui
def execute_python(ui, spacing=0.01, a=0.5, timesteps=500):
"""Execute diffusion stencil using pure Python list indexing."""
nx, ny = ui.shape
dx2, dy2 = spacing**2, spacing**2
dt = dx2 * dy2 / (2 * a * (dx2 + dy2))
u = np.concatenate((ui, np.zeros_like(ui))).reshape((2, nx, ny))
# Execute timestepping loop with alternating buffers
tstart = time.time()
for ti in range(timesteps):
t0 = ti % 2
t1 = (ti + 1) % 2
for i in range(1, nx-1):
for j in range(1, ny-1):
uxx = (u[t0, i+1, j] - 2*u[t0, i, j] + u[t0, i-1, j]) / dx2
uyy = (u[t0, i, j+1] - 2*u[t0, i, j] + u[t0, i, j-1]) / dy2
u[t1, i, j] = u[t0, i, j] + dt * a * (uxx + uyy)
runtime = time.time() - tstart
log("Python: Diffusion with dx=%0.4f, dy=%0.4f, executed %d timesteps in %f seconds"
% (spacing, spacing, timesteps, runtime))
return u[ti % 2, :, :], runtime
def execute_numpy(ui, spacing=0.01, a=0.5, timesteps=500):
"""Execute diffusion stencil using vectorised numpy array accesses."""
nx, ny = ui.shape
dx2, dy2 = spacing**2, spacing**2
dt = dx2 * dy2 / (2 * a * (dx2 + dy2))
u = np.concatenate((ui, np.zeros_like(ui))).reshape((2, nx, ny))
# Execute timestepping loop with alternating buffers
tstart = time.time()
for ti in range(timesteps):
t0 = ti % 2
t1 = (ti + 1) % 2
uxx = (u[t0, 2:, 1:-1] - 2*u[t0, 1:-1, 1:-1] + u[t0, :-2, 1:-1]) / dx2
uyy = (u[t0, 1:-1, 2:] - 2*u[t0, 1:-1, 1:-1] + u[t0, 1:-1, :-2]) / dy2
u[t1, 1:-1, 1:-1] = u[t0, 1:-1, 1:-1] + a * dt * (uxx + uyy)
runtime = time.time() - tstart
log("Numpy: Diffusion with dx=%0.4f, dy=%0.4f, executed %d timesteps in %f seconds"
% (spacing, spacing, timesteps, runtime))
return u[ti % 2, :, :], runtime
def execute_lambdify(ui, spacing=0.01, a=0.5, timesteps=500):
"""Execute diffusion stencil using vectorised numpy array accesses."""
nx, ny = ui.shape
dx2, dy2 = spacing**2, spacing**2
dt = dx2 * dy2 / (2 * a * (dx2 + dy2))
u = np.concatenate((ui, np.zeros_like(ui))).reshape((2, nx, ny))
def diffusion_stencil():
"""Create stencil and substitutions for the diffusion equation"""
p = sympy.Function('p')
x, y, t, h, s = sympy.symbols('x y t h s')
dx2 = p(x, y, t).diff(x, x).as_finite_difference([x - h, x, x + h])
dy2 = p(x, y, t).diff(y, y).as_finite_difference([y - h, y, y + h])
dt = p(x, y, t).diff(t).as_finite_difference([t, t + s])
eqn = Eq(dt, a * (dx2 + dy2))
stencil = solve(eqn, p(x, y, t + s))
return stencil, (p(x, y, t), p(x + h, y, t), p(x - h, y, t),
p(x, y + h, t), p(x, y - h, t), s, h)
stencil, subs = diffusion_stencil()
kernel = sympy.lambdify(subs, stencil, 'numpy')
# Execute timestepping loop with alternating buffers
tstart = time.time()
for ti in range(timesteps):
t0 = ti % 2
t1 = (ti + 1) % 2
u[t1, 1:-1, 1:-1] = kernel(u[t0, 1:-1, 1:-1], u[t0, 2:, 1:-1],
u[t0, :-2, 1:-1], u[t0, 1:-1, 2:],
u[t0, 1:-1, :-2], dt, spacing)
runtime = time.time() - tstart
log("Lambdify: Diffusion with dx=%0.4f, dy=%0.4f, executed %d timesteps in %f seconds"
% (spacing, spacing, timesteps, runtime))
return u[ti % 2, :, :], runtime
def execute_devito(ui, spacing=0.01, a=0.5, timesteps=500):
"""Execute diffusion stencil using the devito Operator API."""
nx, ny = ui.shape
dx2, dy2 = spacing**2, spacing**2
dt = dx2 * dy2 / (2 * a * (dx2 + dy2))
# Allocate the grid and set initial condition
# Note: This should be made simpler through the use of defaults
grid = Grid(shape=(nx, ny))
u = TimeFunction(name='u', grid=grid, time_order=1, space_order=2)
u.data[0, :] = ui[:]
# Derive the stencil according to devito conventions
eqn = Eq(u.dt, a * (u.dx2 + u.dy2))
stencil = solve(eqn, u.forward)
op = Operator(Eq(u.forward, stencil))
# Execute the generated Devito stencil operator
tstart = time.time()
op.apply(u=u, t=timesteps, dt=dt)
runtime = time.time() - tstart
log("Devito: Diffusion with dx=%0.4f, dy=%0.4f, executed %d timesteps in %f seconds"
% (spacing, spacing, timesteps, runtime))
return u.data[1, :], runtime
def animate(field):
"""Animate solution field"""
import matplotlib.pyplot as plt
xvals = np.linspace(0, 1., field.shape[0])
yvals = np.linspace(0, 1., field.shape[1])
levels = np.linspace(-.1, 1.1, 21)
cs = plt.contourf(xvals, yvals, field, levels=levels)
plt.colorbar(cs)
plt.show()
def test_diffusion2d(spacing=0.01, timesteps=1000):
ui = ring_initial(spacing=spacing)
u, _ = execute_devito(ui, spacing=spacing, timesteps=timesteps)
assert(u.max() < 2.4)
assert(np.linalg.norm(u, ord=2) < 13)
if __name__ == "__main__":
description = """
Example script demonstrating several approaches to solving the
2D diffusion equation. The various modes are pure Python,
vectorized numpy, a lambdified SymPy equation and the Devito API.
Please note that the default settings for spacing and timesteps are
chosen to highlight that Devito and "vectorised numpy" are faster than
pure Python or the sympy.lambdify kernels approach. For a fair
performance comparison between Devito and "vectorised numpy" we
recommend using --spacing 0.001 -t 1000.
"""
parser = ArgumentParser(description=description)
parser.add_argument('-m', '--mode', nargs='+', default=['devito'],
choices=['python', 'numpy', 'lambdify', 'devito'],
help="Example modes: python, numpy, lambdify, devito")
parser.add_argument('-s', '--spacing', type=float, nargs='+', default=[0.01],
help='Grid spacing in x and y)')
parser.add_argument('-t', '--timesteps', type=int, default=20,
help='Number of timesteps to run')
parser.add_argument('--show', action='store_true',
help="Show animation of the solution field")
args = parser.parse_args()
if len(args.spacing) > 2:
raise ValueError("Too many arguments encountered for --spacing")
# Map "mode" arg to execution function
exec_func = {'python': execute_python, 'numpy': execute_numpy,
'lambdify': execute_lambdify, 'devito': execute_devito}
# Create parameters dict and remove script-specific args
parameters = vars(args).copy()
del parameters['show']
if len(args.spacing) > 1:
raise ValueError("Run mode only supports a single --spacing value")
# Run and animate example runs
for mode in args.mode:
ui = ring_initial(spacing=args.spacing[0])
if args.show:
animate(ui)
u, _ = exec_func[mode](ui, spacing=args.spacing[0],
timesteps=args.timesteps)
if args.show:
animate(u)