#!/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 argparse import time from numpy import sqrt import ipyparallel as ipp def setup_partitioner(comm, addrs, index, num_procs, gnum_cells, parts): """create a partitioner in the engine namespace""" global partitioner p = ZMQRectPartitioner2D( # noqa: F821 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) # noqa: F821 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.0, help="Wave speed (I think)") paa('-Ly', type=float, default=1.0, help="system size (in y)") paa('-Lx', type=float, default=1.0, help="system size (in x)") paa('-t', '--tstop', type=float, default=1.0, help="Time units to run") paa( '--profile', type=str, default='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 = ipp.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(f"Running {grid} system on {partition} processes until {tstop:f}") # 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) # noqa: F821 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, ipp.Reference('com'), peers, ipp.Reference('my_id'), num_procs, grid, partition, ) time.sleep(1) # convenience lambda to call solver.solve: def _solve(*args, **kwargs): return 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=ipp.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(f'scalar inner-version, Wtime={t1 - t0:g}, norm={norm:g}') # 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=ipp.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(f'vector inner-version, Wtime={t1 - t0:g}, norm={norm:g}') # 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()