|
|
#!/usr/bin/env python
|
|
|
"""
|
|
|
A simple WaveSolver class for evolving the wave equation in 2D.
|
|
|
This works in parallel by using a RectPartitioner object.
|
|
|
|
|
|
Authors
|
|
|
-------
|
|
|
|
|
|
* Xing Cai
|
|
|
* Min Ragan-Kelley
|
|
|
|
|
|
"""
|
|
|
import time
|
|
|
|
|
|
from numpy import exp, zeros, newaxis, sqrt, arange
|
|
|
|
|
|
def iseq(start=0, stop=None, inc=1):
|
|
|
"""
|
|
|
Generate integers from start to (and including!) stop,
|
|
|
with increment of inc. Alternative to range/xrange.
|
|
|
"""
|
|
|
if stop is None: # allow isequence(3) to be 0, 1, 2, 3
|
|
|
# take 1st arg as stop, start as 0, and inc=1
|
|
|
stop = start; start = 0; inc = 1
|
|
|
return arange(start, stop+inc, inc)
|
|
|
|
|
|
class WaveSolver(object):
|
|
|
"""
|
|
|
Solve the 2D wave equation u_tt = u_xx + u_yy + f(x,y,t) with
|
|
|
u = bc(x,y,t) on the boundary and initial condition du/dt = 0.
|
|
|
|
|
|
Parallelization by using a RectPartitioner object 'partitioner'
|
|
|
|
|
|
nx and ny are the total number of global grid cells in the x and y
|
|
|
directions. The global grid points are numbered as (0,0), (1,0), (2,0),
|
|
|
..., (nx,0), (0,1), (1,1), ..., (nx, ny).
|
|
|
|
|
|
dt is the time step. If dt<=0, an optimal time step is used.
|
|
|
|
|
|
tstop is the stop time for the simulation.
|
|
|
|
|
|
I, f are functions: I(x,y), f(x,y,t)
|
|
|
|
|
|
user_action: function of (u, x, y, t) called at each time
|
|
|
level (x and y are one-dimensional coordinate vectors).
|
|
|
This function allows the calling code to plot the solution,
|
|
|
compute errors, etc.
|
|
|
|
|
|
implementation: a dictionary specifying how the initial
|
|
|
condition ('ic'), the scheme over inner points ('inner'),
|
|
|
and the boundary conditions ('bc') are to be implemented.
|
|
|
Normally, values are legal: 'scalar' or 'vectorized'.
|
|
|
'scalar' means straight loops over grid points, while
|
|
|
'vectorized' means special NumPy vectorized operations.
|
|
|
|
|
|
If a key in the implementation dictionary is missing, it
|
|
|
defaults in this function to 'scalar' (the safest strategy).
|
|
|
Note that if 'vectorized' is specified, the functions I, f,
|
|
|
and bc must work in vectorized mode. It is always recommended
|
|
|
to first run the 'scalar' mode and then compare 'vectorized'
|
|
|
results with the 'scalar' results to check that I, f, and bc
|
|
|
work.
|
|
|
|
|
|
verbose: true if a message at each time step is written,
|
|
|
false implies no output during the simulation.
|
|
|
|
|
|
final_test: true means the discrete L2-norm of the final solution is
|
|
|
to be computed.
|
|
|
"""
|
|
|
|
|
|
def __init__(self, I, f, c, bc, Lx, Ly, partitioner=None, dt=-1,
|
|
|
user_action=None,
|
|
|
implementation={'ic': 'vectorized', # or 'scalar'
|
|
|
'inner': 'vectorized',
|
|
|
'bc': 'vectorized'}):
|
|
|
|
|
|
nx = partitioner.global_num_cells[0] # number of global cells in x dir
|
|
|
ny = partitioner.global_num_cells[1] # number of global cells in y dir
|
|
|
dx = Lx/float(nx)
|
|
|
dy = Ly/float(ny)
|
|
|
loc_nx, loc_ny = partitioner.get_num_loc_cells()
|
|
|
nx = loc_nx; ny = loc_ny # now use loc_nx and loc_ny instead
|
|
|
lo_ix0 = partitioner.subd_lo_ix[0]
|
|
|
lo_ix1 = partitioner.subd_lo_ix[1]
|
|
|
hi_ix0 = partitioner.subd_hi_ix[0]
|
|
|
hi_ix1 = partitioner.subd_hi_ix[1]
|
|
|
x = iseq(dx*lo_ix0, dx*hi_ix0, dx) # local grid points in x dir
|
|
|
y = iseq(dy*lo_ix1, dy*hi_ix1, dy) # local grid points in y dir
|
|
|
self.x = x
|
|
|
self.y = y
|
|
|
xv = x[:,newaxis] # for vectorized expressions with f(xv,yv)
|
|
|
yv = y[newaxis,:] # -- " --
|
|
|
if dt <= 0:
|
|
|
dt = (1/float(c))*(1/sqrt(1/dx**2 + 1/dy**2)) # max time step
|
|
|
Cx2 = (c*dt/dx)**2; Cy2 = (c*dt/dy)**2; dt2 = dt**2 # help variables
|
|
|
|
|
|
u = zeros((nx+1,ny+1)) # solution array
|
|
|
u_1 = u.copy() # solution at t-dt
|
|
|
u_2 = u.copy() # solution at t-2*dt
|
|
|
|
|
|
# preserve for self.solve
|
|
|
implementation=dict(implementation) # copy
|
|
|
|
|
|
if 'ic' not in implementation:
|
|
|
implementation['ic'] = 'scalar'
|
|
|
if 'bc' not in implementation:
|
|
|
implementation['bc'] = 'scalar'
|
|
|
if 'inner' not in implementation:
|
|
|
implementation['inner'] = 'scalar'
|
|
|
|
|
|
self.implementation = implementation
|
|
|
self.Lx = Lx
|
|
|
self.Ly = Ly
|
|
|
self.I=I
|
|
|
self.f=f
|
|
|
self.c=c
|
|
|
self.bc=bc
|
|
|
self.user_action = user_action
|
|
|
self.partitioner=partitioner
|
|
|
|
|
|
# set initial condition (pointwise - allows straight if-tests in I(x,y)):
|
|
|
t=0.0
|
|
|
if implementation['ic'] == 'scalar':
|
|
|
for i in xrange(0,nx+1):
|
|
|
for j in xrange(0,ny+1):
|
|
|
u_1[i,j] = I(x[i], y[j])
|
|
|
|
|
|
for i in xrange(1,nx):
|
|
|
for j in xrange(1,ny):
|
|
|
u_2[i,j] = u_1[i,j] + \
|
|
|
0.5*Cx2*(u_1[i-1,j] - 2*u_1[i,j] + u_1[i+1,j]) + \
|
|
|
0.5*Cy2*(u_1[i,j-1] - 2*u_1[i,j] + u_1[i,j+1]) + \
|
|
|
dt2*f(x[i], y[j], 0.0)
|
|
|
|
|
|
# boundary values of u_2 (equals u(t=dt) due to du/dt=0)
|
|
|
i = 0
|
|
|
for j in xrange(0,ny+1):
|
|
|
u_2[i,j] = bc(x[i], y[j], t+dt)
|
|
|
j = 0
|
|
|
for i in xrange(0,nx+1):
|
|
|
u_2[i,j] = bc(x[i], y[j], t+dt)
|
|
|
i = nx
|
|
|
for j in xrange(0,ny+1):
|
|
|
u_2[i,j] = bc(x[i], y[j], t+dt)
|
|
|
j = ny
|
|
|
for i in xrange(0,nx+1):
|
|
|
u_2[i,j] = bc(x[i], y[j], t+dt)
|
|
|
|
|
|
elif implementation['ic'] == 'vectorized':
|
|
|
u_1 = I(xv,yv)
|
|
|
u_2[1:nx,1:ny] = u_1[1:nx,1:ny] + \
|
|
|
0.5*Cx2*(u_1[0:nx-1,1:ny] - 2*u_1[1:nx,1:ny] + u_1[2:nx+1,1:ny]) + \
|
|
|
0.5*Cy2*(u_1[1:nx,0:ny-1] - 2*u_1[1:nx,1:ny] + u_1[1:nx,2:ny+1]) + \
|
|
|
dt2*(f(xv[1:nx,1:ny], yv[1:nx,1:ny], 0.0))
|
|
|
# boundary values (t=dt):
|
|
|
i = 0; u_2[i,:] = bc(x[i], y, t+dt)
|
|
|
j = 0; u_2[:,j] = bc(x, y[j], t+dt)
|
|
|
i = nx; u_2[i,:] = bc(x[i], y, t+dt)
|
|
|
j = ny; u_2[:,j] = bc(x, y[j], t+dt)
|
|
|
|
|
|
if user_action is not None:
|
|
|
user_action(u_1, x, y, t) # allow user to plot etc.
|
|
|
# print(list(self.us[2][2]))
|
|
|
self.us = (u,u_1,u_2)
|
|
|
|
|
|
|
|
|
def solve(self, tstop, dt=-1, user_action=None, verbose=False, final_test=False):
|
|
|
t0=time.time()
|
|
|
f=self.f
|
|
|
c=self.c
|
|
|
bc=self.bc
|
|
|
partitioner = self.partitioner
|
|
|
implementation = self.implementation
|
|
|
nx = partitioner.global_num_cells[0] # number of global cells in x dir
|
|
|
ny = partitioner.global_num_cells[1] # number of global cells in y dir
|
|
|
dx = self.Lx/float(nx)
|
|
|
dy = self.Ly/float(ny)
|
|
|
loc_nx, loc_ny = partitioner.get_num_loc_cells()
|
|
|
nx = loc_nx; ny = loc_ny # now use loc_nx and loc_ny instead
|
|
|
x = self.x
|
|
|
y = self.y
|
|
|
xv = x[:,newaxis] # for vectorized expressions with f(xv,yv)
|
|
|
yv = y[newaxis,:] # -- " --
|
|
|
if dt <= 0:
|
|
|
dt = (1/float(c))*(1/sqrt(1/dx**2 + 1/dy**2)) # max time step
|
|
|
Cx2 = (c*dt/dx)**2; Cy2 = (c*dt/dy)**2; dt2 = dt**2 # help variables
|
|
|
# id for the four possible neighbor subdomains
|
|
|
lower_x_neigh = partitioner.lower_neighbors[0]
|
|
|
upper_x_neigh = partitioner.upper_neighbors[0]
|
|
|
lower_y_neigh = partitioner.lower_neighbors[1]
|
|
|
upper_y_neigh = partitioner.upper_neighbors[1]
|
|
|
u,u_1,u_2 = self.us
|
|
|
# u_1 = self.u_1
|
|
|
|
|
|
t = 0.0
|
|
|
while t <= tstop:
|
|
|
t_old = t; t += dt
|
|
|
if verbose:
|
|
|
print('solving (%s version) at t=%g' % \
|
|
|
(implementation['inner'], t))
|
|
|
# update all inner points:
|
|
|
if implementation['inner'] == 'scalar':
|
|
|
for i in xrange(1, nx):
|
|
|
for j in xrange(1, ny):
|
|
|
u[i,j] = - u_2[i,j] + 2*u_1[i,j] + \
|
|
|
Cx2*(u_1[i-1,j] - 2*u_1[i,j] + u_1[i+1,j]) + \
|
|
|
Cy2*(u_1[i,j-1] - 2*u_1[i,j] + u_1[i,j+1]) + \
|
|
|
dt2*f(x[i], y[j], t_old)
|
|
|
elif implementation['inner'] == 'vectorized':
|
|
|
u[1:nx,1:ny] = - u_2[1:nx,1:ny] + 2*u_1[1:nx,1:ny] + \
|
|
|
Cx2*(u_1[0:nx-1,1:ny] - 2*u_1[1:nx,1:ny] + u_1[2:nx+1,1:ny]) + \
|
|
|
Cy2*(u_1[1:nx,0:ny-1] - 2*u_1[1:nx,1:ny] + u_1[1:nx,2:ny+1]) + \
|
|
|
dt2*f(xv[1:nx,1:ny], yv[1:nx,1:ny], t_old)
|
|
|
|
|
|
# insert boundary conditions (if there's no neighbor):
|
|
|
if lower_x_neigh < 0:
|
|
|
if implementation['bc'] == 'scalar':
|
|
|
i = 0
|
|
|
for j in xrange(0, ny+1):
|
|
|
u[i,j] = bc(x[i], y[j], t)
|
|
|
elif implementation['bc'] == 'vectorized':
|
|
|
u[0,:] = bc(x[0], y, t)
|
|
|
if upper_x_neigh < 0:
|
|
|
if implementation['bc'] == 'scalar':
|
|
|
i = nx
|
|
|
for j in xrange(0, ny+1):
|
|
|
u[i,j] = bc(x[i], y[j], t)
|
|
|
elif implementation['bc'] == 'vectorized':
|
|
|
u[nx,:] = bc(x[nx], y, t)
|
|
|
if lower_y_neigh < 0:
|
|
|
if implementation['bc'] == 'scalar':
|
|
|
j = 0
|
|
|
for i in xrange(0, nx+1):
|
|
|
u[i,j] = bc(x[i], y[j], t)
|
|
|
elif implementation['bc'] == 'vectorized':
|
|
|
u[:,0] = bc(x, y[0], t)
|
|
|
if upper_y_neigh < 0:
|
|
|
if implementation['bc'] == 'scalar':
|
|
|
j = ny
|
|
|
for i in xrange(0, nx+1):
|
|
|
u[i,j] = bc(x[i], y[j], t)
|
|
|
elif implementation['bc'] == 'vectorized':
|
|
|
u[:,ny] = bc(x, y[ny], t)
|
|
|
|
|
|
# communication
|
|
|
partitioner.update_internal_boundary (u)
|
|
|
|
|
|
if user_action is not None:
|
|
|
user_action(u, x, y, t)
|
|
|
# update data structures for next step
|
|
|
u_2, u_1, u = u_1, u, u_2
|
|
|
|
|
|
t1 = time.time()
|
|
|
print('my_id=%2d, dt=%g, %s version, slice_copy=%s, net Wtime=%g'\
|
|
|
%(partitioner.my_id,dt,implementation['inner'],\
|
|
|
partitioner.slice_copy,t1-t0))
|
|
|
# save the us
|
|
|
self.us = u,u_1,u_2
|
|
|
# check final results; compute discrete L2-norm of the solution
|
|
|
if final_test:
|
|
|
loc_res = 0.0
|
|
|
for i in iseq(start=1, stop=nx-1):
|
|
|
for j in iseq(start=1, stop=ny-1):
|
|
|
loc_res += u_1[i,j]**2
|
|
|
return loc_res
|
|
|
return dt
|
|
|
|
|
|
|
