def test_refine():
assert comm.Get_size() < 7
N = (8, 9, 10)
F0 = Basis(8, 'F', dtype='D')
F1 = Basis(9, 'F', dtype='D')
F2 = Basis(10, 'F', dtype='d')
T = TensorProductSpace(comm, (F0, F1, F2), slab=True, collapse_fourier=True)
u_hat = Function(T)
u = Array(T)
u[:] = np.random.random(u.shape)
u_hat = u.forward(u_hat)
Tp = T.get_dealiased(padding_factor=(2, 2, 2))
u_ = Array(Tp)
up_hat = Function(Tp)
assert up_hat.commsizes == u_hat.commsizes
u2 = u_hat.refine(2*np.array(N))
V = VectorTensorProductSpace(T)
u_hat = Function(V)
u = Array(V)
u[:] = np.random.random(u.shape)
u_hat = u.forward(u_hat)
Vp = V.get_dealiased(padding_factor=(2, 2, 2))
u_ = Array(Vp)
up_hat = Function(Vp)
assert up_hat.commsizes == u_hat.commsizes
u3 = u_hat.refine(2*np.array(N))
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_project_hermite(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (20, 19)
funcs = {
(1, 0): (cos(4*y)*sin(2*x))*exp(-x**2/2),
(1, 1): (cos(4*x)*sin(2*y))*exp(-y**2/2),
(2, 0): (sin(3*z)*cos(4*y)*sin(2*x))*exp(-x**2/2),
(2, 1): (sin(2*z)*cos(4*x)*sin(2*y))*exp(-y**2/2),
(2, 2): (sin(2*x)*cos(4*y)*sin(2*z))*exp(-z**2/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))
for axis in range(dim+1):
ST0 = ST(3*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)]
u_h = project(ue, fft)
ul = lambdify(syms[dim], ue, 'numpy')
uq = ul(*X).astype(typecode)
uf = u_h.backward()
assert np.linalg.norm(uq-uf) < 1e-5
bases.pop(axis)
fft.destroy()
def test_assign(fam):
x, y = symbols("x,y")
for bc in (None, 'Dirichlet', 'Biharmonic'):
dtype = 'D' if fam == 'F' else 'd'
bc = 'periodic' if fam == 'F' else bc
if bc == 'Biharmonic' and fam in ('La', 'H'):
continue
tol = 1e-12 if fam in ('C', 'L', 'F') else 1e-5
N = (10, 12)
B0 = Basis(N[0], fam, dtype=dtype, bc=bc)
B1 = Basis(N[1], fam, dtype=dtype, bc=bc)
u_hat = Function(B0)
u_hat[1:4] = 1
ub_hat = Function(B1)
u_hat.assign(ub_hat)
assert abs(inner(1, u_hat)-inner(1, ub_hat)) < tol
T = TensorProductSpace(comm, (B0, B1))
u_hat = Function(T)
u_hat[1:4, 1:4] = 1
Tp = T.get_refined((2*N[0], 2*N[1]))
ub_hat = Function(Tp)
u_hat.assign(ub_hat)
assert abs(inner(1, u_hat)-inner(1, ub_hat)) < tol
VT = VectorTensorProductSpace(T)
u_hat = Function(VT)
u_hat[:, 1:4, 1:4] = 1
Tp = T.get_refined((2*N[0], 2*N[1]))
VTp = VectorTensorProductSpace(Tp)
ub_hat = Function(VTp)
u_hat.assign(ub_hat)
assert abs(inner((1, 1), u_hat)-inner((1, 1), ub_hat)) < tol
def test_eval_expression():
import sympy as sp
from shenfun import div, grad
x, y, z = sp.symbols('x,y,z')
B0 = Basis(16, 'C')
B1 = Basis(17, 'C')
B2 = Basis(20, 'F', dtype='d')
TB = TensorProductSpace(comm, (B0, B1, B2))
f = sp.sin(x)+sp.sin(y)+sp.sin(z)
dfx = f.diff(x, 2) + f.diff(y, 2) + f.diff(z, 2)
fa = Array(TB, buffer=f).forward()
dfe = div(grad(fa))
dfa = project(dfe, TB)
xyz = np.array([[0.25, 0.5, 0.75],
[0.25, 0.5, 0.75],
[0.25, 0.5, 0.75]])
f0 = lambdify((x, y, z), dfx)(*xyz)
f1 = dfe.eval(xyz)
f2 = dfa.eval(xyz)
assert np.allclose(f0, f1, 1e-7)
assert np.allclose(f1, f2, 1e-7)
def test_project_hermite(typecode, dim):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (20, 19)
funcs = {
(1, 0): (cos(4*y)*sin(2*x))*exp(-x**2/2),
(1, 1): (cos(4*x)*sin(2*y))*exp(-y**2/2),
(2, 0): (sin(3*z)*cos(4*y)*sin(2*x))*exp(-x**2/2),
(2, 1): (sin(2*z)*cos(4*x)*sin(2*y))*exp(-y**2/2),
(2, 2): (sin(2*x)*cos(4*y)*sin(2*z))*exp(-z**2/2)
}
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 = hBasis[0](3*shape[-1])
bases.insert(axis, ST0)
fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
due = ue.diff(xs[0], 1)
u_h = project(ue, fft)
du_h = project(due, fft)
du2 = project(Dx(u_h, 0, 1), fft)
assert np.linalg.norm(du_h-du2) < 1e-5
bases.pop(axis)
fft.destroy()
def test_biharmonic2D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (16, 16)
SD = Basis(N[axis], family=family, bc='Biharmonic')
K1 = Basis(N[(axis + 1) % 2], family='F', dtype='d')
subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, allaxes2D[axis])
bases = [K1]
bases.insert(axis, SD)
T = TensorProductSpace(subcomms, bases, axes=allaxes2D[axis])
u = TrialFunction(T)
v = TestFunction(T)
if family == 'chebyshev':
mat = inner(v, div(grad(div(grad(u)))))
else:
mat = inner(div(grad(v)), div(grad(u)))
H = la.Biharmonic(*mat)
u = Function(T)
u[:] = np.random.random(u.shape) + 1j * np.random.random(u.shape)
f = Function(T)
f = H.matvec(u, f)
g0 = Function(T)
g1 = Function(T)
g2 = Function(T)
M = {d.get_key(): d for d in mat}
g0 = M['SBBmat'].matvec(u, g0)
g1 = M['ABBmat'].matvec(u, g1)
g2 = M['BBBmat'].matvec(u, g2)
assert np.linalg.norm(f - (g0 + g1 + g2)) < 1e-8
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_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_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 = 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 interp(self, y, eigvals, eigvectors, eigval=1, verbose=False):
"""Interpolate solution eigenvector and it's derivative onto y
Parameters
----------
y : array
Interpolation points
eigvals : array
All computed eigenvalues
eigvectors : array
All computed eigenvectors
eigval : int, optional
The chosen eigenvalue, ranked with descending imaginary
part. The largest imaginary part is 1, the second
largest is 2, etc.
verbose : bool, optional
Print information or not
"""
nx, eigval = self.get_eigval(eigval, eigvals, verbose)
SB = Basis(self.N, 'C', bc='Biharmonic', quad=self.quad, dtype='D')
phi_hat = Function(SB)
phi_hat[:-4] = np.squeeze(eigvectors[:, nx])
phi = phi_hat.eval(y)
dphidy = Dx(phi_hat, 0, 1).eval(y)
return eigval, phi, dphidy
def assemble(self):
N = self.N
SB = Basis(N, 'C', bc='Biharmonic', quad=self.quad)
SB.plan((N, N), 0, np.float, {})
# (u'', v)
K = inner_product((SB, 0), (SB, 2))
# ((1-x**2)u, v)
x = sp.symbols('x', real=True)
K1 = inner_product((SB, 0), (SB, 0), measure=(1-x**2))
# ((1-x**2)u'', v)
K2 = inner_product((SB, 0), (SB, 2), measure=(1-x**2))
# (u'''', v)
Q = inner_product((SB, 0), (SB, 4))
# (u, v)
M = inner_product((SB, 0), (SB, 0))
Re = self.Re
a = self.alfa
B = -Re*a*1j*(K-a**2*M)
A = Q-2*a**2*K+a**4*M - 2*a*Re*1j*M - 1j*a*Re*(K2-a**2*K1)
return A.diags().toarray(), B.diags().toarray()
def interp(self, y, eigvals, eigvectors, eigval=1, verbose=False):
"""Interpolate solution eigenvector and it's derivative onto y
Parameters
----------
y : array
Interpolation points
eigvals : array
All computed eigenvalues
eigvectors : array
All computed eigenvectors
eigval : int, optional
The chosen eigenvalue, ranked with descending imaginary
part. The largest imaginary part is 1, the second
largest is 2, etc.
verbose : bool, optional
Print information or not
"""
N = self.N
nx, eigval = self.get_eigval(eigval, eigvals, verbose)
phi_hat = np.zeros(N, np.complex)
phi_hat[:-4] = np.squeeze(eigvectors[:, nx])
if not len(self.P4) == len(y):
SB = Basis(N, 'C', bc='Biharmonic', quad=self.quad)
self.P4 = SB.evaluate_basis_all(x=y)
self.T4x = SB.evaluate_basis_derivative_all(x=y, k=1)
phi = np.dot(self.P4, phi_hat)
dphidy = np.dot(self.T4x, phi_hat)
return eigval, phi, dphidy
def test_solve(quad):
SD = Basis(N, 'C', bc=(0, 0), quad=quad, plan=True)
u = TrialFunction(SD)
v = TestFunction(SD)
A = inner(div(grad(u)), v)
b = np.ones(N)
u_hat = Function(SD)
u_hat = A.solve(b, u=u_hat)
w_hat = Function(SD)
B = SparseMatrix(dict(A), (N - 2, N - 2))
w_hat[:-2] = B.solve(b[:-2], w_hat[:-2])
assert np.all(abs(w_hat[:-2] - u_hat[:-2]) < 1e-8)
ww = w_hat[:-2].repeat(N - 2).reshape((N - 2, N - 2))
bb = b[:-2].repeat(N - 2).reshape((N - 2, N - 2))
ww = B.solve(bb, ww, axis=0)
assert np.all(
abs(ww - u_hat[:-2].repeat(N - 2).reshape((N - 2, N - 2))) < 1e-8)
bb = bb.transpose()
ww = B.solve(bb, ww, axis=1)
assert np.all(
abs(ww - u_hat[:-2].repeat(N - 2).reshape(
(N - 2, N - 2)).transpose()) < 1e-8)
def test_eval_fourier(typecode, dim):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (20, 19)
funcs = {
2: cos(4 * x) + sin(6 * y),
3: sin(6 * x) + cos(4 * y) + sin(8 * z)
}
syms = {2: (x, y), 3: (x, y, z)}
points = None
if comm.Get_rank() == 0:
points = np.random.random((dim, 3))
points = comm.bcast(points)
t_0 = 0
t_1 = 0
t_2 = 0
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)
X = fft.local_mesh(True)
ue = funcs[dim]
ul = lambdify(syms[dim], ue, 'numpy')
uu = ul(*X).astype(typecode)
uq = ul(*points).astype(typecode)
u_hat = Function(fft)
u_hat = fft.forward(uu, u_hat)
t0 = time()
result = fft.eval(points, u_hat, method=0)
t_0 += time() - t0
assert allclose(uq, result), print(uq, result)
t0 = time()
result = fft.eval(points, u_hat, method=1)
t_1 += time() - t0
assert allclose(uq, result)
t0 = time()
result = fft.eval(points, u_hat, method=2)
t_2 += time() - t0
assert allclose(uq, result), print(uq, result)
print('method=0', t_0)
print('method=1', t_1)
print('method=2', t_2)
def test_project(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (20, 19)
funcs = {
(1, 0): (cos(1*y)*sin(1*np.pi*x))*(1-x**2),
(1, 1): (cos(1*x)*sin(1*np.pi*y))*(1-y**2),
(2, 0): (sin(1*z)*cos(1*y)*sin(1*np.pi*x))*(1-x**2),
(2, 1): (sin(1*z)*cos(1*x)*sin(1*np.pi*y))*(1-y**2),
(2, 2): (sin(1*x)*cos(1*y)*sin(1*np.pi*z))*(1-z**2)
}
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)
fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
dfft = fft.get_orthogonal()
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
uq = Array(fft, buffer=ue)
uh = Function(fft)
uh = fft.forward(uq, uh)
due = ue.diff(xs[axis], 1)
duq = Array(fft, buffer=due)
uf = project(Dx(uh, axis, 1), dfft).backward()
assert np.linalg.norm(uf-duq) < 1e-5
for ax in (x for x in range(dim+1) if x is not axis):
due = ue.diff(xs[axis], 1, xs[ax], 1)
duq = Array(fft, buffer=due)
uf = project(Dx(Dx(uh, axis, 1), ax, 1), dfft).backward()
assert np.linalg.norm(uf-duq) < 1e-5
bases.pop(axis)
fft.destroy()
dfft.destroy()
def test_helmholtz3D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9, 10)
SD = Basis(N[allaxes3D[axis][0]], family=family, bc=(0, 0))
K1 = Basis(N[allaxes3D[axis][1]], family='F', dtype='D')
K2 = Basis(N[allaxes3D[axis][2]], family='F', dtype='d')
subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, [0, 1, 1])
bases = [0] * 3
bases[allaxes3D[axis][0]] = SD
bases[allaxes3D[axis][1]] = K1
bases[allaxes3D[axis][2]] = K2
T = TensorProductSpace(subcomms, bases, axes=allaxes3D[axis])
u = TrialFunction(T)
v = TestFunction(T)
if family == 'chebyshev':
mat = inner(v, div(grad(u)))
else:
mat = inner(grad(v), grad(u))
H = la.Helmholtz(*mat)
u = Function(T)
s = SD.sl[SD.slice()]
u[s] = np.random.random(u[s].shape) + 1j * np.random.random(u[s].shape)
f = Function(T)
f = H.matvec(u, f)
g0 = Function(T)
g1 = Function(T)
M = {d.get_key(): d for d in mat}
g0 = M['ADDmat'].matvec(u, g0)
g1 = M['BDDmat'].matvec(u, g1)
assert np.linalg.norm(f - (g0 + g1)) < 1e-12, np.linalg.norm(f - (g0 + g1))
uc = Function(T)
uc = H(uc, f)
assert np.linalg.norm(uc - u) < 1e-12
def test_inner(f0, f1):
if f0 == 'F' and f1 == 'F':
B0 = Basis(8, f0, dtype='D', domain=(-2*np.pi, 2*np.pi))
else:
B0 = Basis(8, f0)
c = Array(B0, val=1)
d = inner(1, c)
assert abs(d-(B0.domain[1]-B0.domain[0])) < 1e-7
B1 = Basis(8, f1)
T = TensorProductSpace(comm, (B0, B1))
a0 = Array(T, val=1)
c0 = inner(1, a0)
L = np.array([b.domain[1]-b.domain[0] for b in (B0, B1)])
assert abs(c0-np.prod(L)) < 1e-7
if not (f0 == 'F' or f1 == 'F'):
B2 = Basis(8, f1, domain=(-2, 2))
T = TensorProductSpace(comm, (B0, B1, B2))
a0 = Array(T, val=1)
c0 = inner(1, a0)
L = np.array([b.domain[1]-b.domain[0] for b in (B0, B1, B2)])
assert abs(c0-np.prod(L)) < 1e-7
def test_mul2():
mat = SparseMatrix({0: 1}, (3, 3))
v = np.ones(3)
c = mat * v
assert np.allclose(c, 1)
mat = SparseMatrix({-2: 1, -1: 1, 0: 1, 1: 1, 2: 1}, (3, 3))
c = mat * v
assert np.allclose(c, 3)
SD = Basis(8, "L", bc=(0, 0), scaled=True)
u = TrialFunction(SD)
v = TestFunction(SD)
mat = inner(grad(u), grad(v))
z = Function(SD, val=1)
c = mat * z
assert np.allclose(c[:6], 1)
def assemble(self):
N = self.N
SB = Basis(N, 'C', bc='Biharmonic', quad=self.quad)
SB.plan((N, N), 0, np.float, {})
x, _ = self.x, self.w = SB.points_and_weights(N)
# Trial function
P4 = SB.evaluate_basis_all(x=x)
# Second derivatives
T2x = SB.evaluate_basis_derivative_all(x=x, k=2)
# (u'', v)
K = np.zeros((N, N))
K[:-4, :-4] = inner_product((SB, 0), (SB, 2)).diags().toarray()
# ((1-x**2)u, v)
xx = np.broadcast_to((1 - x**2)[:, np.newaxis], (N, N))
#K1 = np.dot(w*P4.T, xx*P4) # Alternative: K1 = np.dot(w*P4.T, ((1-x**2)*P4.T).T)
K1 = np.zeros((N, N))
K1 = SB.scalar_product(xx * P4, K1)
K1 = extract_diagonal_matrix(
K1).diags().toarray() # For improved roundoff
# ((1-x**2)u'', v)
K2 = np.zeros((N, N))
K2 = SB.scalar_product(xx * T2x, K2)
K2 = extract_diagonal_matrix(
K2).diags().toarray() # For improved roundoff
# (u'''', v)
Q = np.zeros((self.N, self.N))
Q[:-4, :-4] = inner_product((SB, 0), (SB, 4)).diags().toarray()
# (u, v)
M = np.zeros((self.N, self.N))
M[:-4, :-4] = inner_product((SB, 0), (SB, 0)).diags().toarray()
Re = self.Re
a = self.alfa
B = -Re * a * 1j * (K - a**2 * M)
A = Q - 2 * a**2 * K + a**4 * M - 2 * a * Re * 1j * M - 1j * a * Re * (
K2 - a**2 * K1)
return A, B
def test_PDMA(quad):
SB = Basis(N, 'C', bc='Biharmonic', quad=quad, plan=True)
u = TrialFunction(SB)
v = TestFunction(SB)
points, weights = SB.points_and_weights(N)
fj = Array(SB, buffer=np.random.randn(N))
f_hat = Function(SB)
f_hat = inner(v, fj, output_array=f_hat)
A = inner(v, div(grad(u)))
B = inner(v, u)
s = SB.slice()
H = A + B
P = PDMA(A, B, A.scale, B.scale, solver='cython')
u_hat = Function(SB)
u_hat[s] = solve(H.diags().toarray()[s, s], f_hat[s])
u_hat2 = Function(SB)
u_hat2 = P(u_hat2, f_hat)
assert np.allclose(u_hat2, u_hat)
def get_context():
float, complex, mpitype = datatypes(params.precision)
collapse_fourier = False if params.dealias == '3/2-rule' else True
dim = len(params.N)
dtype = lambda d: float if d == dim - 1 else complex
V = [
Basis(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i))
for i in range(dim)
]
kw0 = {
'threads': params.threads,
'planner_effort': params.planner_effort['fft']
}
T = TensorProductSpace(comm,
V,
dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier,
**kw0)
VT = VectorTensorProductSpace(T)
VM = MixedTensorProductSpace([T] * 2 * dim)
mask = T.mask_nyquist() if params.mask_nyquist else None
kw = {
'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'
}
Vp = [
Basis(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i), **kw)
for i in range(dim)
]
Tp = TensorProductSpace(comm,
Vp,
dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier,
**kw0)
VTp = VectorTensorProductSpace(Tp)
VMp = MixedTensorProductSpace([Tp] * 2 * dim)
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
for i in range(dim):
X[i] = X[i].astype(float)
K[i] = K[i].astype(float)
K2 = np.zeros(T.shape(True), dtype=float)
for i in range(dim):
K2 += K[i] * K[i]
# Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
for i in range(dim):
Kx[i] = Kx[i].astype(float)
K_over_K2 = np.zeros(VT.shape(True), dtype=float)
for i in range(dim):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
UB = Array(VM)
P = Array(T)
curl = Array(VT)
UB_hat = Function(VM)
P_hat = Function(T)
dU = Function(VM)
Source = Array(VM)
ub_dealias = Array(VMp)
ZZ_hat = np.zeros((3, 3) + Tp.shape(True), dtype=complex) # Work array
# Create views into large data structures
U = UB[:3]
U_hat = UB_hat[:3]
B = UB[3:]
B_hat = UB_hat[3:]
# Primary variable
u = UB_hat
hdf5file = MHDFile(config.params.solver,
checkpoint={
'space': VM,
'data': {
'0': {
'UB': [UB_hat]
}
}
},
results={
'space': VM,
'data': {
'UB': [UB]
}
})
return config.AttributeDict(locals())
# Collect basis and solver from either Chebyshev or Legendre submodules
family = sys.argv[-1].lower() if len(sys.argv) == 2 else 'chebyshev'
# Use sympy to compute a rhs, given an analytical solution
x, y = symbols("x,y")
ue = (sin(2 * np.pi * x) * sin(4 * np.pi * y)) * (1 - x**2) * (1 - y**2)
fe = ue.diff(x, 4) + ue.diff(y, 4) + 2 * ue.diff(x, 2, y, 2)
# Lambdify for faster evaluation
ul = lambdify((x, y), ue, 'numpy')
fl = lambdify((x, y), fe, 'numpy')
# Size of discretization
N = (36, 36)
S0 = Basis(N[0], family=family, bc='Biharmonic')
S1 = Basis(N[1], family=family, bc='Biharmonic')
T = TensorProductSpace(comm, (S0, S1), axes=(0, 1))
X = T.local_mesh(True)
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fl(*X))
# Compute right hand side of biharmonic equation
f_hat = inner(v, fj)
# Get left hand side of biharmonic equation
if family == 'chebyshev': # No integration by parts due to weights
matrices = inner(v, div(grad(div(grad(u)))))
def test_eval_tensor(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
# Testing for Dirichlet and regular basis
x, y, z = symbols("x,y,z")
sizes = (22, 21)
funcx = {
'': (1 - x**2) * sin(np.pi * x),
'Dirichlet': (1 - x**2) * sin(np.pi * x),
'Neumann': (1 - x**2) * sin(np.pi * x),
'Biharmonic': (1 - x**2) * sin(2 * np.pi * x)
}
funcy = {
'': (1 - y**2) * sin(np.pi * y),
'Dirichlet': (1 - y**2) * sin(np.pi * y),
'Neumann': (1 - y**2) * sin(np.pi * y),
'Biharmonic': (1 - y**2) * sin(2 * np.pi * y)
}
funcz = {
'': (1 - z**2) * sin(np.pi * z),
'Dirichlet': (1 - z**2) * sin(np.pi * z),
'Neumann': (1 - z**2) * sin(np.pi * z),
'Biharmonic': (1 - z**2) * sin(2 * np.pi * z)
}
funcs = {
(1, 0): cos(2 * y) * funcx[ST.boundary_condition()],
(1, 1): cos(2 * x) * funcy[ST.boundary_condition()],
(2, 0): sin(6 * z) * cos(4 * y) * funcx[ST.boundary_condition()],
(2, 1): sin(2 * z) * cos(4 * x) * funcy[ST.boundary_condition()],
(2, 2): sin(2 * x) * cos(4 * y) * funcz[ST.boundary_condition()]
}
syms = {1: (x, y), 2: (x, y, z)}
points = None
if comm.Get_rank() == 0:
points = np.random.random((dim + 1, 4))
points = comm.bcast(points)
t_0 = 0
t_1 = 0
t_2 = 0
for shape in product(*([sizes] * dim)):
#for shape in ((64, 64),):
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):
#for axis in (0,):
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])
print('axes', axes[dim][axis])
print('bases', bases)
#print(bases[0].axis, bases[1].axis)
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
ul = lambdify(syms[dim], ue, 'numpy')
uu = ul(*X).astype(typecode)
uq = ul(*points).astype(typecode)
u_hat = Function(fft)
u_hat = fft.forward(uu, u_hat)
t0 = time()
result = fft.eval(points, u_hat, method=0)
t_0 += time() - t0
assert np.allclose(uq, result, 0, 1e-6)
t0 = time()
result = fft.eval(points, u_hat, method=1)
t_1 += time() - t0
assert np.allclose(uq, result, 0, 1e-6)
t0 = time()
result = fft.eval(points, u_hat, method=2)
t_2 += time() - t0
print(uq)
assert np.allclose(uq, result, 0, 1e-6), uq / result
result = u_hat.eval(points)
assert np.allclose(uq, result, 0, 1e-6)
ua = u_hat.backward()
assert np.allclose(uu, ua, 0, 1e-6)
ua = Array(fft)
ua = u_hat.backward(ua)
assert np.allclose(uu, ua, 0, 1e-6)
bases.pop(axis)
fft.destroy()
print('method=0', t_0)
print('method=1', t_1)
print('method=2', t_2)
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()
domain = (-1., 1.)
alpha = 1.
a = -1.
b = 1.
x = symbols("x")
ue = sin(4*np.pi*x)*(x+domain[0])*(x+domain[1]) + a*(x-domain[0])/2. + b*(domain[1] - x)/2.
fe = ue.diff(x, 2) + alpha*ue.diff(x, 1)
# Lambdify for faster evaluation
ul = lambdify(x, ue, 'numpy')
fl = lambdify(x, fe, 'numpy')
# Size of discretization
N = int(sys.argv[-2])
SD = Basis(N, family=family, bc=(a, b), domain=domain)
X = SD.mesh()
u = TrialFunction(SD)
v = TestFunction(SD)
# Get f on quad points
fj = Array(SD, buffer=fl(X))
# Compute right hand side of Poisson equation
f_hat = Function(SD)
f_hat = inner(v, fj, output_array=f_hat)
if family == 'legendre':
f_hat *= -1.
# Get left hand side of Poisson equation
if family == 'chebyshev':
def get_context():
"""Set up context for solver"""
# Get points and weights for Chebyshev weighted integrals
ST = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad)
SB = Basis(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad)
ST0 = Basis(params.N[0], 'C', bc=(0, 0),
quad=params.Dquad) # For 1D problem
CT = ST.CT # Chebyshev transform
Nu = params.N[0] - 2 # Number of velocity modes in Shen basis
Nb = params.N[0] - 4 # Number of velocity modes in Shen biharmonic basis
u_slice = slice(0, Nu)
v_slice = slice(0, Nb)
FST = SlabShen_R2C(params.N,
params.L,
comm,
threads=params.threads,
communication=params.communication,
planner_effort=params.planner_effort,
dealias_cheb=params.dealias_cheb)
float, complex, mpitype = datatypes("double")
ST.plan(FST.complex_shape(), 0, complex, {
'threads': params.threads,
'planner_effort': params.planner_effort["dct"]
})
SB.plan(FST.complex_shape(), 0, complex, {
'threads': params.threads,
'planner_effort': params.planner_effort["dct"]
})
# Mesh variables
X = FST.get_local_mesh(ST)
x0, x1, x2 = FST.get_mesh_dims(ST)
K = FST.get_local_wavenumbermesh(scaled=True)
K2 = K[1] * K[1] + K[2] * K[2]
K4 = K2**2
# Set Nyquist frequency to zero on K that is used for odd derivatives
Kx = FST.get_local_wavenumbermesh(scaled=True, eliminate_highest_freq=True)
K_over_K2 = zeros((2, ) + FST.complex_shape())
for i in range(2):
K_over_K2[i] = K[i + 1] / np.where(K2 == 0, 1, K2)
# Solution variables
U = zeros((3, ) + FST.real_shape(), dtype=float)
U0 = zeros((3, ) + FST.real_shape(), dtype=float)
U_hat = zeros((3, ) + FST.complex_shape(), dtype=complex)
U_hat0 = zeros((3, ) + FST.complex_shape(), dtype=complex)
g = zeros(FST.complex_shape(), dtype=complex)
# primary variable
u = (U_hat, g)
H_hat = zeros((3, ) + FST.complex_shape(), dtype=complex)
H_hat0 = zeros((3, ) + FST.complex_shape(), dtype=complex)
H_hat1 = zeros((3, ) + FST.complex_shape(), dtype=complex)
dU = zeros((3, ) + FST.complex_shape(), dtype=complex)
hv = zeros(FST.complex_shape(), dtype=complex)
hg = zeros(FST.complex_shape(), dtype=complex)
Source = zeros((3, ) + FST.real_shape(), dtype=float)
Sk = zeros((3, ) + FST.complex_shape(), dtype=complex)
work = work_arrays()
nu, dt, N = params.nu, params.dt, params.N
kx = K[0][:, 0, 0]
alfa = K2 - 2.0 / nu / dt
# Collect all matrices
mat = config.AttributeDict(
dict(
CDD=inner_product((ST, 0), (ST, 1)),
AB=HelmholtzCoeff(N[0], 1.0, -alfa, ST.quad),
AC=BiharmonicCoeff(N[0],
nu * dt / 2., (1. - nu * dt * K2[0]),
-(K2[0] - nu * dt / 2. * K4[0]),
quad=SB.quad),
# Matrices for biharmonic equation
CBD=inner_product((SB, 0), (ST, 1)),
ABB=inner_product((SB, 0), (SB, 2)),
BBB=inner_product((SB, 0), (SB, 0)),
SBB=inner_product((SB, 0), (SB, 4)),
# Matrices for Helmholtz equation
ADD=inner_product((ST, 0), (ST, 2)),
BDD=inner_product((ST, 0), (ST, 0)),
BBD=inner_product((SB, 0), (ST, 0)),
CDB=inner_product((ST, 0), (SB, 1)),
ADD0=inner_product((ST0, 0), (ST0, 2)),
BDD0=inner_product((ST0, 0), (ST0, 0))))
mat.ADD.axis = 0
mat.BDD.axis = 0
mat.SBB.axis = 0
# Collect all linear algebra solvers
la = config.AttributeDict(
dict(HelmholtzSolverG=Helmholtz(mat.ADD, mat.BDD, -np.ones((1, 1, 1)),
(K2 + 2.0 / nu / dt)),
BiharmonicSolverU=Biharmonic(mat.SBB, mat.ABB, mat.BBB,
-nu * dt / 2. * np.ones(
(1, 1, 1)), (1. + nu * dt * K2),
(-(K2 + nu * dt / 2. * K4))),
HelmholtzSolverU0=Helmholtz(mat.ADD0, mat.BDD0, np.array([-1.]),
np.array([2. / nu / dt])),
TDMASolverD=TDMA(inner_product((ST, 0), (ST, 0)))))
hdf5file = KMMWriter({
"U": U[0],
"V": U[1],
"W": U[2]
},
chkpoint={
'current': {
'U': U
},
'previous': {
'U': U0
}
},
filename=params.solver + ".h5",
mesh={
"x": x0,
"y": x1,
"z": x2
})
return config.AttributeDict(locals())
assert len(sys.argv) == 2, 'Call with one command-line argument'
assert isinstance(int(sys.argv[-1]), int)
comm = MPI.COMM_WORLD
# Use sympy to compute a rhs, given an analytical solution
x, y = symbols("x,y")
#ue = sin(4*x)*exp(-x**2)
ue = cos(4 * y) * hermite(4, x) * exp(-x**2 / 2)
fe = ue.diff(x, 2) + ue.diff(y, 2)
# Size of discretization
N = int(sys.argv[-1])
SD = Basis(N, 'Hermite')
K0 = Basis(N, 'Fourier', dtype='d')
T = TensorProductSpace(comm, (SD, K0), axes=(0, 1))
X = T.local_mesh(True)
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side of Poisson equation
f_hat = Function(T)
f_hat = inner(v, -fj, output_array=f_hat)
# Get left hand side of Poisson equation
matrices = inner(grad(v), grad(u))
from sympy import Symbol, cos, sin, lambdify
import numpy as np
from shenfun import inner, grad, TestFunction, TrialFunction, Basis, Function, \
Array
# Use sympy to compute a rhs, given an analytical solution
x = Symbol("x")
ue = cos(4 * x) + 1j * sin(6 * x)
#ue = cos(4*x)
fe = ue.diff(x, 2)
# Size of discretization
N = 40
dtype = {True: np.complex, False: np.float}[ue.has(1j)]
ST = Basis(N, dtype=dtype, domain=(-2 * np.pi, 2 * np.pi))
u = TrialFunction(ST)
v = TestFunction(ST)
X = ST.mesh()
# Get f on quad points and exact solution
fj = Array(ST, buffer=fe)
uj = Array(ST, buffer=ue)
# Compute right hand side
f_hat = Function(ST)
f_hat = inner(v, fj, output_array=f_hat)
# Solve Poisson equation
A = inner(grad(v), grad(u))
