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 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
self.LT.plan(shape, axis, dtype, options)
U, V = self.LT.forward.input_array, self.LT.forward.output_array
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)
self.backward_uniform = Transform(self.backward_uniform, 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.backward_uniform = Transform(self.backward_uniform, 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 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
self.LT.plan(shape, axis, dtype, options)
U, V = self.LT.forward.input_array, self.LT.forward.output_array
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)
self.si = islicedict(axis=self.axis, dimensions=self.dimensions)
self.sl = slicedict(axis=self.axis, dimensions=self.dimensions)
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)
self.si = islicedict(axis=self.axis, dimensions=self.dimensions())
self.sl = slicedict(axis=self.axis, dimensions=self.dimensions())
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
self.CT.plan(shape, axis, dtype, options)
self.CT.tensorproductspace = self.tensorproductspace
xfftn_fwd = self.CT.forward.xfftn
xfftn_bck = self.CT.backward.xfftn
U = xfftn_fwd.input_array
V = xfftn_fwd.output_array
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 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 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())
