def get_fft_plan(a, shape=None, axes=None, value_type='C2C'):
""" Generate a CUDA FFT plan for transforming up to three axes.
This is a convenient handle to cupy.fft.fft._get_cufft_plan_nd.
Args:
a (cupy.ndarray): Array to be transform, assumed to be either C- or
F- contiguous.
shape (None or tuple of ints): Shape of the transformed axes of the
output. If ``shape`` is not given, the lengths of the input along
the axes specified by ``axes`` are used.
axes (None or int or tuple of int): The axes of the array to
transform. Currently, these must be a set of up to three adjacent
axes and must include either the first or the last axis of the
array. If `None`, it is assumed that all axes are transformed.
value_type ('C2C'): The FFT type to perform.
Currently only complex-to-complex transforms are supported.
Returns:
plan (cupy.cuda.cufft.PlanNd): The cuFFT Plan.
"""
fft_type = _convert_fft_type(a, value_type)
if fft_type not in [cufft.CUFFT_C2C, cufft.CUFFT_Z2Z]:
raise NotImplementedError("Only C2C and Z2Z are supported.")
if a.flags.c_contiguous:
order = 'C'
elif a.flags.f_contiguous:
order = 'F'
else:
raise ValueError("Input array a must be contiguous")
if (shape is not None) and (axes is not None) and len(shape) != len(axes):
raise ValueError("Shape and axes have different lengths.")
if (shape is not None) and (axes is None) and len(shape) != a.ndim:
raise ValueError("Shape and axes have different lengths.")
if (shape is None) and (axes is not None) and (len(axes) > a.ndim):
raise ValueError("The number of axes exceeds a.ndim.")
# let _get_cufft_plan_nd() check (shape is None and axes is None)
# Note that "shape" here refers to the shape along trasformed axes, not
# the shape of the output array, and we need to convert it to the latter.
# The result is as if "a=_cook_shape(a); return a.shape" is called.
transformed_shape = shape
shape = list(a.shape)
if transformed_shape is not None:
if axes is None:
axes = [i for i in range(a.ndim)]
for s, axis in zip(transformed_shape, axes):
shape[axis] = s
shape = tuple(shape)
plan = _get_cufft_plan_nd(shape, fft_type, axes=axes, order=order)
return plan
def get_fft_plan(a, shape=None, axes=None, value_type='C2C'):
""" Generate a CUDA FFT plan for transforming up to three axes.
Args:
a (cupy.ndarray): Array to be transform, assumed to be either C- or
F- contiguous.
shape (None or tuple of ints): Shape of the transformed axes of the
output. If ``shape`` is not given, the lengths of the input along
the axes specified by ``axes`` are used.
axes (None or int or tuple of int): The axes of the array to
transform. If `None`, it is assumed that all axes are transformed.
Currently, for performing N-D transform these must be a set of up
to three adjacent axes, and must include either the first or the
last axis of the array.
value_type (str): The FFT type to perform. Acceptable values are:
* 'C2C': complex-to-complex transform (default)
* 'R2C': real-to-complex transform
* 'C2R': complex-to-real transform
Returns:
a cuFFT plan for either 1D transform (``cupy.cuda.cufft.Plan1d``) or
N-D transform (``cupy.cuda.cufft.PlanNd``).
.. note::
The returned plan can not only be passed as one of the arguments of
the functions in ``cupyx.scipy.fftpack``, but also be used as a
context manager for both ``cupy.fft`` and ``cupyx.scipy.fftpack``
functions:
.. code-block:: python
x = cupy.random.random(16).reshape(4, 4).astype(cupy.complex)
plan = cupyx.scipy.fftpack.get_fft_plan(x)
with plan:
y = cupy.fft.fftn(x)
# alternatively:
y = cupyx.scipy.fftpack.fftn(x) # no explicit plan is given!
# alternatively:
y = cupyx.scipy.fftpack.fftn(x, plan=plan) # pass plan explicitly
In the first case, no cuFFT plan will be generated automatically,
even if ``cupy.fft.config.enable_nd_planning = True`` is set.
.. warning::
This API is a deviation from SciPy's, is currently experimental, and
may be changed in the future version.
"""
# check input array
if a.flags.c_contiguous:
order = 'C'
elif a.flags.f_contiguous:
order = 'F'
else:
raise ValueError('Input array a must be contiguous')
if isinstance(shape, int):
shape = (shape, )
if isinstance(axes, int):
axes = (axes, )
if (shape is not None) and (axes is not None) and len(shape) != len(axes):
raise ValueError('Shape and axes have different lengths.')
# check axes
# n=1: 1d (need axis1D); n>1: Nd
if axes is None:
n = a.ndim if shape is None else len(shape)
axes = tuple(i for i in range(-n, 0))
if n == 1:
axis1D = 0
else: # axes is a tuple
n = len(axes)
if n == 1:
axis1D = axes[0]
if axis1D >= a.ndim or axis1D < -a.ndim:
err = 'The chosen axis ({0}) exceeds the number of '\
'dimensions of a ({1})'.format(axis1D, a.ndim)
raise ValueError(err)
elif n > 3:
raise ValueError('Only up to three axes is supported')
# Note that "shape" here refers to the shape along trasformed axes, not
# the shape of the output array, and we need to convert it to the latter.
# The result is as if "a=_cook_shape(a); return a.shape" is called.
# Because of this, we need to use (possibly unsorted) axes.
transformed_shape = shape
shape = list(a.shape)
if transformed_shape is not None:
for s, axis in zip(transformed_shape, axes):
if s is not None:
if axis == axes[-1] and value_type == 'C2R':
s = s // 2 + 1
shape[axis] = s
shape = tuple(shape)
# check value_type
out_dtype = _output_dtype(a.dtype, value_type)
fft_type = _convert_fft_type(out_dtype, value_type)
# TODO(leofang): figure out if we really have to skip F-order?
if n > 1 and value_type != 'C2C' and a.flags.f_contiguous:
raise ValueError('C2R/R2C PlanNd for F-order arrays is not supported')
# generate plan
if n > 1: # ND transform
out_size = _get_fftn_out_size(shape, transformed_shape, axes[-1],
value_type)
plan = _get_cufft_plan_nd(shape,
fft_type,
axes=axes,
order=order,
out_size=out_size)
else: # 1D transform
if value_type != 'C2R':
out_size = shape[axis1D]
else:
out_size = _get_fftn_out_size(shape, transformed_shape, axis1D,
value_type)
batch = prod(shape) // shape[axis1D]
devices = None if not config.use_multi_gpus else config._devices
plan = cufft.Plan1d(out_size, fft_type, batch, devices=devices)
return plan
def get_fft_plan(a, shape=None, axes=None, value_type='C2C'):
""" Generate a CUDA FFT plan for transforming up to three axes.
Args:
a (cupy.ndarray): Array to be transform, assumed to be either C- or
F- contiguous.
shape (None or tuple of ints): Shape of the transformed axes of the
output. If ``shape`` is not given, the lengths of the input along
the axes specified by ``axes`` are used.
axes (None or int or tuple of int): The axes of the array to
transform. If `None`, it is assumed that all axes are transformed.
Currently, for performing N-D transform these must be a set of up
to three adjacent axes, and must include either the first or the
last axis of the array.
value_type ('C2C'): The FFT type to perform.
Currently only complex-to-complex transforms are supported.
Returns:
a cuFFT plan for either 1D transform (``cupy.cuda.cufft.Plan1d``) or
N-D transform (``cupy.cuda.cufft.PlanNd``).
.. note::
The returned plan can not only be passed as one of the arguments of
the functions in ``cupyx.scipy.fftpack``, but also be used as a
context manager for both ``cupy.fft`` and ``cupyx.scipy.fftpack``
functions:
.. code-block:: python
x = cupy.random.random(16).reshape(4, 4).astype(cupy.complex)
plan = cupyx.scipy.fftpack.get_fft_plan(x)
with plan:
y = cupy.fft.fftn(x)
# alternatively:
y = cupyx.scipy.fftpack.fftn(x) # no explicit plan is given!
# alternatively:
y = cupyx.scipy.fftpack.fftn(x, plan=plan) # pass plan explicitly
In the first case, no cuFFT plan will be generated automatically,
even if ``cupy.fft.config.enable_nd_planning = True`` is set.
.. warning::
This API is a deviation from SciPy's, is currently experimental, and
may be changed in the future version.
"""
# check input array
if a.flags.c_contiguous:
order = 'C'
elif a.flags.f_contiguous:
order = 'F'
else:
raise ValueError('Input array a must be contiguous')
# check axes
# n=1: 1d (need axis1D); n>1: Nd
if axes is None:
n = a.ndim
axes = tuple(i for i in range(n))
if n == 1:
axis1D = 0
elif isinstance(axes, int):
n = 1
axis1D = axes
axes = (axes, )
if axis1D >= a.ndim or axis1D < -a.ndim:
raise ValueError('The chosen axis ({0}) exceeds the number of '
'dimensions of a ({1})'.format(axis1D, a.ndim))
else: # axes is a tuple
n = len(axes)
if n == 1:
axis1D = axes[0]
elif n > 3:
raise ValueError('Only up to three axes is supported')
# check shape
if isinstance(shape, int):
shape = (shape, )
if (shape is not None) and len(shape) != n:
raise ValueError('Shape and axes have different lengths.')
# Note that "shape" here refers to the shape along trasformed axes, not
# the shape of the output array, and we need to convert it to the latter.
# The result is as if "a=_cook_shape(a); return a.shape" is called.
transformed_shape = shape
shape = list(a.shape)
if transformed_shape is not None:
for s, axis in zip(transformed_shape, axes):
shape[axis] = s
shape = tuple(shape)
# check value_type
fft_type = _convert_fft_type(a, value_type)
if n > 1 and fft_type not in [cufft.CUFFT_C2C, cufft.CUFFT_Z2Z]:
raise NotImplementedError('Only C2C and Z2Z are supported for N-dim'
' transform.')
# generate plan
if n > 1: # ND transform
plan = _get_cufft_plan_nd(shape, fft_type, axes=axes, order=order)
else: # 1D transform
out_size = shape[axis1D]
batch = prod(shape) // out_size
plan = cufft.Plan1d(out_size, fft_type, batch)
return plan
def get_fft_plan(a, shape=None, axes=None, value_type='C2C'):
""" Generate a CUDA FFT plan for transforming up to three axes.
Args:
a (cupy.ndarray): Array to be transform, assumed to be either C- or
F- contiguous.
shape (None or tuple of ints): Shape of the transformed axes of the
output. If ``shape`` is not given, the lengths of the input along
the axes specified by ``axes`` are used.
axes (None or int or tuple of int): The axes of the array to
transform. If `None`, it is assumed that all axes are transformed.
Currently, for performing N-D transform these must be a set of up
to three adjacent axes, and must include either the first or the
last axis of the array.
value_type ('C2C'): The FFT type to perform.
Currently only complex-to-complex transforms are supported.
Returns:
plan: a cuFFT plan for either 1D transform (cupy.cuda.cufft.Plan1d)
or N-D transform (cupy.cuda.cufft.PlanNd).
"""
# check input array
if a.flags.c_contiguous:
order = 'C'
elif a.flags.f_contiguous:
order = 'F'
else:
raise ValueError('Input array a must be contiguous')
# check axes
# n=1: 1d (need axis1D); n>1: Nd
if axes is None:
n = a.ndim
axes = tuple(i for i in range(n))
if n == 1:
axis1D = 0
elif isinstance(axes, int):
n = 1
axis1D = axes
axes = (axes, )
if axis1D >= a.ndim or axis1D < -a.ndim:
raise ValueError('The chosen axis ({0}) exceeds the number of '
'dimensions of a ({1})'.format(axis1D, a.ndim))
else: # axes is a tuple
n = len(axes)
if n == 1:
axis1D = axes[0]
elif n > 3:
raise ValueError('Only up to three axes is supported')
# check shape
if isinstance(shape, int):
shape = (shape, )
if (shape is not None) and len(shape) != n:
raise ValueError('Shape and axes have different lengths.')
# Note that "shape" here refers to the shape along trasformed axes, not
# the shape of the output array, and we need to convert it to the latter.
# The result is as if "a=_cook_shape(a); return a.shape" is called.
transformed_shape = shape
shape = list(a.shape)
if transformed_shape is not None:
for s, axis in zip(transformed_shape, axes):
shape[axis] = s
shape = tuple(shape)
# check value_type
fft_type = _convert_fft_type(a, value_type)
if n > 1 and fft_type not in [cufft.CUFFT_C2C, cufft.CUFFT_Z2Z]:
raise NotImplementedError('Only C2C and Z2Z are supported for N-dim'
' transform.')
# generate plan
if n > 1: # ND transform
plan = _get_cufft_plan_nd(shape, fft_type, axes=axes, order=order)
else: # 1D transform
out_size = shape[axis1D]
batch = prod(shape) // out_size
plan = cufft.Plan1d(out_size, fft_type, batch)
return plan