/
_fft.py
493 lines (396 loc) · 19 KB
/
_fft.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
from numpy import prod
import cupy
from cupy.cuda import cufft
from cupy.fft import config
from cupy.fft._fft import (_convert_fft_type, _default_fft_func, _fft,
_get_cufft_plan_nd, _get_fftn_out_size,
_output_dtype)
from cupy.fft._cache import get_plan_cache
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(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.
.. note::
If this function is called under the context of
:func:`~cupy.fft.config.set_cufft_callbacks`, the generated plan will
have callbacks enabled.
.. 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
# (load from cache if it exists, otherwise create one but don't cache it)
if n > 1: # ND transform
if cupy.cuda.runtime.is_hip and value_type == 'C2R':
raise RuntimeError("hipFFT's C2R PlanNd is buggy and unsupported")
out_size = _get_fftn_out_size(
shape, transformed_shape, axes[-1], value_type)
# _get_cufft_plan_nd interacts with plan cache and callback
plan = _get_cufft_plan_nd(
shape, fft_type, axes=axes, order=order, out_size=out_size,
to_cache=False)
else: # 1D transform
# prepare plan arguments
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
keys = (out_size, fft_type, batch, devices)
mgr = config.get_current_callback_manager()
if mgr is not None:
# to avoid a weird segfault, we generate and cache distinct plans
# for every possible (load_aux, store_aux) pairs; the plans are
# still generated from the same external Python module
load_aux = mgr.cb_load_aux_arr
store_aux = mgr.cb_store_aux_arr
keys += (mgr.cb_load, mgr.cb_store,
0 if load_aux is None else load_aux.data.ptr,
0 if store_aux is None else store_aux.data.ptr)
cache = get_plan_cache()
cached_plan = cache.get(keys)
if cached_plan is not None:
plan = cached_plan
elif mgr is None:
plan = cufft.Plan1d(out_size, fft_type, batch, devices=devices)
else: # has callback
# TODO(leofang): support multi-GPU callback (devices is ignored)
if devices:
raise NotImplementedError('multi-GPU cuFFT callbacks are not '
'yet supported')
plan = mgr.create_plan(('Plan1d', keys[:-3]))
mgr.set_callbacks(plan)
return plan
def fft(x, n=None, axis=-1, overwrite_x=False, plan=None):
"""Compute the one-dimensional FFT.
Args:
x (cupy.ndarray): Array to be transformed.
n (None or int): Length of the transformed axis of the output. If ``n``
is not given, the length of the input along the axis specified by
``axis`` is used.
axis (int): Axis over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for
transforming ``x`` over ``axis``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(x, axis)
Note that `plan` is defaulted to None, meaning CuPy will use an
auto-generated plan behind the scene.
Returns:
cupy.ndarray:
The transformed array which shape is specified by ``n`` and type
will convert to complex if that of the input is another.
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
.. seealso:: :func:`scipy.fftpack.fft`
"""
return _fft(x, (n,), (axis,), None, cufft.CUFFT_FORWARD,
overwrite_x=overwrite_x, plan=plan)
def ifft(x, n=None, axis=-1, overwrite_x=False, plan=None):
"""Compute the one-dimensional inverse FFT.
Args:
x (cupy.ndarray): Array to be transformed.
n (None or int): Length of the transformed axis of the output. If ``n``
is not given, the length of the input along the axis specified by
``axis`` is used.
axis (int): Axis over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for
transforming ``x`` over ``axis``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(x, axis)
Note that `plan` is defaulted to None, meaning CuPy will use an
auto-generated plan behind the scene.
Returns:
cupy.ndarray:
The transformed array which shape is specified by ``n`` and type
will convert to complex if that of the input is another.
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
.. seealso:: :func:`scipy.fftpack.ifft`
"""
return _fft(x, (n,), (axis,), None, cufft.CUFFT_INVERSE,
overwrite_x=overwrite_x, plan=plan)
def fft2(x, shape=None, axes=(-2, -1), overwrite_x=False, plan=None):
"""Compute the two-dimensional FFT.
Args:
x (cupy.ndarray): Array to be transformed.
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 (tuple of ints): Axes over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for
transforming ``x`` over ``axes``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(x, axes)
Note that `plan` is defaulted to None, meaning CuPy will either
use an auto-generated plan behind the scene if cupy.fft.config.
enable_nd_planning = True, or use no cuFFT plan if it is set to
False.
Returns:
cupy.ndarray:
The transformed array which shape is specified by ``shape`` and
type will convert to complex if that of the input is another.
.. seealso:: :func:`scipy.fftpack.fft2`
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
"""
func = _default_fft_func(x, shape, axes, plan)
return func(x, shape, axes, None, cufft.CUFFT_FORWARD,
overwrite_x=overwrite_x, plan=plan)
def ifft2(x, shape=None, axes=(-2, -1), overwrite_x=False, plan=None):
"""Compute the two-dimensional inverse FFT.
Args:
x (cupy.ndarray): Array to be transformed.
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 (tuple of ints): Axes over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for
transforming ``x`` over ``axes``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(x, axes)
Note that `plan` is defaulted to None, meaning CuPy will either
use an auto-generated plan behind the scene if cupy.fft.config.
enable_nd_planning = True, or use no cuFFT plan if it is set to
False.
Returns:
cupy.ndarray:
The transformed array which shape is specified by ``shape`` and
type will convert to complex if that of the input is another.
.. seealso:: :func:`scipy.fftpack.ifft2`
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
"""
func = _default_fft_func(x, shape, axes, plan)
return func(x, shape, axes, None, cufft.CUFFT_INVERSE,
overwrite_x=overwrite_x, plan=plan)
def fftn(x, shape=None, axes=None, overwrite_x=False, plan=None):
"""Compute the N-dimensional FFT.
Args:
x (cupy.ndarray): Array to be transformed.
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 (tuple of ints): Axes over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for
transforming ``x`` over ``axes``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(x, axes)
Note that `plan` is defaulted to None, meaning CuPy will either
use an auto-generated plan behind the scene if cupy.fft.config.
enable_nd_planning = True, or use no cuFFT plan if it is set to
False.
Returns:
cupy.ndarray:
The transformed array which shape is specified by ``shape`` and
type will convert to complex if that of the input is another.
.. seealso:: :func:`scipy.fftpack.fftn`
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
"""
func = _default_fft_func(x, shape, axes, plan)
return func(x, shape, axes, None, cufft.CUFFT_FORWARD,
overwrite_x=overwrite_x, plan=plan)
def ifftn(x, shape=None, axes=None, overwrite_x=False, plan=None):
"""Compute the N-dimensional inverse FFT.
Args:
x (cupy.ndarray): Array to be transformed.
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 (tuple of ints): Axes over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.PlanNd` or ``None``): a cuFFT plan for
transforming ``x`` over ``axes``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(x, axes)
Note that `plan` is defaulted to None, meaning CuPy will either
use an auto-generated plan behind the scene if cupy.fft.config.
enable_nd_planning = True, or use no cuFFT plan if it is set to
False.
Returns:
cupy.ndarray:
The transformed array which shape is specified by ``shape`` and
type will convert to complex if that of the input is another.
.. seealso:: :func:`scipy.fftpack.ifftn`
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
"""
func = _default_fft_func(x, shape, axes, plan)
return func(x, shape, axes, None, cufft.CUFFT_INVERSE,
overwrite_x=overwrite_x, plan=plan)
def rfft(x, n=None, axis=-1, overwrite_x=False, plan=None):
"""Compute the one-dimensional FFT for real input.
The returned real array contains
.. code-block:: python
[y(0),Re(y(1)),Im(y(1)),...,Re(y(n/2))] # if n is even
[y(0),Re(y(1)),Im(y(1)),...,Re(y(n/2)),Im(y(n/2))] # if n is odd
Args:
x (cupy.ndarray): Array to be transformed.
n (None or int): Length of the transformed axis of the output. If ``n``
is not given, the length of the input along the axis specified by
``axis`` is used.
axis (int): Axis over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
plan (:class:`cupy.cuda.cufft.Plan1d` or ``None``): a cuFFT plan for
transforming ``x`` over ``axis``, which can be obtained using::
plan = cupyx.scipy.fftpack.get_fft_plan(
x, axes, value_type='R2C')
Note that `plan` is defaulted to None, meaning CuPy will either
use an auto-generated plan behind the scene if cupy.fft.config.
enable_nd_planning = True, or use no cuFFT plan if it is set to
False.
Returns:
cupy.ndarray:
The transformed array.
.. seealso:: :func:`scipy.fftpack.rfft`
.. note::
The argument `plan` is currently experimental and the interface may be
changed in the future version.
"""
if n is None:
n = x.shape[axis]
shape = list(x.shape)
shape[axis] = n
f = _fft(x, (n,), (axis,), None, cufft.CUFFT_FORWARD, 'R2C',
overwrite_x=overwrite_x, plan=plan)
z = cupy.empty(shape, f.real.dtype)
slice_z = [slice(None)] * x.ndim
slice_f = [slice(None)] * x.ndim
slice_z[axis] = slice(1)
slice_f[axis] = slice(1)
z[tuple(slice_z)] = f[tuple(slice_f)].real
slice_z[axis] = slice(1, None, 2)
slice_f[axis] = slice(1, None)
z[tuple(slice_z)] = f[tuple(slice_f)].real
slice_z[axis] = slice(2, None, 2)
slice_f[axis] = slice(1, n - f.shape[axis] + 1)
z[tuple(slice_z)] = f[tuple(slice_f)].imag
return z
def irfft(x, n=None, axis=-1, overwrite_x=False):
"""Compute the one-dimensional inverse FFT for real input.
Args:
x (cupy.ndarray): Array to be transformed.
n (None or int): Length of the transformed axis of the output. If ``n``
is not given, the length of the input along the axis specified by
``axis`` is used.
axis (int): Axis over which to compute the FFT.
overwrite_x (bool): If True, the contents of ``x`` can be destroyed.
Returns:
cupy.ndarray:
The transformed array.
.. seealso:: :func:`scipy.fftpack.irfft`
.. note::
This function does not support a precomputed `plan`. If you need this
capability, please consider using :func:`cupy.fft.irfft` or :func:`
cupyx.scipy.fft.irfft`.
"""
if n is None:
n = x.shape[axis]
m = min(n, x.shape[axis])
shape = list(x.shape)
shape[axis] = n // 2 + 1
if x.dtype in (cupy.float16, cupy.float32):
z = cupy.zeros(shape, dtype=cupy.complex64)
else:
z = cupy.zeros(shape, dtype=cupy.complex128)
slice_x = [slice(None)] * x.ndim
slice_z = [slice(None)] * x.ndim
slice_x[axis] = slice(1)
slice_z[axis] = slice(1)
z[tuple(slice_z)].real = x[tuple(slice_x)]
slice_x[axis] = slice(1, m, 2)
slice_z[axis] = slice(1, m // 2 + 1)
z[tuple(slice_z)].real = x[tuple(slice_x)]
slice_x[axis] = slice(2, m, 2)
slice_z[axis] = slice(1, (m + 1) // 2)
z[tuple(slice_z)].imag = x[tuple(slice_x)]
return _fft(z, (n,), (axis,), None, cufft.CUFFT_INVERSE, 'C2R',
overwrite_x=overwrite_x)