def test_assign(fam):
x, y = symbols("x,y")
for bc in (None, (0, 0), (0, 0, 0, 0)):
dtype = 'D' if fam == 'F' else 'd'
bc = 'periodic' if fam == 'F' else bc
if bc == (0, 0, 0, 0) and fam in ('La', 'H'):
continue
tol = 1e-12 if fam in ('C', 'L', 'F') else 1e-5
N = (10, 12)
B0 = FunctionSpace(N[0], fam, dtype=dtype, bc=bc)
B1 = FunctionSpace(N[1], fam, dtype=dtype, bc=bc)
u_hat = Function(B0)
u_hat[1:4] = 1
ub_hat = Function(B1)
u_hat.assign(ub_hat)
assert abs(inner(1, u_hat)-inner(1, ub_hat)) < tol
T = TensorProductSpace(comm, (B0, B1))
u_hat = Function(T)
u_hat[1:4, 1:4] = 1
Tp = T.get_refined((2*N[0], 2*N[1]))
ub_hat = Function(Tp)
u_hat.assign(ub_hat)
assert abs(inner(1, u_hat)-inner(1, ub_hat)) < tol
VT = VectorSpace(T)
u_hat = Function(VT)
u_hat[:, 1:4, 1:4] = 1
Tp = T.get_refined((2*N[0], 2*N[1]))
VTp = VectorSpace(Tp)
ub_hat = Function(VTp)
u_hat.assign(ub_hat)
assert abs(inner((1, 1), u_hat)-inner((1, 1), ub_hat)) < tol
def test_refine():
assert comm.Get_size() < 7
N = (8, 9, 10)
F0 = FunctionSpace(8, 'F', dtype='D')
F1 = FunctionSpace(9, 'F', dtype='D')
F2 = FunctionSpace(10, 'F', dtype='d')
T = TensorProductSpace(comm, (F0, F1, F2), slab=True, collapse_fourier=True)
u_hat = Function(T)
u = Array(T)
u[:] = np.random.random(u.shape)
u_hat = u.forward(u_hat)
Tp = T.get_dealiased(padding_factor=(2, 2, 2))
u_ = Array(Tp)
up_hat = Function(Tp)
assert up_hat.commsizes == u_hat.commsizes
u2 = u_hat.refine(2*np.array(N))
V = VectorSpace(T)
u_hat = Function(V)
u = Array(V)
u[:] = np.random.random(u.shape)
u_hat = u.forward(u_hat)
Vp = V.get_dealiased(padding_factor=(2, 2, 2))
u_ = Array(Vp)
up_hat = Function(Vp)
assert up_hat.commsizes == u_hat.commsizes
u3 = u_hat.refine(2*np.array(N))
def test_curl_cc():
theta, phi = sp.symbols('x,y', real=True, positive=True)
psi = (theta, phi)
r = 1
rv = (r * sp.sin(theta) * sp.cos(phi), r * sp.sin(theta) * sp.sin(phi),
r * sp.cos(theta))
# Manufactured solution
sph = sp.functions.special.spherical_harmonics.Ynm
ue = sph(6, 3, theta, phi)
N, M = 16, 12
L0 = FunctionSpace(N, 'C', domain=(0, np.pi))
F1 = FunctionSpace(M, 'F', dtype='D')
T = TensorProductSpace(comm, (L0, F1), coordinates=(psi, rv))
u_hat = Function(T, buffer=ue)
du = curl(grad(u_hat))
du.terms() == [[]]
r, theta, z = psi = sp.symbols('x,y,z', real=True, positive=True)
rv = (r * sp.cos(theta), r * sp.sin(theta), z)
# Manufactured solution
ue = (r * (1 - r) * sp.cos(4 * theta) - 1 * (r - 1)) * sp.cos(4 * z)
N = 12
F0 = FunctionSpace(N, 'F', dtype='D')
F1 = FunctionSpace(N, 'F', dtype='d')
L = FunctionSpace(N, 'L', bc='Dirichlet', domain=(0, 1))
T = TensorProductSpace(comm, (L, F0, F1), coordinates=(psi, rv))
T1 = T.get_orthogonal()
V = VectorSpace(T1)
u_hat = Function(T, buffer=ue)
du = project(curl(grad(u_hat)), V)
assert np.linalg.norm(du) < 1e-10
def test_mixed_3D(backend, forward_output, as_scalar):
if (backend == 'netcdf4' and forward_output is True) or skip[backend]:
return
K0 = FunctionSpace(N[0], 'F', dtype='D', domain=(0, np.pi))
K1 = FunctionSpace(N[1], 'F', dtype='d', domain=(0, 2 * np.pi))
K2 = FunctionSpace(N[2], 'C')
T = TensorProductSpace(comm, (K0, K1, K2))
TT = VectorSpace(T)
filename = 'test3Dm_{}'.format(ex[forward_output])
hfile = writer(filename, TT, backend=backend)
uf = Function(TT, val=2) if forward_output else Array(TT, val=2)
uf[0] = 1
data = {
'ux': (uf[0], (uf[0], [slice(None), 4,
slice(None)]), (uf[0], [slice(None), 4, 4])),
'uy': (uf[1], (uf[1], [slice(None), 4,
slice(None)]), (uf[1], [slice(None), 4, 4])),
'u': [uf, (uf, [slice(None), 4, slice(None)])]
}
hfile.write(0, data, as_scalar=as_scalar)
hfile.write(1, data, as_scalar=as_scalar)
if not forward_output and backend == 'hdf5' and comm.Get_rank() == 0:
generate_xdmf(filename + '.h5')
if as_scalar is False:
u0 = Function(TT) if forward_output else Array(TT)
read = reader(filename, TT, backend=backend)
read.read(u0, 'u', step=1)
assert np.allclose(u0, uf)
else:
u0 = Function(T) if forward_output else Array(T)
read = reader(filename, T, backend=backend)
read.read(u0, 'u0', step=1)
assert np.allclose(u0, uf[0])
cleanup()
def test_curl(typecode):
K0 = FunctionSpace(N[0], 'F', dtype=typecode.upper())
K1 = FunctionSpace(N[1], 'F', dtype=typecode.upper())
K2 = FunctionSpace(N[2], 'F', dtype=typecode)
T = TensorProductSpace(comm, (K0, K1, K2), dtype=typecode)
X = T.local_mesh(True)
K = T.local_wavenumbers()
Tk = VectorSpace(T)
u = shenfun.TrialFunction(Tk)
v = shenfun.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)
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_cylinder():
T = get_function_space('cylinder')
u = TrialFunction(T)
du = div(grad(u))
assert du.tolatex() == '\\frac{\\partial^2 u}{\\partial x^2 }+\\frac{1}{x}\\frac{\\partial u}{\\partial x }+\\frac{1}{x^{2}}\\frac{\\partial^2 u}{\\partial y^2 }+\\frac{\\partial^2 u}{\\partial z^2 }'
V = VectorSpace(T)
u = TrialFunction(V)
du = div(grad(u))
assert du.tolatex() == '\\left( \\frac{\\partial^2 u^{x}}{\\partial x^2 }+\\frac{1}{x}\\frac{\\partial u^{x}}{\\partial x }+\\frac{1}{x^{2}}\\frac{\\partial^2 u^{x}}{\\partial y^2 }- \\frac{2}{x}\\frac{\\partial u^{y}}{\\partial y }- \\frac{1}{x^{2}}u^{x}+\\frac{\\partial^2 u^{x}}{\\partial z^2 }\\right) \\mathbf{b}_{x} \\\\+\\left( \\frac{\\partial^2 u^{y}}{\\partial x^2 }+\\frac{3}{x}\\frac{\\partial u^{y}}{\\partial x }+\\frac{2}{x^{3}}\\frac{\\partial u^{x}}{\\partial y }+\\frac{1}{x^{2}}\\frac{\\partial^2 u^{y}}{\\partial y^2 }+\\frac{\\partial^2 u^{y}}{\\partial z^2 }\\right) \\mathbf{b}_{y} \\\\+\\left( \\frac{\\partial^2 u^{z}}{\\partial x^2 }+\\frac{1}{x}\\frac{\\partial u^{z}}{\\partial x }+\\frac{1}{x^{2}}\\frac{\\partial^2 u^{z}}{\\partial y^2 }+\\frac{\\partial^2 u^{z}}{\\partial z^2 }\\right) \\mathbf{b}_{z} \\\\'
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 test_vector_laplace(space):
"""Test that
div(grad(u)) = grad(div(u)) - curl(curl(u))
"""
T = get_function_space(space)
V = VectorSpace(T)
u = TrialFunction(V)
v = _TestFunction(V)
du = div(grad(u))
dv = grad(div(u)) - curl(curl(u))
u_hat = Function(V)
u_hat[:] = np.random.random(u_hat.shape) + np.random.random(u_hat.shape)*1j
A0 = inner(v, du)
A1 = inner(v, dv)
a0 = BlockMatrix(A0)
a1 = BlockMatrix(A1)
b0 = Function(V)
b1 = Function(V)
b0 = a0.matvec(u_hat, b0)
b1 = a1.matvec(u_hat, b1)
assert np.linalg.norm(b0-b1) < 1e-8
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())
N = (200, 200) K0 = FunctionSpace(N[0], 'F', dtype='D', domain=(-1., 1.)) K1 = FunctionSpace(N[1], 'F', dtype='d', domain=(-1., 1.)) T = TensorProductSpace(comm, (K0, K1)) u = TrialFunction(T) v = TestFunction(T) # For nonlinear term we can use the 3/2-rule with padding Tp = T.get_dealiased((1.5, 1.5)) # Turn on padding by commenting #Tp = T # Create vector spaces and a test function for the regular vector space TV = VectorSpace(T) TVp = VectorSpace(Tp) vv = TestFunction(TV) uu = TrialFunction(TV) # Declare solution arrays and work arrays UV = Array(TV, buffer=(u0, v0)) UVp = Array(TVp) U, V = UV # views into vector components UV_hat = Function(TV) w0 = Function(TV) # Work array spectral space w1 = Array(TVp) # Work array physical space e1 = 0.00002 e2 = 0.00001 b0 = 0.03
def test_transform(typecode, dim):
s = (True,)
if comm.Get_size() > 2 and dim > 2:
s = (True, False)
for slab in s:
for shape in product(*([sizes]*dim)):
bases = []
for n in shape[:-1]:
bases.append(FunctionSpace(n, 'F', dtype=typecode.upper()))
bases.append(FunctionSpace(shape[-1], 'F', dtype=typecode))
fft = TensorProductSpace(comm, bases, dtype=typecode, slab=slab)
if comm.rank == 0:
grid = [c.size for c in fft.subcomm]
print('grid:{} shape:{} typecode:{}'
.format(grid, shape, typecode))
U = random_like(fft.forward.input_array)
F = fft.forward(U)
V = fft.backward(F)
assert allclose(V, U)
# Alternative method
fft.forward.input_array[...] = U
fft.forward(fast_transform=False)
fft.backward(fast_transform=False)
V = fft.backward.output_array
assert allclose(V, U)
TT = VectorSpace(fft)
U = Array(TT)
V = Array(TT)
F = Function(TT)
U[:] = random_like(U)
F = TT.forward(U, F)
V = TT.backward(F, V)
assert allclose(V, U)
TM = CompositeSpace([fft, fft])
U = Array(TM)
V = Array(TM)
F = Function(TM)
U[:] = random_like(U)
F = TM.forward(U, F)
V = TM.backward(F, V)
assert allclose(V, U)
fftp = fft.get_dealiased(padding_factor=1.5)
#fft.destroy()
#padding = 1.5
#bases = []
#for n in shape[:-1]:
# bases.append(FunctionSpace(n, 'F', dtype=typecode.upper(), padding_factor=padding))
#bases.append(FunctionSpace(shape[-1], 'F', dtype=typecode, padding_factor=padding))
#fft = TensorProductSpace(comm, bases, dtype=typecode)
if comm.rank == 0:
grid = [c.size for c in fftp.subcomm]
print('grid:{} shape:{} typecode:{}'
.format(grid, shape, typecode))
U = random_like(fftp.forward.input_array)
F = fftp.forward(U)
Fc = F.copy()
V = fftp.backward(F)
F = fftp.forward(V)
assert allclose(F, Fc)
# Alternative method
fftp.backward.input_array[...] = F
fftp.backward()
fftp.forward()
V = fftp.forward.output_array
assert allclose(F, V)
fftp.destroy()
fft.destroy()
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 = [
FunctionSpace(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 = VectorSpace(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 = [
FunctionSpace(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 = VectorSpace(Tp)
mask = T.get_mask_nyquist() if params.mask_nyquist else None
# 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]
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())
ua = ((x - 1)**2 * (y - 1) * (y + 1), 5 * (y + 1)**2 * (x - 1) * (x + 1))
f = (-ua[0].diff(x, 2) - ua[0].diff(y, 2),
-ua[1].diff(x, 2) - ua[1].diff(y, 2))
neumann_condition_x = ua[0].diff(x).evalf(subs={x: -1})
neumann_condition_y = ua[1].diff(y).evalf(subs={y: 1})
FXX = FunctionSpace(20, family='legendre', bc=(None, 0))
FXY = FunctionSpace(20, family='legendre', bc=(0, 0))
FYX = FunctionSpace(20, family='legendre', bc=(0, 0))
FYY = FunctionSpace(20, family='legendre', bc=(0, None))
TX = TensorProductSpace(comm, (FXX, FXY))
TY = TensorProductSpace(comm, (FYX, FYY))
V = VectorSpace([TX, TY])
u = TrialFunction(V)
v = TestFunction(V)
mat = inner(grad(u), grad(v))
fj = Array(V, buffer=f)
rhs = inner(v, fj)
# boundary integrals
# x - component
v_bndry_x = TestFunction(FXY)
gn_x = Array(FXY, buffer=neumann_condition_x)
evaluate_bndry_x = FXX.evaluate_basis_all(-1)
project_gn_x = inner(gn_x, v_bndry_x)
bndry_integral_x = -np.outer(evaluate_bndry_x, project_gn_x)
# y - component
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 = [
FunctionSpace(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 = VectorSpace(T)
VM = CompositeSpace([T] * 2 * dim)
mask = T.get_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 = [
FunctionSpace(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 = VectorSpace(Tp)
VMp = CompositeSpace([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]
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 = FunctionSpace(params.N[0], 'C', bc=(0, 0), quad=params.Dquad)
CT = FunctionSpace(params.N[0], 'C', quad=params.Dquad)
CP = FunctionSpace(params.N[0], 'C', 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, 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 = 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], 'C', bc=(0, 0), quad=params.Dquad)
CTp = FunctionSpace(params.N[0], 'C', 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")
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 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 = FunctionSpace(params.N[0], 'C', bc=(0, 0), quad=params.Dquad)
SB = FunctionSpace(params.N[0], 'C', bc='Biharmonic', quad=params.Bquad)
CT = FunctionSpace(params.N[0], 'C', quad=params.Dquad)
ST0 = FunctionSpace(params.N[0], 'C', bc=(0, 0),
quad=params.Dquad) # For 1D problem
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')
kw0 = {
'threads': params.threads,
'planner_effort': params.planner_effort["dct"],
'slab': (params.decomposition == 'slab'),
'collapse_fourier': collapse_fourier,
'modify_spaces_inplace': True
}
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 = VectorSpace([FSB, FST, FST])
VFST = VectorSpace([FST, FST, FST])
VUG = CompositeSpace([FSB, FST])
VCT = VectorSpace(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 = FunctionSpace(params.N[0], 'C', bc=(0, 0), quad=params.Dquad)
SBp = FunctionSpace(params.N[0],
'C',
bc='Biharmonic',
quad=params.Bquad)
CTp = FunctionSpace(params.N[0], 'C', quad=params.Dquad)
else:
STp, SBp, CTp = ST, SB, 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)
FSBp = TensorProductSpace(comm, (SBp, K0p, K1p), **kw0)
FCTp = TensorProductSpace(comm, (CTp, K0p, K1p), **kw0)
VFSp = VectorSpace([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 - 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())