def test_helmholtz3D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9, 10)
SD = FunctionSpace(N[allaxes3D[axis][0]], family=family, bc=(0, 0))
K1 = FunctionSpace(N[allaxes3D[axis][1]], family='F', dtype='D')
K2 = FunctionSpace(N[allaxes3D[axis][2]], family='F', dtype='d')
subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, [0, 1, 1])
bases = [0]*3
bases[allaxes3D[axis][0]] = SD
bases[allaxes3D[axis][1]] = K1
bases[allaxes3D[axis][2]] = K2
T = TensorProductSpace(subcomms, bases, axes=allaxes3D[axis], modify_spaces_inplace=True)
u = shenfun.TrialFunction(T)
v = shenfun.TestFunction(T)
if family == 'chebyshev':
mat = inner(v, div(grad(u)))
else:
mat = inner(grad(v), grad(u))
H = la.Helmholtz(*mat)
u = Function(T)
s = SD.sl[SD.slice()]
u[s] = np.random.random(u[s].shape) + 1j*np.random.random(u[s].shape)
f = Function(T)
f = H.matvec(u, f)
g0 = Function(T)
g1 = Function(T)
M = {d.get_key(): d for d in mat}
g0 = M['ADDmat'].matvec(u, g0)
g1 = M['BDDmat'].matvec(u, g1)
assert np.linalg.norm(f-(g0+g1)) < 1e-12, np.linalg.norm(f-(g0+g1))
uc = Function(T)
uc = H(uc, f)
assert np.linalg.norm(uc-u) < 1e-12
def test_biharmonic3D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (16, 16, 16)
SD = FunctionSpace(N[allaxes3D[axis][0]], family=family, bc='Biharmonic')
K1 = FunctionSpace(N[allaxes3D[axis][1]], family='F', dtype='D')
K2 = FunctionSpace(N[allaxes3D[axis][2]], family='F', dtype='d')
subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, [0, 1, 1])
bases = [0] * 3
bases[allaxes3D[axis][0]] = SD
bases[allaxes3D[axis][1]] = K1
bases[allaxes3D[axis][2]] = K2
T = TensorProductSpace(subcomms, bases, axes=allaxes3D[axis])
u = shenfun.TrialFunction(T)
v = shenfun.TestFunction(T)
if family == 'chebyshev':
mat = inner(v, div(grad(div(grad(u)))))
else:
mat = inner(div(grad(v)), div(grad(u)))
H = la.Biharmonic(*mat)
u = Function(T)
u[:] = np.random.random(u.shape) + 1j * np.random.random(u.shape)
f = Function(T)
f = H.matvec(u, f)
g0 = Function(T)
g1 = Function(T)
g2 = Function(T)
M = {d.get_key(): d for d in mat}
g0 = M['SBBmat'].matvec(u, g0)
g1 = M['ABBmat'].matvec(u, g1)
g2 = M['BBBmat'].matvec(u, g2)
assert np.linalg.norm(f - (g0 + g1 + g2)) < 1e-8, np.linalg.norm(f -
(g0 + g1 +
g2))
def test_eval_expression2(fam):
import sympy as sp
from shenfun import div, grad
x, y = sp.symbols('x,y')
B0 = FunctionSpace(16, fam, domain=(-2, 2))
B1 = FunctionSpace(17, fam, domain=(-1, 3))
T = TensorProductSpace(comm, (B0, B1))
f = sp.sin(x)+sp.sin(y)
dfx = f.diff(x, 2) + f.diff(y, 2)
fa = Function(T, buffer=f)
dfe = div(grad(fa))
dfa = project(dfe, T)
xy = np.array([[0.25, 0.5, 0.75],
[0.25, 0.5, 0.75]])
f0 = lambdify((x, y), dfx)(*xy)
f1 = dfe.eval(xy)
f2 = dfa.eval(xy)
assert np.allclose(f0, f1, 1e-7)
assert np.allclose(f1, f2, 1e-7)
def test_project_hermite(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (24, 23)
funcs = {
(1, 0): (cos(4 * y) * sin(2 * x)) * exp(-x**2 / 2),
(1, 1): (cos(4 * x) * sin(2 * y)) * exp(-y**2 / 2),
(2, 0): (sin(6 * z) * cos(4 * y) * sin(2 * x)) * exp(-x**2 / 2),
(2, 1): (sin(2 * z) * cos(4 * x) * sin(2 * y)) * exp(-y**2 / 2),
(2, 2): (sin(2 * x) * cos(4 * y) * sin(2 * z)) * exp(-z**2 / 2)
}
syms = {1: (x, y), 2: (x, y, z)}
xs = {0: x, 1: y, 2: z}
for shape in product(*([sizes] * dim)):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
for axis in range(dim + 1):
ST0 = ST(3 * shape[-1], quad=quad)
bases.insert(axis, ST0)
fft = TensorProductSpace(comm,
bases,
dtype=typecode,
axes=axes[dim][axis])
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
u_h = project(ue, fft)
ul = lambdify(syms[dim], ue, 'numpy')
uq = ul(*X).astype(typecode)
uf = u_h.backward()
assert np.linalg.norm(uq - uf) < 1e-5
bases.pop(axis)
fft.destroy()
def test_helmholtz2D(family, axis):
la = lla
if family == 'chebyshev':
la = cla
N = (8, 9)
SD = FunctionSpace(N[axis], family=family, bc=(0, 0))
K1 = FunctionSpace(N[(axis + 1) % 2], family='F', dtype='d')
subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, allaxes2D[axis])
bases = [K1]
bases.insert(axis, SD)
T = TensorProductSpace(subcomms, bases, axes=allaxes2D[axis])
u = shenfun.TrialFunction(T)
v = shenfun.TestFunction(T)
if family == 'chebyshev':
mat = inner(v, div(grad(u)))
else:
mat = inner(grad(v), grad(u))
H = la.Helmholtz(*mat)
u = Function(T)
s = SD.sl[SD.slice()]
u[s] = np.random.random(u[s].shape) + 1j * np.random.random(u[s].shape)
f = Function(T)
f = H.matvec(u, f)
g0 = Function(T)
g1 = Function(T)
M = {d.get_key(): d for d in mat}
g0 = M['ADDmat'].matvec(u, g0)
g1 = M['BDDmat'].matvec(u, g1)
assert np.linalg.norm(f - (g0 + g1)) < 1e-12, np.linalg.norm(f - (g0 + g1))
uc = Function(T)
uc = H(uc, f)
assert np.linalg.norm(uc - u) < 1e-12
def test_project_lag(typecode, dim):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (20, 17)
funcs = {
(1, 0): (cos(4*y)*sin(2*x))*exp(-x),
(1, 1): (cos(4*x)*sin(2*y))*exp(-y),
(2, 0): (sin(3*z)*cos(4*y)*sin(2*x))*exp(-x),
(2, 1): (sin(2*z)*cos(4*x)*sin(2*y))*exp(-y),
(2, 2): (sin(2*x)*cos(4*y)*sin(2*z))*exp(-z)
}
xs = {0:x, 1:y, 2:z}
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))
for axis in range(dim+1):
ST1 = lagBasis[1](3*shape[-1])
bases.insert(axis, ST1)
fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
dfft = fft.get_orthogonal()
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
due = ue.diff(xs[0], 1)
u_h = project(ue, fft)
du_h = project(due, dfft)
du2 = project(Dx(u_h, 0, 1), dfft)
uf = u_h.backward()
assert np.linalg.norm(du2-du_h) < 1e-3
bases.pop(axis)
fft.destroy()
dfft.destroy()
def test_Mult_CTD_3D(quad):
SD = FunctionSpace(N, 'C', bc=(0, 0))
F0 = FunctionSpace(4, 'F', dtype='D')
F1 = FunctionSpace(4, 'F', dtype='d')
T = TensorProductSpace(comm, (SD, F0, F1))
TO = T.get_orthogonal()
CT = TO.bases[0]
C = inner_product((CT, 0), (SD, 1))
B = inner_product((CT, 0), (CT, 0))
vk = Array(T)
wk = Array(T)
vk[:] = np.random.random(vk.shape)
wk[:] = np.random.random(vk.shape)
bv = Function(T)
bw = Function(T)
vk0 = vk.forward()
vk = vk0.backward()
wk0 = wk.forward()
wk = wk0.backward()
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 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())
import matplotlib.pyplot as plt
from mpi4py_fft import generate_xdmf, fftw
from shenfun import inner, div, grad, TestFunction, TrialFunction, \
TensorProductSpace, Array, Function, ETDRK4, HDF5File, FunctionSpace, comm
# Use sympy to set up initial condition
x, y = symbols("x,y", real=True)
#ue = (1j*x + y)*exp(-0.03*(x**2+y**2))
ue = (x + y) * exp(-0.03 * (x**2 + y**2))
# Size of discretization
N = (129, 129)
K0 = FunctionSpace(N[0], 'F', dtype='D', domain=(-50, 50))
K1 = FunctionSpace(N[1], 'F', dtype='D', domain=(-50, 50))
T = TensorProductSpace(comm, (K0, K1), **{'planner_effort': 'FFTW_MEASURE'})
Tp = T.get_dealiased((1.5, 1.5))
u = TrialFunction(T)
v = TestFunction(T)
# Try to import wisdom. Note that wisdom must be imported after creating the Bases (that initializes the wisdom somehow?)
try:
fftw.import_wisdom('GL.wisdom')
print('Importing wisdom')
except:
print('No wisdom imported')
# Turn on padding by commenting:
#Tp = T
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 test_eval_tensor(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
# Testing for Dirichlet and regular basis
x, y, z = symbols("x,y,z")
sizes = (22, 21)
funcx = {
'': (1 - x**2) * sin(np.pi * x),
'Dirichlet': (1 - x**2) * sin(np.pi * x),
'Neumann': (1 - x**2) * sin(np.pi * x),
'Biharmonic': (1 - x**2) * sin(2 * np.pi * x)
}
funcy = {
'': (1 - y**2) * sin(np.pi * y),
'Dirichlet': (1 - y**2) * sin(np.pi * y),
'Neumann': (1 - y**2) * sin(np.pi * y),
'Biharmonic': (1 - y**2) * sin(2 * np.pi * y)
}
funcz = {
'': (1 - z**2) * sin(np.pi * z),
'Dirichlet': (1 - z**2) * sin(np.pi * z),
'Neumann': (1 - z**2) * sin(np.pi * z),
'Biharmonic': (1 - z**2) * sin(2 * np.pi * z)
}
funcs = {
(1, 0): cos(2 * y) * funcx[ST.boundary_condition()],
(1, 1): cos(2 * x) * funcy[ST.boundary_condition()],
(2, 0): sin(6 * z) * cos(4 * y) * funcx[ST.boundary_condition()],
(2, 1): sin(2 * z) * cos(4 * x) * funcy[ST.boundary_condition()],
(2, 2): sin(2 * x) * cos(4 * y) * funcz[ST.boundary_condition()]
}
syms = {1: (x, y), 2: (x, y, z)}
points = None
if comm.Get_rank() == 0:
points = np.random.random((dim + 1, 4))
points = comm.bcast(points)
t_0 = 0
t_1 = 0
t_2 = 0
for shape in product(*([sizes] * dim)):
#for shape in ((64, 64),):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
for axis in range(dim + 1):
#for axis in (0,):
ST0 = ST(shape[-1], quad=quad)
bases.insert(axis, ST0)
# Spectral space must be aligned in nonperiodic direction, hence axes
fft = TensorProductSpace(comm,
bases,
dtype=typecode,
axes=axes[dim][axis])
print('axes', axes[dim][axis])
print('bases', bases)
#print(bases[0].axis, bases[1].axis)
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
ul = lambdify(syms[dim], ue, 'numpy')
uu = ul(*X).astype(typecode)
uq = ul(*points).astype(typecode)
u_hat = Function(fft)
u_hat = fft.forward(uu, u_hat)
t0 = time()
result = fft.eval(points, u_hat, method=0)
t_0 += time() - t0
assert np.allclose(uq, result, 0, 1e-6)
t0 = time()
result = fft.eval(points, u_hat, method=1)
t_1 += time() - t0
assert np.allclose(uq, result, 0, 1e-6)
t0 = time()
result = fft.eval(points, u_hat, method=2)
t_2 += time() - t0
print(uq)
assert np.allclose(uq, result, 0, 1e-6), uq / result
result = u_hat.eval(points)
assert np.allclose(uq, result, 0, 1e-6)
ua = u_hat.backward()
assert np.allclose(uu, ua, 0, 1e-6)
ua = Array(fft)
ua = u_hat.backward(ua)
assert np.allclose(uu, ua, 0, 1e-6)
bases.pop(axis)
fft.destroy()
print('method=0', t_0)
print('method=1', t_1)
print('method=2', t_2)
def test_project2(typecode, dim, ST, quad):
# Using sympy to compute an analytical solution
x, y, z = symbols("x,y,z")
sizes = (22, 21)
funcx = ((2 * np.pi**2 * (x**2 - 1) - 1) * cos(2 * np.pi * x) -
2 * np.pi * x * sin(2 * np.pi * x)) / (4 * np.pi**3)
funcy = ((2 * np.pi**2 * (y**2 - 1) - 1) * cos(2 * np.pi * y) -
2 * np.pi * y * sin(2 * np.pi * y)) / (4 * np.pi**3)
funcz = ((2 * np.pi**2 * (z**2 - 1) - 1) * cos(2 * np.pi * z) -
2 * np.pi * z * sin(2 * np.pi * z)) / (4 * np.pi**3)
funcs = {
(1, 0): cos(4 * y) * funcx,
(1, 1): cos(4 * x) * funcy,
(2, 0): sin(6 * z) * cos(4 * y) * funcx,
(2, 1): sin(2 * z) * cos(4 * x) * funcy,
(2, 2): sin(2 * x) * cos(4 * y) * funcz
}
syms = {1: (x, y), 2: (x, y, z)}
xs = {0: x, 1: y, 2: z}
for shape in product(*([sizes] * dim)):
bases = []
for n in shape[:-1]:
bases.append(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(shape[-1], 'F', dtype=typecode))
for axis in range(dim + 1):
ST0 = ST(shape[-1], quad=quad)
bases.insert(axis, ST0)
# Spectral space must be aligned in nonperiodic direction, hence axes
fft = TensorProductSpace(comm,
bases,
dtype=typecode,
axes=axes[dim][axis])
X = fft.local_mesh(True)
ue = funcs[(dim, axis)]
ul = lambdify(syms[dim], ue, 'numpy')
uq = ul(*X).astype(typecode)
uh = Function(fft)
uh = fft.forward(uq, uh)
due = ue.diff(xs[axis], 1)
dul = lambdify(syms[dim], due, 'numpy')
duq = dul(*X).astype(typecode)
uf = project(Dx(uh, axis, 1), fft)
uy = Array(fft)
uy = fft.backward(uf, uy)
assert np.allclose(uy, duq, 0, 1e-5)
# Test also several derivatives
for ax in (x for x in range(dim + 1) if x is not axis):
due = ue.diff(xs[ax], 1, xs[axis], 1)
dul = lambdify(syms[dim], due, 'numpy')
duq = dul(*X).astype(typecode)
uf = project(Dx(Dx(uh, ax, 1), axis, 1), fft)
uy = Array(fft)
uy = fft.backward(uf, uy)
assert np.allclose(uy, duq, 0, 1e-5)
bases.pop(axis)
fft.destroy()
# Use sympy to compute a rhs, given an analytical solution
a = 1.
b = -1.
if family == 'jacobi':
a = 0
b = 0
x, y = symbols("x,y")
ue = (cos(4*x) + sin(2*y))*(1 - x**2) + a*(1 - x)/2. + b*(1 + x)/2.
fe = ue.diff(x, 2) + ue.diff(y, 2)
# Size of discretization
N = (int(sys.argv[-2]), int(sys.argv[-2])+1)
SD = FunctionSpace(N[0], family=family, scaled=True, bc=(a, b))
K1 = FunctionSpace(N[1], family='F', dtype='d', domain=(-2*np.pi, 2*np.pi))
T = TensorProductSpace(comm, (SD, K1), axes=(0, 1))
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side of Poisson equation
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
# Get left hand side of Poisson equation
matrices = inner(v, div(grad(u)))
# Create Helmholtz linear algebra solver
H = Solver(*matrices)
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())
# Use sympy to compute a rhs, given an analytical solution
a = 1.
x, y = symbols("x,y")
ue = (cos(4*y)*sin(2*x))*(1-x**2)*(1-y**2)
fe = a*ue - ue.diff(x, 2) - ue.diff(y, 2)
# Lambdify for faster evaluation
ul = lambdify((x, y), ue, 'numpy')
fl = lambdify((x, y), fe, 'numpy')
# Size of discretization
N = (int(sys.argv[-3]), int(sys.argv[-2]))
SD0 = Basis(N[0], family, bc=(0, 0), scaled=True)
SD1 = Basis(N[1], family, bc=(0, 0), scaled=True)
T = TensorProductSpace(comm, (SD0, SD1))
X = T.local_mesh(True)
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fl(*X))
# Compute right hand side of Poisson equation
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
# Get left hand side of Poisson equation
if family == 'legendre':
matrices = inner(grad(v), grad(u))
else:
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():
float, complex, mpitype = datatypes(params.precision)
"""Set up context for classical (NS) solver"""
V0 = Basis(params.N[0], 'F', domain=(0, params.L[0]), dtype=complex)
V1 = Basis(params.N[1], 'F', domain=(0, params.L[1]), dtype=complex)
V2 = Basis(params.N[2], 'F', domain=(0, params.L[2]), dtype=float)
kw0 = {'threads': params.threads,
'planner_effort': params.planner_effort['fft']}
T = TensorProductSpace(comm, (V0, V1, V2), dtype=float,
slab=params.decomposition=='slab', **kw0)
VT = VectorTensorProductSpace(T)
kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
'dealias_direct': params.dealias == '2/3-rule'}
V0p = Basis(params.N[0], 'F', domain=(0, params.L[0]), dtype=complex, **kw)
V1p = Basis(params.N[1], 'F', domain=(0, params.L[1]), dtype=complex, **kw)
V2p = Basis(params.N[2], 'F', domain=(0, params.L[2]), dtype=float, **kw)
Tp = TensorProductSpace(comm, (V0p, V1p, V2p), dtype=float,
slab=params.decomposition=='slab', **kw0)
VTp = VectorTensorProductSpace(Tp)
FFT = T # For compatibility - to be removed
# Mesh variables
X = T.local_mesh(True)
K = T.local_wavenumbers(scaled=True)
for i in range(3):
X[i] = X[i].astype(float)
K[i] = K[i].astype(float)
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)
for i in range(3):
Kx[i] = Kx[i].astype(float)
K_over_K2 = np.zeros((3,)+VT.local_shape(), dtype=float)
for i in range(3):
K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)
# Velocity and pressure
U = Array(VT)
U_hat = Function(VT)
P = Array(T)
P_hat = Function(T)
# Primary variable
u = U_hat
# RHS array
dU = Function(VT)
curl = Array(VT)
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())
a = 0
b = 0
x, y, z = symbols("x,y,z", real=True)
ue = (cos(4 * x) + sin(2 * y) +
sin(4 * z)) * (1 - y**2) + a * (1 - y) / 2. + b * (1 + y) / 2.
# Size of discretization
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
# Use sympy to compute a rhs, given an analytical solution
x, y, z = symbols("x,y,z")
ue = (sin(2 * np.pi * z) * sin(4 * np.pi * y) *
cos(4 * x)) * (1 - y**2) * (1 - z**2)
fe = ue.diff(x, 4) + ue.diff(y, 4) + ue.diff(z, 4) + 2 * ue.diff(
x, 2, y, 2) + 2 * ue.diff(x, 2, z, 2) + 2 * ue.diff(y, 2, z, 2)
# Size of discretization
N = (36, 36, 36)
K0 = FunctionSpace(N[0], 'Fourier', dtype='d')
S0 = FunctionSpace(N[1], family=family, bc='Biharmonic')
S1 = FunctionSpace(N[2], family=family, bc='Biharmonic')
T = TensorProductSpace(comm, (K0, S0, S1), axes=(1, 0, 2), slab=True)
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side of biharmonic equation
f_hat = inner(v, fj)
# Get left hand side of biharmonic equation
matrices = inner(v, div(grad(div(grad(u)))))
# Create linear algebra solver
H = SolverGeneric2ND(matrices)
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())
# Use sympy to compute a rhs, given an analytical solution
x, y, z = symbols("x,y,z", real=True)
ue = (sin(4 * np.pi * x) * sin(6 * z) * cos(4 * y)) * (1 - x**2)
fe = ue.diff(x, 4) + ue.diff(y, 4) + ue.diff(z, 4) + 2 * ue.diff(
x, 2, y, 2) + 2 * ue.diff(x, 2, z, 2) + 2 * ue.diff(y, 2, z, 2)
# Size of discretization
N = (36, 36, 36)
if family == 'chebyshev':
assert N[0] % 2 == 0, "Biharmonic solver only implemented for even numbers"
SD = FunctionSpace(N[0], family=family, bc=(0, 0, 0, 0))
K1 = FunctionSpace(N[1], family='F', dtype='D')
K2 = FunctionSpace(N[2], family='F', dtype='d')
T = TensorProductSpace(comm, (SD, K1, K2), axes=(0, 1, 2))
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side of biharmonic equation
f_hat = inner(v, fj)
# Get left hand side of biharmonic equation
matrices = inner(v, div(grad(div(grad(u)))))
# Create linear algebra solver
H = BiharmonicSolver(*matrices)
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 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(Basis(n, 'F', dtype=typecode.upper()))
bases.append(Basis(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 = VectorTensorProductSpace(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 = MixedTensorProductSpace([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)
fft.destroy()
padding = 1.5
bases = []
for n in shape[:-1]:
bases.append(
Basis(n,
'F',
dtype=typecode.upper(),
padding_factor=padding))
bases.append(
Basis(shape[-1], 'F', dtype=typecode, padding_factor=padding))
fft = TensorProductSpace(comm, bases, dtype=typecode)
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)
Fc = F.copy()
V = fft.backward(F)
F = fft.forward(V)
assert allclose(F, Fc)
# Alternative method
fft.backward.input_array[...] = F
fft.backward()
fft.forward()
V = fft.forward.output_array
assert allclose(F, V)
fft.destroy()
u0 = 0.5 * (1 - ((0.5 * (erf((x - 0.04) / a) + 1) - 0.5 * (erf(
(x + 0.04) / a) + 1)) * (0.5 * (erf((y - 0.04) / a) + 1) - 0.5 * (erf(
(y + 0.04) / a) + 1)))) + 0.5
v0 = 0.25 * (0.5 * (erf((x - 0.04) / a) + 1) - 0.5 * (erf(
(x + 0.04) / a) + 1)) * (0.5 * (erf((y - 0.04) / a) + 1) - 0.5 * (erf(
(y + 0.04) / a) + 1))
ul = lambdify((x, y), u0, modules=['numpy', {'erf': scipy.special.erf}])
vl = lambdify((x, y), v0, modules=['numpy', {'erf': scipy.special.erf}])
# Size of discretization
N = (200, 200)
K0 = Basis(N[0], 'F', dtype='D', domain=(-1., 1.))
K1 = Basis(N[1], 'F', dtype='d', domain=(-1., 1.))
T = TensorProductSpace(comm, (K0, K1))
X = T.local_mesh(True)
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 = VectorTensorProductSpace(T)
TVp = VectorTensorProductSpace(Tp)
vv = TestFunction(TV)
uu = TrialFunction(TV)
# Use sympy to compute a rhs, given an analytical solution
x, y = symbols("x,y")
ue = (sin(2 * np.pi * x) * sin(4 * np.pi * y)) * (1 - x**2) * (1 - y**2)
fe = ue.diff(x, 4) + ue.diff(y, 4) + 2 * ue.diff(x, 2, y, 2)
# Lambdify for faster evaluation
ul = lambdify((x, y), ue, 'numpy')
fl = lambdify((x, y), fe, 'numpy')
# Size of discretization
N = (36, 36)
S0 = Basis(N[0], family=family, bc='Biharmonic')
S1 = Basis(N[1], family=family, bc='Biharmonic')
T = TensorProductSpace(comm, (S0, S1), axes=(0, 1))
X = T.local_mesh(True)
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fl(*X))
# Compute right hand side of biharmonic equation
f_hat = inner(v, fj)
# Get left hand side of biharmonic equation
if family == 'chebyshev': # No integration by parts due to weights
matrices = inner(v, div(grad(div(grad(u)))))
else: # Use form with integration by parts.
matrices = inner(div(grad(v)), div(grad(u)))
TensorProductSpace, Array, Function, dx
comm = MPI.COMM_WORLD
# Use sympy to compute a rhs, given an analytical solution
x, y, z = symbols("x,y,z", real=True)
ue = cos(4*x) + sin(4*y) + sin(6*z)
fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2) + ue
# Size of discretization
N = 16
K0 = FunctionSpace(N, 'F', dtype='D')
K1 = FunctionSpace(N, 'F', dtype='D')
K2 = FunctionSpace(N, 'F', dtype='d')
T = TensorProductSpace(comm, (K0, K1, K2), slab=True)
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
# Solve Poisson equation
A = inner(v, u+div(grad(u)))
f_hat = A.solve(f_hat)
uq = T.backward(f_hat, fast_transform=True)
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())
# Use sympy to compute a rhs, given an analytical solution
x, y = symbols("x,y")
ue = cos(4*y)*sin(2*np.pi*x)*(1-x**2)
fe = ue.diff(x, 2) + ue.diff(y, 2)
# Lambdify for faster evaluation
ul = lambdify((x, y), ue, 'numpy')
fl = lambdify((x, y), fe, 'numpy')
# Size of discretization
N = (31, 32)
SD = Basis(N[0], family=family, bc='Neumann')
K1 = Basis(N[1], family='F', dtype='d')
T = TensorProductSpace(comm, (SD, K1))
X = T.local_mesh(True) # With broadcasting=True the shape of X is local_shape, even though the number of datapoints are still the same as in 1D
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fl(*X))
# Compute right hand side of Poisson equation
f_hat = inner(v, fj)
if family == 'legendre':
f_hat *= -1.
# Get left hand side of Poisson equation
if family == 'chebyshev':
matrices = inner(v, div(grad(u)))
comm = MPI.COMM_WORLD
# Use sympy to compute a rhs, given an analytical solution
x, y, z, r = symbols("x,y,z,r")
ue = cos(4 * x) + sin(4 * y) + sin(6 * z) + cos(6 * r)
fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2) + ue.diff(r, 2)
# Size of discretization
N = (8, 10, 12, 14)
K0 = Basis(N[0], 'F', dtype='D')
K1 = Basis(N[1], 'F', dtype='D')
K2 = Basis(N[2], 'F', dtype='D')
K3 = Basis(N[3], 'F', dtype='d')
T = TensorProductSpace(comm, (K0, K1, K2, K3))
X = T.local_mesh(
True
) # With broadcasting=True the shape of X is local_shape, even though the number of datapoints are still the same as in 1D
u = TrialFunction(T)
v = TestFunction(T)
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
# Solve Poisson equation
A = inner(v, div(grad(u)))
y = Symbol("y")
# Initial conditions
a = 0.0001
u0 = 0.5*(1-((0.5*(erf((x-0.04)/a)+1) - 0.5*(erf((x+0.04)/a)+1))*(0.5*(erf((y-0.04)/a)+1) - 0.5*(erf((y+0.04)/a)+1))))+0.5
v0 = 0.25*(0.5*(erf((x-0.04)/a)+1) - 0.5*(erf((x+0.04)/a)+1))*(0.5*(erf((y-0.04)/a)+1) - 0.5*(erf((y+0.04)/a)+1))
ul = lambdify((x, y), u0, modules=['numpy', {'erf': scipy.special.erf}])
vl = lambdify((x, y), v0, modules=['numpy', {'erf': scipy.special.erf}])
# Size of discretization
N = (200, 200)
K0 = Basis(N[0], 'F', dtype='D', domain=(-1., 1.))
K1 = Basis(N[1], 'F', dtype='d', domain=(-1., 1.))
T = TensorProductSpace(comm, (K0, K1))
X = T.local_mesh(True)
u = TrialFunction(T)
v = TestFunction(T)
# For nonlinear term we can use the 3/2-rule with padding
Kp0 = Basis(N[0], 'F', dtype='D', domain=(-1., 1.), padding_factor=1.5)
Kp1 = Basis(N[1], 'F', dtype='d', domain=(-1., 1.), padding_factor=1.5)
Tp = TensorProductSpace(comm, (Kp0, Kp1))
# Turn on padding by commenting
Tp = T
# Create vector spaces and a test function for the regular vector space
TV = VectorTensorProductSpace(T)
TVp = VectorTensorProductSpace(Tp)
