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 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 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_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 = 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 get_context():
"""Set up context for solver"""
# Get points and weights for Chebyshev weighted integrals
ST = ShenDirichletBasis(params.N[0], quad=params.Dquad)
SB = ShenBiharmonicBasis(params.N[0], quad=params.Bquad)
CT = ST.CT # Chebyshev transform
Nu = params.N[0] - 2 # Number of velocity modes in Shen basis
Nb = params.N[0] - 4 # Number of velocity modes in Shen biharmonic basis
u_slice = slice(0, Nu)
v_slice = slice(0, Nb)
FST = SlabShen_R2C(params.N,
params.L,
comm,
threads=params.threads,
communication=params.communication,
planner_effort=params.planner_effort,
dealias_cheb=params.dealias_cheb)
float, complex, mpitype = datatypes("double")
ST.plan(FST.complex_shape(), 0, complex, {
'threads': params.threads,
'planner_effort': params.planner_effort["dct"]
})
SB.plan(FST.complex_shape(), 0, complex, {
'threads': params.threads,
'planner_effort': params.planner_effort["dct"]
})
# Mesh variables
X = FST.get_local_mesh(ST)
x0, x1, x2 = FST.get_mesh_dims(ST)
K = FST.get_local_wavenumbermesh(scaled=True)
K2 = K[1] * K[1] + K[2] * K[2]
K_over_K2 = zeros((2, ) + FST.complex_shape())
for i in range(2):
K_over_K2[i] = K[i + 1] / np.where(K2 == 0, 1, K2)
# Solution variables
U = zeros((3, ) + FST.real_shape(), dtype=float)
U0 = zeros((3, ) + FST.real_shape(), dtype=float)
U_hat = zeros((3, ) + FST.complex_shape(), dtype=complex)
U_hat0 = zeros((3, ) + FST.complex_shape(), dtype=complex)
g = zeros(FST.complex_shape(), dtype=complex)
# primary variable
u = (U_hat, g)
H_hat = zeros((3, ) + FST.complex_shape(), dtype=complex)
H_hat0 = zeros((3, ) + FST.complex_shape(), dtype=complex)
H_hat1 = zeros((3, ) + FST.complex_shape(), dtype=complex)
dU = zeros((3, ) + FST.complex_shape(), dtype=complex)
hv = zeros(FST.complex_shape(), dtype=complex)
hg = zeros(FST.complex_shape(), dtype=complex)
Source = zeros((3, ) + FST.real_shape(), dtype=float)
Sk = zeros((3, ) + FST.complex_shape(), dtype=complex)
work = work_arrays()
nu, dt, N = params.nu, params.dt, params.N
K4 = K2**2
kx = K[0][:, 0, 0]
# Collect all linear algebra solvers
la = config.AttributeDict(
dict(HelmholtzSolverG=Helmholtz(N[0], np.sqrt(K2[0] + 2.0 / nu / dt),
ST),
BiharmonicSolverU=Biharmonic(N[0],
-nu * dt / 2.,
1. + nu * dt * K2[0],
-(K2[0] + nu * dt / 2. * K4[0]),
quad=SB.quad,
solver="cython"),
HelmholtzSolverU0=Helmholtz(N[0], np.sqrt(2. / nu / dt), ST),
TDMASolverD=TDMA(inner_product((ST, 0), (ST, 0)))))
alfa = K2[0] - 2.0 / nu / dt
# Collect all matrices
mat = config.AttributeDict(
dict(
CDD=inner_product((ST, 0), (ST, 1)),
AB=HelmholtzCoeff(N[0], 1.0, -alfa, ST.quad),
AC=BiharmonicCoeff(N[0],
nu * dt / 2., (1. - nu * dt * K2[0]),
-(K2[0] - nu * dt / 2. * K4[0]),
quad=SB.quad),
# Matrices for biharmonic equation
CBD=inner_product((SB, 0), (ST, 1)),
ABB=inner_product((SB, 0), (SB, 2)),
BBB=inner_product((SB, 0), (SB, 0)),
SBB=inner_product((SB, 0), (SB, 4)),
# Matrices for Helmholtz equation
ADD=inner_product((ST, 0), (ST, 2)),
BDD=inner_product((ST, 0), (ST, 0)),
BBD=inner_product((SB, 0), (ST, 0)),
CDB=inner_product((ST, 0), (SB, 1))))
hdf5file = KMMWriter({
"U": U[0],
"V": U[1],
"W": U[2]
},
chkpoint={
'current': {
'U': U
},
'previous': {
'U': U0
}
},
filename=params.solver + ".h5",
mesh={
"x": x0,
"y": x1,
"z": x2
})
return config.AttributeDict(locals())
