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_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_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()
