def dx(u, FST, axis=0):
"""Compute integral of u over domain for channel solvers"""
sx = list(range(3))
sx.pop(axis)
uu = np.sum(u, axis=tuple(sx))
sl = FST.local_slice(False)[axis]
M = FST.shape()[axis]
c = aligned(M, fill=0)
cc = aligned(M, fill=0)
cc[sl] = uu
FST.comm.Reduce(cc, c, op=MPI.SUM, root=0)
quad = FST.bases[axis].quad
if FST.comm.Get_rank() == 0:
if quad == 'GL':
ak = aligned_like(c)
dct = dctn(aligned_like(c), axes=(0, ), type=1)
ak = dct(c, ak)
ak /= (M - 1)
w = np.arange(0, M, 1, dtype=float)
w[2:] = 2. / (1 - w[2:]**2)
w[0] = 1
w[1::2] = 0
return sum(ak * w) * np.prod(
np.take(config.params.L / config.params.N, sx))
assert quad == 'GC'
d = aligned(M, fill=0)
k = 2 * (1 + np.arange((M - 1) // 2))
d[::2] = (2. / M) / np.hstack((1., 1. - k * k))
w = aligned_like(d)
dct = dctn(w, axes=(0, ), type=3)
w = dct(d, w)
return np.sum(c * w) * np.prod(
np.take(config.params.L / config.params.N, sx))
return 0
def dx(u, FST, axis=0):
"""Compute integral of u over domain for channel solvers"""
sx = list(range(3))
sx.pop(axis)
uu = np.sum(u, axis=tuple(sx))
sl = FST.local_slice(False)[axis]
M = FST.shape()[axis]
c = aligned(M, fill=0)
cc = aligned(M, fill=0)
cc[sl] = uu
FST.comm.Reduce(cc, c, op=MPI.SUM, root=0)
quad = FST.bases[axis].quad
if FST.comm.Get_rank() == 0:
if quad == 'GL':
ak = aligned_like(c)
dct = dctn(aligned_like(c), axes=(0,), type=1)
ak = dct(c, ak)
ak /= (M-1)
w = np.arange(0, M, 1, dtype=float)
w[2:] = 2./(1-w[2:]**2)
w[0] = 1
w[1::2] = 0
return sum(ak*w)*np.prod(np.take(config.params.L/config.params.N, sx))
assert quad == 'GC'
d = aligned(M, fill=0)
k = 2*(1 + np.arange((M-1)//2))
d[::2] = (2./M)/np.hstack((1., 1.-k*k))
w = aligned_like(d)
dct = dctn(w, axes=(0,), type=3)
w = dct(d, w)
return np.sum(c*w)*np.prod(np.take(config.params.L/config.params.N, sx))
return 0
def test_r2r():
N = (5, 6, 7, 8, 9)
assert MPI.COMM_WORLD.Get_size() < 6
dctn = functools.partial(fftw.dctn, type=3)
idctn = functools.partial(fftw.idctn, type=3)
dstn = functools.partial(fftw.dstn, type=3)
idstn = functools.partial(fftw.idstn, type=3)
fft = PFFT(MPI.COMM_WORLD, N, axes=((0,), (1, 2), (3, 4)), grid=(-1,),
transforms={(1, 2): (dctn, idctn), (3, 4): (dstn, idstn)})
A = newDistArray(fft, forward_output=False)
A[:] = np.random.random(A.shape)
C = fftw.aligned_like(A)
B = fft.forward(A)
C = fft.backward(B, C)
assert np.allclose(A, C)
def main(N, alpha, method, tau):
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
xb = basis.ChebyshevT(c, size=N, bounds=(-1, 1))
x = xb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb,), dtype=np.float64)
ux = field.Field(name='ux', dist=d, bases=(xb,), dtype=np.float64)
f = field.Field(name='f', dist=d, bases=(xb,), dtype=np.float64)
t1 = field.Field(name='t1', dist=d, dtype=np.float64)
t2 = field.Field(name='t2', dist=d, dtype=np.float64)
pi, sin, cos = np.pi, np.sin, np.cos
f['g'] = 64*pi**3*(4*pi*sin(4*pi*x)**2*sin(4*pi*cos(4*pi*x)) + cos(4*pi*x)*cos(4*pi*cos(4*pi*x))) + alpha*sin(4*pi*cos(4*pi*x))
# Tau polynomials
xb1 = xb._new_a_b(xb.a+1, xb.b+1)
xb2 = xb._new_a_b(xb.a+2, xb.b+2)
if tau == 0:
# First-order classical Chebyshev tau: T[-1], dx(T[-1])
p1 = field.Field(name='p1', dist=d, bases=(xb,), dtype=np.float64)
p2 = field.Field(name='p2', dist=d, bases=(xb,), dtype=np.float64)
p1['c'][-1] = 1
p2['c'][-1] = 1
elif tau == 1:
# First-order ultraspherical tau: U[-1], dx(U[-1])
p1 = field.Field(name='p1', dist=d, bases=(xb1,), dtype=np.float64)
p2 = field.Field(name='p2', dist=d, bases=(xb1,), dtype=np.float64)
p1['c'][-1] = 1
p2['c'][-1] = 1
# Problem
dx = lambda A: operators.Differentiate(A, c)
problem = problems.LBVP([u, ux, t1, t2])
problem.add_equation((alpha*u - dx(ux) + t1*p1, f))
problem.add_equation((ux - dx(u) + t2*p2, 0))
problem.add_equation((u(x=-1), 0))
problem.add_equation((u(x=+1), 0))
solver = solvers.LinearBoundaryValueSolver(problem)
# Methods
if method == 0:
# Condition number
L_exp = solver.subproblems[0].L_exp
result = np.linalg.cond(L_exp.A)
elif method == 1:
# Roundtrip roundoff with uniform u
A = solver.subproblems[0].L_exp
v = np.random.rand(A.shape[1])
f = A * v
u = solver.subproblem_matsolvers[solver.subproblems[0]].solve(f)
result = np.max(np.abs(u-v))
elif method == 2:
# Manufactured solution
from mpi4py_fft import fftw as mpi4py_fftw
solver.solve()
ue = np.sin(4*np.pi*np.cos(4*np.pi*x))
d = mpi4py_fftw.aligned(N, fill=0)
k = 2*(1 + np.arange((N-1)//2))
d[::2] = (2./N)/np.hstack((1., 1.-k*k))
w = mpi4py_fftw.aligned_like(d)
dct = mpi4py_fftw.dctn(w, axes=(0,), type=3)
weights = dct(d, w)
result = np.sqrt(np.sum(weights*(u['g']-ue)**2))
elif method == 3:
# Roundtrip roundoff with uniform f
A = solver.subproblems[0].L_exp
g = np.random.rand(A.shape[0])
v = solver.subproblem_matsolvers[solver.subproblems[0]].solve(g)
f = A * v
u = solver.subproblem_matsolvers[solver.subproblems[0]].solve(f)
result = np.max(np.abs(u-v))
return result