def get_context():
"""Set up context for classical (NS) solver"""
V0 = C2CBasis(params.N[0], domain=(0, params.L[0]))
V1 = C2CBasis(params.N[1], domain=(0, params.L[1]))
V2 = R2CBasis(params.N[2], domain=(0, params.L[2]))
T = TensorProductSpace(comm, (V0, V1, V2), **{'threads': params.threads})
VT = VectorTensorProductSpace([T]*3)
kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'}
V0p = C2CBasis(params.N[0], domain=(0, params.L[0]), **kw)
V1p = C2CBasis(params.N[1], domain=(0, params.L[1]), **kw)
V2p = R2CBasis(params.N[2], domain=(0, params.L[2]), **kw)
Tp = TensorProductSpace(comm, (V0p, V1p, V2p), **{'threads': params.threads})
VTp = VectorTensorProductSpace([Tp]*3)
float, complex, mpitype = datatypes(params.precision)
FFT = T # For compatibility - to be removed
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
K2 = K[0]*K[0] + K[1]*K[1] + K[2]*K[2]
# Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
K_over_K2 = np.zeros((3,)+VT.local_shape())
for i in range(3):
K_over_K2[i] = K[i] / np.where(K2==0, 1, K2)
# Velocity and pressure
U = Array(VT, False)
U_hat = Array(VT)
P = Array(T, False)
P_hat = Array(T)
# Primary variable
u = U_hat
# RHS array
dU = Array(VT)
curl = Array(VT, False)
Source = Array(VT) # Possible source term initialized to zero
work = work_arrays()
hdf5file = NSWriter({"U":U[0], "V":U[1], "W":U[2], "P":P},
chkpoint={"current":{"U":U, "P":P}, "previous":{}},
filename=params.h5filename+".h5")
return config.AttributeDict(locals())
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_curl2():
# Test projection of curl
K0 = FunctionSpace(N[0], 'C', bc=(0, 0))
K1 = FunctionSpace(N[1], 'F', dtype='D')
K2 = FunctionSpace(N[2], 'F', dtype='d')
K3 = FunctionSpace(N[0], 'C')
T = TensorProductSpace(comm, (K0, K1, K2))
TT = TensorProductSpace(comm, (K3, K1, K2))
X = T.local_mesh(True)
K = T.local_wavenumbers(False)
Tk = VectorSpace(T)
TTk = VectorSpace([T, T, TT])
U = Array(Tk)
U_hat = Function(Tk)
curl_hat = Function(TTk)
curl_ = Array(TTk)
# Initialize a Taylor Green vortex
U[0] = np.sin(X[0]) * np.cos(X[1]) * np.cos(X[2]) * (1 - X[0]**2)
U[1] = -np.cos(X[0]) * np.sin(X[1]) * np.cos(X[2]) * (1 - X[0]**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)
# Compute curl first by computing each term individually
curl_hat[0] = 1j * (K[1] * U_hat[2] - K[2] * U_hat[1])
curl_[0] = T.backward(
curl_hat[0], curl_[0]) # No x-derivatives, still in Dirichlet space
dwdx_hat = project(Dx(U_hat[2], 0, 1), TT) # Need to use space without bc
dvdx_hat = project(Dx(U_hat[1], 0, 1), TT) # Need to use space without bc
dwdx = Array(TT)
dvdx = Array(TT)
dwdx = TT.backward(dwdx_hat, dwdx)
dvdx = TT.backward(dvdx_hat, dvdx)
curl_hat[1] = 1j * K[2] * U_hat[0]
curl_hat[2] = -1j * K[1] * U_hat[0]
curl_[1] = T.backward(curl_hat[1], curl_[1])
curl_[2] = T.backward(curl_hat[2], curl_[2])
curl_[1] -= dwdx
curl_[2] += dvdx
# Now do it with project
w_hat = project(curl(U_hat), TTk)
w = Array(TTk)
w = TTk.backward(w_hat, w)
assert allclose(w, curl_)
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())
def get_context():
float, complex, mpitype = datatypes(params.precision)
collapse_fourier = False if params.dealias == '3/2-rule' else True
dim = len(params.N)
dtype = lambda d: float if d == dim - 1 else complex
V = [
Basis(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i))
for i in range(dim)
]
kw0 = {
'threads': params.threads,
'planner_effort': params.planner_effort['fft']
}
T = TensorProductSpace(comm,
V,
dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier,
**kw0)
VT = VectorTensorProductSpace(T)
VM = MixedTensorProductSpace([T] * 2 * dim)
mask = T.mask_nyquist() if params.mask_nyquist else None
kw = {
'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'
}
Vp = [
Basis(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i), **kw)
for i in range(dim)
]
Tp = TensorProductSpace(comm,
Vp,
dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier,
**kw0)
VTp = VectorTensorProductSpace(Tp)
VMp = MixedTensorProductSpace([Tp] * 2 * dim)
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
for i in range(dim):
X[i] = X[i].astype(float)
K[i] = K[i].astype(float)
K2 = np.zeros(T.shape(True), dtype=float)
for i in range(dim):
K2 += K[i] * K[i]
# Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
for i in range(dim):
Kx[i] = Kx[i].astype(float)
K_over_K2 = np.zeros(VT.shape(True), dtype=float)
for i in range(dim):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
UB = Array(VM)
P = Array(T)
curl = Array(VT)
UB_hat = Function(VM)
P_hat = Function(T)
dU = Function(VM)
Source = Array(VM)
ub_dealias = Array(VMp)
ZZ_hat = np.zeros((3, 3) + Tp.shape(True), dtype=complex) # Work array
# Create views into large data structures
U = UB[:3]
U_hat = UB_hat[:3]
B = UB[3:]
B_hat = UB_hat[3:]
# Primary variable
u = UB_hat
hdf5file = MHDFile(config.params.solver,
checkpoint={
'space': VM,
'data': {
'0': {
'UB': [UB_hat]
}
}
},
results={
'space': VM,
'data': {
'UB': [UB]
}
})
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
collapse_fourier = False if params.dealias == '3/2-rule' else True
ST = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad)
CT = Basis(params.N[0], 'C', quad=params.Dquad)
CP = Basis(params.N[0], 'C', quad=params.Dquad)
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')
CP.slice = lambda: slice(0, CT.N)
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 = VectorTensorProductSpace(FST)
VCT = VectorTensorProductSpace(FCT)
VQ = MixedTensorProductSpace([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 = Basis(params.N[0], 'C', bc=(0, 0), quad=params.Dquad)
CTp = Basis(params.N[0], 'C', quad=params.Dquad)
else:
STp, CTp = ST, 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)
FCTp = TensorProductSpace(comm, (CTp, K0p, K1p), **kw0)
VFSp = VectorTensorProductSpace(FSTp)
VCp = MixedTensorProductSpace([FSTp, FCTp, FCTp])
float, complex, mpitype = datatypes("double")
constraints = ((3, 0, 0),
(3, params.N[0]-1, 0))
# 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
# primary variable
u = UP_hat
H_hat = Function(VFS)
H_hat0 = Function(VFS)
H_hat1 = Function(VFS)
dU = Function(VQ)
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
alfa = 2./nu/dt
a0 = inner(vt, (2./nu/dt)*ut-div(grad(ut)))
a1 = inner(vt, (2./nu)*grad(pt))
a2 = inner(qt, (2./nu)*div(ut))
M = BlockMatrix(a0+a1+a2)
# Collect all matrices
mat = config.AttributeDict(
dict(CDD=inner_product((ST, 0), (ST, 1)),
AB=HelmholtzCoeff(N[0], 1., alfa-K2, 0, ST.quad),))
la = None
hdf5file = CoupledFile(config.params.solver,
checkpoint={'space': VQ,
'data': {'0': {'UP': [UP_hat]},
'1': {'UP': [UP_hat0]}}},
results={'space': VFS,
'data': {'U': [U]}})
return config.AttributeDict(locals())
def get_context():
"""Set up context for classical (NS) solver"""
float, complex, mpitype = datatypes(params.precision)
collapse_fourier = False if params.dealias == '3/2-rule' else True
dim = len(params.N)
dtype = lambda d: float if d == dim - 1 else complex
V = [
Basis(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i))
for i in range(dim)
]
kw0 = {
'threads': params.threads,
'planner_effort': params.planner_effort['fft']
}
T = TensorProductSpace(comm,
V,
dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier,
**kw0)
VT = VectorTensorProductSpace(T)
# Different bases for nonlinear term, either 2/3-rule or 3/2-rule
kw = {
'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'
}
Vp = [
Basis(params.N[i], 'F', domain=(0, params.L[i]), dtype=dtype(i), **kw)
for i in range(dim)
]
Tp = TensorProductSpace(comm,
Vp,
dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier,
**kw0)
VTp = VectorTensorProductSpace(Tp)
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
for i in range(dim):
X[i] = X[i].astype(float)
K[i] = K[i].astype(float)
K2 = np.zeros(T.shape(True), dtype=float)
for i in range(dim):
K2 += K[i] * K[i]
# Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
for i in range(dim):
Kx[i] = Kx[i].astype(float)
K_over_K2 = np.zeros(VT.shape(True), dtype=float)
for i in range(dim):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
# Velocity and pressure. Use ndarray view for efficiency
U = Array(VT)
U_hat = Function(VT)
P = Array(T)
P_hat = Function(T)
u_dealias = Array(VTp)
# Primary variable
u = U_hat
# RHS array
dU = Function(VT)
curl = Array(VT)
Source = Function(VT) # Possible source term initialized to zero
work = work_arrays()
hdf5file = NSFile(config.params.solver,
checkpoint={
'space': VT,
'data': {
'0': {
'U': [U_hat]
}
}
},
results={
'space': VT,
'data': {
'U': [U],
'P': [P]
}
})
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 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())
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())
N = int(sys.argv[-2]) N = [N, N + 1, N + 2] #N = (14, 15, 16) SD = FunctionSpace(N[1], family=family, bc=(a, b)) K1 = FunctionSpace(N[0], family='F', dtype='D') K2 = FunctionSpace(N[2], family='F', dtype='d') subcomms = Subcomm(MPI.COMM_WORLD, [0, 0, 1]) T = TensorProductSpace(subcomms, (K1, SD, K2), axes=(1, 0, 2)) u = TrialFunction(T) v = TestFunction(T) # Get manufactured right hand side fe = div(grad(u)).tosympy(basis=ue) K = T.local_wavenumbers() # Get f on quad points fj = Array(T, buffer=fe) # Compute right hand side of Poisson equation f_hat = inner(v, fj) # 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
def get_context():
"""Set up context for classical (NS) solver"""
float, complex, mpitype = datatypes(params.precision)
collapse_fourier = False if params.dealias == '3/2-rule' else True
dim = len(params.N)
dtype = lambda d: float if d == dim-1 else complex
V = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
dtype=dtype(i)) for i in range(dim)]
kw0 = {'threads': params.threads,
'planner_effort': params.planner_effort['fft']}
T = TensorProductSpace(comm, V, dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier, **kw0)
VT = VectorTensorProductSpace(T)
# Different bases for nonlinear term, either 2/3-rule or 3/2-rule
kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'}
Vp = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
dtype=dtype(i), **kw) for i in range(dim)]
Tp = TensorProductSpace(comm, Vp, dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier, **kw0)
VTp = VectorTensorProductSpace(Tp)
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
for i in range(dim):
X[i] = X[i].astype(float)
K[i] = K[i].astype(float)
K2 = np.zeros(T.shape(True), dtype=float)
for i in range(dim):
K2 += K[i]*K[i]
# Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
for i in range(dim):
Kx[i] = Kx[i].astype(float)
K_over_K2 = np.zeros(VT.shape(True), dtype=float)
for i in range(dim):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
# Velocity and pressure. Use ndarray view for efficiency
U = Array(VT)
U_hat = Function(VT)
P = Array(T)
P_hat = Function(T)
u_dealias = Array(VTp)
# Primary variable
u = U_hat
# RHS array
dU = Function(VT)
curl = Array(VT)
Source = Function(VT) # Possible source term initialized to zero
work = work_arrays()
hdf5file = NSFile(config.params.solver,
checkpoint={'space': VT,
'data': {'0': {'U': [U_hat]}}},
results={'space': VT,
'data': {'U': [U], 'P': [P]}})
return config.AttributeDict(locals())
def get_context():
float, complex, mpitype = datatypes(params.precision)
collapse_fourier = False if params.dealias == '3/2-rule' else True
dim = len(params.N)
dtype = lambda d: float if d == dim-1 else complex
V = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
dtype=dtype(i)) for i in range(dim)]
kw0 = {'threads': params.threads,
'planner_effort': params.planner_effort['fft']}
T = TensorProductSpace(comm, V, dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier, **kw0)
VT = VectorTensorProductSpace(T)
VM = MixedTensorProductSpace([T]*2*dim)
kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'}
Vp = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
dtype=dtype(i), **kw) for i in range(dim)]
Tp = TensorProductSpace(comm, Vp, dtype=float,
slab=(params.decomposition == 'slab'),
collapse_fourier=collapse_fourier, **kw0)
VTp = VectorTensorProductSpace(Tp)
VMp = MixedTensorProductSpace([Tp]*2*dim)
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
for i in range(dim):
X[i] = X[i].astype(float)
K[i] = K[i].astype(float)
K2 = np.zeros(T.shape(True), dtype=float)
for i in range(dim):
K2 += K[i]*K[i]
# Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
for i in range(dim):
Kx[i] = Kx[i].astype(float)
K_over_K2 = np.zeros(VT.shape(True), dtype=float)
for i in range(dim):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
UB = Array(VM)
P = Array(T)
curl = Array(VT)
UB_hat = Function(VM)
P_hat = Function(T)
dU = Function(VM)
Source = Array(VM)
ub_dealias = Array(VMp)
ZZ_hat = np.zeros((3, 3) + Tp.shape(True), dtype=complex) # Work array
# Create views into large data structures
U = UB[:3]
U_hat = UB_hat[:3]
B = UB[3:]
B_hat = UB_hat[3:]
# Primary variable
u = UB_hat
hdf5file = MHDFile(config.params.solver,
checkpoint={'space': VM,
'data': {'0': {'UB': [UB_hat]}}},
results={'space': VM,
'data': {'UB': [UB]}})
return config.AttributeDict(locals())
