def test_curl2(): # Test projection of curl K0 = Basis(N[0], 'C', bc=(0, 0)) K1 = Basis(N[1], 'F', dtype='D') K2 = Basis(N[2], 'F', dtype='d') K3 = Basis(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 = VectorTensorProductSpace(T) TTk = MixedTensorProductSpace([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()