Wrap 1D, 2D, and ND real and complex FFT functions Takes a function that behaves like ``numpy.fft`` functions and a specified kind to match it to that are named after the functions in the ``numpy.fft`` API. Supported kinds include: * fft * fft2 * fftn
(fft_func, kind=None, dtype=None, allow_fftpack=False)
| 116 | |
| 117 | |
| 118 | def fft_wrap(fft_func, kind=None, dtype=None, allow_fftpack=False): |
| 119 | """Wrap 1D, 2D, and ND real and complex FFT functions |
| 120 | |
| 121 | Takes a function that behaves like ``numpy.fft`` functions and |
| 122 | a specified kind to match it to that are named after the functions |
| 123 | in the ``numpy.fft`` API. |
| 124 | |
| 125 | Supported kinds include: |
| 126 | |
| 127 | * fft |
| 128 | * fft2 |
| 129 | * fftn |
| 130 | * ifft |
| 131 | * ifft2 |
| 132 | * ifftn |
| 133 | * rfft |
| 134 | * rfft2 |
| 135 | * rfftn |
| 136 | * irfft |
| 137 | * irfft2 |
| 138 | * irfftn |
| 139 | * hfft |
| 140 | * ihfft |
| 141 | |
| 142 | Examples |
| 143 | -------- |
| 144 | >>> import dask.array.fft as dff |
| 145 | >>> parallel_fft = dff.fft_wrap(np.fft.fft) |
| 146 | >>> parallel_ifft = dff.fft_wrap(np.fft.ifft) |
| 147 | """ |
| 148 | if scipy is not None: |
| 149 | if fft_func.__module__.startswith("scipy.fftpack"): |
| 150 | if not allow_fftpack: |
| 151 | warnings.warn( |
| 152 | f"Function {fft_func.__name__} from `scipy.fftpack` does not " |
| 153 | "match NumPy's API and is considered legacy. Please use " |
| 154 | "`scipy.fft` instead. To suppress this warning and allow usage" |
| 155 | ", set `allow_fftpack=True`. Support for `scipy.fftpack` will " |
| 156 | "be deprecated in future releases.", |
| 157 | FutureWarning, |
| 158 | ) |
| 159 | # If allow_fftpack is True, we proceed but we skip passing the norm |
| 160 | # argument. |
| 161 | |
| 162 | if kind is None: |
| 163 | kind = fft_func.__name__ |
| 164 | try: |
| 165 | out_chunk_fn = _out_chunk_fns[kind.rstrip("2n")] |
| 166 | except KeyError: |
| 167 | raise ValueError(f"Given unknown `kind` {kind}.") |
| 168 | |
| 169 | def func(a, s=None, axes=None, norm=None): |
| 170 | a = asarray(a) |
| 171 | if axes is None: |
| 172 | if kind.endswith("2"): |
| 173 | axes = (-2, -1) |
| 174 | elif kind.endswith("n"): |
| 175 | if s is None: |
searching dependent graphs…