Repository: dask/dask

Issue:
NPY support
It would be nice to have a method `array.to_npy('file.npy')`



Discussion:
Sure.  Just to be clear, my understanding is that npy files do not support partial reads so we should be able to assume that the result fits in memory.  Is this correct?

If this is the case then the implementation will probably be something only marginally more complex than the following:

``` python
def save(fn, x):
    np.save(fn, x.compute()
```

Actually I would need this for huge out-of-memory arrays. AFAIU [npy does support memmapping](http://docs.scipy.org/doc/numpy/reference/generated/numpy.load.html)?

To write a big array on disk maybe you could first write an empty array in a file, and fill it chunk by chunk? Otherwise there might be a way to just append data to an existing file, since [npy](http://docs.scipy.org/doc/numpy-dev/neps/npy-format.html) is just a flat binary file with a header.

Ah, I had forgotten about memmap.

Your suggested approach should work well.  Here is a small test:

``` python
def test_memmap():
    with tmpfile('npy') as fn_1:
        with tmpfile('npy') as fn_2:
            x = da.arange(100, chunks=15)
            target = np.memmap(fn_1, shape=x.shape, mode='w+', dtype=x.dtype)

            x.store(target)
            assert eq(target, x)

            np.save(fn_2, target)
            assert eq(np.load(fn_2, mmap_mode='r'), x)
```

Hrm, this approach unfortunately involves two complete writes to disk.

From the format spec, we could probably do an append approach if we wanted (https://github.com/numpy/numpy/blob/master/doc/neps/npy-format.rst). Would you want the arrays written as chunks (format looks like multiple arrays on disk to numpy), or as a single array (looks like 1 array on disk to numpy)?

@jcrist probably multiple arrays on disk. Here's what I'm trying to do:

Initially I have a dask array `D=(D1, D2, ..., Dn)` (2D arrays chunked along axis=1) loaded from a big flat binary array. Then a function `f` is applied to every chunk in parallel, returning array `O=(O1, ..., On)`. I want to persist this big output array on disk, using the `npy` format, in a sensible way that I can control. For example, something like `.dask_cache/output_array_1.npy` etc. AFAIK we can't write to a single big file because of file concurrency issues with multiprocessing.

There would actually be several tasks like this, caching results in different subfolders.

These files would be used in two ways:
- if you relaunch the task, or if the process crashed in the middle, the existing results are directly reused
- the user may want to export the output to a big `npy` array, concatenation of the chunked `npy` files

Ah, ok, so you don't want a single npy file, you're looking for a more custom/chunked solution.  Implementing this is probably easier than figuring out the right API/spec.  What were you thinking?  Perhaps the following?

``` python
x.to_npy(dirname, axis=0)
```

Yes, something like that. I guess there would be one directory per dask array, with one file per chunk?

Conversely, how would I load this set of files into a dask array without loading all files into memory? This would be to export this multi-npy-files array into a single npy array.

I was thinking one file per slice rather than one file per chunk.  

Regarding loading, this work would accompany a second function to load.

So @rossant just to be clear on things before I start on this.  Does the following look right to you?

``` python
>>> da.to_npy_stack('/path/to/data/', my_dask_array, axis=0)
>>> y = da.from_npy_stack('/path/to/data/')

$ ls /path/to/data/
0.npy
1.npy
2.npy
```

Just a note -- reading mmapped arrays with dask is already highly performant with the multi-threaded engine. We do this all the time for reading netcdf3 files in xray via scipy (which uses mmap for reads).

@mrocklin to be clear can you precise the structure of `my_dask_array` and the contents of `0.npy` etc?

``` python
>>> my_dask_array.shape
(5, 1000, 2000)

>>> np.load('/path/to/data/0.npy').shape
(1000, 2000)
```

``` python
>>> (my_dask_array[0].compute() == np.load('/path/to/data/0.npy')).all()
True
```

ah, okay, so that's not exactly what I had in mind. In my case I have `my_dask_array.shape == (5000, 2000)` with, say, 100 chunks of size 50, and I'd expect `np.load('/path/to/data/0.npy').shape == (50, 2000)`


Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -4,7 +4,7 @@
 from .core import (Array, stack, concatenate, take, tensordot, transpose,
         from_array, choose, where, coarsen, insert, broadcast_to, fromfunction,
         unique, store, squeeze, topk, bincount, histogram, map_blocks, atop,
-        to_hdf5, dot, cov, array)
+        to_hdf5, dot, cov, array, to_npy_stack, from_npy_stack)
 from .core import (logaddexp, logaddexp2, conj, exp, log, log2, log10, log1p,
         expm1, sqrt, square, sin, cos, tan, arcsin, arccos, arctan, arctan2,
         hypot, sinh, cosh, tanh, arcsinh, arccosh, arctanh, deg2rad, rad2deg,
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -14,6 +14,9 @@
                            memoize, map, groupby, valmap, accumulate, merge,
                            curry, reduce, interleave, sliding_window, partial)
 import numpy as np
+import os
+import pickle
+import uuid
 
 from threading import Lock
 from . import chunk
@@ -1089,6 +1092,7 @@ def squeeze(self):
         return squeeze(self)
 
     def rechunk(self, chunks):
+        """ See da.rechunk for docstring """
         from .rechunk import rechunk
         return rechunk(self, chunks)
 
@@ -2580,3 +2584,80 @@ def cov(m, y=None, rowvar=1, bias=0, ddof=None):
         return (dot(X.T, X.conj()) / fact).squeeze()
     else:
         return (dot(X, X.T.conj()) / fact).squeeze()
+
+
+def to_npy_stack(dirname, x, axis=0):
+    """ Write dask array to a stack of .npy files
+
+    This partitions the dask.array along one axis and stores each block along
+    that axis as a single .npy file in the specified directory
+
+    Example
+    -------
+
+    >>> x = da.ones((5, 10, 10), chunks=(2, 4, 4))  # doctest: +SKIP
+    >>> da.to_npy_stack('data/', x, axis=0)  # doctest: +SKIP
+
+    $ tree data/
+    data/
+    |-- 0.npy
+    |-- 1.npy
+    |-- 2.npy
+    |-- info
+
+    The ``.npy`` files store numpy arrays for ``x[0:2], x[2:4], and x[4:5]``
+    respectively, as is specified by the chunk size along the zeroth axis.  The
+    info file stores the dtype, chunks, and axis information of the array.
+
+    You can load these stacks with the ``da.from_npy_stack`` function.
+
+    >>> y = da.from_npy_stack('data/')  # doctest: +SKIP
+
+    See also:
+        from_npy_stack
+    """
+
+    chunks = tuple((c if i == axis else (sum(c),))
+                   for i, c in enumerate(x.chunks))
+    xx = x.rechunk(chunks)
+
+    if not os.path.exists(dirname):
+        os.path.mkdir(dirname)
+
+    meta = {'chunks': chunks, 'dtype': x.dtype, 'axis': axis}
+
+    with open(os.path.join(dirname, 'info'), 'wb') as f:
+        pickle.dump(meta, f)
+
+    name = 'to-npy-stack-' + str(uuid.uuid1())
+    dsk = dict(((name, i), (np.save, os.path.join(dirname, '%d.npy' % i), key))
+              for i, key in enumerate(core.flatten(xx._keys())))
+
+    Array._get(merge(dsk, xx.dask), list(dsk))
+
+def from_npy_stack(dirname, mmap_mode='r'):
+    """ Load dask array from stack of npy files
+
+    See ``da.to_npy_stack`` for docstring
+
+    Parameters
+    ----------
+    dirname: string
+        Directory of .npy files
+    mmap_mode: (None or 'r')
+        Read data in memory map mode
+    """
+    with open(os.path.join(dirname, 'info'), 'rb') as f:
+        info = pickle.load(f)
+
+    dtype = info['dtype']
+    chunks = info['chunks']
+    axis = info['axis']
+
+    name = 'from-npy-stack-%s' % dirname
+    keys = list(product([name], *[range(len(c)) for c in chunks]))
+    values = [(np.load, os.path.join(dirname, '%d.npy' % i), mmap_mode)
+              for i in range(len(chunks[axis]))]
+    dsk = dict(zip(keys, values))
+
+    return Array(dsk, name, chunks, dtype)


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -4,6 +4,9 @@
 pytest.importorskip('numpy')
 
 from operator import add
+from tempfile import mkdtemp
+import shutil
+import os
 
 from toolz import merge
 from toolz.curried import identity
@@ -1401,3 +1404,34 @@ def test_cov():
     assert eq(da.cov(e, d), np.cov(y, x))
 
     assert raises(ValueError, lambda: da.cov(d, ddof=1.5))
+
+
+def test_memmap():
+    with tmpfile('npy') as fn_1:
+        with tmpfile('npy') as fn_2:
+            x = da.arange(100, chunks=15)
+            target = np.memmap(fn_1, shape=x.shape, mode='w+', dtype=x.dtype)
+
+            x.store(target)
+
+            assert eq(target, x)
+
+            np.save(fn_2, target)
+
+            assert eq(np.load(fn_2, mmap_mode='r'), x)
+
+
+def test_to_npy_stack():
+    x = np.arange(5*10*10).reshape((5, 10, 10))
+    d = da.from_array(x, chunks=(2, 4, 4))
+
+    dirname = mkdtemp()
+    try:
+        da.to_npy_stack(dirname, d, axis=0)
+        assert os.path.exists(os.path.join(dirname, '0.npy'))
+        assert (np.load(os.path.join(dirname, '1.npy')) == x[2:4]).all()
+
+        e = da.from_npy_stack(dirname)
+        assert eq(d, e)
+    finally:
+        shutil.rmtree(dirname)