x = y = np.linspace(-10., 10., 41)
xv, yv = plt.ndgrid(x, y) # Grid for x, y values (km)
hv = h0 / (1 + (xv**2 + yv**2) / (R**2)) # Compute height (m)
x = y = np.linspace(-10., 10., 11)
x2v, y2v = plt.ndgrid(x, y) # Define a coarser grid for the vector field
h2v = h0 / (1 + (x2v**2 + y2v**2) / (R**2)) # Compute height for new grid
dhdx, dhdy = np.gradient(h2v) # Compute the gradient vector (dh/dx,dh/dy)
# Draw contours and gradient field of h
plt.figure(9)
plt.quiver(x2v, y2v, dhdx, dhdy, 0, 'r')
plt.hold('on')
plt.contour(xv, yv, hv)
plt.axis('equal')
# end draw contours and gradient field of h
x = y = np.linspace(-5, 5, 11)
xv, yv = plt.ndgrid(x, y)
u = xv**2 + 2 * yv - .5 * xv * yv
v = -3 * yv
# Draw 2D-field
plt.figure(10)
plt.quiver(xv, yv, u, v, 200, 'b')
plt.axis('equal')
# end draw 2D-field
plt.figure(9)
plt.savefig('images/quiver_scitools_advanced.pdf')
def solve(self,T=10.4,version="scalar",quiet_mode=True,display=False,\
save_figs=False,filename="figure",destination=None):
"""
The solver solves the 2D wave equation with varying coefficients inside the
spatial derivatives. By default the solver uses the scalar version of
the implementation. Other implementations includes a vectorized version, and
soon to be implemented Cython version.
Variables :
T : End time for simulation
version : version of implementation
Choices include : - scalar (default)
- vectorized
- cython
quiet_mode : If True (default), solver does not display any progress message.
And if False, solver gives detailed information what it is doing.
save_figs : If True, saves the plots at certain time steps. These figures can
later be preprocessed to a single gif file. If False (defualt),
no figures are saved.
"""
if save_figs or display:
viz = self.Visualize(destination=destination,filename=filename)
destination = viz.destination # in case None was given as destination
if version=="scalar" or version=="vectorized":
self.version = version
else:
import sys
print "Only following versions have been implemented for the solver : "
print "scalar and vectorized"
print "versions must be case-sensitive"
sys.exit(1)
import numpy as np
self.T = T
dt = self.physprm.gridprm.dt
Lx = self.physprm.gridprm.Lx
Ly = self.physprm.gridprm.Ly
dt = self.physprm.gridprm.dt
dx = self.physprm.gridprm.dx
dy = self.physprm.gridprm.dy
Nx = self.physprm.gridprm.Nx
Ny = self.physprm.gridprm.Ny
Nt = int(round(T/float(dt)))
t = np.linspace(0,T,Nt+1)
#t = np.arange(0,T+dt,dt)
x = np.linspace(0,Lx,Nx+3)
y = np.linspace(0,Ly,Ny+3)
self.x = x
self.y = y
if self.physprm.gridprm.problem3D :
if version=="scalar":
print "Only vectorized version available, setting version to vectorized."
version="vectorized"
Lz = self.physprm.gridprm.Lz
dz = self.physprm.gridprm.dz
Nz = self.physprm.gridprm.Nz
z = np.linspace(0,Lz,Nz+3)
self.z = z
# Set bcs (either Dirichlet BCs or Neumann BCs) :
BC = self.BC
# Define the arrays to store the three time steps :
# time step u^{n+1} :
if self.physprm.gridprm.problem3D :
u = np.zeros([Nx+3,Ny+3,Nz+3])
self.u = u
# time step u^{n} :
u_1 = np.zeros([Nx+3,Ny+3,Nz+3])
self.u_1 = u_1
# time step u^{n-1} :
u_2 = np.zeros([Nx+3,Ny+3,Nz+3])
self.u_2 = u_2
else :
u = np.zeros([Nx+3,Ny+3])
self.u = u
# time step u^{n} :
u_1 = np.zeros([Nx+3,Ny+3])
self.u_1 = u_1
# time step u^{n-1} :
u_2 = np.zeros([Nx+3,Ny+3])
self.u_2 = u_2
if version=="vectorized":
xv = x[:,np.newaxis] # for vectorized function evaluations
yv = y[np.newaxis,:]
if self.physprm.gridprm.problem3D :
import scitools.easyviz as plt
# what do we do here ?
xx,yy,zz = plt.ndgrid(x,y,z)
# These are the real physical points on our grid :
Ix = range(1,u.shape[0]-1)
self.Ix = Ix
Iy = range(1,u.shape[1]-1)
self.Iy = Iy
if self.physprm.gridprm.problem3D :
Iz = range(1,u.shape[2]-1)
self.Iz = Iz
b = self.physprm.b
# Assuming the following are functions
f = self.physprm.f
I = self.physprm.I
V = self.physprm.V
q = self.physprm.q
if self.physprm.gridprm.problem3D :
def q_i_half(x,y,z,x_dx):
return 0.5*(q(x,y,z)+q(x_dx,y,z))
def q_j_half(x,y,z,y_dy):
return 0.5*(q(x,y,z)+q(x,y_dy,z))
def q_k_half(x,y,z,z_dz):
return 0.5*(q(x,y,z)+q(x,y,z_dz))
q_i_phalf_array = np.zeros([Nx+3,Ny+3,Nz+3])
q_i_mhalf_array = np.zeros([Nx+3,Ny+3,Nz+3])
q_j_phalf_array = np.zeros([Nx+3,Ny+3,Nz+3])
q_j_mhalf_array = np.zeros([Nx+3,Ny+3,Nz+3])
q_k_phalf_array = np.zeros([Nx+3,Ny+3,Nz+3])
q_k_mhalf_array = np.zeros([Nx+3,Ny+3,Nz+3])
q_i_phalf_array[1:-1,1:-1,1:-1] = q_i_half(x[1:-1],y[1:-1],z[1:-1],x[2:])
q_i_mhalf_array[1:-1,1:-1,1:-1] = q_i_half(x[1:-1],y[1:-1],z[1:-1],x[:-2])
q_j_phalf_array[1:-1,1:-1,1:-1] = q_j_half(x[1:-1],y[1:-1],z[1:-1],y[2:])
q_j_mhalf_array[1:-1,1:-1,1:-1] = q_j_half(x[1:-1],y[1:-1],z[1:-1],y[:-2])
q_k_phalf_array[1:-1,1:-1,1:-1] = q_j_half(x[1:-1],y[1:-1],z[1:-1],z[2:])
q_k_mhalf_array[1:-1,1:-1,1:-1] = q_j_half(x[1:-1],y[1:-1],z[1:-1],z[:-2])
else :
# The following functions will be used as well later through-out
def q_i_half(x,y,x_dx):
# used arithmetic mean for q(i+1/2,j) and the other corresponding terms
return 0.5*(q(x,y)+q(x_dx,y)) # x_dx should either be x + dx or x - dx
def q_j_half(x,y,y_dy):
return 0.5*(q(x,y)+q(x,y_dy)) # y_dy should either be y + dy or y - dy
if self.outlet:
q_a = np.zeros([Nx+3,Ny+3])
C5 = np.zeros([Nx+3,Ny+3])
q_i_phalf_array = np.zeros([Nx+3,Ny+3])
q_i_mhalf_array = np.zeros([Nx+3,Ny+3])
q_j_phalf_array = np.zeros([Nx+3,Ny+3])
q_j_mhalf_array = np.zeros([Nx+3,Ny+3])
# Pre-compute q, f and V.
if version=="scalar":
# Fill all q_***_arrays with their respective values.
if self.outlet:
for i in Ix:
for j in Iy:
q_a[i,j] = q(x[i],y[j]) # used for outlet bcs
C5[i,j] = (dx/dt)/np.sqrt(q_a[i,j])
for i in Ix:
for j in Iy:
q_i_phalf_array[i,j] = q_i_half(x[i],y[j],x[i+1])
q_i_mhalf_array[i,j] = q_i_half(x[i],y[j],x[i-1])
q_j_phalf_array[i,j] = q_j_half(x[i],y[j],y[j+1])
q_j_mhalf_array[i,j] = q_j_half(x[i],y[j],y[j-1])
elif version=="vectorized":
q_i_phalf_array[1:-1,1:-1] = q_i_half(x[1:-1],y[1:-1],x[2:])
q_i_mhalf_array[1:-1,1:-1] = q_i_half(x[1:-1],y[1:-1],x[:-2])
q_j_phalf_array[1:-1,1:-1] = q_j_half(x[1:-1],y[1:-1],y[2:])
q_j_mhalf_array[1:-1,1:-1] = q_j_half(x[1:-1],y[1:-1],y[:-2])
if self.outlet:
q_a[1:-1,1:-1] = q(x[1:-1],y[1:-1]) # used for outlet bcs
C5[1:-1,1:-1] = (dx/dt)/np.sqrt(q_a[1:-1,1:-1])
if not self.bypass_stability:
if self.physprm.gridprm.problem3D :
q_val = (np.fabs(q(x[:],y[:],z[:]))).max()
stab_coeff = (1/float(np.sqrt(q_val)))*(1/np.sqrt(1/dx**2 + 1/dy**2 + 1/dz**2))
else:
q_val = (np.fabs(q(x[:],y[:]))).max()
stab_coeff = (1/float(np.sqrt(q_val)))*(1/np.sqrt(1/dx**2 + 1/dy**2))
if stab_coeff < t[1] :
import sys
error_msg = "Time step dt=%f exceeds the stability limit %f."
sys.exit(error_msg%(t[1],stab_coeff))
# Following constants will be used throughout the solver
C1 = 2.0/(2.0 + b*dt)
Cb = b*dt*0.5 - 1.0
C11 = dt*(1 - 0.5*b*dt)
dtdt = dt*dt
C2 = dtdt/(dx*dx)
C3 = dtdt/(dy*dy)
if self.physprm.gridprm.problem3D :
C4 = dtdt/(dz*dz)
########## Step 1 : Initialize for the zeroth time step ##########
if not quiet_mode:
print "Initializing zeroth time step"
if version=="scalar":
for i in Ix:
for j in Iy:
u_2[i,j] = I(x[i-Ix[0]],y[j-Iy[0]])
elif version=="vectorized":
if self.physprm.gridprm.problem3D :
u_2[:,:] = I(xv,yv,z[:])
else :
u_2[:,:] = I(xv,yv)
if save_figs : # store graphs -> create animation at end
import time
if self.physprm.gridprm.problem3D :
plt.setp(interactive=False)
h = plt.isosurface(xx,yy,zz,u_2,-3)
h.setp(opacity=0.5)
plt.shading("interp")
plt.daspect([1,1,1])
plt.view(3)
plt.axis("tight")
plt.show()
raw_input("enter")
else :
viz.store_fig(t[0],Ix,Iy,u_2,u_1,u)
time.sleep(1.0)
elif display:
import time
viz.display(t[0],Ix,Iy,u_2,u_1,u)
time.sleep(1.0)
############ Update ghost points for u_2 ############
if not quiet_mode:
print "Updating ghost points"
if self.physprm.gridprm.problem3D :
BC(self.u_2,Ix,Iy,Iz,version)
else :
BC(self.u_2,Ix,Iy,version)
########## Step 2 : Implement special case for t[n]=0 ##########
"""
Here we use a leap frog scheme for u_t (short notation for
du/dt) and insert the unknown u(x,y,-dt) into the main scheme
to find an expression for u_1, i.e u^{n}_{i,j}.
"""
if not quiet_mode:
print "Initializing time step dt"
# Inserting u^{n=-1} = u_1 - 2*dt*V(x,y) into the main scheme and solving for u_1 :
# Insert n = 0 in main scheme (see while loop) and insert expression for u^-1
if version=="scalar":
for i in Ix:
for j in Iy:
q_i_phalf = q_i_phalf_array[i,j]
q_i_mhalf = q_i_mhalf_array[i,j]
q_j_phalf = q_j_phalf_array[i,j]
q_j_mhalf = q_j_mhalf_array[i,j]
u_1[i,j] = u_2[i,j] + C11*V(x[i],y[j])\
+ 0.5*C2*(q_i_phalf*(u_2[i+1,j]-u_2[i,j])-\
q_i_mhalf*(u_2[i,j]-u_2[i-1,j]))\
+ 0.5*C3*(q_j_phalf*(u_2[i,j+1]-u_2[i,j])-\
q_j_mhalf*(u_2[i,j]-u_2[i,j-1]))\
+ 0.5*dtdt*f(x[i],y[j],t[0])
elif version=="vectorized":
if self.physprm.gridprm.problem3D :
u_1[1:-1,1:-1,1:-1] = u_2[1:-1,1:-1,1:-1] + C11*V(x[1:-1],y[1:-1],z[1:-1])\
+ 0.5*C2*(q_i_phalf_array[1:-1,1:-1,1:-1]*\
(u_2[2:,1:-1,1:-1]-u_2[1:-1,1:-1,1:-1])-\
q_i_mhalf_array[1:-1,1:-1,1:-1]*\
(u_2[1:-1,1:-1,1:-1]-u_2[:-2,1:-1,1:-1]))\
+ 0.5*C3*(q_j_phalf_array[1:-1,1:-1,1:-1]*\
(u_2[1:-1,2:,1:-1]-u_2[1:-1,1:-1,1:-1])-\
q_j_mhalf_array[1:-1,1:-1,1:-1]*\
(u_2[1:-1,1:-1,1:-1]-u_2[1:-1,:-2,1:-1]))\
+ 0.5*C4*(q_k_phalf_array[1:-1,1:-1,1:-1]*\
(u_2[1:-1,1:-1,2:]-u_2[1:-1,1:-1,1:-1])-\
q_k_mhalf_array[1:-1,1:-1,1:-1]*\
(u_2[1:-1,1:-1,1:-1]-u_2[1:-1,1:-1,:-2]))\
+ 0.5*dtdt*f(x[1:-1],y[1:-1],z[1:-1],t[0])
else :
u_1[1:-1,1:-1] = u_2[1:-1,1:-1] + C11*V(x[1:-1],y[1:-1])\
+ 0.5*C2*(q_i_phalf_array[1:-1,1:-1]*\
(u_2[2:,1:-1]-u_2[1:-1,1:-1])-\
q_i_mhalf_array[1:-1,1:-1]*\
(u_2[1:-1,1:-1]-u_2[:-2,1:-1]))\
+ 0.5*C3*(q_j_phalf_array[1:-1,1:-1]*\
(u_2[1:-1,2:]-u_2[1:-1,1:-1])-\
q_j_mhalf_array[1:-1,1:-1]*\
(u_2[1:-1,1:-1]-u_2[1:-1,:-2]))\
+ 0.5*dtdt*f(x[1:-1],y[1:-1],t[0])
if save_figs :
if self.physprm.gridprm.problem3D :
plt.setp(interactive=False)
h = plt.isosurface(xx,yy,zz,u_2,-3)
h.setp(opacity=0.5)
plt.shading("interp")
plt.daspect([1,1,1])
plt.view(3)
plt.axis("tight")
plt.show()
else :
viz.store_fig(t[1],Ix,Iy,u_2,u_1,u)
time.sleep(0.5)
elif display:
viz.display(t[1],Ix,Iy,u_2,u_1,u)
time.sleep(0.5)
############ Update ghost points for u_1 ############
if not quiet_mode:
print "Updating ghost points"
if self.physprm.gridprm.problem3D :
BC(self.u_1,Ix,Iy,Iz,version)
else :
BC(self.u_1,Ix,Iy,version)
if self.outlet:
term = np.zeros([Nx+3,Ny+3])
############ Run scheme for all time steps ############
n = 1
if not quiet_mode:
print "Starting scheme"
while t[n]<self.T-1e-9:
if not quiet_mode:
print "Starting time iteration for t=%.3f, n=%g"%(t[n],n)
if version=="scalar":
for i in Ix:
for j in Iy: # use arithmetic mean for q(i+1/2,j) etc.
q_i_phalf = q_i_phalf_array[i,j]
q_i_mhalf = q_i_mhalf_array[i,j]
q_j_phalf = q_j_phalf_array[i,j]
q_j_mhalf = q_j_mhalf_array[i,j]
u[i,j] = C1*(2*u_1[i,j] + Cb*u_2[i,j]\
+ C2*(q_i_phalf*(u_1[i+1,j]-u_1[i,j])-\
q_i_mhalf*(u_1[i,j]-u_1[i-1,j]))\
+ C3*(q_j_phalf*(u_1[i,j+1]-u_1[i,j])-\
q_j_mhalf*(u_1[i,j]-u_1[i,j-1]))\
+ dtdt*f(x[i],y[j],t[n]))
elif version=="vectorized":
if self.physprm.gridprm.problem3D :
u[1:-1,1:-1,1:-1] = C1*(2*u_1[1:-1,1:-1,1:-1] \
+ Cb*u_2[1:-1,1:-1,1:-1]\
+ C2*(q_i_phalf_array[1:-1,1:-1,1:-1]*\
(u_1[2:,1:-1,1:-1]-u_1[1:-1,1:-1,1:-1])-\
q_i_mhalf_array[1:-1,1:-1,1:-1]*\
(u_1[1:-1,1:-1,1:-1]-u_1[:-2,1:-1,1:-1]))\
+ C3*(q_j_phalf_array[1:-1,1:-1,1:-1]*\
(u_1[1:-1,2:,1:-1]-u_1[1:-1,1:-1,1:-1])-\
q_j_mhalf_array[1:-1,1:-1,1:-1]*\
(u_1[1:-1,1:-1,1:-1]-u_1[1:-1,:-2,1:-1]))\
+ C4*(q_k_phalf_array[1:-1,1:-1,1:-1]*\
(u_1[1:-1,1:-1,2:]-u_1[1:-1,1:-1,1:-1])-\
q_k_mhalf_array[1:-1,1:-1,1:-1]*\
(u_1[1:-1,1:-1,1:-1]-u_1[1:-1,1:-1,:-2]))\
+ dtdt*f(x[1:-1],y[1:-1],z[1:-1],t[n]))
else :
u[1:-1,1:-1] = C1*(2*u_1[1:-1,1:-1] + Cb*u_2[1:-1,1:-1]\
+ C2*(q_i_phalf_array[1:-1,1:-1]*(u_1[2:,1:-1]-u_1[1:-1,1:-1])-\
q_i_mhalf_array[1:-1,1:-1]*(u_1[1:-1,1:-1]-u_1[:-2,1:-1]))\
+ C3*(q_j_phalf_array[1:-1,1:-1]*(u_1[1:-1,2:]-u_1[1:-1,1:-1])-\
q_j_mhalf_array[1:-1,1:-1]*(u_1[1:-1,1:-1]-u_1[1:-1,:-2]))\
+ dtdt*f(x[1:-1],y[1:-1],t[n]))
############ Update the ghost points for u ############
if not quiet_mode:
print "Updating ghost points"
if self.physprm.gridprm.problem3D :
BC(self.u,Ix,Iy,Iz,version)
else :
if self.outlet:
if version=="scalar":
for i in Ix:
for j in Iy:
term[i,j] = C5[i,j]*(u_2[i,j] - u[i,j])
elif version=="vectorized":
term[1:-1,1:-1] = C5[1:-1,1:-1]*(u_2[1:-1,1:-1] - u[1:-1,1:-1])
BC(self.u,Ix,Iy,version,term)
else:
BC(self.u,Ix,Iy,version)
############ Update all values for next time iteration ############
if self.physprm.gridprm.problem3D :
u_2[:,:,:], u_1[:,:,:] = u_1, u
else :
u_2[:,:], u_1[:,:] = u_1, u
n += 1 # Update time level
if save_figs :
if self.physprm.gridprm.problem3D :
plt.setp(interactive=False)
h = plt.isosurface(xx,yy,zz,u_2,-3)
h.setp(opacity=0.5)
plt.shading("interp")
plt.daspect([1,1,1])
plt.view(3)
plt.axis("tight")
plt.show()
else :
viz.store_fig(t[n],Ix,Iy,u_2,u_1,u)
elif display:
viz.display(t[n],Ix,Iy,u_2,u_1,u)
if save_figs:
viz.movie()
return u
x = y = np.linspace(-10., 10., 41)
xv, yv = plt.ndgrid(x, y) # Grid for x, y values (km)
hv = h0/(1+(xv**2+yv**2)/(R**2)) # Compute height (m)
x = y = np.linspace(-10.,10.,11)
x2v, y2v = plt.ndgrid(x, y) # Define a coarser grid for the vector field
h2v = h0/(1+(x2v**2+y2v**2)/(R**2)) # Compute height for new grid
dhdx, dhdy = np.gradient(h2v) # Compute the gradient vector (dh/dx,dh/dy)
# Draw contours and gradient field of h
plt.figure(9)
plt.quiver(x2v, y2v, dhdx, dhdy, 0, 'r')
plt.hold('on')
plt.contour(xv, yv, hv)
plt.axis('equal')
# end draw contours and gradient field of h
x = y = np.linspace(-5, 5, 11)
xv, yv = plt.ndgrid(x, y)
u = xv**2 + 2*yv - .5*xv*yv
v = -3*yv
# Draw 2D-field
plt.figure(10)
plt.quiver(xv, yv, u, v, 200, 'b')
plt.axis('equal')
# end draw 2D-field
