def get_context(): """Set up context for solver""" # Get points and weights for Chebyshev weighted integrals ST = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad) SB = Basis(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad) ST0 = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad) # For 1D problem 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] K4 = K2**2 # Set Nyquist frequency to zero on K that is used for odd derivatives Kx = FST.get_local_wavenumbermesh(scaled=True, eliminate_highest_freq=True) 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 kx = K[0][:, 0, 0] alfa = K2 - 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)), ADD0=inner_product((ST0, 0), (ST0, 2)), BDD0=inner_product((ST0, 0), (ST0, 0)))) mat.ADD.axis = 0 mat.BDD.axis = 0 mat.SBB.axis = 0 # Collect all linear algebra solvers 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 = 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())
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[0], 'C', bc=(0, 0), quad=params.Dquad) SB = Basis(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad) CT = Basis(params.N[0], 'C', quad=params.Dquad) ST0 = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad) # For 1D problem K0 = Basis(params.N[1], 'F', domain=(0, params.L[1]), dtype='D') K1 = Basis(params.N[2], 'F', domain=(0, params.L[2]), dtype='d') 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 FSB = TensorProductSpace(comm, (SB, K0, K1), **kw0) # Biharmonic FCT = TensorProductSpace(comm, (CT, K0, K1), **kw0) # Regular Chebyshev VFS = VectorTensorProductSpace([FSB, FST, FST]) VFST = VectorTensorProductSpace([FST, FST, FST]) VUG = MixedTensorProductSpace([FSB, FST]) VCT = VectorTensorProductSpace(FCT) 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 = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad) SBp = Basis(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad) CTp = Basis(params.N[0], 'C', quad=params.Dquad) else: STp, SBp, CTp = ST, SB, CT K0p = Basis(params.N[1], 'F', dtype='D', domain=(0, params.L[1]), **kw) K1p = Basis(params.N[2], 'F', dtype='d', domain=(0, params.L[2]), **kw) FSTp = TensorProductSpace(comm, (STp, K0p, K1p), **kw0) FSBp = TensorProductSpace(comm, (SBp, K0p, K1p), **kw0) FCTp = TensorProductSpace(comm, (CTp, K0p, K1p), **kw0) VFSp = VectorTensorProductSpace([FSBp, FSTp, FSTp]) 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(VFS) hv = Function(FSB) hg = Function(FST) Source = Array(VFS) Sk = Function(VFS) K2 = K[1]*K[1]+K[2]*K[2] K4 = K2**2 K_over_K2 = np.zeros((2,)+g.shape) for i in range(2): K_over_K2[i] = K[i+1] / np.where(K2 == 0, 1, K2) for i in range(3): K[i] = K[i].astype(float) work = work_arrays() u_dealias = Array(VFSp) u0_hat = np.zeros((2, params.N[0]), dtype=complex) h0_hat = np.zeros((2, params.N[0]), dtype=complex) w = np.zeros((params.N[0], ), dtype=complex) w1 = np.zeros((params.N[0], ), dtype=complex) nu, dt, N = params.nu, params.dt, params.N 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., -(K2 - 2.0/nu/dt), 0, ST.quad), AC=BiharmonicCoeff(N[0], nu*dt/2., (1. - nu*dt*K2), -(K2 - nu*dt/2.*K4), 0, 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())
def get_context(): """Set up context for solver""" # Get points and weights for Chebyshev weighted integrals assert params.Dquad == params.Bquad ST = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad) SB = Basis(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad) CT = Basis(params.N[0], 'C', quad=params.Dquad) ST0 = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad) # For 1D problem K0 = Basis(params.N[1], 'F', domain=(0, params.L[1]), dtype='D') K1 = Basis(params.N[2], 'F', domain=(0, params.L[2]), dtype='d') kw0 = { 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] } FST = TensorProductSpace(comm, (ST, K0, K1), axes=(0, 1, 2), collapse_fourier=False, **kw0) # Dirichlet FSB = TensorProductSpace(comm, (SB, K0, K1), axes=(0, 1, 2), collapse_fourier=False, **kw0) # Biharmonic FCT = TensorProductSpace(comm, (CT, K0, K1), axes=(0, 1, 2), collapse_fourier=False, **kw0) # Regular Chebyshev VFS = MixedTensorProductSpace([FSB, FST, FST]) VUG = MixedTensorProductSpace([FSB, FST]) # 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[0], 'C', bc=(0, 0), quad=params.Dquad) SBp = Basis(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad) CTp = Basis(params.N[0], 'C', quad=params.Dquad) else: STp, SBp, CTp = ST, SB, CT K0p = Basis(params.N[1], 'F', dtype='D', domain=(0, params.L[1]), **kw) K1p = Basis(params.N[2], 'F', dtype='d', domain=(0, params.L[2]), **kw) FSTp = TensorProductSpace(comm, (STp, K0p, K1p), axes=(0, 1, 2), collapse_fourier=False, **kw0) FSBp = TensorProductSpace(comm, (SBp, K0p, K1p), axes=(0, 1, 2), collapse_fourier=False, **kw0) FCTp = TensorProductSpace(comm, (CTp, K0p, K1p), axes=(0, 1, 2), collapse_fourier=False, **kw0) VFSp = MixedTensorProductSpace([FSBp, FSTp, FSTp]) 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) 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(VFS) H_hat0 = Function(VFS) H_hat1 = Function(VFS) dU = Function(VUG) hv = Function(FST) hg = Function(FST) Source = Array(VFS) Sk = Function(VFS) K2 = K[1] * K[1] + K[2] * K[2] 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 + 1] / np.where(K2 == 0, 1, K2) work = work_arrays() nu, dt, N = params.nu, params.dt, params.N 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, -(K2 - 2.0 / nu / dt), ST.quad), AC=BiharmonicCoeff(N[0], nu * dt / 2., (1. - nu * dt * K2), -(K2 - nu * dt / 2. * K4), 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)), ADD0=inner_product((ST0, 0), (ST0, 2)), BDD0=inner_product((ST0, 0), (ST0, 0)), )) ## Collect all linear algebra solvers #la = config.AttributeDict(dict( #HelmholtzSolverG = old_Helmholtz(N[0], np.sqrt(K2[0]+2.0/nu/dt), ST), #BiharmonicSolverU = old_Biharmonic(N[0], -nu*dt/2., 1.+nu*dt*K2[0], #-(K2[0] + nu*dt/2.*K4[0]), quad=SB.quad, #solver="cython"), #HelmholtzSolverU0 = old_Helmholtz(N[0], np.sqrt(2./nu/dt), ST), #TDMASolverD = TDMA(inner_product((ST, 0), (ST, 0))) #) #) mat.ADD.axis = 0 mat.BDD.axis = 0 mat.SBB.axis = 0 la = config.AttributeDict( dict(HelmholtzSolverG=Helmholtz(mat.ADD, mat.BDD, -np.ones( (1, 1, 1)), (K2[0] + 2.0 / nu / dt)[np.newaxis, :, :]), BiharmonicSolverU=Biharmonic( mat.SBB, mat.ABB, mat.BBB, -nu * dt / 2. * np.ones( (1, 1, 1)), (1. + nu * dt * K2[0])[np.newaxis, :, :], (-(K2[0] + nu * dt / 2. * K4[0]))[np.newaxis, :, :]), HelmholtzSolverU0=Helmholtz(mat.ADD0, mat.BDD0, np.array([-1.]), np.array([2. / nu / dt])), TDMASolverD=TDMA(inner_product((ST, 0), (ST, 0))))) 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())
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 = Basis(params.N[0], quad=params.Dquad) ST0 = ShenDirichletBasis(params.N[0], quad=params.Dquad, plan=True) # For 1D problem K0 = C2CBasis(params.N[1], domain=(0, params.L[1])) K1 = R2CBasis(params.N[2], domain=(0, params.L[2])) #CT = ST.CT # Chebyshev transform FST = TensorProductSpace(comm, (ST, K0, K1), **{ 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] }) # Dirichlet FSB = TensorProductSpace(comm, (SB, K0, K1), **{ 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] }) # Biharmonic FCT = TensorProductSpace(comm, (CT, K0, K1), **{ 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] }) # Regular Chebyshev VFS = VectorTensorProductSpace([FSB, FST, FST]) # Padded STp = ShenDirichletBasis(params.N[0], quad=params.Dquad) SBp = ShenBiharmonicBasis(params.N[0], quad=params.Bquad) CTp = Basis(params.N[0], quad=params.Dquad) K0p = C2CBasis(params.N[1], padding_factor=1.5, domain=(0, params.L[1])) K1p = R2CBasis(params.N[2], padding_factor=1.5, domain=(0, params.L[2])) FSTp = TensorProductSpace( comm, (STp, K0p, K1p), **{ 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] }) FSBp = TensorProductSpace( comm, (SBp, K0p, K1p), **{ 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] }) FCTp = TensorProductSpace( comm, (CTp, K0p, K1p), **{ 'threads': params.threads, 'planner_effort': params.planner_effort["dct"] }) VFSp = VectorTensorProductSpace([FSBp, FSTp, FSTp]) VFSp = VFS FCTp = FCT FSTp = FST FSBp = FSB 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) 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, False) U0 = Array(VFS, False) U_hat = Array(VFS) U_hat0 = Array(VFS) g = Array(FST) # primary variable u = (U_hat, g) H_hat = Array(VFS) H_hat0 = Array(VFS) H_hat1 = Array(VFS) dU = Array(VFS) hv = Array(FST) hg = Array(FST) Source = Array(VFS, False) Sk = Array(VFS) K2 = K[1] * K[1] + K[2] * K[2] K_over_K2 = np.zeros((2, ) + g.shape) for i in range(2): K_over_K2[i] = K[i + 1] / np.where(K2 == 0, 1, K2) work = work_arrays() nu, dt, N = params.nu, params.dt, params.N K4 = K2**2 kx = K[0][:, 0, 0] alfa = K2[0] - 2.0 / nu / dt # Collect all matrices mat = config.AttributeDict( dict( CDD=inner_product((ST, 0), (ST, 1)), AB=HelmholtzCoeff(kx, -1.0, -alfa, ST.quad), AC=BiharmonicCoeff(kx, 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)), ADD0=inner_product((ST0, 0), (ST0, 2)), BDD0=inner_product((ST0, 0), (ST0, 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))) #) #) mat.ADD.axis = 0 mat.BDD.axis = 0 mat.SBB.axis = 0 la = config.AttributeDict( dict(HelmholtzSolverG=Helmholtz(mat.ADD, mat.BDD, -np.ones( (1, 1, 1)), (K2[0] + 2.0 / nu / dt)[np.newaxis, :, :]), BiharmonicSolverU=Biharmonic( mat.SBB, mat.ABB, mat.BBB, -nu * dt / 2. * np.ones( (1, 1, 1)), (1. + nu * dt * K2[0])[np.newaxis, :, :], (-(K2[0] + nu * dt / 2. * K4[0]))[np.newaxis, :, :]), HelmholtzSolverU0=old_Helmholtz(N[0], np.sqrt(2. / nu / dt), ST), TDMASolverD=TDMA(inner_product((ST, 0), (ST, 0))))) 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())
def initialize(solver, context): # Initialize with pertubation ala perturbU (https://github.com/wyldckat/perturbU) for openfoam U = context.U X = context.X U_hat = context.U_hat params = config.params Y = np.where(X[0] < 0, 1+X[0], 1-X[0]) utau = params.nu*params.Re_tau Um = 46.9091*utau # For Re_tau=180 #Um = 56.*utau Xplus = Y*params.Re_tau Yplus = X[1]*params.Re_tau Zplus = X[2]*params.Re_tau duplus = Um*0.2/utau #Um*0.25/utau alfaplus = params.L[1]/200. # May have to adjust these two for different Re betaplus = params.L[2]/100. # sigma = 0.00055 # 0.00055 epsilon = Um/200. #Um/200. U[:] = 0 U[1] = Um*(Y-0.5*Y**2) dev = 1+0.0000001*np.random.randn(Y.shape[0], Y.shape[1], Y.shape[2]) #dev = np.fromfile('dev.dat').reshape((64, 64, 64)) dd = utau*duplus/2.0*Xplus/40.*np.exp(-sigma*Xplus**2+0.5)*np.cos(betaplus*Zplus)*dev[:, slice(0, 1), :] U[1] += dd U[2] += epsilon*np.sin(alfaplus*Yplus)*Xplus*np.exp(-sigma*Xplus**2)*dev[:, :, slice(0, 1)] U_hat = U.forward(U_hat) U = U_hat.backward(U) U_hat = U.forward(U_hat) if "KMM" in params.solver: context.g[:] = 1j*context.K[1]*U_hat[2] - 1j*context.K[2]*U_hat[1] # Set the flux #flux[0] = context.FST.dx(U[1], context.ST.quad) #solver.comm.Bcast(flux) # project to zero divergence div_u = solver.get_divergence(**context) print('div0 ', dx(div_u**2, context.FST)) u = TrialFunction(context.FST) v = TestFunction(context.FST) A = inner(v, div(grad(u))) b = inner(v, div(U_hat)) from shenfun.chebyshev.la import Helmholtz sol = Helmholtz(*A) phi = Function(context.FST) phi = sol(phi, b) U_hat -= project(grad(phi), context.VFS) U = U_hat.backward(U) div_u = solver.get_divergence(**context) print('div1 ', dx(div_u**2, context.FST)) if solver.rank == 0: print("Flux {}".format(flux[0])) if not 'KMM' in params.solver: P_hat = solver.compute_pressure(**context) P = P_hat.backward(context.P) if not 'RK3' in params.solver: context.U_hat0[:] = context.U_hat[:] context.H_hat1[:] = solver.get_convection(**context)