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())