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 plan(self, shape, axis, dtype, options):
if shape in (0, (0, )):
return
if isinstance(axis, tuple):
assert len(axis) == 1
axis = axis[0]
if isinstance(self.forward, Transform):
if self.forward.input_array.shape == shape and self.axis == axis:
# Already planned
return
U = fftw.aligned(shape, dtype=dtype)
V = fftw.aligned(shape, dtype=dtype)
U.fill(0)
V.fill(0)
self.axis = axis
if self.padding_factor > 1. + 1e-8:
trunc_array = self._get_truncarray(shape, V.dtype)
self.forward = Transform(self.forward, None, U, V, trunc_array)
self.backward = Transform(self.backward, None, trunc_array, V, U)
else:
self.forward = Transform(self.forward, None, U, V, V)
self.backward = Transform(self.backward, None, V, V, U)
self.scalar_product = Transform(self.scalar_product, None, U, V, V)
self.si = islicedict(axis=self.axis, dimensions=self.dimensions)
self.sl = slicedict(axis=self.axis, dimensions=self.dimensions)
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 plan(self, shape, axis, dtype, options):
if isinstance(axis, tuple):
assert len(axis) == 1
axis = axis[0]
if isinstance(self.forward, Transform):
if self.forward.input_array.shape == shape and self.axis == axis:
# Already planned
return
U = fftw.aligned(shape, dtype=dtype)
V = fftw.aligned(shape, dtype=dtype)
U.fill(0)
V.fill(0)
self.axis = axis
self.forward = Transform(self.forward, None, U, V, V)
self.backward = Transform(self.backward, None, V, V, U)
self.scalar_product = Transform(self.scalar_product, None, U, V, V)
def _get_truncarray(self, shape, dtype):
shape = list(shape)
shape[self.axis] = int(shape[self.axis] / self.padding_factor)
shape[self.axis] = shape[self.axis]//2 + 1
return fftw.aligned(shape, dtype=dtype)
def plan(self, shape, axis, dtype, options):
if isinstance(axis, int):
axis = [axis]
s = tuple(np.take(shape, axis))
if isinstance(self.forward, Transform):
if self.forward.input_array.shape == shape and axis == self._planned_axes:
# Already planned
return
plan_fwd = self._xfftn_fwd
plan_bck = self._xfftn_bck
if 'builders' in self._xfftn_fwd.__module__:
opts = dict(
avoid_copy=True,
overwrite_input=True,
auto_align_input=True,
auto_contiguous=True,
planner_effort='FFTW_MEASURE',
threads=1,
)
opts.update(options)
U = fftw.aligned(shape, dtype=dtype)
xfftn_fwd = plan_fwd(U, s=s, axes=axis, **opts)
V = xfftn_fwd.output_array
xfftn_bck = plan_bck(V, s=s, axes=axis, **opts)
V.fill(0)
U.fill(0)
xfftn_fwd.update_arrays(U, V)
xfftn_bck.update_arrays(V, U)
self._M = 1./np.prod(np.take(shape, axis))
else:
opts = dict(
overwrite_input='FFTW_DESTROY_INPUT',
planner_effort='FFTW_MEASURE',
threads=1,
)
opts.update(options)
flags = (fftw.flag_dict[opts['planner_effort']],
fftw.flag_dict[opts['overwrite_input']])
threads = opts['threads']
U = fftw.aligned(shape, dtype=dtype)
xfftn_fwd = plan_fwd(U, s=s, axes=axis, threads=threads, flags=flags)
V = xfftn_fwd.output_array
if np.issubdtype(dtype, np.floating):
flags = (fftw.flag_dict[opts['planner_effort']],)
xfftn_bck = plan_bck(V, s=s, axes=axis, threads=threads, flags=flags, output_array=U)
V.fill(0)
U.fill(0)
self._M = xfftn_fwd.get_normalization()
self.axis = axis[-1]
self._planned_axes = axis
if self.padding_factor > 1.+1e-8:
trunc_array = self._get_truncarray(shape, V.dtype)
self.forward = Transform(self.forward, xfftn_fwd, U, V, trunc_array)
self.backward = Transform(self.backward, xfftn_bck, trunc_array, V, U)
else:
self.forward = Transform(self.forward, xfftn_fwd, U, V, V)
self.backward = Transform(self.backward, xfftn_bck, V, V, U)
# scalar_product is not padded, just the forward/backward
self.scalar_product = Transform(self.scalar_product, xfftn_fwd, U, V, V)
self.si = islicedict(axis=self.axis, dimensions=self.dimensions)
self.sl = slicedict(axis=self.axis, dimensions=self.dimensions)
def _get_truncarray(self, shape, dtype):
shape = list(shape) if np.ndim(shape) else [shape]
shape[self.axis] = int(np.round(shape[self.axis] /
self.padding_factor))
return fftw.aligned(shape, dtype=dtype)
def plan(self, shape, axis, dtype, options):
"""Plan transform
Allocate work arrays for transforms and set up methods `forward`,
`backward` and `scalar_product` with or without padding
Parameters
----------
shape : array
Local shape of global array
axis : int
This base's axis in global TensorProductSpace
dtype : numpy.dtype
Type of array
options : dict
Options for planning transforms
"""
if shape in (0, (0, )):
return
if isinstance(axis, tuple):
axis = axis[0]
if isinstance(self.forward, Transform):
if self.forward.input_array.shape == shape and self.axis == axis:
# Already planned
return
plan_fwd = self._xfftn_fwd
plan_bck = self._xfftn_bck
if 'builders' in self._xfftn_fwd.__module__: #pragma: no cover
opts = dict(
avoid_copy=True,
overwrite_input=True,
auto_align_input=True,
auto_contiguous=True,
planner_effort='FFTW_MEASURE',
threads=1,
backward=False,
)
opts.update(options)
n = shape[axis]
U = fftw.aligned(shape, dtype=dtype)
xfftn_fwd = plan_fwd(U, n=n, axis=axis, **opts)
V = xfftn_fwd.output_array
xfftn_bck = plan_bck(V, n=n, axis=axis, **opts)
V.fill(0)
U.fill(0)
xfftn_fwd.update_arrays(U, V)
xfftn_bck.update_arrays(V, U)
self._M = 1. / np.prod(np.take(shape, axis))
else:
opts = dict(
overwrite_input='FFTW_DESTROY_INPUT',
planner_effort='FFTW_MEASURE',
threads=1,
)
opts.update(options)
flags = (fftw.flag_dict[opts['planner_effort']],
fftw.flag_dict[opts['overwrite_input']])
threads = opts['threads']
n = (shape[axis], )
U = fftw.aligned(shape, dtype=dtype)
xfftn_fwd = plan_fwd(U, n, (axis, ), threads=threads, flags=flags)
V = xfftn_fwd.output_array
if np.issubdtype(dtype, np.floating):
flags = (fftw.flag_dict[opts['planner_effort']], )
xfftn_bck = plan_bck(V,
n, (axis, ),
threads=threads,
flags=flags,
output_array=U)
V.fill(0)
U.fill(0)
self._M = xfftn_fwd.get_normalization()
self.axis = axis
if self.padding_factor != 1:
trunc_array = self._get_truncarray(shape, V.dtype)
self.forward = Transform(self.forward, xfftn_fwd, U, V,
trunc_array)
self.backward = Transform(self.backward, xfftn_bck, trunc_array, V,
U)
else:
self.forward = Transform(self.forward, xfftn_fwd, U, V, V)
self.backward = Transform(self.backward, xfftn_bck, V, V, U)
# scalar_product is not padded, just the forward/backward
self.scalar_product = Transform(self.scalar_product, xfftn_fwd, U, V,
V)
self.si = islicedict(axis=self.axis, dimensions=self.dimensions)
self.sl = slicedict(axis=self.axis, dimensions=self.dimensions)
try:
fftw.import_wisdom('wisdom.dat')
except AssertionError:
pass
N = (64, 64, 64)
loops = 50
axis = 1
threads = 4
implicit = True
flags = (fftw.FFTW_PATIENT, fftw.FFTW_DESTROY_INPUT)
# Transform complex to complex
#A = pyfftw.byte_align(np.random.random(N).astype('D'))
#A = np.random.random(N).astype(np.dtype('D'))
A = fftw.aligned(N, n=8, dtype=np.dtype('D'))
A[:] = np.random.random(N).astype(np.dtype('D'))
#print(A.ctypes.data % 32)
input_array = fftw.aligned(A.shape, n=32, dtype=A.dtype)
output_array = fftw.aligned(A.shape, n=32, dtype=A.dtype)
AC = A.copy()
ptime = [[], []]
ftime = [[], []]
stime = [[], []]
for axis in ((1, 2), 0, 1, 2):
axes = axis if np.ndim(axis) else [axis]
# pyfftw
def plan(self, shape, axis, dtype, options):
if isinstance(axis, tuple):
assert len(axis) == 1
axis = axis[-1]
if isinstance(self.forward, Transform):
if self.forward.input_array.shape == shape and self.axis == axis:
# Already planned
return
plan_fwd = self._xfftn_fwd
plan_bck = self._xfftn_bck
if 'builders' in self._xfftn_fwd.func.__module__:
opts = dict(
avoid_copy=True,
overwrite_input=True,
auto_align_input=True,
auto_contiguous=True,
planner_effort='FFTW_MEASURE',
threads=1,
)
opts.update(options)
U = fftw.aligned(shape, dtype=np.float)
xfftn_fwd = plan_fwd(U, axis=axis, **opts)
V = xfftn_fwd.output_array
xfftn_bck = plan_bck(V, axis=axis, **opts)
V.fill(0)
U.fill(0)
xfftn_fwd.update_arrays(U, V)
xfftn_bck.update_arrays(V, U)
else: # fftw wrapped with mpi4py-fft
opts = dict(
overwrite_input='FFTW_DESTROY_INPUT',
planner_effort='FFTW_MEASURE',
threads=1,
)
opts.update(options)
flags = (fftw.flag_dict[opts['planner_effort']],
fftw.flag_dict[opts['overwrite_input']])
threads = opts['threads']
U = fftw.aligned(shape, dtype=np.float)
xfftn_fwd = plan_fwd(U, axes=(axis,), threads=threads, flags=flags)
V = xfftn_fwd.output_array
xfftn_bck = plan_bck(V, axes=(axis,), threads=threads, flags=flags, output_array=U)
V.fill(0)
U.fill(0)
if np.dtype(dtype) is np.dtype('complex'):
# dct only works on real data, so need to wrap it
U = fftw.aligned(shape, dtype=np.complex)
V = fftw.aligned(shape, dtype=np.complex)
U.fill(0)
V.fill(0)
xfftn_fwd = DCTWrap(xfftn_fwd, U, V)
xfftn_bck = DCTWrap(xfftn_bck, V, U)
self.axis = axis
self.forward = Transform(self.forward, xfftn_fwd, U, V, V)
self.backward = Transform(self.backward, xfftn_bck, V, V, U)
self.scalar_product = Transform(self.scalar_product, xfftn_fwd, U, V, V)
self.si = islicedict(axis=self.axis, dimensions=self.dimensions())
self.sl = slicedict(axis=self.axis, dimensions=self.dimensions())
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
