forked from ipython/ipython
-
Notifications
You must be signed in to change notification settings - Fork 4
/
parallelwave.py
executable file
·209 lines (175 loc) · 6.65 KB
/
parallelwave.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
205
206
207
208
209
#!/usr/bin/env python
"""
A simple python program of solving a 2D wave equation in parallel.
Domain partitioning and inter-processor communication
are done by an object of class ZMQRectPartitioner2D
(which is a subclass of RectPartitioner2D and uses 0MQ via pyzmq)
An example of running the program is (8 processors, 4x2 partition,
200x200 grid cells)::
$ ipcluster start -n 8 # start 8 engines
$ python parallelwave.py --grid 200 200 --partition 4 2
See also parallelwave-mpi, which runs the same program, but uses MPI
(via mpi4py) for the inter-engine communication.
Authors
-------
* Xing Cai
* Min Ragan-Kelley
"""
#
import sys
import time
from numpy import exp, zeros, newaxis, sqrt
from IPython.external import argparse
from IPython.parallel import Client, Reference
def setup_partitioner(comm, addrs, index, num_procs, gnum_cells, parts):
"""create a partitioner in the engine namespace"""
global partitioner
p = ZMQRectPartitioner2D(comm, addrs, my_id=index, num_procs=num_procs)
p.redim(global_num_cells=gnum_cells, num_parts=parts)
p.prepare_communication()
# put the partitioner into the global namespace:
partitioner=p
def setup_solver(*args, **kwargs):
"""create a WaveSolver in the engine namespace."""
global solver
solver = WaveSolver(*args, **kwargs)
def wave_saver(u, x, y, t):
"""save the wave state for each timestep."""
global u_hist
global t_hist
t_hist.append(t)
u_hist.append(1.0*u)
# main program:
if __name__ == '__main__':
parser = argparse.ArgumentParser()
paa = parser.add_argument
paa('--grid', '-g',
type=int, nargs=2, default=[100,100], dest='grid',
help="Cells in the grid, e.g. --grid 100 200")
paa('--partition', '-p',
type=int, nargs=2, default=None,
help="Process partition grid, e.g. --partition 4 2 for 4x2")
paa('-c',
type=float, default=1.,
help="Wave speed (I think)")
paa('-Ly',
type=float, default=1.,
help="system size (in y)")
paa('-Lx',
type=float, default=1.,
help="system size (in x)")
paa('-t', '--tstop',
type=float, default=1.,
help="Time units to run")
paa('--profile',
type=unicode, default=u'default',
help="Specify the ipcluster profile for the client to connect to.")
paa('--save',
action='store_true',
help="Add this flag to save the time/wave history during the run.")
paa('--scalar',
action='store_true',
help="Also run with scalar interior implementation, to see vector speedup.")
ns = parser.parse_args()
# set up arguments
grid = ns.grid
partition = ns.partition
Lx = ns.Lx
Ly = ns.Ly
c = ns.c
tstop = ns.tstop
if ns.save:
user_action = wave_saver
else:
user_action = None
num_cells = 1.0*(grid[0]-1)*(grid[1]-1)
final_test = True
# create the Client
rc = Client(profile=ns.profile)
num_procs = len(rc.ids)
if partition is None:
partition = [num_procs,1]
else:
num_procs = min(num_procs, partition[0]*partition[1])
assert partition[0]*partition[1] == num_procs, "can't map partition %s to %i engines"%(partition, num_procs)
# construct the View:
view = rc[:num_procs]
print("Running %s system on %s processes until %f"%(grid, partition, tstop))
# functions defining initial/boundary/source conditions
def I(x,y):
from numpy import exp
return 1.5*exp(-100*((x-0.5)**2+(y-0.5)**2))
def f(x,y,t):
return 0.0
# from numpy import exp,sin
# return 10*exp(-(x - sin(100*t))**2)
def bc(x,y,t):
return 0.0
# initialize t_hist/u_hist for saving the state at each step (optional)
view['t_hist'] = []
view['u_hist'] = []
# set vector/scalar implementation details
impl = {}
impl['ic'] = 'vectorized'
impl['inner'] = 'scalar'
impl['bc'] = 'vectorized'
# execute some files so that the classes we need will be defined on the engines:
view.execute('import numpy')
view.run('communicator.py')
view.run('RectPartitioner.py')
view.run('wavesolver.py')
# scatter engine IDs
view.scatter('my_id', range(num_procs), flatten=True)
# create the engine connectors
view.execute('com = EngineCommunicator()')
# gather the connection information into a single dict
ar = view.apply_async(lambda : com.info)
peers = ar.get_dict()
# print peers
# this is a dict, keyed by engine ID, of the connection info for the EngineCommunicators
# setup remote partitioner
# note that Reference means that the argument passed to setup_partitioner will be the
# object named 'com' in the engine's namespace
view.apply_sync(setup_partitioner, Reference('com'), peers, Reference('my_id'), num_procs, grid, partition)
time.sleep(1)
# convenience lambda to call solver.solve:
_solve = lambda *args, **kwargs: solver.solve(*args, **kwargs)
if ns.scalar:
impl['inner'] = 'scalar'
# setup remote solvers
view.apply_sync(setup_solver, I,f,c,bc,Lx,Ly, partitioner=Reference('partitioner'), dt=0,implementation=impl)
# run first with element-wise Python operations for each cell
t0 = time.time()
ar = view.apply_async(_solve, tstop, dt=0, verbose=True, final_test=final_test, user_action=user_action)
if final_test:
# this sum is performed element-wise as results finish
s = sum(ar)
# the L2 norm (RMS) of the result:
norm = sqrt(s/num_cells)
else:
norm = -1
t1 = time.time()
print('scalar inner-version, Wtime=%g, norm=%g'%(t1-t0, norm))
# run again with faster numpy-vectorized inner implementation:
impl['inner'] = 'vectorized'
# setup remote solvers
view.apply_sync(setup_solver, I,f,c,bc,Lx,Ly,partitioner=Reference('partitioner'), dt=0,implementation=impl)
t0 = time.time()
ar = view.apply_async(_solve, tstop, dt=0, verbose=True, final_test=final_test, user_action=user_action)
if final_test:
# this sum is performed element-wise as results finish
s = sum(ar)
# the L2 norm (RMS) of the result:
norm = sqrt(s/num_cells)
else:
norm = -1
t1 = time.time()
print('vector inner-version, Wtime=%g, norm=%g'%(t1-t0, norm))
# if ns.save is True, then u_hist stores the history of u as a list
# If the partion scheme is Nx1, then u can be reconstructed via 'gather':
if ns.save and partition[-1] == 1:
import matplotlib.pyplot as plt
view.execute('u_last=u_hist[-1]')
u_last = view.gather('u_last', block=True)
plt.pcolor(u_last)
plt.show()