def test_biharmonic2D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (16, 16)
SD = FunctionSpace(N[axis], family=family, bc='Biharmonic')
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(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_cylinder():
T = get_function_space('cylinder')
u = TrialFunction(T)
du = div(grad(u))
assert du.tolatex() == '\\frac{\\partial^2 u}{\\partial x^2 }+\\frac{1}{x}\\frac{\\partial u}{\\partial x }+\\frac{1}{x^{2}}\\frac{\\partial^2 u}{\\partial y^2 }+\\frac{\\partial^2 u}{\\partial z^2 }'
V = VectorSpace(T)
u = TrialFunction(V)
du = div(grad(u))
assert du.tolatex() == '\\left( \\frac{\\partial^2 u^{x}}{\\partial x^2 }+\\frac{1}{x}\\frac{\\partial u^{x}}{\\partial x }+\\frac{1}{x^{2}}\\frac{\\partial^2 u^{x}}{\\partial y^2 }- \\frac{2}{x}\\frac{\\partial u^{y}}{\\partial y }- \\frac{1}{x^{2}}u^{x}+\\frac{\\partial^2 u^{x}}{\\partial z^2 }\\right) \\mathbf{b}_{x} \\\\+\\left( \\frac{\\partial^2 u^{y}}{\\partial x^2 }+\\frac{3}{x}\\frac{\\partial u^{y}}{\\partial x }+\\frac{2}{x^{3}}\\frac{\\partial u^{x}}{\\partial y }+\\frac{1}{x^{2}}\\frac{\\partial^2 u^{y}}{\\partial y^2 }+\\frac{\\partial^2 u^{y}}{\\partial z^2 }\\right) \\mathbf{b}_{y} \\\\+\\left( \\frac{\\partial^2 u^{z}}{\\partial x^2 }+\\frac{1}{x}\\frac{\\partial u^{z}}{\\partial x }+\\frac{1}{x^{2}}\\frac{\\partial^2 u^{z}}{\\partial y^2 }+\\frac{\\partial^2 u^{z}}{\\partial z^2 }\\right) \\mathbf{b}_{z} \\\\'
def test_quasiGalerkin(basis):
N = 40
T = FunctionSpace(N, 'C')
S = FunctionSpace(N, 'C', basis=basis)
u = TrialFunction(S)
v = TestFunction(T)
A = inner(v, div(grad(u)))
B = inner(v, u)
Q = chebyshev.quasi.QIGmat(N)
A = Q*A
B = Q*B
M = B-A
sol = la.Solve(M, S)
f_hat = inner(v, Array(T, buffer=fe))
f_hat[:-2] = Q.diags('csc')*f_hat[:-2]
u_hat = Function(S)
u_hat = sol(f_hat, u_hat)
uj = u_hat.backward()
ua = Array(S, buffer=ue)
bb = Q*np.ones(N)
assert abs(np.sum(bb[:8])) < 1e-8
if S.boundary_condition().lower() == 'neumann':
xj, wj = S.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_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_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 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_helmholtz2D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9)
SD = shenfun.Basis(N[axis], family=family, bc=(0, 0))
K1 = shenfun.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 = shenfun.TensorProductSpace(subcomms, bases, axes=allaxes2D[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 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], 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['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_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_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 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_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_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 test_inner():
v = shenfun.TestFunction(TT)
u = shenfun.TrialFunction(TT)
p = shenfun.TrialFunction(T)
q = shenfun.TestFunction(T)
A = inner(div(u), div(v))
B = inner(grad(u), grad(v))
C = inner(q, div(u))
D = inner(curl(v), curl(u))
E = inner(grad(q), grad(p))
F = inner(v, grad(div(u)))
wq = shenfun.TrialFunction(VT)
w, q = wq
hf = shenfun.TestFunction(VT)
h, f = hf
G = inner(h, div(grad(w)))
H = inner(f, div(div(grad(w))))
I = inner(h, curl(w))
J = inner(curl(h), curl(w))
K = inner(h, grad(div(w)))
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 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 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 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])
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_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_eval_expression2(fam):
import sympy as sp
from shenfun import div, grad
x, y = sp.symbols('x,y')
B0 = FunctionSpace(16, fam, domain=(-2, 2))
B1 = FunctionSpace(17, fam, domain=(-1, 3))
T = TensorProductSpace(comm, (B0, B1))
f = sp.sin(x)+sp.sin(y)
dfx = f.diff(x, 2) + f.diff(y, 2)
fa = Function(T, buffer=f)
dfe = div(grad(fa))
dfa = project(dfe, T)
xy = np.array([[0.25, 0.5, 0.75],
[0.25, 0.5, 0.75]])
f0 = lambdify((x, y), dfx)(*xy)
f1 = dfe.eval(xy)
f2 = dfa.eval(xy)
assert np.allclose(f0, f1, 1e-7)
assert np.allclose(f1, f2, 1e-7)
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())
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()
# Compare with analytical solution
uj = ul(*X)
print(abs(uj - uq).max())
assert np.allclose(uj, uq)
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
f_hat[0] += 0.5*np.pi*alpha*b
elif family == 'legendre':
f_hat[0] += alpha*a
f_hat[0] -= alpha*b
f_hat[:-2] = A.solve(f_hat[:-2])
f_hat[-2] = a
f_hat[-1] = b
uj = SD.backward(f_hat)
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() uh = uq.forward() # Compare with analytical solution uj = Array(T, buffer=ue) assert np.allclose(uj, uq) if 'pytest' not in os.environ:
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)
K1 = R2CBasis(N[1])
T = TensorProductSpace(comm, (SD, K1), axes=(0, 1))
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)
# Get left hand side of Helmholtz equation
if basis == 'chebyshev':
matrices = inner(v, alpha*u - div(grad(u)))
else:
matrices = inner(grad(v), grad(u)) # Both ADDmat and BDDmat
B = inner(v, alpha*u)
matrices['BDDmat'] += B
# 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 = Function(T, False)
uq = T.backward(u_hat, uq)
# Compare with analytical solution
# Size of discretization
N = int(sys.argv[-2])
SD = FunctionSpace(N, family=family, bc='NeumannDirichlet', 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)
u_hat = A.solve(f_hat, u_hat)
uj = u_hat.backward()
uh = uj.forward()
# Compare with analytical solution
ua = Array(SD, buffer=ue)
print("Error=%2.16e" % (np.sqrt(dx((uj - ua)**2))))
assert np.allclose(uj, ua)
if 'pytest' not in os.environ:
import matplotlib.pyplot as plt
plt.plot(SD.mesh(), uj, 'b', SD.mesh(), ua, 'r')
plt.show()
