def test_curl(typecode):
K0 = FunctionSpace(N[0], 'F', dtype=typecode.upper())
K1 = FunctionSpace(N[1], 'F', dtype=typecode.upper())
K2 = FunctionSpace(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 = shenfun.TrialFunction(Tk)
v = shenfun.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)
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_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 _setup_variational_problem(self):
self._setup_function_space()
u = sf.TrialFunction(self._V)
v = sf.TestFunction(self._V)
self._elastic_law.set_material_parameters(self._material_parameters)
self._dw_int = self._elastic_law.dw_int(u, v)
self._dw_ext = inner(
v, sf.Array(self._V.get_orthogonal(), buffer=(0, ) * self._dim))
if self._body_forces is not None:
V_body_forces = self._V.get_orthogonal()
body_forces_quad = sf.Array(V_body_forces,
buffer=self._body_forces)
self._dw_ext = inner(v, body_forces_quad)
def test_project_1D(basis):
ue = sin(2*np.pi*x)*(1-x**2)
T = basis(12)
u = shenfun.TrialFunction(T)
v = shenfun.TestFunction(T)
u_tilde = shenfun.Function(T)
X = T.mesh()
ua = shenfun.Array(T, buffer=ue)
u_tilde = shenfun.inner(v, ua, output_array=u_tilde)
M = shenfun.inner(u, v)
u_p = shenfun.Function(T)
u_p = M.solve(u_tilde, u=u_p)
u_0 = shenfun.Function(T)
u_0 = shenfun.project(ua, T)
assert np.allclose(u_0, u_p)
u_1 = shenfun.project(ue, T)
assert np.allclose(u_1, u_p)
def ComputeRHS(rhs, u_hat, rk, solver, H_hat, VFSp, FSTp, FCTp, VCp, work, K,
K2, u_dealias, curl_dealias, curl_hat, mat, la, vt, Sk, hv, a,
b, mask, **context):
"""Compute right hand side of Navier Stokes
Parameters
----------
rhs : array
The right hand side to be returned
u_hat : array
The velocity vector at current time
solver : module
The current solver module
Remaining args are extracted from context
"""
w0 = work[(u_hat[0], 0, False)]
# Nonlinear convection term at current u_hat
H_hat = solver.conv(H_hat, u_hat, K, VFSp, VCp, FSTp, FCTp, work,
u_dealias, curl_dealias, curl_hat, mat, la)
if mask is not None:
H_hat.mask_nyquist(mask)
# Assemble rhs
rhs[:] = 0
rhs_u, rhs_p = rhs
if context['ST'].family() == 'chebyshev':
rhs_u[0] = mat.AB[rk].matvec(u_hat[0], rhs_u[0])
rhs_u[1] = mat.AB[rk].matvec(u_hat[1], rhs_u[1])
rhs_u[2] = mat.AB[rk].matvec(u_hat[2], rhs_u[2])
else:
#rhs_u[:] = inner(vt, (2./params.nu/params.dt/(a[rk]+b[rk]))*u_hat)
#rhs_u += inner(vt, div(grad(u_hat)))
rhs_u[0] = -(K2 - 2. / params.nu / params.dt /
(a[rk] + b[rk])) * mat.BDD.matvec(u_hat[0], w0)
rhs_u[1] = -(K2 - 2. / params.nu / params.dt /
(a[rk] + b[rk])) * mat.BDD.matvec(u_hat[1], w0)
rhs_u[2] = -(K2 - 2. / params.nu / params.dt /
(a[rk] + b[rk])) * mat.BDD.matvec(u_hat[2], w0)
rhs_u[0] += mat.ADD.matvec(u_hat[0], w0)
rhs_u[1] += mat.ADD.matvec(u_hat[1], w0)
rhs_u[2] += mat.ADD.matvec(u_hat[2], w0)
# Convection
hv[1] = inner(vt, H_hat)
# Source
hv[1, 1] -= Sk[1]
rhs_u[:] += (hv[1] * a[rk] + hv[0] * b[rk]) * 2. / params.nu / (a[rk] +
b[rk])
hv[0] = hv[1]
return rhs
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 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 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_biharmonic3D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (16, 16, 16)
SD = shenfun.Basis(N[allaxes3D[axis][0]], family=family, bc='Biharmonic')
K1 = shenfun.Basis(N[allaxes3D[axis][1]], family='F', dtype='D')
K2 = shenfun.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 = shenfun.TensorProductSpace(subcomms, bases, axes=allaxes3D[axis])
u = shenfun.TrialFunction(T)
v = shenfun.TestFunction(T)
if family == 'chebyshev':
mat = shenfun.inner(
v, shenfun.div(shenfun.grad(shenfun.div(shenfun.grad(u)))))
else:
mat = shenfun.inner(shenfun.div(shenfun.grad(v)),
shenfun.div(shenfun.grad(u)))
H = la.Biharmonic(*mat)
H = la.Biharmonic(*mat)
u = shenfun.Function(T)
u[:] = np.random.random(u.shape) + 1j * np.random.random(u.shape)
f = shenfun.Function(T)
f = H.matvec(u, f)
f = H.matvec(u, f)
g0 = shenfun.Function(T)
g1 = shenfun.Function(T)
g2 = shenfun.Function(T)
M = {d.get_key(): d for d in mat}
amat = 'ABBmat' if family == 'chebyshev' else 'PBBmat'
g0 = M['SBBmat'].matvec(u, g0)
g1 = M[amat].matvec(u, g1)
g2 = M['BBBmat'].matvec(u, g2)
assert np.linalg.norm(f - (g0 + g1 + g2)) < 1e-8, np.linalg.norm(f -
(g0 + g1 +
g2))
def check_with_pde(self, u_hat, material_parameters, body_forces):
assert isinstance(u_hat, sf.Function)
assert len(material_parameters) == self._n_material_parameters
for comp in body_forces:
assert isinstance(comp, (sp.Expr, float, int))
lambd, mu = material_parameters
V = u_hat.function_space().get_orthogonal()
# left hand side of pde
lhs = (lambd + mu) * grad(div(u_hat)) + mu * div(grad(u_hat))
error_array = sf.Array(V, buffer=body_forces)
error_array += sf.project(lhs, V).backward()
error = np.sqrt(inner((1, 1), error_array ** 2))
# scale by magnitude of solution
scale = np.sqrt(inner((1, 1), u_hat.backward() ** 2))
return error / scale
def test_div2(basis, quad):
B = basis(8, quad=quad)
u = shenfun.TrialFunction(B)
v = shenfun.TestFunction(B)
m = inner(u, v)
z = Function(B, val=1)
c = m / z
m2 = m.diags('csr')
c2 = spsolve(m2, z[B.slice()])
assert np.allclose(c2, c[B.slice()])
def test_helmholtz3D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9, 10)
SD = FunctionSpace(N[allaxes3D[axis][0]], family=family, bc=(0, 0))
K1 = FunctionSpace(N[allaxes3D[axis][1]], family='F', dtype='D')
K2 = FunctionSpace(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],
modify_spaces_inplace=True)
u = shenfun.TrialFunction(T)
v = shenfun.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['ASDSDmat'].matvec(u, g0)
g1 = M['BSDSDmat'].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 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 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 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 compute_numerical_error(u_ana, u_hat):
assert isinstance(u_hat, sf.Function)
V = u_hat.function_space()
# evaluate u_ana at quadrature points
error_array = sf.Array(V, buffer=u_ana)
# subtract numerical solution
error_array -= u_hat.backward()
# compute integral error
error = np.sqrt(sf.inner((1, 1), error_array**2))
return error
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 dw_int(self, u, v):
assert isinstance(u, sf.TrialFunction)
assert isinstance(v, sf.TestFunction)
assert hasattr(self, "_material_parameters")
assert len(self._material_parameters) == self._n_material_parameters
self.dim = u.dimensions
lmbda, mu, c1, c2, c3, c4, c5 = self._material_parameters
dw_int = []
if c1 != 0.0:
dw_int += inner(c1 * div(grad(u)), div(grad(v)))
if c2 != 0.0:
dw_int += inner(c2 * div(grad(u)), grad(div(v)))
if c3 != 0.0:
dw_int += inner(c3 * grad(div(u)), grad(div(v)))
if c4 != 0.0:
for i in range(self.dim):
for j in range(self.dim):
for k in range(self.dim):
mat = inner(c4 * Dx(Dx(u[i], j), k),
Dx(Dx(v[i], j), k))
if isinstance(mat, list):
dw_int += mat
else:
dw_int += [mat]
if c5 != 0.0:
for i in range(self.dim):
for j in range(self.dim):
for k in range(self.dim):
mat = inner(c5 * Dx(Dx(u[j], i), k),
Dx(Dx(v[i], j), k))
if isinstance(mat, list):
dw_int += mat
else:
dw_int += [mat]
# add the classical cauchy-terms
dw_int += inner(mu * grad(u), grad(v))
for i in range(self.dim):
for j in range(self.dim):
mat = inner(mu * Dx(u[i], j), Dx(v[j], i))
if isinstance(mat, list):
dw_int += mat
else:
dw_int += [mat]
dw_int += inner(lmbda * div(u), div(v))
return dw_int
def dw_int(self, u, v):
assert isinstance(u, sf.TrialFunction)
assert isinstance(v, sf.TestFunction)
assert hasattr(self, "_material_parameters")
assert len(self._material_parameters) == self._n_material_parameters
self.dim = u.dimensions
lmbda, mu = self._material_parameters
A = inner(mu * grad(u), grad(v))
B = []
for i in range(self.dim):
for j in range(self.dim):
mat = inner(mu * Dx(u[i], j), Dx(v[j], i))
if isinstance(mat, list):
B += mat
else:
B += [mat]
C = inner(lmbda * div(u), div(v))
dw_int = A + B + C
return dw_int
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_vector_laplace(space):
"""Test that
div(grad(u)) = grad(div(u)) - curl(curl(u))
"""
T = get_function_space(space)
V = VectorSpace(T)
u = TrialFunction(V)
v = _TestFunction(V)
du = div(grad(u))
dv = grad(div(u)) - curl(curl(u))
u_hat = Function(V)
u_hat[:] = np.random.random(u_hat.shape) + np.random.random(u_hat.shape)*1j
A0 = inner(v, du)
A1 = inner(v, dv)
a0 = BlockMatrix(A0)
a1 = BlockMatrix(A1)
b0 = Function(V)
b1 = Function(V)
b0 = a0.matvec(u_hat, b0)
b1 = a1.matvec(u_hat, b1)
assert np.linalg.norm(b0-b1) < 1e-8
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_helmholtz2D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9)
SD = FunctionSpace(N[axis], family=family, bc=(0, 0))
K1 = FunctionSpace(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 = shenfun.TrialFunction(T)
v = shenfun.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_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 = FunctionSpace(8, "L", bc=(0, 0), scaled=True)
u = shenfun.TrialFunction(SD)
v = shenfun.TestFunction(SD)
mat = inner(grad(u), grad(v))
z = Function(SD, val=1)
c = mat * z
assert np.allclose(c[:6], 1)
def test_helmholtz3D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9, 10)
SD = shenfun.Basis(N[allaxes3D[axis][0]], family=family, bc=(0, 0))
K1 = shenfun.Basis(N[allaxes3D[axis][1]], family='F', dtype='D')
K2 = shenfun.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 = shenfun.TensorProductSpace(subcomms, bases, axes=allaxes3D[axis])
u = shenfun.TrialFunction(T)
v = shenfun.TestFunction(T)
if family == 'chebyshev':
mat = shenfun.inner(v, shenfun.div(shenfun.grad(u)))
else:
mat = shenfun.inner(shenfun.grad(v), shenfun.grad(u))
H = la.Helmholtz(*mat)
H = la.Helmholtz(*mat)
u = shenfun.Function(T)
u[:] = np.random.random(u.shape) + 1j * np.random.random(u.shape)
f = shenfun.Function(T)
f = H.matvec(u, f)
f = H.matvec(u, f)
g0 = shenfun.Function(T)
g1 = shenfun.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))
def solve(self):
if self._dim == 2:
map_boundary_to_component_index = {
'right': 0,
'top': 1,
'left': 0,
'bottom': 1
}
map_boundary_to_start_end_index = {
'right': 1,
'top': 1,
'left': -1,
'bottom': -1
}
if self._traction_bcs:
boundary_traction_term = Function(self._V.get_orthogonal())
# b: boundary, c: component
for b, c, value in self._traction_bcs:
if self._dim == 2:
side_index = map_boundary_to_component_index[b]
bdry_basis_index = 1 if side_index == 0 else 0
boundary_basis = self._function_spaces[c][bdry_basis_index]
start_or_end_index = map_boundary_to_start_end_index[b]
# test function for boundary integral
v_boundary = sf.TestFunction(boundary_basis)
if isinstance(value, (float, int)):
trac = Array(boundary_basis, val=value)
else:
trac = Array(boundary_basis, buffer=value)
evaluate_on_boundary = self._function_spaces[c][
side_index].evaluate_basis_all(start_or_end_index)
project_traction = inner(trac, v_boundary)
if side_index == 0:
boundary_traction_term[c] += np.outer(
evaluate_on_boundary, project_traction)
elif side_index == 1:
boundary_traction_term[c] += np.outer(
project_traction, evaluate_on_boundary)
else:
raise ValueError()
M = sf.BlockMatrix(self._dw_int)
if self._traction_bcs:
self._dw_ext += boundary_traction_term
u_hat = M.solve(self._dw_ext)
return u_hat
def ComputeRHS(rhs, u_hat, solver,
H_hat, H_hat1, H_hat0, VFSp, FSTp, FCTp, VCp, work, K, K2,
u_dealias, curl_dealias, curl_hat, mat, la, vt, Sk, mask, **context):
"""Compute right hand side of Navier Stokes
Parameters
----------
rhs : array
The right hand side to be returned
u_hat : array
The velocity vector at current time
solver : module
The current solver module
Remaining args are extracted from context
"""
# Nonlinear convection term at current u_hat
H_hat = solver.conv(H_hat, u_hat, K, VFSp, VCp, FSTp, FCTp, work,
u_dealias, curl_dealias, curl_hat, mat, la)
# Assemble convection with Adams-Bashforth at time = n+1/2
H_hat0 = solver.assembleAB(H_hat0, H_hat, H_hat1)
if mask is not None:
H_hat0.mask_nyquist(mask)
# Assemble rhs
rhs_u, rhs_p = rhs
rhs_u[0] = mat.AB.matvec(u_hat[0], rhs_u[0])
rhs_u[1] = mat.AB.matvec(u_hat[1], rhs_u[1])
rhs_u[2] = mat.AB.matvec(u_hat[2], rhs_u[2])
# Convection
rhs_u[:] += 2./params.nu*inner(vt, H_hat0)
# Source
rhs_u[1] -= 2./params.nu*Sk[1]
return rhs
err += " & {:2.2e} & {:2.2e} ".format(errb/M, errs/M)
err += " \\\ "
print(err)
print("\hline")
for z in Z:
err = str(z)
for n in N:
errh = 0
vh = zeros(n)
sh = zeros(n)
alfa = sqrt(z**2+2.0/nu/dt)
basis = chebyshev.bases.ShenDirichletBasis(n, plan=True, quad='GC')
u = TrialFunction(basis)
v = TestFunction(basis)
A = inner(v, div(grad(u)))
A.axis = 0
B = inner(v, u)
B.axis = 0
HS = chebyshev.la.Helmholtz(A, B, array([1.]), array([alfa]))
for m in range(M):
u = random.randn(n)
u[-2:] = 0
vh = HS.matvec(u, vh)
sh = HS(sh, vh)
errh += max(abs(sh-u)) / max(abs(u))
err += " & {:2.2e} ".format(errh/M)
err += " \\\ "
print(err)
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())
# 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
u_hat = H(f_hat, u_hat) # Solve
uq = u_hat.backward()
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)))
assert np.allclose(uj, ua, atol=1e-5)
point = np.array([0.1, 0.2])
p = SD.eval(point, f_hat)
