def test_project(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (25, 24)
funcs = {
(1, 0): (cos(4 * y) * sin(2 * np.pi * x)) * (1 - x**2),
(1, 1): (cos(4 * x) * sin(2 * np.pi * y)) * (1 - y**2),
(2, 0): (sin(6 * z) * cos(4 * y) * sin(2 * np.pi * x)) * (1 - x**2),
(2, 1): (sin(2 * z) * cos(4 * x) * sin(2 * np.pi * y)) * (1 - y**2),
(2, 2): (sin(2 * x) * cos(4 * y) * sin(2 * np.pi * z)) * (1 - z**2)
}
syms = {1: (x, y), 2: (x, y, z)}
xs = {0: x, 1: y, 2: z}
for shape in product(*([sizes] * dim)):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
if dim < 3:
n = min(shape)
if typecode in 'fdg':
n //= 2
n += 1
if n < comm.size:
continue
for axis in range(dim + 1):
ST0 = ST(shape[-1], quad=quad)
bases.insert(axis, ST0)
fft = TensorProductSpace(comm,
bases,
dtype=typecode,
axes=axes[dim][axis])
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
ul = lambdify(syms[dim], ue, 'numpy')
uq = ul(*X).astype(typecode)
uh = Function(fft)
uh = fft.forward(uq, uh)
due = ue.diff(xs[axis], 1)
dul = lambdify(syms[dim], due, 'numpy')
duq = dul(*X).astype(typecode)
uf = project(Dx(uh, axis, 1), fft)
uy = Array(fft)
uy = fft.backward(uf, uy)
assert np.allclose(uy, duq, 0, 1e-6)
for ax in (x for x in range(dim + 1) if x is not axis):
due = ue.diff(xs[axis], 1, xs[ax], 1)
dul = lambdify(syms[dim], due, 'numpy')
duq = dul(*X).astype(typecode)
uf = project(Dx(Dx(uh, axis, 1), ax, 1), fft)
uy = Array(fft)
uy = fft.backward(uf, uy)
assert np.allclose(uy, duq, 0, 1e-6)
bases.pop(axis)
fft.destroy()
def test_project2(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (18, 17)
funcx = ((2*np.pi**2*(x**2 - 1) - 1)* cos(2*np.pi*x) - 2*np.pi*x*sin(2*np.pi*x))/(4*np.pi**3)
funcy = ((2*np.pi**2*(y**2 - 1) - 1)* cos(2*np.pi*y) - 2*np.pi*y*sin(2*np.pi*y))/(4*np.pi**3)
funcz = ((2*np.pi**2*(z**2 - 1) - 1)* cos(2*np.pi*z) - 2*np.pi*z*sin(2*np.pi*z))/(4*np.pi**3)
funcs = {
(1, 0): cos(4*y)*funcx,
(1, 1): cos(4*x)*funcy,
(2, 0): sin(3*z)*cos(4*y)*funcx,
(2, 1): sin(2*z)*cos(4*x)*funcy,
(2, 2): sin(2*x)*cos(4*y)*funcz
}
syms = {1: (x, y), 2:(x, y, z)}
xs = {0:x, 1:y, 2:z}
for shape in product(*([sizes]*dim)):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
for axis in range(dim+1):
ST0 = ST(shape[-1], quad=quad)
bases.insert(axis, ST0)
# Spectral space must be aligned in nonperiodic direction, hence axes
fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
ul = lambdify(syms[dim], ue, 'numpy')
uq = ul(*X).astype(typecode)
uh = Function(fft)
uh = fft.forward(uq, uh)
due = ue.diff(xs[axis], 1)
dul = lambdify(syms[dim], due, 'numpy')
duq = dul(*X).astype(typecode)
uf = project(Dx(uh, axis, 1), fft)
uy = Array(fft)
uy = fft.backward(uf, uy)
assert np.allclose(uy, duq, 0, 1e-3)
# Test also several derivatives
for ax in (x for x in range(dim+1) if x is not axis):
due = ue.diff(xs[ax], 1, xs[axis], 1)
dul = lambdify(syms[dim], due, 'numpy')
duq = dul(*X).astype(typecode)
uf = project(Dx(Dx(uh, ax, 1), axis, 1), fft)
uy = Array(fft)
uy = fft.backward(uf, uy)
assert np.allclose(uy, duq, 0, 1e-3)
bases.pop(axis)
fft.destroy()
def test_curl2():
# Test projection of curl
K0 = FunctionSpace(N[0], 'C', bc=(0, 0))
K1 = FunctionSpace(N[1], 'F', dtype='D')
K2 = FunctionSpace(N[2], 'F', dtype='d')
K3 = FunctionSpace(N[0], 'C')
T = TensorProductSpace(comm, (K0, K1, K2))
TT = TensorProductSpace(comm, (K3, K1, K2))
X = T.local_mesh(True)
K = T.local_wavenumbers(False)
Tk = VectorSpace(T)
TTk = VectorSpace([T, T, TT])
U = Array(Tk)
U_hat = Function(Tk)
curl_hat = Function(TTk)
curl_ = Array(TTk)
# Initialize a Taylor Green vortex
U[0] = np.sin(X[0]) * np.cos(X[1]) * np.cos(X[2]) * (1 - X[0]**2)
U[1] = -np.cos(X[0]) * np.sin(X[1]) * np.cos(X[2]) * (1 - X[0]**2)
U[2] = 0
U_hat = Tk.forward(U, U_hat)
Uc = U_hat.copy()
U = Tk.backward(U_hat, U)
U_hat = Tk.forward(U, U_hat)
assert allclose(U_hat, Uc)
# Compute curl first by computing each term individually
curl_hat[0] = 1j * (K[1] * U_hat[2] - K[2] * U_hat[1])
curl_[0] = T.backward(
curl_hat[0], curl_[0]) # No x-derivatives, still in Dirichlet space
dwdx_hat = project(Dx(U_hat[2], 0, 1), TT) # Need to use space without bc
dvdx_hat = project(Dx(U_hat[1], 0, 1), TT) # Need to use space without bc
dwdx = Array(TT)
dvdx = Array(TT)
dwdx = TT.backward(dwdx_hat, dwdx)
dvdx = TT.backward(dvdx_hat, dvdx)
curl_hat[1] = 1j * K[2] * U_hat[0]
curl_hat[2] = -1j * K[1] * U_hat[0]
curl_[1] = T.backward(curl_hat[1], curl_[1])
curl_[2] = T.backward(curl_hat[2], curl_[2])
curl_[1] -= dwdx
curl_[2] += dvdx
# Now do it with project
w_hat = project(curl(U_hat), TTk)
w = Array(TTk)
w = TTk.backward(w_hat, w)
assert allclose(w, curl_)
def test_project_2dirichlet(quad):
x, y = symbols("x,y")
ue = (cos(4 * y) * sin(2 * x)) * (1 - x**2) * (1 - y**2)
sizes = (25, 24)
D0 = lbases.ShenDirichletBasis(sizes[0], quad=quad)
D1 = lbases.ShenDirichletBasis(sizes[1], quad=quad)
B0 = lbases.Basis(sizes[0], quad=quad)
B1 = lbases.Basis(sizes[1], quad=quad)
DD = TensorProductSpace(comm, (D0, D1))
BD = TensorProductSpace(comm, (B0, D1))
DB = TensorProductSpace(comm, (D0, B1))
BB = TensorProductSpace(comm, (B0, B1))
X = DD.local_mesh(True)
ul = lambdify((x, y), ue, 'numpy')
uq = Array(DD)
uq[:] = ul(*X)
uh = Function(DD)
uh = DD.forward(uq, uh)
dudx_hat = project(Dx(uh, 0, 1), BD)
dudx = Array(BD)
dudx = BD.backward(dudx_hat, dudx)
duedx = ue.diff(x, 1)
duxl = lambdify((x, y), duedx, 'numpy')
dx = duxl(*X)
assert np.allclose(dx, dudx)
dudy_hat = project(Dx(uh, 1, 1), DB)
dudy = Array(DB)
dudy = DB.backward(dudy_hat, dudy)
duedy = ue.diff(y, 1)
duyl = lambdify((x, y), duedy, 'numpy')
dy = duyl(*X)
assert np.allclose(dy, dudy)
us_hat = Function(BB)
us_hat = project(uq, BB, output_array=us_hat)
us = Array(BB)
us = BB.backward(us_hat, us)
assert np.allclose(us, uq)
dudxy_hat = project(Dx(us_hat, 0, 1) + Dx(us_hat, 1, 1), BB)
dudxy = Array(BB)
dudxy = BB.backward(dudxy_hat, dudxy)
duedxy = ue.diff(x, 1) + ue.diff(y, 1)
duxyl = lambdify((x, y), duedxy, 'numpy')
dxy = duxyl(*X)
assert np.allclose(dxy, dudxy)
def test_shentransform(typecode, dim, ST, quad):
for shape in product(*([sizes] * dim)):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
if dim < 3:
n = min(shape)
if typecode in 'fdg':
n //= 2
n += 1
if n < comm.size:
continue
for axis in range(dim + 1):
ST0 = ST(shape[-1], quad=quad)
bases.insert(axis, ST0)
fft = TensorProductSpace(comm, bases, dtype=typecode)
U = random_like(fft.forward.input_array)
F = fft.forward(U)
Fc = F.copy()
V = fft.backward(F)
F = fft.forward(U)
assert allclose(F, Fc)
bases.pop(axis)
fft.destroy()
def test_curl(typecode):
K0 = Basis(N[0], 'F', dtype=typecode.upper())
K1 = Basis(N[1], 'F', dtype=typecode.upper())
K2 = Basis(N[2], 'F', dtype=typecode)
T = TensorProductSpace(comm, (K0, K1, K2), dtype=typecode)
X = T.local_mesh(True)
K = T.local_wavenumbers()
Tk = VectorTensorProductSpace(T)
u = TrialFunction(Tk)
v = TestFunction(Tk)
U = Array(Tk)
U_hat = Function(Tk)
curl_hat = Function(Tk)
curl_ = Array(Tk)
# Initialize a Taylor Green vortex
U[0] = np.sin(X[0]) * np.cos(X[1]) * np.cos(X[2])
U[1] = -np.cos(X[0]) * np.sin(X[1]) * np.cos(X[2])
U[2] = 0
U_hat = Tk.forward(U, U_hat)
Uc = U_hat.copy()
U = Tk.backward(U_hat, U)
U_hat = Tk.forward(U, U_hat)
assert allclose(U_hat, Uc)
divu_hat = project(div(U_hat), T)
divu = Array(T)
divu = T.backward(divu_hat, divu)
assert allclose(divu, 0)
curl_hat[0] = 1j * (K[1] * U_hat[2] - K[2] * U_hat[1])
curl_hat[1] = 1j * (K[2] * U_hat[0] - K[0] * U_hat[2])
curl_hat[2] = 1j * (K[0] * U_hat[1] - K[1] * U_hat[0])
curl_ = Tk.backward(curl_hat, curl_)
w_hat = Function(Tk)
w_hat = inner(v, curl(U_hat), output_array=w_hat)
A = inner(v, u)
for i in range(3):
w_hat[i] = A[i].solve(w_hat[i])
w = Array(Tk)
w = Tk.backward(w_hat, w)
#from IPython import embed; embed()
assert allclose(w, curl_)
u_hat = Function(Tk)
u_hat = inner(v, U, output_array=u_hat)
for i in range(3):
u_hat[i] = A[i].solve(u_hat[i])
uu = Array(Tk)
uu = Tk.backward(u_hat, uu)
assert allclose(u_hat, U_hat)
def test_shentransform(typecode, dim, ST, quad):
for shape in product(*([sizes]*dim)):
bases = []
for n in shape[:-1]:
bases.append(FunctionSpace(n, 'F', dtype=typecode.upper()))
bases.append(FunctionSpace(shape[-1], 'F', dtype=typecode))
for axis in range(dim+1):
ST0 = ST(shape[-1], quad=quad)
bases.insert(axis, ST0)
fft = TensorProductSpace(comm, bases, dtype=typecode)
U = random_like(fft.forward.input_array)
F = fft.forward(U)
Fc = F.copy()
V = fft.backward(F)
F = fft.forward(U)
assert allclose(F, Fc)
bases.pop(axis)
fft.destroy()
def test_transform(typecode, dim):
s = (True, )
if comm.Get_size() > 2 and dim > 2:
s = (True, False)
for slab in s:
for shape in product(*([sizes] * dim)):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
fft = TensorProductSpace(comm, bases, dtype=typecode, slab=slab)
if comm.rank == 0:
grid = [c.size for c in fft.subcomm]
print('grid:{} shape:{} typecode:{}'.format(
grid, shape, typecode))
U = random_like(fft.forward.input_array)
F = fft.forward(U)
V = fft.backward(F)
assert allclose(V, U)
# Alternative method
fft.forward.input_array[...] = U
fft.forward(fast_transform=False)
fft.backward(fast_transform=False)
V = fft.backward.output_array
assert allclose(V, U)
TT = VectorTensorProductSpace(fft)
U = Array(TT)
V = Array(TT)
F = Function(TT)
U[:] = random_like(U)
F = TT.forward(U, F)
V = TT.backward(F, V)
assert allclose(V, U)
TM = MixedTensorProductSpace([fft, fft])
U = Array(TM)
V = Array(TM)
F = Function(TM)
U[:] = random_like(U)
F = TM.forward(U, F)
V = TM.backward(F, V)
assert allclose(V, U)
fft.destroy()
padding = 1.5
bases = []
for n in shape[:-1]:
bases.append(
Basis(n,
'F',
dtype=typecode.upper(),
padding_factor=padding))
bases.append(
Basis(shape[-1], 'F', dtype=typecode, padding_factor=padding))
fft = TensorProductSpace(comm, bases, dtype=typecode)
if comm.rank == 0:
grid = [c.size for c in fft.subcomm]
print('grid:{} shape:{} typecode:{}'.format(
grid, shape, typecode))
U = random_like(fft.forward.input_array)
F = fft.forward(U)
Fc = F.copy()
V = fft.backward(F)
F = fft.forward(V)
assert allclose(F, Fc)
# Alternative method
fft.backward.input_array[...] = F
fft.backward()
fft.forward()
V = fft.forward.output_array
assert allclose(F, V)
fft.destroy()
if family == 'legendre':
f_hat *= -1.
# Get left hand side of Poisson equation
if family == 'chebyshev':
matrices = inner(v, div(grad(u)))
else:
matrices = inner(grad(v), grad(u))
# Create Helmholtz linear algebra solver
H = Solver(**matrices)
# Solve and transform to real space
u_hat = Function(T) # Solution spectral space
u_hat = H(u_hat, f_hat) # Solve
uq = T.backward(u_hat)
# Compare with analytical solution
uj = ul(*X)
print(abs(uj-uq).max())
assert np.allclose(uj, uq)
if plt is not None and not 'pytest' in os.environ:
plt.figure()
plt.contourf(X[0], X[1], uq)
plt.colorbar()
plt.figure()
plt.contourf(X[0], X[1], uj)
plt.colorbar()
) # With broadcasting=True the shape of X is local_shape, even though the number of datapoints are still the same as in 1D
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
# Solve Poisson equation
A = inner(v, div(grad(u)))
f_hat = A.solve(f_hat)
uq = T.backward(f_hat)
uj = Array(T, buffer=ue)
print(np.sqrt(dx((uj - uq)**2)))
assert np.allclose(uj, uq)
print(f_hat.commsizes, fj.commsizes)
if 'pytest' not in os.environ and comm.Get_size() == 1:
import matplotlib.pyplot as plt
plt.figure()
plt.contourf(X[0][:, :, 0, 0], X[1][:, :, 0, 0], uq[:, :, 0, 0])
plt.colorbar()
plt.figure()
plt.contourf(X[0][:, :, 0, 0], X[1][:, :, 0, 0], uj[:, :, 0, 0])
# Get f on quad points fj = Array(T, buffer=fl(*X)) # Compute right hand side f_hat = Function(T) f_hat = inner(v, fj, output_array=f_hat) # Solve Poisson equation u_hat = Function(T) #A = inner(v, div(grad(u))) A = inner(grad(v), grad(u)) u_hat = A.solve(-f_hat, u_hat) uq = Array(T) uq = T.backward(u_hat, uq, fast_transform=True) uj = ul(*X) assert np.allclose(uj, uq) #from shenfun.tensorproductspace import Convolve #S0 = Basis(N[0], family='F', dtype='D', padding_factor=2.0) #S1 = Basis(N[1], family='F', dtype='d', padding_factor=2.0) #Tp = TensorProductSpace(comm, (S0, S1), axes=(0, 1)) #C0 = Convolve(Tp) #ff_hat = C0(f_hat, f_hat) # Test eval at point point = np.array([[0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.1, 0.3]]) p = T.eval(point, u_hat) assert np.allclose(p, ul(*point.T))
) # With broadcasting=True the shape of X is local_shape, even though the number of datapoints are still the same as in 1D
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fl(*X))
# Compute right hand side
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
# Solve Poisson equation
A = inner(v, div(grad(u)))
f_hat = A.solve(f_hat)
uq = T.backward(f_hat, fast_transform=True)
uj = ul(*X)
print(abs(uj - uq).max())
assert np.allclose(uj, uq)
# Test eval at point
point = np.array([[0.1, 0.2, 0.3], [0.2, 0.3, 0.3], [0.3, 0.4, 0.1]])
p = T.eval(point, f_hat)
assert np.allclose(p, ul(*point.T))
p2 = f_hat.eval(point)
assert np.allclose(p2, ul(*point.T))
if plt is not None and not 'pytest' in os.environ:
plt.figure()
plt.contourf(X[0][:, :, 0], X[1][:, :, 0], uq[:, :, 0])
# Get left hand side of Poisson equation
if family == 'legendre':
matrices = inner(grad(v), grad(u))
else:
matrices = inner(v, -div(grad(u)))
matrices += inner(v, a*u)
# Create Helmholtz linear algebra solver
H = Solver(matrices)
# Solve and transform to real space
u_hat = Function(T) # Solution spectral space
u_hat = H(f_hat, u_hat) # Solve
uq = Array(T)
uq = T.backward(u_hat, uq)
# Compare with analytical solution
uj = ul(*X)
print(abs(uj-uq).max())
assert np.allclose(uj, uq)
if 'pytest' not in os.environ:
import matplotlib.pyplot as plt
plt.figure()
plt.contourf(X[0], X[1], uq)
plt.colorbar()
plt.figure()
plt.contourf(X[0], X[1], uj)
f_hat *= -1.
# Get left hand side of Poisson equation
if family == 'chebyshev':
matrices = inner(v, div(grad(u)))
else:
matrices = inner(grad(v), grad(u))
# Create Helmholtz linear algebra solver
H = Solver(**matrices)
# Solve and transform to real space
u_hat = Function(T) # Solution spectral space
t0 = time.time()
u_hat = H(u_hat, f_hat) # Solve
uq = T.backward(u_hat, fast_transform=False)
# Compare with analytical solution
uj = ul(*X)
error = comm.reduce(np.linalg.norm(uj - uq)**2)
if comm.Get_rank() == 0 and regtest == True:
print("Error=%2.16e" % (np.sqrt(error)))
#assert np.allclose(uj, uq)
if plt is not None and not 'pytest' in os.environ:
plt.figure()
plt.contourf(X[2][0, 0, :], X[0][:, 0, 0], uq[:, 2, :])
plt.colorbar()
plt.figure()
plt.contourf(X[2][0, 0, :], X[0][:, 0, 0], uj[:, 2, :])
comm = MPI.COMM_WORLD N = (32, 33, 34) K0 = ShenBiharmonicBasis(N[0]) K1 = C2CBasis(N[1]) K2 = R2CBasis(N[2]) W = TensorProductSpace(comm, (K0, K1, K2)) u = TrialFunction(W) v = TestFunction(W) matrices = inner(v, div(grad(div(grad(u))))) fj = Function(W, False) fj[:] = np.random.random(fj.shape) f_hat = inner(v, fj) # Some right hand side B = Solver(**matrices) # Solve and transform to real space u_hat = Function(W) # Solution spectral space u_hat = B(u_hat, f_hat) # Solve u = Function(W, False) u = W.backward(u_hat, u) # compute dudx of the solution K0 = Basis(N[0]) W0 = TensorProductSpace(comm, (K0, K1, K2)) du_hat = project(Dx(u, 0, 1), W0, uh_hat=u_hat) du = Function(W0, False) du = W0.backward(du_hat, du)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side of Poisson equation
f_hat = inner(v, fj)
# Get left hand side of Poisson equation
matrices = inner(v, u - div(grad(u)))
# Create Helmholtz linear algebra solver
H = Solver(*matrices)
# Solve and transform to real space
u_hat = Function(T).set_boundary_dofs() # Solution spectral space
u_hat = H(u_hat, f_hat) # Solve
uq = T.backward(u_hat).copy()
# Compare with analytical solution
uj = Array(T, buffer=ue)
print(abs(uj - uq).max())
assert np.allclose(uj, uq)
c = H.matvec(u_hat, Function(T))
f_hat = inner(v, fj)
assert np.allclose(c, f_hat)
if 'pytest' not in os.environ:
import matplotlib.pyplot as plt
plt.figure()
X = T.local_mesh(
True
) # With broadcasting=True the shape of X is local_shape, even though the number of datapoints are still the same as in 1D
