Repository: dask/dask

Issue:
Implement numpy.fft functionality
FFTs are used in a variety of fields for everything from data compression to spectral analysis.  Using the [Cooley-Tukey FFT](https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm) it should be possible to used a chunked algorithm to do the computation since the algorithm is already inherently chunked*.  Both FFT and inverse FFT support would be needed for the most common use cases.

*(Note: Would need a good way to handle arrays with length that is a prime number and a good way to specify a minimum sized chunk to use)



Discussion:
FFT would be doable if the direction of FFT is not chunked (this might be the common case).  If it is chunked however then we might find this a bit hard because FFT is a fairly non-local algorithm.  I would be quite happy to learn otherwise if you can find a blocked algorithm.  I'm not an expert here and would love to be proven incorrect.

On the other hand, something like the discrete wavelet transform would be very doable.

Out of curiousity are you more interested in the multi-core aspects of dasks for fast FFTs or the out-of-core nature?  Do you need to block within the direction of transformation or are you in the "stack of images" case?

a common technique in signal processing is known as overlap-add. It allows
for example to do out of core filtering using FFTs.

@mrocklin Keep in mind that my use case is only one of many for FFTs.  It would be best to make this functionality as general as possible.

My specific use case is usually performing lots of FFTs on lots of relatively small data sets, usually on the scale of (2000, 1) to (100000, 1) shaped arrays.  For that upper end of 100000 points it would be useful to have some splitting over cores for a single FFT, but on the lower end there wouldn't be much gained (if anything) from parallelization.

One of my most common use cases would be a graph similar to

```
def build_graph(x, y):
    return {
        'x': x,
        'y': y,

        'x1': (f1, 'x'),
        'x2': (f2, 'x'),
        'X1': (np.fft.rfft, 'x1'),
        'X2': (np.fft.rfft, 'x2'),

        'y1': (f1, 'y'),
        'y2': (f2, 'y'),
        'Y1': (np.fft.rfft, 'y1'),
        'Y2': (np.fft.rfft, 'y2'),

        'X': (f3, 'X1', 'X2'),
        'Y': (f3, 'Y1', 'Y2'),

        'output': (operator.div, 'X', 'Y')
    }
```

Where `f1`, `f2`, and `f3` are simple functions that don't really need multicore (mostly array slicing and a small amount of arithmetic).  This could definitely benefit from some parallelization, especially considering I might call this function several dozen times in one run of my software.  It's definitely not the biggest data set out there, but we're still limited to one core unless we want to mangle the code with a lot of multiprocessing logic.

The Cooley-Tukey FFT algorithm is really a divide, compute, and re-combine operation.  It depends on how well you can factor the length of the array, so it works best for powers of 2 length arrays and works worst for a prime number of points (I've actually run into that problem in the past with numpy's FFT, it becomes the world's slowest prime number finder).  I don't know how well dask can handle differently sized chunks, but it is essentially a chunk-then-combine algorithm.

I'll take a closer look at Cooley Tukey, probably sometime next week.

Regarding your particular computation it looks like you could just run that through a scheduler directly and bypass dask.array entirely.

``` python
from dask.threaded import get
dsk = build_graph(x, y)
output = get(dsk, 'output')  # executes with multiple threads
```

@bheklilr if you are not already using the MKL linked version of NumPy that should give you a significant speed boost. Also I think `scipy.fftpack.fft` is faster than numpy's default.

In Cooley-Tukey you "chunk" an FFT of `n` samples into an multiple FFTs of `n1`, `n2`, `n2`, ... samples. Where these `n`s are the prime factors of the original `n`. These chunks can be computed in parallel.

The first problem I notice, is that these are not "chunks" in the dask array sense, they are not contiguous, but cyclic.

A few use cases I can think of:
- The input data fits comfortably in memory (considering an FFT operation will make a few copies). I'm going to assume that NumPy/SciPy will beat us here. 
- Data does not fit in memory. Having the input data be a dask array is probably bad because each CT chunk may touch many dask.array chunks. So we would probably have to compute the whole dask.array before we FFT it. So having the data on disk would probably be a better application for dask. But we would still be doing non-contiguous reads which would probably hurt. The algorithm would look like:
  - We first compute CT chunking (prime factors in each dimension).
  - then grab the chunks we want with indexing, and FFT each one with np.fft. This can be parallized.
  - Write the data to disk? This can also be parallelized. Having this output a dask array seems problematic since each chunk is so non-contiguous

The thing I notice here is that this does not seem like a good fit for dask.array. But rather a custom dask graph. Like @bheklilr has, but also considering data size.

Jake Vanderplas has a nice blog post about Cooley-Tukey [here](https://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/). It implements Cooley-Tookey in NumPy but does not handle cases where the input data is not a power of 2.


Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -21,7 +21,7 @@
 from .percentile import percentile
 with ignoring(ImportError):
     from .reductions import nanprod
-from . import random, linalg, ghost, learn
+from . import random, linalg, ghost, learn, fft
 from .wrap import ones, zeros, empty, full
 from .rechunk import rechunk
 from ..context import set_options
diff --git a/dask/array/conftest.py b/dask/array/conftest.py
new file mode 100644
--- /dev/null
+++ b/dask/array/conftest.py
@@ -0,0 +1,6 @@
+import pytest
+
+
+def pytest_ignore_collect(path, config):
+    if 'fft.py' in str(path):
+        return True
diff --git a/dask/array/fft.py b/dask/array/fft.py
new file mode 100644
--- /dev/null
+++ b/dask/array/fft.py
@@ -0,0 +1,101 @@
+from __future__ import absolute_import
+
+import numpy as np
+import numpy.fft as npfft
+
+from .core import map_blocks
+
+
+chunk_error = ("Dask array only supports taking an FFT along an axis that \n"
+               "has a single chunk. An FFT operation was tried on axis %s \n"
+               "which has chunks %s. To change the array's chunks use "
+               "dask.Array.rechunk.")
+
+fft_preamble = """
+    Wrapping of numpy.fft.%s
+
+    The axis along which the FFT is applied must have a one chunk. To change
+    the array's chunking use dask.Array.rechunk.
+
+    The numpy.fft.%s docstring follows below:
+
+    """
+
+
+def _fft_wrap(fft_func, dtype, out_chunk_fn):
+    def func(a, n=None, axis=-1):
+        if len(a.chunks[axis]) != 1:
+            raise ValueError(chunk_error % (axis, a.chunks[axis]))
+
+        chunks = out_chunk_fn(a, n, axis)
+
+        return map_blocks(lambda x: fft_func(x, n=n, axis=axis), a, dtype=dtype,
+                          chunks=chunks)
+
+    np_name = fft_func.__name__
+    func.__doc__ = (fft_preamble % (np_name, np_name)) + fft_func.__doc__
+    func.__name__ = np_name
+    return func
+
+
+def _fft_out_chunks(a, n, axis):
+    """ For computing the output chunks of fft and ifft"""
+    if n is None:
+        return a.chunks
+    chunks = list(a.chunks)
+    chunks[axis] = (n,)
+    return chunks
+
+
+def _rfft_out_chunks(a, n, axis):
+    if n is None:
+        n = a.chunks[axis][0]
+    chunks = list(a.chunks)
+    chunks[axis] = (n//2 + 1,)
+    return chunks
+
+
+def _irfft_out_chunks(a, n, axis):
+    if n is None:
+        n = 2 * (a.chunks[axis][0] - 1)
+    chunks = list(a.chunks)
+    chunks[axis] = (n,)
+    return chunks
+
+
+def _hfft_out_chunks(a, n, axis):
+    if n is None:
+        n = 2 * (a.chunks[axis][0] - 1)
+    chunks = list(a.chunks)
+    chunks[axis] = (n,)
+    return chunks
+
+
+def _ihfft_out_chunks(a, n, axis):
+    if n is None:
+        n = a.chunks[axis][0]
+    chunks = list(a.chunks)
+    if n % 2 == 0:
+        m = (n//2) + 1
+    else:
+        m = (n + 1)//2
+    chunks[axis] = (m,)
+    return chunks
+
+
+fft = _fft_wrap(npfft.fft, np.complex_, _fft_out_chunks)
+
+
+ifft = _fft_wrap(npfft.ifft, np.complex_, _fft_out_chunks)
+
+
+rfft = _fft_wrap(npfft.rfft, np.complex_, _rfft_out_chunks)
+
+
+irfft = _fft_wrap(npfft.irfft, np.float_, _irfft_out_chunks)
+
+
+hfft = _fft_wrap(npfft.hfft, np.float_, _hfft_out_chunks)
+
+
+ihfft = _fft_wrap(npfft.ihfft, np.complex_, _ihfft_out_chunks)


Test patch:
diff --git a/dask/array/tests/test_fft.py b/dask/array/tests/test_fft.py
new file mode 100644
--- /dev/null
+++ b/dask/array/tests/test_fft.py
@@ -0,0 +1,128 @@
+import numpy as np
+import numpy.fft as npfft
+
+import dask
+from dask.array.core import Array
+from dask.utils import raises
+import dask.array as da
+from dask.array.fft import fft, ifft, rfft, irfft, hfft, ihfft
+
+
+def mismatch_err(mismatch_type, got, expected):
+    return '%s mismatch, got %s, expected %s' % (mismatch_type, got, expected)
+
+
+def eq(a, b):
+    assert a.shape == b.shape, mismatch_err('shape', a.shape, b.shape)
+
+    if isinstance(a, Array):
+        adt = a._dtype
+        a = a.compute(get=dask.get)
+    else:
+        adt = getattr(a, 'dtype', None)
+    if isinstance(b, Array):
+        bdt = b._dtype
+        b = b.compute(get=dask.get)
+    else:
+        bdt = getattr(b, 'dtype', None)
+
+    assert str(adt) == str(bdt), mismatch_err('dtype', str(adt), str(bdt))
+
+    try:
+        return np.allclose(a, b)
+    except TypeError:
+        pass
+
+    c = a == b
+
+    if isinstance(c, np.ndarray):
+        return c.all()
+    else:
+        return c
+
+
+nparr = np.arange(100).reshape(10, 10)
+darr = da.from_array(nparr, chunks=(1, 10))
+darr2 = da.from_array(nparr, chunks=(10, 1))
+
+
+def test_cant_fft_chunked_axis():
+    bad_darr = da.from_array(nparr, chunks=(5, 5))
+    assert raises(ValueError, lambda: fft(bad_darr))
+    assert raises(ValueError, lambda: fft(bad_darr, axis=0))
+
+
+def test_fft():
+    assert eq(fft(darr), npfft.fft(nparr))
+    assert eq(fft(darr2, axis=0), npfft.fft(nparr, axis=0))
+
+
+def test_fft_n_kwarg():
+    assert eq(fft(darr, 5), npfft.fft(nparr, 5))
+    assert eq(fft(darr, 13), npfft.fft(nparr, 13))
+    assert eq(fft(darr2, 5, axis=0), npfft.fft(nparr, 5, axis=0))
+    assert eq(fft(darr2, 13, axis=0), npfft.fft(nparr, 13, axis=0))
+
+
+def test_ifft():
+    assert eq(ifft(darr), npfft.ifft(nparr))
+    assert eq(ifft(darr2, axis=0), npfft.ifft(nparr, axis=0))
+
+
+def test_ifft_n_kwarg():
+    assert eq(ifft(darr, 5), npfft.ifft(nparr, 5))
+    assert eq(ifft(darr, 13), npfft.ifft(nparr, 13))
+    assert eq(ifft(darr2, 5, axis=0), npfft.ifft(nparr, 5, axis=0))
+    assert eq(ifft(darr2, 13, axis=0), npfft.ifft(nparr, 13, axis=0))
+
+
+def test_rfft():
+    assert eq(rfft(darr), npfft.rfft(nparr))
+    assert eq(rfft(darr2, axis=0), npfft.rfft(nparr, axis=0))
+
+
+def test_rfft_n_kwarg():
+    assert eq(rfft(darr, 5), npfft.rfft(nparr, 5))
+    assert eq(rfft(darr, 13), npfft.rfft(nparr, 13))
+    assert eq(rfft(darr2, 5, axis=0), npfft.rfft(nparr, 5, axis=0))
+    assert eq(rfft(darr2, 13, axis=0), npfft.rfft(nparr, 13, axis=0))
+    assert eq(rfft(darr2, 12, axis=0), npfft.rfft(nparr, 12, axis=0))
+
+
+def test_irfft():
+    assert eq(irfft(darr), npfft.irfft(nparr))
+    assert eq(irfft(darr2, axis=0), npfft.irfft(nparr, axis=0))
+
+
+def test_irfft_n_kwarg():
+    assert eq(irfft(darr, 5), npfft.irfft(nparr, 5))
+    assert eq(irfft(darr, 13), npfft.irfft(nparr, 13))
+    assert eq(irfft(darr2, 5, axis=0), npfft.irfft(nparr, 5, axis=0))
+    assert eq(irfft(darr2, 13, axis=0), npfft.irfft(nparr, 13, axis=0))
+    assert eq(irfft(darr2, 12, axis=0), npfft.irfft(nparr, 12, axis=0))
+
+
+def test_hfft():
+    assert eq(hfft(darr), npfft.hfft(nparr))
+    assert eq(hfft(darr2, axis=0), npfft.hfft(nparr, axis=0))
+
+
+def test_hfft_nkwarg():
+    assert eq(hfft(darr, 5), npfft.hfft(nparr, 5))
+    assert eq(hfft(darr, 13), npfft.hfft(nparr, 13))
+    assert eq(hfft(darr2, 5, axis=0), npfft.hfft(nparr, 5, axis=0))
+    assert eq(hfft(darr2, 13, axis=0), npfft.hfft(nparr, 13, axis=0))
+    assert eq(hfft(darr2, 12, axis=0), npfft.hfft(nparr, 12, axis=0))
+
+
+def test_ihfft():
+    assert eq(ihfft(darr), npfft.ihfft(nparr))
+    assert eq(ihfft(darr2, axis=0), npfft.ihfft(nparr, axis=0))
+
+
+def test_ihfft_n_kwarg():
+    assert eq(ihfft(darr, 5), npfft.ihfft(nparr, 5))
+    assert eq(ihfft(darr, 13), npfft.ihfft(nparr, 13))
+    assert eq(ihfft(darr2, 5, axis=0), npfft.ihfft(nparr, 5, axis=0))
+    assert eq(ihfft(darr2, 13, axis=0), npfft.ihfft(nparr, 13, axis=0))
+    assert eq(ihfft(darr2, 12, axis=0), npfft.ihfft(nparr, 12, axis=0))