M = 5
N = array([2**M, 2**(M-1), 2**(M-1)])
L = array([2, 2*pi, 2*pi])
dx = (L / N).astype(float)
comm = MPI.COMM_WORLD
num_processes = comm.Get_size()
rank = comm.Get_rank()
Np = N / num_processes
# Get points and weights for Chebyshev weighted integrals
quad = "GC"
BC1 = array([1,0,0, 1,0,0])
BC2 = array([0,1,0, 0,1,0])
BC3 = array([0,1,0, 1,0,0])
SC = ChebyshevTransform(quad)
ST = ShenBasis(BC1, quad)
SN = ShenBasis(BC2, quad, Neumann = True)
SR = ShenBasis(BC3, quad)
SB = ShenBiharmonicBasis(quad, fast_transform=False)
points, weights = ST.points_and_weights(N[0])
pointsN, weightsN = SN.points_and_weights(N[0])
x1 = arange(N[1], dtype=float)*L[1]/N[1]
x2 = arange(N[2], dtype=float)*L[2]/N[2]
# Get grid for velocity points
X = array(meshgrid(points[rank*Np[0]:(rank+1)*Np[0]], x1, x2, indexing='ij'), dtype=float)
Y = array(meshgrid(pointsN[rank*Np[0]:(rank+1)*Np[0]], x1, x2, indexing='ij'), dtype=float)
Nf = N[2]/2+1 # Number of independent complex wavenumbers in z-direction
print " Menu "
print "==============================="
print "0 - Poisson Neumann BCs"
print "1 - Poisson Dirichlet BCs"
print "2 - Helmholtz Dirichlet BCs"
print "3 - Biharmonic Dirichlet BCs"
print "==============================="
test = raw_input("Your choice: ")
M = 2 ** 3
quad = "GL"
BC1 = array([1, 0, 0, 1, 0, 0])
BC2 = array([0, 1, 0, 0, 1, 0])
BC3 = array([0, 1, 0, 1, 0, 0])
SC = ChebyshevTransform(quad)
ST = ShenBasis(BC1, quad)
SN = ShenBasis(BC2, quad, Neumann=True)
SR = ShenBasis(BC3, quad)
SB = ShenBiharmonicBasis(quad, fast_transform=False)
if test == "0":
v_exact = zeros(M)
f = zeros(M)
points, weights = SN.points_and_weights(M)
f[:] = (-8.0 / 3.0) * points # -((pi/2.)**2)*sin(pi*points/2.)
v_exact = -(4.0 / 9.0) * (points ** 3 - 3.0 * points) # sin(pi*points/2.)
v = Poisson1DNeumann(M, f, quad, SC, SN, BC2)
print "Error: ", linalg.norm(v - v_exact, inf)
assert allclose(v, v_exact)
pl.plot(points, v_exact, points, v)