def some_basic_tests():
import pyfftw #pylint: disable=import-error
from mpi4py import MPI
comm = MPI.COMM_WORLD
N = 8
K0 = C2CBasis(N)
K1 = C2CBasis(N)
K2 = C2CBasis(N)
K3 = R2CBasis(N)
T = TensorProductSpace(comm, (K0, K1, K2, K3))
# Create data on rank 0 for testing
if comm.Get_rank() == 0:
f_g = np.random.random(T.shape())
f_g_hat = pyfftw.interfaces.numpy_fft.rfftn(f_g, axes=(0, 1, 2, 3))
else:
f_g = np.zeros(T.shape())
f_g_hat = np.zeros(T.shape(), dtype=np.complex)
# Distribute test data to all ranks
comm.Bcast(f_g, root=0)
comm.Bcast(f_g_hat, root=0)
# Create a function in real space to hold the test data
fj = Array(T)
fj[:] = f_g[T.local_slice(False)]
# Perform forward transformation
f_hat = T.forward(fj)
assert np.allclose(f_g_hat[T.local_slice(True)], f_hat * N**4)
# Perform backward transformation
fj2 = Array(T)
fj2 = T.backward(f_hat)
assert np.allclose(fj, fj2)
f_hat = T.scalar_product(fj)
# Padding
# Needs new instances of bases because arrays have new sizes
Kp0 = C2CBasis(N, padding_factor=1.5)
Kp1 = C2CBasis(N, padding_factor=1.5)
Kp2 = C2CBasis(N, padding_factor=1.5)
Kp3 = R2CBasis(N, padding_factor=1.5)
Tp = TensorProductSpace(comm, (Kp0, Kp1, Kp2, Kp3))
f_g_pad = Tp.backward(f_hat)
f_hat2 = Tp.forward(f_g_pad)
assert np.allclose(f_hat2, f_hat)
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 __init__(self, comm, N):
method_common.__init__(self, comm, N)
# initial spectral spaces
K0 = C2CBasis(N[0], domain=(-2 * np.pi, 2 * np.pi))
K1 = R2CBasis(N[1], domain=(-2 * np.pi, 2 * np.pi))
self.T = TensorProductSpace(comm, (K0, K1),
**{'planner_effort': 'FFTW_MEASURE'})
# Use sympy to set up initial condition
x, y = symbols("x,y")
Le=500.*1e3
he = 1.0*exp(-(x**2 + y**2)/Le**2)
ue = 0.
ve = 0.
ul = lambdify((x, y), ue, 'numpy')
vl = lambdify((x, y), ve, 'numpy')
hl = lambdify((x, y), he, 'numpy')
# Size of discretization
N = (128, 128)
# initial spectral spaces
K0 = C2CBasis(N[0], domain=(-2*np.pi*L, 2*np.pi*L))
K1 = R2CBasis(N[1], domain=(-2*np.pi*L, 2*np.pi*L))
T = TensorProductSpace(comm, (K0, K1), **{'planner_effort': 'FFTW_MEASURE'})
# Turn on padding by commenting out:
Tp = T
#
TTT = MixedTensorProductSpace([T, T, T])
# init vectors
X = T.local_mesh(True) # physical grid
uvh = Array(TTT, False) # in physical space
u, v, h = uvh[:]
up = Array(Tp, False)
vp = Array(Tp, False)
# for time stepping, everything is in physical space
a = -2
b = 2
x, y = symbols("x,y")
ue = (cos(4 * np.pi * y) +
sin(2 * x)) * (1 - y**2) + a * (1 + y) / 2. + b * (1 - y) / 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 = (eval(sys.argv[-2]), eval(sys.argv[-2]))
SD = Basis(N[1], scaled=True, bc=(a, b))
K1 = R2CBasis(N[0])
T = TensorProductSpace(comm, (K1, SD), axes=(1, 0))
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 = fl(*X)
# Compute right hand side of Poisson equation
f_hat = Array(T)
f_hat = inner(v, fj, output_array=f_hat)
if basis == 'legendre':
f_hat *= -1.
# Use sympy to compute a rhs, given an analytical solution
alpha = 2.
x, y = symbols("x,y")
ue = (cos(4*np.pi*x) + sin(2*y))*(1-x**2)
fe = alpha*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 = (eval(sys.argv[-2]),)*2
SD = Basis(N[0], scaled=True)
K1 = R2CBasis(N[1])
T = TensorProductSpace(comm, (SD, K1), axes=(0, 1))
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 = fl(*X)
# Compute right hand side of Poisson equation
f_hat = Array(T)
f_hat = inner(v, fj, output_array=f_hat)
# Get left hand side of Helmholtz equation
if basis == 'chebyshev':
matrices = inner(v, alpha*u - div(grad(u)))
import shenfun from shenfun.operators import div, grad, Dx from shenfun import inner_product from shenfun.fourier.bases import R2CBasis from mpl_toolkits.mplot3d import axes3d from matplotlib.collections import PolyCollection from time import time N = 256 t = 0 dt = 0.1 / N**2 tstep = 0 tmax = 0.006 tsteps = int(tmax / dt) nplt = int(tmax / 25 / dt) SD = R2CBasis(N, True) Nh = SD.spectral_shape() x = SD.points_and_weights()[0] u = np.zeros(N) u_1 = np.zeros(N) u_2 = np.zeros(N) u_ab = np.zeros(N) u_hat = np.zeros(Nh, np.complex) u_hat_1 = np.zeros(Nh, np.complex) u_hat_2 = np.zeros(Nh, np.complex) uu_hat = np.zeros(Nh, np.complex) k = SD.wavenumbers(N) A = 25. B = 16.
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())
rank = comm.Get_rank()
# Use sympy to set up initial condition
x, y, z = symbols("x,y,z")
ue = 0.1 * exp(-(x**2 + y**2 + z**2))
ul = lambdify((x, y, z), ue, 'numpy')
# Size of discretization
N = (32, 32, 32)
# Defocusing or focusing
gamma = 1
K0 = C2CBasis(N[0], domain=(-2 * np.pi, 2 * np.pi))
K1 = C2CBasis(N[1], domain=(-2 * np.pi, 2 * np.pi))
K2 = R2CBasis(N[2], domain=(-2 * np.pi, 2 * np.pi))
T = TensorProductSpace(comm, (K0, K1, K2),
**{'planner_effort': 'FFTW_MEASURE'})
TT = MixedTensorProductSpace([T, T])
TV = VectorTensorProductSpace([T, T, T])
X = T.local_mesh(True)
fu = Array(TT, False) # solution real space
f, u = fu[:] # split into views
up = Array(T, False) # work array real space
w0 = Array(T) # work array spectral space
dfu = Array(TT) # solution spectral space
df, du = dfu[:] # views
comm = MPI.COMM_WORLD
# Use sympy to compute a rhs, given an analytical solution
x, y, z = symbols("x,y,z")
ue = cos(4 * x) + sin(4 * y) + sin(6 * z)
fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2)
ul = lambdify((x, y, z), ue, 'numpy')
fl = lambdify((x, y, z), fe, 'numpy')
# Size of discretization
N = (14, 15, 16)
K0 = C2CBasis(N[0])
K1 = C2CBasis(N[1])
K2 = R2CBasis(N[2])
T = TensorProductSpace(comm, (K0, K1, K2))
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 = fl(*X)
# Compute right hand side
f_hat = inner(v, fj)
# Solve Poisson equation
A = inner(v, div(grad(u)))
timer = Timer()
# Use sympy to set up initial condition
x, y, z = symbols("x,y,z")
ue = 0.1*exp(-(x**2 + y**2 + z**2))
ul = lambdify((x, y, z), ue, 'numpy')
# Size of discretization
N = (32, 32, 32)
# Defocusing or focusing
gamma = 1
K0 = C2CBasis(N[0], domain=(-2*np.pi, 2*np.pi))
K1 = C2CBasis(N[1], domain=(-2*np.pi, 2*np.pi))
K2 = R2CBasis(N[2], domain=(-2*np.pi, 2*np.pi))
T = TensorProductSpace(comm, (K0, K1, K2), slab=False, **{'planner_effort': 'FFTW_MEASURE'})
TT = MixedTensorProductSpace([T, T])
Kp0 = C2CBasis(N[0], domain=(-2*np.pi, 2*np.pi), padding_factor=1.5)
Kp1 = C2CBasis(N[1], domain=(-2*np.pi, 2*np.pi), padding_factor=1.5)
Kp2 = R2CBasis(N[2], domain=(-2*np.pi, 2*np.pi), padding_factor=1.5)
Tp = TensorProductSpace(comm, (Kp0, Kp1, Kp2), slab=False, **{'planner_effort': 'FFTW_MEASURE'})
# Turn on padding by commenting out:
Tp = T
file0 = HDF5Writer("KleinGordon{}.h5".format(N[0]), ['u', 'f'], TT)
X = T.local_mesh(True)
from mpi4py import MPI
from shenfun.fourier.bases import R2CBasis, C2CBasis
from shenfun import *
comm = MPI.COMM_WORLD
# Use sympy to set up initial condition
x, y = symbols("x,y")
ue = exp(-0.01*(x**2+y**2)) # + exp(-0.02*((x-15*np.pi)**2+(y)**2))
ul = lambdify((x, y), ue, 'numpy')
# Size of discretization
N = (128, 128)
K0 = C2CBasis(N[0], domain=(-30*np.pi, 30*np.pi))
K1 = R2CBasis(N[1], domain=(-30*np.pi, 30*np.pi))
T = TensorProductSpace(comm, (K0, K1), **{'planner_effort': 'FFTW_MEASURE'})
TV = VectorTensorProductSpace([T, T])
u = TrialFunction(T)
v = TestFunction(T)
# Create solution and work arrays
U = Array(T, False)
U_hat = Array(T)
gradu = Array(TV, False)
K = np.array(T.local_wavenumbers(True, True, True))
def LinearRHS():
# Assemble diagonal bilinear forms
A = inner(u, v)
comm = MPI.COMM_WORLD
# Use sympy to compute a rhs, given an analytical solution
x, y, z = symbols("x,y,z")
ue = cos(4 * x) + sin(4 * y) + sin(6 * z)
fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2) + ue
ul = lambdify((x, y, z), ue, 'numpy')
fl = lambdify((x, y, z), fe, 'numpy')
# Size of discretization
N = 16
K0 = C2CBasis(N)
K1 = C2CBasis(N)
K2 = R2CBasis(N)
T = TensorProductSpace(comm, (K0, K1, K2), slab=True)
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 = fl(*X)
# Compute right hand side
f_hat = inner(v, fj)
# Solve Poisson equation
A = inner(v, u + div(grad(u)))