def test_eval_expression():
import sympy as sp
from shenfun import div, grad
x, y, z = sp.symbols('x,y,z')
B0 = FunctionSpace(16, 'C')
B1 = FunctionSpace(17, 'C')
B2 = FunctionSpace(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_padding_biharmonic(family):
N = 8
B = FunctionSpace(N, family, bc=(0, 0, 0, 0))
Bp = B.get_dealiased(1.5)
u = Function(B)
u[:(N - 4)] = np.random.random(N - 4)
up = Array(Bp)
up = Bp.backward(u, fast_transform=False)
uf = Bp.forward(up, fast_transform=False)
assert np.linalg.norm(uf - u) < 1e-12
if family == 'C':
up = Bp.backward(u)
uf = Bp.forward(up)
assert np.linalg.norm(uf - u) < 1e-12
# Test padding 2D
F = FunctionSpace(N, 'F', dtype='d')
T = TensorProductSpace(comm, (B, F))
Tp = T.get_dealiased(1.5)
u = Function(T)
u[:-4, :-1] = np.random.random(u[:-4, :-1].shape)
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
# Test padding 3D
F1 = FunctionSpace(N, 'F', dtype='D')
T = TensorProductSpace(comm, (F1, F, B))
Tp = T.get_dealiased(1.5)
u = Function(T)
u[:, :, :-4] = np.random.random(u[:, :, :-4].shape)
u = u.backward().forward() # Clean
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
def get_divergence(U_hat, FST, mask, **context):
div_hat = project(div(U_hat), FST)
if mask is not None:
div_hat.mask_nyquist(mask)
div_ = Array(FST)
div_ = div_hat.backward(div_)
return div_
def test_padding_neumann(family):
N = 8
B = FunctionSpace(N, family, bc={'left': ('N', 0), 'right': ('N', 0)})
Bp = B.get_dealiased(1.5)
u = Function(B)
u[1:-2] = np.random.random(N - 3)
up = Array(Bp)
up = Bp.backward(u, fast_transform=False)
uf = Bp.forward(up, fast_transform=False)
assert np.linalg.norm(uf - u) < 1e-12
if family == 'C':
up = Bp.backward(u)
uf = Bp.forward(up)
assert np.linalg.norm(uf - u) < 1e-12
# Test padding 2D
F = FunctionSpace(N, 'F', dtype='d')
T = TensorProductSpace(comm, (B, F))
Tp = T.get_dealiased(1.5)
u = Function(T)
u[1:-2, :-1] = np.random.random(u[1:-2, :-1].shape)
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
# Test padding 3D
F1 = FunctionSpace(N, 'F', dtype='D')
T = TensorProductSpace(comm, (F1, F, B))
Tp = T.get_dealiased(1.5)
u = Function(T)
u[:, :, 1:-2] = np.random.random(u[:, :, 1:-2].shape)
u = u.backward().forward() # Clean
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
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(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)
# Spectral space must be aligned in nonperiodic direction, hence axes
fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
dfft = fft.get_orthogonal()
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
uh = Function(fft, buffer=ue)
due = ue.diff(xs[axis], 1)
duy = Array(fft, buffer=due)
duf = project(Dx(uh, axis, 1), fft).backward()
assert np.allclose(duy, duf, 0, 1e-3), np.linalg.norm(duy-duf)
# 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)
duq = Array(fft, buffer=due)
uf = project(Dx(Dx(uh, ax, 1), axis, 1), fft).backward()
assert np.allclose(uf, duq, 0, 1e-3)
bases.pop(axis)
fft.destroy()
def main(N, family, bci, bcj, plotting=False):
global fe, ue
BX = FunctionSpace(N, family=family, bc=bcx[bci], domain=xdomain)
BY = FunctionSpace(N, family=family, bc=bcy[bcj], domain=ydomain)
T = TensorProductSpace(comm, (BX, BY))
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compare with analytical solution
ua = Array(T, buffer=ue)
if T.use_fixed_gauge:
mean = dx(ua, weighted=True) / inner(1, Array(T, val=1))
# 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
A = inner(v, -div(grad(u)))
u_hat = Function(T)
sol = la.Solver2D(A, fixed_gauge=mean if T.use_fixed_gauge else None)
u_hat = sol(f_hat, u_hat)
uj = u_hat.backward()
assert np.allclose(uj, ua), np.linalg.norm(uj - ua)
print("Error=%2.16e" % (np.sqrt(dx((uj - ua)**2))))
if 'pytest' not in os.environ and plotting is True:
import matplotlib.pyplot as plt
X, Y = T.local_mesh(True)
plt.contourf(X, Y, uj, 100)
plt.colorbar()
plt.figure()
plt.contourf(X, Y, ua, 100)
plt.colorbar()
plt.figure()
plt.contourf(X, Y, ua - uj, 100)
plt.colorbar()
def test_regular_2D(backend, forward_output):
if (backend == 'netcdf4' and forward_output is True) or skip[backend]:
return
K0 = FunctionSpace(N[0], 'F')
K1 = FunctionSpace(N[1], 'C', bc=(0, 0))
T = TensorProductSpace(comm, (K0, K1))
filename = 'test2Dr_{}'.format(ex[forward_output])
hfile = writer(filename, T, backend=backend)
u = Function(T, val=1) if forward_output else Array(T, val=1)
hfile.write(0, {'u': [u]}, forward_output=forward_output)
hfile.write(1, {'u': [u]}, forward_output=forward_output)
if not forward_output and backend == 'hdf5' and comm.Get_rank() == 0:
generate_xdmf(filename + '.h5')
generate_xdmf(filename + '.h5', order='visit')
u0 = Function(T) if forward_output else Array(T)
read = reader(filename, T, backend=backend)
read.read(u0, 'u', forward_output=forward_output, step=1)
assert np.allclose(u0, u)
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(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, 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 main(N, family, bci):
bc = bcs[bci]
if bci == 0:
SD = FunctionSpace(N,
family=family,
bc=bc,
domain=domain,
mean=mean[family.lower()])
else:
SD = FunctionSpace(N, family=family, bc=bc, domain=domain)
u = TrialFunction(SD)
v = TestFunction(SD)
# Get f on quad points
fj = Array(SD, buffer=fe)
# Compute right hand side of Poisson equation
f_hat = Function(SD)
f_hat = inner(v, fj, output_array=f_hat)
# Get left hand side of Poisson equation
A = inner(v, div(grad(u)))
u_hat = Function(SD).set_boundary_dofs()
if isinstance(A, list):
bc_mat = extract_bc_matrices([A])
A = A[0]
f_hat -= bc_mat[0].matvec(u_hat, Function(SD))
u_hat = A.solve(f_hat, u_hat)
uj = u_hat.backward()
uh = uj.forward()
# Compare with analytical solution
ua = Array(SD, buffer=ue)
assert np.allclose(uj, ua), np.linalg.norm(uj - ua)
if 'pytest' not in os.environ:
print("Error=%2.16e" % (np.sqrt(dx((uj - ua)**2))))
import matplotlib.pyplot as plt
plt.plot(SD.mesh(), uj, 'b', SD.mesh(), ua, 'r')
def get_pressure(context, solver):
FCT = context.FCT
FST = context.FST
U = solver.get_velocity(**context)
U0 = context.VFS.backward(context.U_hat0, context.U0)
dt = solver.params.dt
H_hat = solver.get_convection(**context)
Hx = Array(FST)
Hx = FST.backward(H_hat[0], Hx)
v = TestFunction(FCT)
p = TrialFunction(FCT)
U = U.as_function()
U0 = U0.as_function()
rhs_hat = inner((0.5 * context.nu) * div(grad(U[0] + U0[0])), v)
Hx -= 1. / dt * (U[0] - U0[0])
rhs_hat += inner(Hx, v)
CT = inner(Dx(p, 0), v)
# Should implement fast solver. Just a backwards substitution
A = CT.diags().toarray() * CT.scale[0]
A[-1, 0] = 1
a_i = np.linalg.inv(A)
p_hat = Function(context.FCT)
for j in range(p_hat.shape[1]):
for k in range(p_hat.shape[2]):
p_hat[:, j, k] = np.dot(a_i, rhs_hat[:, j, k])
p = Array(FCT)
p = FCT.backward(p_hat, p)
uu = np.sum((0.5 * (U + U0))**2, 0)
uu *= 0.5
return p - uu + 3. / 16.
def test_quasiTau(bc):
N = 40
T = FunctionSpace(N, 'C')
u = TrialFunction(T)
v = TestFunction(T)
A = inner(v, div(grad(u)))
B = inner(v, u)
Q = chebyshev.quasi.QITmat(N)
A = Q*A
B = Q*B
bb = Q*np.ones(N)
assert abs(np.sum(bb[:8])) < 1e-8
M = B-A
M0 = M.diags().tolil()
if bc == 'Dirichlet':
M0[0] = np.ones(N)
nn = np.ones(N)
nn[1::2] = -1
M0[1] = nn
elif bc == 'Neumann':
nn = np.arange(N)**2
M0[0] = nn.copy()
nn[1::2] *= -1
M0[1] = nn
M0 = M0.tocsc()
f_hat = inner(v, Array(T, buffer=fe))
gh = Q.diags('csc')*f_hat
gh[:2] = 0
u_hat = Function(T)
u_hat[:] = scp.linalg.spsolve(M0, gh)
uj = u_hat.backward()
ua = Array(T, buffer=ue)
if bc == 'Neumann':
xj, wj = T.points_and_weights()
ua -= np.sum(ua*wj)/np.pi # normalize
uj -= np.sum(uj*wj)/np.pi # normalize
assert np.sqrt(inner(1, (uj-ua)**2)) < 1e-5
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 get_context():
"""Set up context for classical (NS) solver"""
V0 = C2CBasis(params.N[0], domain=(0, params.L[0]))
V1 = C2CBasis(params.N[1], domain=(0, params.L[1]))
V2 = R2CBasis(params.N[2], domain=(0, params.L[2]))
T = TensorProductSpace(comm, (V0, V1, V2), **{'threads': params.threads})
VT = VectorTensorProductSpace([T]*3)
kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'}
V0p = C2CBasis(params.N[0], domain=(0, params.L[0]), **kw)
V1p = C2CBasis(params.N[1], domain=(0, params.L[1]), **kw)
V2p = R2CBasis(params.N[2], domain=(0, params.L[2]), **kw)
Tp = TensorProductSpace(comm, (V0p, V1p, V2p), **{'threads': params.threads})
VTp = VectorTensorProductSpace([Tp]*3)
float, complex, mpitype = datatypes(params.precision)
FFT = T # For compatibility - to be removed
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
K2 = K[0]*K[0] + K[1]*K[1] + K[2]*K[2]
# 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)
K_over_K2 = np.zeros((3,)+VT.local_shape())
for i in range(3):
K_over_K2[i] = K[i] / np.where(K2==0, 1, K2)
# Velocity and pressure
U = Array(VT, False)
U_hat = Array(VT)
P = Array(T, False)
P_hat = Array(T)
# Primary variable
u = U_hat
# RHS array
dU = Array(VT)
curl = Array(VT, False)
Source = Array(VT) # Possible source term initialized to zero
work = work_arrays()
hdf5file = NSWriter({"U":U[0], "V":U[1], "W":U[2], "P":P},
chkpoint={"current":{"U":U, "P":P}, "previous":{}},
filename=params.h5filename+".h5")
return config.AttributeDict(locals())
def test_PDMA(quad):
SB = FunctionSpace(N, 'C', bc=(0, 0, 0, 0), quad=quad)
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 test_TwoDMA():
N = 12
SD = FunctionSpace(N, 'C', basis='ShenDirichlet')
HH = FunctionSpace(N, 'C', basis='Heinrichs')
u = TrialFunction(HH)
v = TestFunction(SD)
points, weights = SD.points_and_weights(N)
fj = Array(SD, buffer=np.random.randn(N))
f_hat = Function(SD)
f_hat = inner(v, fj, output_array=f_hat)
A = inner(v, div(grad(u)))
sol = TwoDMA(A)
u_hat = Function(HH)
u_hat = sol(f_hat, u_hat)
sol2 = la.Solve(A, HH)
u_hat2 = Function(HH)
u_hat2 = sol2(f_hat, u_hat2)
assert np.allclose(u_hat2, u_hat)
def get_pressure(context, solver):
FCT = context.FCT
FST = context.FST
U_hat = context.U_hat
U_hat0 = context.U_hat0
Um = Function(context.FST)
Um[:] = 0.5*(U_hat[0] + U_hat0[0])
U = U_hat.backward(context.U)
U0 = U_hat0.backward(context.U0)
dt = solver.params.dt
Hx = Function(context.FST)
Hx[:] = solver.get_convection(**context)[0]
v = TestFunction(FCT)
p = TrialFunction(FCT)
rhs_hat = inner(context.nu*div(grad(Um)), v)
Hx -= 1./dt*(U_hat[0]-U_hat0[0])
rhs_hat += inner(Hx, v)
CT = inner(Dx(p, 0, 1), v)
# Should implement fast solver. Just a backwards substitution
# Apply integral constraint
A = CT.mats[0]
N = A.shape[0]
A[-(N-1)] = 1
p_hat = Function(context.FCT)
p_hat = CT.solve(rhs_hat, p_hat)
p = Array(FCT)
p = FCT.backward(p_hat, p)
uu = 0.
if params.convection == 'Vortex':
uu = np.sum((0.5*(U+U0))**2, 0)
uu *= 0.5
return p-uu
def test_padding_orthogonal(family):
N = 8
B = FunctionSpace(N, family)
Bp = B.get_dealiased(1.5)
u = Function(B)
u[:] = np.random.random(u.shape)
up = Array(Bp)
if family != 'F':
up = Bp.backward(u, fast_transform=False)
uf = Bp.forward(up, fast_transform=False)
assert np.linalg.norm(uf - u) < 1e-12
if family in ('C', 'F'):
up = Bp.backward(u, fast_transform=True)
uf = Bp.forward(up, fast_transform=True)
assert np.linalg.norm(uf - u) < 1e-12
# Test padding 2D
dtype = 'D' if family == 'F' else 'd'
F = FunctionSpace(N, 'F', dtype=dtype)
T = TensorProductSpace(comm, (F, B))
Tp = T.get_dealiased(1.5)
u = Function(T)
u[:] = np.random.random(u.shape)
u = u.backward().forward()
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
# Test padding 3D
F1 = FunctionSpace(N, 'F', dtype='D')
T = TensorProductSpace(comm, (F1, F, B))
Tp = T.get_dealiased(1.5)
u = Function(T)
u[:] = np.random.random(u.shape)
u = u.backward().forward() # Clean
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
def test_padding(family):
N = 8
B = FunctionSpace(N, family, bc=(-1, 1), domain=(-2, 2))
Bp = B.get_dealiased(1.5)
u = Function(B).set_boundary_dofs()
#u[:(N-2)] = np.random.random(N-2)
u[:(N - 2)] = 1
up = Array(Bp)
up = Bp.backward(u, fast_transform=False)
uf = Bp.forward(up, fast_transform=False)
assert np.linalg.norm(uf - u) < 1e-12
if family == 'C':
up = Bp.backward(u)
uf = Bp.forward(up)
assert np.linalg.norm(uf - u) < 1e-12
# Test padding 2D
F = FunctionSpace(N, 'F', dtype='d')
T = TensorProductSpace(comm, (B, F))
Tp = T.get_dealiased(1.5)
u = Function(T).set_boundary_dofs()
u[:-2, :-1] = np.random.random(u[:-2, :-1].shape)
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
# Test padding 3D
F1 = FunctionSpace(N, 'F', dtype='D')
T = TensorProductSpace(comm, (F1, F, B))
Tp = T.get_dealiased(1.5)
u = Function(T).set_boundary_dofs()
u[:, :, :-2] = np.random.random(u[:, :, :-2].shape)
u = u.backward().forward() # Clean
up = Tp.backward(u)
uc = Tp.forward(up)
assert np.linalg.norm(u - uc) < 1e-8
def test_Mult_CTD_3D(quad):
SD = FunctionSpace(N, 'C', bc=(0, 0))
F0 = FunctionSpace(4, 'F', dtype='D')
F1 = FunctionSpace(4, 'F', dtype='d')
T = TensorProductSpace(comm, (SD, F0, F1))
TO = T.get_orthogonal()
CT = TO.bases[0]
C = inner_product((CT, 0), (SD, 1))
B = inner_product((CT, 0), (CT, 0))
vk = Array(T)
wk = Array(T)
vk[:] = np.random.random(vk.shape)
wk[:] = np.random.random(vk.shape)
bv = Function(T)
bw = Function(T)
vk0 = vk.forward()
vk = vk0.backward()
wk0 = wk.forward()
wk = wk0.backward()
LUsolve.Mult_CTD_3D_ptr(N, vk0, wk0, bv, bw, 0)
cv = np.zeros_like(vk0)
cw = np.zeros_like(wk0)
cv = C.matvec(vk0, cv)
cw = C.matvec(wk0, cw)
cv /= B[0].repeat(np.array(bv.shape[1:]).prod()).reshape(bv.shape)
cw /= B[0].repeat(np.array(bv.shape[1:]).prod()).reshape(bv.shape)
assert np.allclose(cv, bv)
assert np.allclose(cw, bw)
def get_context():
"""Set up context for solver"""
collapse_fourier = False if params.dealias == '3/2-rule' else True
family = 'C'
ST = FunctionSpace(params.N[0], family, bc=(0, 0), quad=params.Dquad)
CT = FunctionSpace(params.N[0], family, quad=params.Dquad)
CP = FunctionSpace(params.N[0], family, quad=params.Dquad)
K0 = FunctionSpace(params.N[1], 'F', domain=(0, params.L[1]), dtype='D')
K1 = FunctionSpace(params.N[2], 'F', domain=(0, params.L[2]), dtype='d')
#CP.slice = lambda: slice(0, CP.N-2)
constraints = ((3, 0, 0), (3, params.N[0] - 1, 0))
kw0 = {
'threads': params.threads,
'planner_effort': params.planner_effort["dct"],
'slab': (params.decomposition == 'slab'),
'collapse_fourier': collapse_fourier
}
FST = TensorProductSpace(comm, (ST, K0, K1), **kw0) # Dirichlet
FCT = TensorProductSpace(comm, (CT, K0, K1), **kw0) # Regular Chebyshev N
FCP = TensorProductSpace(comm, (CP, K0, K1),
**kw0) # Regular Chebyshev N-2
VFS = VectorSpace(FST)
VCT = VectorSpace(FCT)
VQ = CompositeSpace([VFS, FCP])
mask = FST.get_mask_nyquist() if params.mask_nyquist else None
# Padded
kw = {
'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'
}
if params.dealias == '3/2-rule':
# Requires new bases due to planning and transforms on different size arrays
STp = FunctionSpace(params.N[0], family, bc=(0, 0), quad=params.Dquad)
CTp = FunctionSpace(params.N[0], family, quad=params.Dquad)
else:
STp, CTp = ST, CT
K0p = FunctionSpace(params.N[1],
'F',
dtype='D',
domain=(0, params.L[1]),
**kw)
K1p = FunctionSpace(params.N[2],
'F',
dtype='d',
domain=(0, params.L[2]),
**kw)
FSTp = TensorProductSpace(comm, (STp, K0p, K1p), **kw0)
FCTp = TensorProductSpace(comm, (CTp, K0p, K1p), **kw0)
VFSp = VectorSpace(FSTp)
VCp = CompositeSpace([FSTp, FCTp, FCTp])
float, complex, mpitype = datatypes("double")
# Mesh variables
X = FST.local_mesh(True)
x0, x1, x2 = FST.mesh()
K = FST.local_wavenumbers(scaled=True)
# Solution variables
UP_hat = Function(VQ)
UP_hat0 = Function(VQ)
U_hat, P_hat = UP_hat
U_hat0, P_hat0 = UP_hat0
UP = Array(VQ)
UP0 = Array(VQ)
U, P = UP
U0, P0 = UP0
# RK parameters
a = (8. / 15., 5. / 12., 3. / 4.)
b = (0.0, -17. / 60., -5. / 12.)
# primary variable
u = UP_hat
H_hat = Function(VFS)
dU = Function(VQ)
hv = np.zeros((2, ) + H_hat.shape, dtype=np.complex)
Source = Array(
VFS) # Note - not using VQ. Only used for constant pressure gradient
Sk = Function(VFS)
K2 = K[1] * K[1] + K[2] * K[2]
for i in range(3):
K[i] = K[i].astype(float)
work = work_arrays()
u_dealias = Array(VFSp)
curl_hat = Function(VCp)
curl_dealias = Array(VCp)
nu, dt, N = params.nu, params.dt, params.N
up = TrialFunction(VQ)
vq = TestFunction(VQ)
ut, pt = up
vt, qt = vq
M = []
for rk in range(3):
a0 = inner(vt, (2. / nu / dt / (a[rk] + b[rk])) * ut - div(grad(ut)))
a1 = inner(vt, (2. / nu / (a[rk] + b[rk])) * grad(pt))
a2 = inner(qt, (2. / nu / (a[rk] + b[rk])) * div(ut))
M.append(BlockMatrix(a0 + a1 + a2))
# Collect all matrices
if ST.family() == 'chebyshev':
mat = config.AttributeDict(
dict(AB=[
HelmholtzCoeff(N[0], 1.,
-(K2 - 2. / nu / dt / (a[rk] + b[rk])), 0,
ST.quad) for rk in range(3)
], ))
else:
mat = config.AttributeDict(
dict(ADD=inner_product((ST, 0), (ST, 2)),
BDD=inner_product((ST, 0), (ST, 0))))
la = None
hdf5file = CoupledRK3File(config.params.solver,
checkpoint={
'space': VQ,
'data': {
'0': {
'UP': [UP_hat]
}
}
},
results={
'space': VFS,
'data': {
'U': [U]
}
})
del rk
return config.AttributeDict(locals())
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())
def get_divergence(T, Kx, U_hat, **context):
div_u = Array(T)
div_u = T.backward(
1j * (Kx[0] * U_hat[0] + Kx[1] * U_hat[1] + Kx[2] * U_hat[2]), div_u)
return div_u
# 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)))))
else: # Use form with integration by parts.
matrices = inner(div(grad(v)), div(grad(u)))
# Create linear algebra solver
H = SolverGeneric2NP(matrices)
# Solve and transform to real space
u_hat = Function(T) # Solution spectral space
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()
def test_project2(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (22, 21)
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(6 * 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-5)
# 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-5)
bases.pop(axis)
fft.destroy()
# Use sympy to compute a rhs, given an analytical solution
x = symbols("x", real=True)
#ue = sin(4*x)*exp(-x**2)
ue = hermite(4, x) * exp(-x**2 / 2)
fe = ue.diff(x, 2)
# Size of discretization
N = int(sys.argv[-1])
SD = FunctionSpace(N, 'Hermite')
u = TrialFunction(SD)
v = TestFunction(SD)
# Get f on quad points
fj = Array(SD, buffer=fe)
# Compute right hand side of Poisson equation
f_hat = Function(SD)
f_hat = inner(v, -fj, output_array=f_hat)
# Get left hand side of Poisson equation
A = inner(grad(v), grad(u))
f_hat = A / f_hat
uj = f_hat.backward()
uh = uj.forward()
# Compare with analytical solution
ua = Array(SD, buffer=ue)
print("Error=%2.16e" % (np.linalg.norm(uj - ua)))
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':
A = inner(v, div(grad(u))) + alpha*inner(v, grad(u))
else:
A = inner(grad(v), grad(u)) - alpha*inner(v, grad(u))
if family == 'chebyshev':
f_hat[0] -= 0.5*np.pi*alpha*a
b = 0
x, y = symbols("x,y")
ue = (cos(4*x) + sin(2*y))*(1 - x**2) + a*(1 - x)/2. + b*(1 + x)/2.
fe = ue.diff(x, 2) + ue.diff(y, 2)
# Size of discretization
N = (int(sys.argv[-2]), int(sys.argv[-2])+1)
SD = FunctionSpace(N[0], family=family, scaled=True, bc=(a, b))
K1 = FunctionSpace(N[1], family='F', dtype='d', domain=(-2*np.pi, 2*np.pi))
T = TensorProductSpace(comm, (SD, K1), axes=(0, 1))
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(v, div(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 = u_hat.backward()
N = (eval(sys.argv[-2]), eval(sys.argv[-2]))
SD = Basis(N[1], scaled=True, bc=(a, b))
K1 = R2CBasis(N[0])
T = TensorProductSpace(comm, (K1, SD), axes=(1, 0))
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
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = fl(*X)
# Compute right hand side of Poisson equation
f_hat = Array(T)
f_hat = inner(v, fj, output_array=f_hat)
if basis == 'legendre':
f_hat *= -1.
#from IPython import embed; embed()
# Get left hand side of Poisson equation
if basis == '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
