def test_Helmholtz_matvec(quad):
M = 2*N
SD = ShenDirichletBasis(M, quad=quad)
kx = 11
uj = np.random.randn(M)
u_hat = np.zeros(M)
u_hat = SD.forward(uj, u_hat)
uj = SD.backward(u_hat, uj)
B = inner_product((SD, 0), (SD, 0))
A = inner_product((SD, 0), (SD, 2))
AB = HelmholtzCoeff(M, 1, kx**2, SD.quad)
u1 = np.zeros(M)
u1 = SD.forward(uj, u1)
c0 = np.zeros_like(u1)
c1 = np.zeros_like(u1)
c = A.matvec(u1, c0)+kx**2*B.matvec(u1, c1)
b = np.zeros(M)
#LUsolve.Mult_Helmholtz_1D(M, SD.quad=="GL", 1, kx**2, u1, b)
b = AB.matvec(u1, b)
#from IPython import embed; embed()
assert np.allclose(c, b)
b = np.zeros((M, 4, 4), dtype=np.complex)
u1 = u1.repeat(16).reshape((M, 4, 4)) +1j*u1.repeat(16).reshape((M, 4, 4))
kx = np.zeros((1, 4, 4))+kx
#LUsolve.Mult_Helmholtz_3D_complex(M, SD.quad=="GL", 1.0, kx**2, u1, b)
AB = HelmholtzCoeff(M, 1, kx**2, SD.quad)
b = AB.matvec(u1, b)
assert np.linalg.norm(b[:, 2, 2].real - c)/(M*16) < 1e-12
assert np.linalg.norm(b[:, 2, 2].imag - c)/(M*16) < 1e-12
def test_Biharmonic(quad):
M = 128
SB = ShenBiharmonicBasis(M, quad=quad)
x = Symbol("x")
u = sin(6*pi*x)**2
a = 1.0
b = 1.0
f = -u.diff(x, 4) + a*u.diff(x, 2) + b*u
ul = lambdify(x, u, 'numpy')
fl = lambdify(x, f, 'numpy')
points, _ = SB.points_and_weights(M)
uj = ul(points)
fj = fl(points)
A = inner_product((SB, 0), (SB, 4))
B = inner_product((SB, 0), (SB, 0))
C = inner_product((SB, 0), (SB, 2))
AA = -A.diags() + C.diags() + B.diags()
f_hat = np.zeros(M)
f_hat = SB.scalar_product(fj, f_hat)
u_hat = np.zeros(M)
u_hat[:-4] = la.spsolve(AA, f_hat[:-4])
u1 = np.zeros(M)
u1 = SB.backward(u_hat, u1)
#from IPython import embed; embed()
assert np.allclose(u1, uj)
def test_Mult_CTD_3D(quad):
SD = ShenDirichlet(N, quad=quad)
SD.plan((N, 4, 4), 0, np.complex, {})
C = inner_product((SD.CT, 0), (SD, 1))
B = inner_product((SD.CT, 0), (SD.CT, 0))
vk = np.random.random((N, 4, 4)) + np.random.random((N, 4, 4)) * 1j
wk = np.random.random((N, 4, 4)) + np.random.random((N, 4, 4)) * 1j
bv = np.zeros((N, 4, 4), dtype=np.complex)
bw = np.zeros((N, 4, 4), dtype=np.complex)
vk0 = np.zeros((N, 4, 4), dtype=np.complex)
wk0 = np.zeros((N, 4, 4), dtype=np.complex)
cv = np.zeros((N, 4, 4), dtype=np.complex)
cw = np.zeros((N, 4, 4), dtype=np.complex)
vk0 = SD.forward(vk, vk0)
vk = SD.backward(vk0, vk)
vk0 = SD.forward(vk, vk0)
wk0 = SD.forward(wk, wk0)
wk = SD.backward(wk0, wk)
wk0 = SD.forward(wk, wk0)
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 test_Mult_CTD_3D(quad):
SD = ShenDirichletBasis(N, quad=quad)
SD.plan((N, 4, 4), 0, np.complex, {})
C = inner_product((SD.CT, 0), (SD, 1))
B = inner_product((SD.CT, 0), (SD.CT, 0))
vk = np.random.random((N, 4, 4))+np.random.random((N, 4, 4))*1j
wk = np.random.random((N, 4, 4))+np.random.random((N, 4, 4))*1j
bv = np.zeros((N, 4, 4), dtype=np.complex)
bw = np.zeros((N, 4, 4), dtype=np.complex)
vk0 = np.zeros((N, 4, 4), dtype=np.complex)
wk0 = np.zeros((N, 4, 4), dtype=np.complex)
cv = np.zeros((N, 4, 4), dtype=np.complex)
cw = np.zeros((N, 4, 4), dtype=np.complex)
vk0 = SD.forward(vk, vk0)
vk = SD.backward(vk0, vk)
vk0 = SD.forward(vk, vk0)
wk0 = SD.forward(wk, wk0)
wk = SD.backward(wk0, wk)
wk0 = SD.forward(wk, wk0)
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 test_Mult_Div():
SD = ShenDirichlet(N, "GC")
SN = ShenNeumann(N, "GC")
SD.plan(N, 0, np.complex, {})
SN.plan(N, 0, np.complex, {})
Cm = inner_product((SN, 0), (SD, 1))
Bm = inner_product((SN, 0), (SD, 0))
uk = np.random.randn((N)) + np.random.randn((N)) * 1j
vk = np.random.randn((N)) + np.random.randn((N)) * 1j
wk = np.random.randn((N)) + np.random.randn((N)) * 1j
b = np.zeros(N, dtype=np.complex)
uk0 = np.zeros(N, dtype=np.complex)
vk0 = np.zeros(N, dtype=np.complex)
wk0 = np.zeros(N, dtype=np.complex)
uk0 = SD.forward(uk, uk0)
uk = SD.backward(uk0, uk)
uk0 = SD.forward(uk, uk0)
vk0 = SD.forward(vk, vk0)
vk = SD.backward(vk0, vk)
vk0 = SD.forward(vk, vk0)
wk0 = SD.forward(wk, wk0)
wk = SD.backward(wk0, wk)
wk0 = SD.forward(wk, wk0)
LUsolve.Mult_Div_1D(N, 7, 7, uk0[:N - 2], vk0[:N - 2], wk0[:N - 2],
b[1:N - 2])
uu = np.zeros_like(uk0)
v0 = np.zeros_like(vk0)
w0 = np.zeros_like(wk0)
uu = Cm.matvec(uk0, uu)
uu += 1j * 7 * Bm.matvec(vk0, v0) + 1j * 7 * Bm.matvec(wk0, w0)
#from IPython import embed; embed()
assert np.allclose(uu, b)
uk0 = uk0.repeat(4 * 4).reshape(
(N, 4, 4)) + 1j * uk0.repeat(4 * 4).reshape((N, 4, 4))
vk0 = vk0.repeat(4 * 4).reshape(
(N, 4, 4)) + 1j * vk0.repeat(4 * 4).reshape((N, 4, 4))
wk0 = wk0.repeat(4 * 4).reshape(
(N, 4, 4)) + 1j * wk0.repeat(4 * 4).reshape((N, 4, 4))
b = np.zeros((N, 4, 4), dtype=np.complex)
m = np.zeros((4, 4)) + 7
n = np.zeros((4, 4)) + 7
LUsolve.Mult_Div_3D(N, m, n, uk0[:N - 2], vk0[:N - 2], wk0[:N - 2],
b[1:N - 2])
uu = np.zeros_like(uk0)
v0 = np.zeros_like(vk0)
w0 = np.zeros_like(wk0)
uu = Cm.matvec(uk0, uu)
uu += 1j * 7 * Bm.matvec(vk0, v0) + 1j * 7 * Bm.matvec(wk0, w0)
assert np.allclose(uu, b)
def __init__(self, N, alfa, basis):
# Prepare LU Helmholtz solver for velocity
self.N = N
self.alfa = alfa
self.basis = basis
self.neumann = True if isinstance(basis, bases.ShenNeumannBasis) else False
quad = basis.quad
M = (N-4)//2 if self.neumann else (N-3)//2
self.s = basis.slice()
if hasattr(alfa, "__len__"):
Ny, Nz = alfa.shape
self.u0 = zeros((2, M+1, Ny, Nz), float) # Diagonal entries of U
self.u1 = zeros((2, M, Ny, Nz), float) # Diagonal+1 entries of U
self.u2 = zeros((2, M-1, Ny, Nz), float) # Diagonal+2 entries of U
self.L = zeros((2, M, Ny, Nz), float) # The single nonzero row of L
LUsolve.LU_Helmholtz_3D(N, self.neumann, quad=="GL", self.alfa, self.u0, self.u1, self.u2, self.L)
else:
self.u0 = zeros((2, M+1), float) # Diagonal entries of U
self.u1 = zeros((2, M), float) # Diagonal+1 entries of U
self.u2 = zeros((2, M-1), float) # Diagonal+2 entries of U
self.L = zeros((2, M), float) # The single nonzero row of L
LUsolve.LU_Helmholtz_1D(N, self.neumann, quad=="GL", self.alfa, self.u0, self.u1, self.u2, self.L)
if not self.neumann:
self.B = inner_product((basis, 0), (basis, 0))
self.A = inner_product((basis, 0), (basis, 2))
def test_Mult_CTD(quad):
SD = ShenDirichletBasis(N, quad=quad)
SD.plan(N, 0, np.complex, {})
C = inner_product((SD.CT, 0), (SD, 1))
B = inner_product((SD.CT, 0), (SD.CT, 0))
vk = np.random.randn((N))+np.random.randn((N))*1j
wk = np.random.randn((N))+np.random.randn((N))*1j
bv = np.zeros(N, dtype=np.complex)
bw = np.zeros(N, dtype=np.complex)
vk0 = np.zeros(N, dtype=np.complex)
wk0 = np.zeros(N, dtype=np.complex)
cv = np.zeros(N, dtype=np.complex)
cw = np.zeros(N, dtype=np.complex)
vk0 = SD.forward(vk, vk0)
vk = SD.backward(vk0, vk)
vk0 = SD.forward(vk, vk0)
wk0 = SD.forward(wk, wk0)
wk = SD.backward(wk0, wk)
wk0 = SD.forward(wk, wk0)
LUsolve.Mult_CTD_1D(N, vk0, wk0, bv, bw)
cv = np.zeros_like(vk0)
cw = np.zeros_like(wk0)
cv = C.matvec(vk0, cv)
cw = C.matvec(wk0, cw)
cv /= B[0]
cw /= B[0]
assert np.allclose(cv, bv)
assert np.allclose(cw, bw)
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 test_CDDmat(quad):
M = 128
SD = cbases.ShenDirichletBasis(M, quad=quad)
u = (1 - x**2) * sin(np.pi * 6 * x)
dudx = u.diff(x, 1)
points, weights = SD.points_and_weights(M)
ul = lambdify(x, u, 'numpy')
dudx_l = lambdify(x, dudx, 'numpy')
dudx_j = dudx_l(points)
uj = ul(points)
u_hat = shenfun.Function(SD)
u_hat = SD.forward(uj, u_hat)
uj = SD.backward(u_hat, uj)
u_hat = SD.forward(uj, u_hat)
uc_hat = shenfun.Function(SD)
uc_hat = SD.CT.forward(uj, uc_hat)
dudx_j = SD.CT.fast_derivative(uj)
Cm = inner_product((SD, 0), (SD, 1))
B = inner_product((SD, 0), (SD, 0))
TDMASolver = TDMA(B)
cs = np.zeros_like(u_hat)
cs = Cm.matvec(u_hat, cs)
# Should equal (but not exact so use extra resolution)
cs2 = np.zeros(M)
cs2 = SD.scalar_product(dudx_j, cs2)
s = SD.slice()
assert np.allclose(cs[s], cs2[s])
cs = TDMASolver(cs)
du = np.zeros(M)
du = SD.backward(cs, du)
assert np.linalg.norm(du[s] - dudx_j[s]) / M < 1e-10
# Multidimensional version
u3_hat = u_hat.repeat(4 * 4).reshape(
(M, 4, 4)) + 1j * u_hat.repeat(4 * 4).reshape((M, 4, 4))
cs = np.zeros_like(u3_hat)
cs = Cm.matvec(u3_hat, cs)
cs2 = np.zeros((M, 4, 4), dtype=np.complex)
du3 = dudx_j.repeat(4 * 4).reshape(
(M, 4, 4)) + 1j * dudx_j.repeat(4 * 4).reshape((M, 4, 4))
SD.plan((M, 4, 4), 0, np.complex, {})
cs2 = SD.scalar_product(du3, cs2)
assert np.allclose(cs[s], cs2[s], 1e-10)
cs = TDMASolver(cs)
d3 = np.zeros((M, 4, 4), dtype=np.complex)
d3 = SD.backward(cs, d3)
assert np.linalg.norm(du3[s] - d3[s]) / (M * 16) < 1e-10
def __init__(self, N, alfa, beta, quad="GL"):
"""alfa*ADD + beta*BDD
"""
self.quad = quad
self.shape = (N - 2, N - 2)
SD = bases.ShenDirichletBasis(N, quad)
self.B = inner_product((SD, 0), (SD, 0))
self.A = inner_product((SD, 0), (SD, 2))
self.alfa = alfa
self.beta = beta
def __init__(self, N, a0, alfa, beta, quad="GL"):
self.quad = quad
self.shape = (N - 4, N - 4)
SB = bases.ShenBiharmonicBasis(N, quad)
self.S = inner_product((SB, 0), (SB, 4))
self.B = inner_product((SB, 0), (SB, 0))
self.A = inner_product((SB, 0), (SB, 2))
self.a0 = a0
self.alfa = alfa
self.beta = beta
def test_Mult_Div():
SD = ShenDirichletBasis(N, "GC")
SN = ShenNeumannBasis(N, "GC")
SD.plan(N, 0, np.complex, {})
SN.plan(N, 0, np.complex, {})
Cm = inner_product((SN, 0), (SD, 1))
Bm = inner_product((SN, 0), (SD, 0))
uk = np.random.randn((N))+np.random.randn((N))*1j
vk = np.random.randn((N))+np.random.randn((N))*1j
wk = np.random.randn((N))+np.random.randn((N))*1j
b = np.zeros(N, dtype=np.complex)
uk0 = np.zeros(N, dtype=np.complex)
vk0 = np.zeros(N, dtype=np.complex)
wk0 = np.zeros(N, dtype=np.complex)
uk0 = SD.forward(uk, uk0)
uk = SD.backward(uk0, uk)
uk0 = SD.forward(uk, uk0)
vk0 = SD.forward(vk, vk0)
vk = SD.backward(vk0, vk)
vk0 = SD.forward(vk, vk0)
wk0 = SD.forward(wk, wk0)
wk = SD.backward(wk0, wk)
wk0 = SD.forward(wk, wk0)
LUsolve.Mult_Div_1D(N, 7, 7, uk0[:N-2], vk0[:N-2], wk0[:N-2], b[1:N-2])
uu = np.zeros_like(uk0)
v0 = np.zeros_like(vk0)
w0 = np.zeros_like(wk0)
uu = Cm.matvec(uk0, uu)
uu += 1j*7*Bm.matvec(vk0, v0) + 1j*7*Bm.matvec(wk0, w0)
#from IPython import embed; embed()
assert np.allclose(uu, b)
uk0 = uk0.repeat(4*4).reshape((N, 4, 4)) + 1j*uk0.repeat(4*4).reshape((N, 4, 4))
vk0 = vk0.repeat(4*4).reshape((N, 4, 4)) + 1j*vk0.repeat(4*4).reshape((N, 4, 4))
wk0 = wk0.repeat(4*4).reshape((N, 4, 4)) + 1j*wk0.repeat(4*4).reshape((N, 4, 4))
b = np.zeros((N, 4, 4), dtype=np.complex)
m = np.zeros((4, 4))+7
n = np.zeros((4, 4))+7
LUsolve.Mult_Div_3D(N, m, n, uk0[:N-2], vk0[:N-2], wk0[:N-2], b[1:N-2])
uu = np.zeros_like(uk0)
v0 = np.zeros_like(vk0)
w0 = np.zeros_like(wk0)
uu = Cm.matvec(uk0, uu)
uu += 1j*7*Bm.matvec(vk0, v0) + 1j*7*Bm.matvec(wk0, w0)
assert np.allclose(uu, b)
def __init__(self, K, alfa, beta, quad="GL"):
"""alfa*ADD + beta*BDD
"""
self.quad = quad
N = self.N = K.shape[0] - 2
self.shape = (N, N)
SD = bases.ShenDirichletBasis(N + 2, quad)
self.B = inner_product((SD, 0), (SD, 0))
self.A = inner_product((SD, 0), (SD, 2))
self.alfa = alfa
self.beta = beta
def __init__(self, N, alfa, beta, axis, quad="GL"):
"""alfa*ADD + beta*BDD
"""
self.quad = quad
self.shape = (N-2, N-2)
SD = bases.ShenDirichletBasis(N, quad)
self.B = inner_product((SD, 0), (SD, 0))
self.A = inner_product((SD, 0), (SD, 2))
self.axis = axis
self.alfa = np.broadcast_to(alfa, beta.shape).copy()
self.beta = beta
def __init__(self, N, alfa, beta, axis, quad="GL"):
"""alfa*ADD + beta*BDD
"""
self.quad = quad
self.shape = (N - 2, N - 2)
SD = bases.ShenDirichletBasis(N, quad)
self.B = inner_product((SD, 0), (SD, 0))
self.A = inner_product((SD, 0), (SD, 2))
self.axis = axis
self.alfa = np.broadcast_to(alfa, beta.shape).copy()
self.beta = beta
def __init__(self, N, a0, alfa, beta, axis, quad="GL"):
self.quad = quad
self.shape = (N-4, N-4)
SB = bases.ShenBiharmonicBasis(N, quad)
self.S = inner_product((SB, 0), (SB, 4))
self.B = inner_product((SB, 0), (SB, 0))
self.A = inner_product((SB, 0), (SB, 2))
self.axis = axis
self.a0 = a0
self.alfa = alfa
self.beta = beta
def test_massmatrices(test, trial, quad):
test = test(N, quad=quad)
trial = trial(N, quad=quad)
f_hat = np.zeros(N)
fj = np.random.random(N)
f_hat = trial.forward(fj, f_hat)
fj = trial.backward(f_hat, fj)
BBD = inner_product((test, 0), (trial, 0))
f_hat = trial.forward(fj, f_hat)
u2 = np.zeros_like(f_hat)
u2 = BBD.matvec(f_hat, u2)
u0 = np.zeros(N)
u0 = test.scalar_product(fj, u0)
s = test.slice()
assert np.allclose(u0[s], u2[s], rtol=1e-5, atol=1e-6)
# Multidimensional version
fj = fj.repeat(N * N).reshape((N, N, N)) + 1j * fj.repeat(N * N).reshape(
(N, N, N))
f_hat = f_hat.repeat(N * N).reshape(
(N, N, N)) + 1j * f_hat.repeat(N * N).reshape((N, N, N))
test.plan((N, ) * 3, 0, np.complex, {})
test.tensorproductspace = ABC(3, test.coors)
u0 = np.zeros((N, ) * 3, dtype=np.complex)
u0 = test.scalar_product(fj, u0)
u2 = np.zeros_like(f_hat)
u2 = BBD.matvec(f_hat, u2)
assert np.linalg.norm(u2[s] - u0[s]) / (N * N * N) < 1e-8
del BBD
def test_axis(ST, quad, axis):
kwargs = {}
if not ST.family() == 'fourier':
kwargs['quad'] = quad
ST = ST(N, **kwargs)
points, weights = ST.points_and_weights(N)
f_hat = shenfun.Function(ST)
f_hat[:] = np.random.random(f_hat.shape[0])
B = inner_product((ST, 0), (ST, 0))
c = shenfun.Function(ST)
c = B.solve(f_hat, c)
# Multidimensional version
f0 = shenfun.Array(ST)
bc = [np.newaxis,]*3
bc[axis] = slice(None)
ST.tensorproductspace = ABC(3, ST.coors)
ST.plan((N,)*3, axis, f0.dtype, {})
if ST.has_nonhomogeneous_bcs:
ST.bc.set_tensor_bcs(ST, ST) # To set Dirichlet boundary conditions on multidimensional array
ck = shenfun.Function(ST)
fk = np.broadcast_to(f_hat[tuple(bc)], ck.shape).copy()
ck = B.solve(fk, ck, axis=axis)
cc = [1,]*3
cc[axis] = slice(None)
assert np.allclose(ck[tuple(cc)], c, rtol=1e-5, atol=1e-6)
def test_jmatvec(b0, b1, quad, format, dim, k0, k1):
"""Testq matrix-vector product"""
global c, c1
b0 = b0(N, quad=quad)
b1 = b1(N, quad=quad)
mat = inner_product((b0, k0), (b1, k1))
c = mat.matvec(a, c, format='csr')
c1 = mat.matvec(a, c1, format=format)
assert np.allclose(c, c1)
for axis in range(0, dim):
b, d, d1 = work[dim]
d.fill(0)
d1.fill(0)
d = mat.matvec(b, d, format='csr', axis=axis)
d1 = mat.matvec(b, d1, format=format, axis=axis)
assert np.allclose(d, d1)
# Test multidimensional with axis equals 1D case
d1.fill(0)
bc = [
np.newaxis,
] * dim
bc[axis] = slice(None)
fj = np.broadcast_to(a[tuple(bc)], (N, ) * dim).copy()
d1 = mat.matvec(fj, d1, format=format, axis=axis)
cc = [
0,
] * dim
cc[axis] = slice(None)
assert np.allclose(c, d1[tuple(cc)])
def test_massmatrices(test, trial, quad):
test = test(N, quad=quad, plan=True)
trial = trial(N, quad=quad, plan=True)
f_hat = np.zeros(N)
fj = np.random.random(N)
f_hat = trial.forward(fj, f_hat)
fj = trial.backward(f_hat, fj)
BBD = inner_product((test, 0), (trial, 0))
f_hat = trial.forward(fj, f_hat)
u2 = np.zeros_like(f_hat)
u2 = BBD.matvec(f_hat, u2)
u0 = np.zeros(N)
u0 = test.scalar_product(fj, u0)
s = test.slice()
#from IPython import embed; embed()
assert np.allclose(u0[s], u2[s])
# Multidimensional version
fj = fj.repeat(N * N).reshape((N, N, N)) + 1j * fj.repeat(N * N).reshape(
(N, N, N))
f_hat = f_hat.repeat(N * N).reshape(
(N, N, N)) + 1j * f_hat.repeat(N * N).reshape((N, N, N))
test.plan((N, ) * 3, 0, np.complex, {})
u0 = np.zeros((N, ) * 3, dtype=np.complex)
u0 = test.scalar_product(fj, u0)
u2 = np.zeros_like(f_hat)
u2 = BBD.matvec(f_hat, u2)
assert np.linalg.norm(u2[s] - u0[s]) / (N * N * N) < 1e-12
del BBD
def interp(self, y, eigval=1, same_mesh=False, verbose=False):
"""Interpolate solution eigenvector and it's derivative onto y
"""
N = self.N
nx, eigval = self.get_eigval(eigval, verbose)
phi_hat = np.zeros(N, np.complex)
phi_hat[:-4] = np.squeeze(self.eigvectors[:, nx])
if same_mesh:
phi = np.zeros_like(phi_hat)
dphidy = np.zeros_like(phi_hat)
if self.SB is None:
self.SB = ShenBiharmonicBasis(N, quad=self.quad, plan=True)
self.SD = ShenDirichletBasis(N, quad=self.quad, plan=True)
self.CDB = inner_product((self.SD, 0), (self.SB, 1))
phi = self.SB.ifst(phi_hat, phi)
dphidy_hat = self.CDB.matvec(phi_hat)
dphidy_hat = self.SD.apply_inverse_mass(dphidy_hat)
dphidy = self.SD.backward(dphidy_hat, dphidy)
else:
# Recompute interpolation matrices if necessary
if not len(self.P4) == len(y):
SB = ShenBiharmonicBasis(N, quad=self.quad)
V = self.V = SB.vandermonde(y)
P4 = self.P4 = SB.get_vandermonde_basis(V)
T4x = self.T4x = SB.get_vandermonde_basis_derivative(V, 1)
phi = np.dot(self.P4, phi_hat)
dphidy = np.dot(self.T4x, phi_hat)
return phi, dphidy
def test_axis(ST, quad, axis):
kwargs = {'plan': True}
if not ST.family() == 'fourier':
kwargs['quad'] = quad
ST = ST(N, **kwargs)
points, weights = ST.points_and_weights(N)
f_hat = shenfun.Function(ST)
f_hat[:] = np.random.random(f_hat.shape[0])
B = inner_product((ST, 0), (ST, 0))
c = shenfun.Function(ST)
c = B.solve(f_hat, c)
# Multidimensional version
f0 = shenfun.Array(ST)
bc = [
np.newaxis,
] * 3
bc[axis] = slice(None)
ST.plan((N, ) * 3, axis, f0.dtype, {})
if hasattr(ST, 'bc'):
ST.bc.set_tensor_bcs(
ST
) # To set Dirichlet boundary conditions on multidimensional array
ck = shenfun.Function(ST)
fk = np.broadcast_to(f_hat[bc], ck.shape).copy()
ck = B.solve(fk, ck, axis=axis)
cc = [
0,
] * 3
cc[axis] = slice(None)
assert np.allclose(ck[cc], c)
def test_ASDSDmat(ST, quad):
M = 2*N
ST = ST(M, quad=quad)
u = (1-x**2)*sin(np.pi*x)
f = u.diff(x, 2)
ul = lambdify(x, u, 'numpy')
fl = lambdify(x, f, 'numpy')
points, weights = ST.points_and_weights(M)
uj = ul(points)
fj = fl(points)
s = ST.slice()
if ST.family() == 'chebyshev':
A = inner_product((ST, 0), (ST, 2))
else:
A = inner_product((ST, 1), (ST, 1))
f_hat = np.zeros(M)
f_hat = ST.scalar_product(fj, f_hat)
if ST.family() == 'legendre':
f_hat *= -1
# Test both solve interfaces
c_hat = f_hat.copy()
c_hat = A.solve(c_hat)
u_hat = np.zeros_like(f_hat)
u_hat = A.solve(f_hat, u_hat)
assert np.allclose(c_hat[s], u_hat[s], rtol=1e-5, atol=1e-6)
u0 = np.zeros(M)
u0 = ST.backward(u_hat, u0)
assert np.allclose(u0, uj)
u1 = np.zeros(M)
u1 = ST.forward(uj, u1)
c = np.zeros_like(u1)
c = A.matvec(u1, c)
s = ST.slice()
assert np.allclose(c[s], f_hat[s], rtol=1e-5, atol=1e-6)
# Multidimensional
c_hat = f_hat.copy()
c_hat = c_hat.repeat(M).reshape((M, M)).transpose().copy()
c_hat = A.solve(c_hat, axis=1)
assert np.allclose(c_hat[0, s], u_hat[s], rtol=1e-5, atol=1e-6)
def get_pressure(context, solver):
c = context
dt = solver.params.dt
FST = c.FST
U = solver.get_velocity(**c)
c.U0[0] = FST.backward(c.U_hat0[0], c.U0[0], c.SB)
for i in range(1, 3):
c.U0[i] = FST.backward(c.U_hat0[i], c.U0[i], c.ST)
H_hat = solver.get_convection(**c)
Hx = zeros(FST.real_shape(), dtype=float)
Hx = FST.backward(H_hat[0], Hx, c.ST)
Hx -= 1. / dt * (c.U[0] - c.U0[0])
rhs_hat = zeros(FST.complex_shape(), dtype=complex)
w0 = zeros(FST.complex_shape(), dtype=complex)
ATB = inner_product((c.CT, 0), (c.SB, 2))
BTB = inner_product((c.CT, 0), (c.SB, 0))
rhs_hat = ATB.matvec(c.U_hat[0] + c.U_hat0[0], rhs_hat)
rhs_hat -= c.K2 * BTB.matvec(c.U_hat[0] + c.U_hat0[0], w0)
rhs_hat *= 0.5 * context.nu
rhs_hat += FST.scalar_product(Hx, w0, c.CT)
CT = inner_product((c.CT, 0), (c.CT, 1))
# Should implement fast solver. Just a backwards substitution
A = CT.diags().toarray()
A[-1, 0] = 1
a_i = np.linalg.inv(A)
p_hat = zeros(FST.complex_shape(), dtype=complex)
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 = zeros(FST.real_shape(), dtype=float)
p = FST.backward(p_hat, p, c.CT)
uu = np.sum((0.5 * (c.U + c.U0))**2, 0)
uu *= 0.5
return p - uu + 3. / 16.
def assemble(self):
N = self.N
SB = ShenBiharmonicBasis(N, quad=self.quad)
SB.plan((N, N), 0, np.float, {})
x, w = self.x, self.w = SB.points_and_weights(N)
V = SB.vandermonde(x)
# Trial function
P4 = SB.get_vandermonde_basis(V)
# Second derivatives
T2x = SB.get_vandermonde_basis_derivative(V, 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_SBBmat(SB, quad):
M = 72
SB = SB(M, quad=quad)
u = sin(4 * pi * x)**2
f = u.diff(x, 4)
ul = lambdify(x, u, 'numpy')
fl = lambdify(x, f, 'numpy')
points, weights = SB.points_and_weights(M)
uj = ul(points)
fj = fl(points)
if isinstance(SB, shenfun.chebyshev.bases.ChebyshevBase):
A = inner_product((SB, 0), (SB, 4))
else:
A = inner_product((SB, 2), (SB, 2))
f_hat = np.zeros(M)
f_hat = SB.scalar_product(fj, f_hat)
u_hat = np.zeros(M)
u_hat = A.solve(f_hat, u_hat)
u0 = np.zeros(M)
u0 = SB.backward(u_hat, u0)
assert np.allclose(u0, uj, rtol=1e-5, atol=1e-6)
u1 = np.zeros(M)
u1 = SB.forward(uj, u1)
c = np.zeros_like(u1)
c = A.matvec(u1, c)
assert np.all(abs(c - f_hat) / c.max() < 1e-10)
# Multidimensional
c2 = (c.repeat(16).reshape((M, 4, 4)) + 1j * c.repeat(16).reshape(
(M, 4, 4)))
u1 = (u1.repeat(16).reshape((M, 4, 4)) + 1j * u1.repeat(16).reshape(
(M, 4, 4)))
c = np.zeros_like(u1)
c = A.matvec(u1, c)
assert np.allclose(c, c2)
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 interp(self,
y,
eigvals,
eigvectors,
eigval=1,
same_mesh=False,
verbose=False):
"""Interpolate solution eigenvector and it's derivative onto y
args:
y Interpolation points
eigvals All computed eigenvalues
eigvectors All computed eigenvectors
kwargs:
eigval The chosen eigenvalue, ranked with descending imaginary
part. The largest imaginary part is 1, the second
largest is 2, etc.
same_mesh Boolean. Whether or not to interpolate to the same
quadrature points as used for computing the eigenvectors
verbose Boolean. 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 same_mesh:
phi = np.zeros_like(phi_hat)
dphidy = np.zeros_like(phi_hat)
if self.SB is None:
self.SB = ShenBiharmonicBasis(N, quad=self.quad, plan=True)
self.SD = ShenDirichletBasis(N, quad=self.quad, plan=True)
self.CDB = inner_product((self.SD, 0), (self.SB, 1))
phi = self.SB.ifst(phi_hat, phi)
dphidy_hat = self.CDB.matvec(phi_hat)
dphidy_hat = self.SD.apply_inverse_mass(dphidy_hat)
dphidy = self.SD.backward(dphidy_hat, dphidy)
else:
# Recompute interpolation matrices if necessary
if not len(self.P4) == len(y):
SB = ShenBiharmonicBasis(N, quad=self.quad)
V = SB.vandermonde(y)
self.P4 = SB.get_vandermonde_basis(V)
self.T4x = SB.get_vandermonde_basis_derivative(V, 1)
phi = np.dot(self.P4, phi_hat)
dphidy = np.dot(self.T4x, phi_hat)
return eigval, phi, dphidy
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 test_Helmholtz2(quad):
M = 2 * N
SD = ShenDirichletBasis(M, quad=quad, plan=True)
kx = 12
points, weights = SD.points_and_weights(M)
uj = np.random.randn(M)
u_hat = np.zeros(M)
u_hat = SD.forward(uj, u_hat)
uj = SD.backward(u_hat, uj)
#from IPython import embed; embed()
A = inner_product((SD, 0), (SD, 2))
B = inner_product((SD, 0), (SD, 0))
s = SD.slice()
u1 = np.zeros(M)
u1 = SD.forward(uj, u1)
c0 = np.zeros_like(u1)
c1 = np.zeros_like(u1)
c = A.matvec(u1, c0) + kx**2 * B.matvec(u1, c1)
b = np.zeros(M)
H = Helmholtz(M, kx, SD)
b = H.matvec(u1, b)
#LUsolve.Mult_Helmholtz_1D(M, SD.quad=="GL", 1, kx**2, u1, b)
assert np.allclose(c, b)
b = np.zeros((M, 4, 4), dtype=np.complex)
u1 = u1.repeat(16).reshape((M, 4, 4)) + 1j * u1.repeat(16).reshape(
(M, 4, 4))
kx = np.zeros((4, 4)) + kx
H = Helmholtz(M, kx, SD)
b = H.matvec(u1, b)
#LUsolve.Mult_Helmholtz_3D_complex(M, SD.quad=="GL", 1.0, kx**2, u1, b)
assert np.linalg.norm(b[:, 2, 2].real - c) / (M * 16) < 1e-12
assert np.linalg.norm(b[:, 2, 2].imag - c) / (M * 16) < 1e-12
def test_jmatvec(b0, b1, quad, format, k0, k1):
"""Testq matrix-vector product"""
global c, c1
b0 = b0(N, quad=quad)
b1 = b1(N, quad=quad)
mat = inner_product((b0, k0), (b1, k1))
c = mat.matvec(a, c, format='csr')
formats = mat._matvec_methods + ['python', 'csr']
for format in formats:
c1 = mat.matvec(a, c1, format=format)
assert np.allclose(c, c1)
dim = 2
b, d, d1 = work[dim]
for axis in range(0, dim):
d = mat.matvec(b, d, format='csr', axis=axis)
for format in formats:
d1 = mat.matvec(b, d1, format=format, axis=axis)
assert np.allclose(d, d1)
def test_cmatvec(b0, b1, quad, k):
"""Test matrix-vector product"""
global c, c1
if quad == 'GL' and (b0 in cBasisGC or b1 in cBasisGC):
return
b0 = b0(N, quad=quad)
b1 = b1(N, quad=quad)
mat = inner_product((b0, 0), (b1, k))
formats = mat._matvec_methods + ['python', 'csr']
c = mat.matvec(a, c, format='csr')
for format in formats:
c1 = mat.matvec(a, c1, format=format)
assert np.allclose(c, c1)
for dim in (2, 3):
b, d, d1 = work[dim]
for axis in range(0, dim):
d = mat.matvec(b, d, format='csr', axis=axis)
for format in formats:
d1 = mat.matvec(b, d1, format=format, axis=axis)
assert np.allclose(d, d1)
def test_massmatrices(test, trial, quad):
test = test(N, quad=quad)
trial = trial(N, quad=quad)
f_hat = np.zeros(N)
fj = np.random.random(N)
f_hat = trial.forward(fj, f_hat)
fj = trial.backward(f_hat, fj)
BBD = inner_product((test, 0), (trial, 0))
f_hat = trial.forward(fj, f_hat)
u2 = np.zeros_like(f_hat)
u2 = BBD.matvec(f_hat, u2)
u0 = np.zeros(N)
u0 = test.scalar_product(fj, u0)
s = test.slice()
assert np.allclose(u0[s], u2[s], rtol=1e-5, atol=1e-6)
del BBD
def test_CXXmat(test, trial):
test = test(N)
trial = trial(N)
CT = cBasis[0](N)
Cm = inner_product((test, 0), (trial, 1))
S2 = Cm.trialfunction[0]
S1 = Cm.testfunction[0]
fj = np.random.randn(N)
# project to S2
f_hat = np.zeros(N)
f_hat = S2.forward(fj, f_hat)
fj = S2.backward(f_hat, fj)
# Check S1.scalar_product(f) equals Cm*S2.forward(f)
f_hat = S2.forward(fj, f_hat)
cs = np.zeros_like(f_hat)
cs = Cm.matvec(f_hat, cs)
df = CT.fast_derivative(fj)
cs2 = np.zeros(N)
cs2 = S1.scalar_product(df, cs2)
s = S1.slice()
assert np.allclose(cs[s], cs2[s])
# Multidimensional version
f_hat = f_hat.repeat(4 * 4).reshape(
(N, 4, 4)) + 1j * f_hat.repeat(4 * 4).reshape((N, 4, 4))
df = df.repeat(4 * 4).reshape((N, 4, 4)) + 1j * df.repeat(4 * 4).reshape(
(N, 4, 4))
cs = np.zeros_like(f_hat)
cs = Cm.matvec(f_hat, cs)
cs2 = np.zeros((N, 4, 4), dtype=np.complex)
S1.tensorproductspace = ABC(3)
S1.plan((N, 4, 4), 0, np.complex, {})
cs2 = S1.scalar_product(df, cs2)
assert np.allclose(cs[s], cs2[s])
def test_CXXmat(test, trial):
test = test(N)
trial = trial(N)
CT = cBasis[0](N)
Cm = inner_product((test, 0), (trial, 1))
S2 = Cm.trialfunction[0]
S1 = Cm.testfunction[0]
fj = shenfun.Array(S2, buffer=np.random.randn(N))
# project to S2
f_hat = fj.forward()
fj = f_hat.backward(fj)
# Check S1.scalar_product(f) equals Cm*S2.forward(f)
f_hat = S2.forward(fj, f_hat)
cs = np.zeros_like(f_hat)
cs = Cm.matvec(f_hat, cs)
df = shenfun.project(shenfun.grad(f_hat), CT).backward()
cs2 = np.zeros(N)
cs2 = S1.scalar_product(df, cs2)
s = S1.slice()
assert np.allclose(cs[s], cs2[s], rtol=1e-5, atol=1e-6)
# Multidimensional version
f_hat = f_hat.repeat(4 * 4).reshape(
(N, 4, 4)) + 1j * f_hat.repeat(4 * 4).reshape((N, 4, 4))
df = df.repeat(4 * 4).reshape((N, 4, 4)) + 1j * df.repeat(4 * 4).reshape(
(N, 4, 4))
cs = np.zeros_like(f_hat)
cs = Cm.matvec(f_hat, cs)
cs2 = np.zeros((N, 4, 4), dtype=np.complex)
S1.tensorproductspace = ABC(3, S1.coors)
S1.plan((N, 4, 4), 0, np.complex, {})
cs2 = S1.scalar_product(df, cs2)
assert np.allclose(cs[s], cs2[s], rtol=1e-5, atol=1e-6)
def get_context():
"""Set up context for solver"""
# Get points and weights for Chebyshev weighted integrals
assert params.Dquad == params.Bquad
collapse_fourier = False if params.dealias == '3/2-rule' else True
ST = Basis(params.N[2], 'C', bc=(0, 0), quad=params.Dquad)
SB = Basis(params.N[2], 'C', bc='Biharmonic', quad=params.Bquad)
CT = Basis(params.N[2], 'C', quad=params.Dquad)
ST0 = Basis(params.N[2], 'C', bc=(0, 0), quad=params.Dquad) # For 1D problem
K0 = Basis(params.N[0], 'F', domain=(0, params.L[0]), dtype='D')
K1 = Basis(params.N[1], 'F', domain=(0, params.L[1]), dtype='d')
kw0 = {'threads':params.threads,
'planner_effort':params.planner_effort["dct"],
'slab': (params.decomposition == 'slab'),
'collapse_fourier': collapse_fourier}
FST = TensorProductSpace(comm, (K0, K1, ST), axes=(2, 0, 1), **kw0) # Dirichlet
FSB = TensorProductSpace(comm, (K0, K1, SB), axes=(2, 0, 1), **kw0) # Biharmonic
FCT = TensorProductSpace(comm, (K0, K1, CT), axes=(2, 0, 1), **kw0) # Regular Chebyshev
VFS = MixedTensorProductSpace([FST, FST, FSB])
VFST = MixedTensorProductSpace([FST, FST, FST])
VUG = MixedTensorProductSpace([FST, FSB])
# 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 = Basis(params.N[2], 'C', bc=(0, 0), quad=params.Dquad)
SBp = Basis(params.N[2], 'C', bc='Biharmonic', quad=params.Bquad)
CTp = Basis(params.N[2], 'C', quad=params.Dquad)
else:
STp, SBp, CTp = ST, SB, CT
K0p = Basis(params.N[0], 'F', dtype='D', domain=(0, params.L[0]), **kw)
K1p = Basis(params.N[1], 'F', dtype='d', domain=(0, params.L[1]), **kw)
FSTp = TensorProductSpace(comm, (K0p, K1p, STp), axes=(2, 0, 1), **kw0)
FSBp = TensorProductSpace(comm, (K0p, K1p, SBp), axes=(2, 0, 1), **kw0)
FCTp = TensorProductSpace(comm, (K0p, K1p, CTp), axes=(2, 0, 1), **kw0)
VFSp = MixedTensorProductSpace([FSTp, FSTp, FSBp])
Nu = params.N[2]-2 # Number of velocity modes in Shen basis
Nb = params.N[2]-4 # Number of velocity modes in Shen biharmonic basis
u_slice = slice(0, Nu)
v_slice = slice(0, Nb)
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
U = Array(VFS)
U0 = Array(VFS)
U_hat = Function(VFS)
U_hat0 = Function(VFS)
g = Function(FST)
# primary variable
u = (U_hat, g)
H_hat = Function(VFST)
H_hat0 = Function(VFST)
H_hat1 = Function(VFST)
dU = Function(VUG)
hv = Function(FST)
hg = Function(FST)
Source = Array(VFS)
Sk = Function(VFS)
K2 = K[0]*K[0]+K[1]*K[1]
K4 = K2**2
# Set Nyquist frequency to zero on K that is used for odd derivatives in nonlinear terms
Kx = FST.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
K_over_K2 = np.zeros((2,)+g.shape)
for i in range(2):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
work = work_arrays()
u_dealias = Array(VFSp)
u0_hat = np.zeros((2, params.N[2]), dtype=complex)
h0_hat = np.zeros((2, params.N[2]), dtype=complex)
w = np.zeros((params.N[2], ), dtype=complex)
w1 = np.zeros((params.N[2], ), dtype=complex)
nu, dt, N = params.nu, params.dt, params.N
# Collect all matrices
mat = config.AttributeDict(
dict(CDD=inner_product((ST, 0), (ST, 1)),
CTD=inner_product((CT, 0), (ST, 1)),
BTT=inner_product((CT, 0), (CT, 0)),
AB=HelmholtzCoeff(N[2], 1.0, -(K2 - 2.0/nu/dt), 2, ST.quad),
AC=BiharmonicCoeff(N[2], nu*dt/2., (1. - nu*dt*K2), -(K2 - nu*dt/2.*K4), 2, 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))))
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 = KMMFile(config.params.solver,
checkpoint={'space': VFS,
'data': {'0': {'U': [U_hat]},
'1': {'U': [U_hat0]}}},
results={'space': VFS,
'data': {'U': [U]}})
return config.AttributeDict(locals())