Adding support for saving matrices to files in npz format

docs/generated/sparse.load_npz.rst

@@ -0,0 +1,6 @@
+load\_npz
+=========
+
+.. currentmodule:: sparse
+
+.. autofunction:: load_npz

docs/generated/sparse.rst

@@ -31,6 +31,8 @@ API
 
     elemwise
 
+    load_npz
+
     nanmax
 
     nanmin
@@ -43,6 +45,8 @@ API
 
     random
 
+    save_npz
+
     stack
 
     tensordot
@@ -52,4 +56,3 @@ API
     triu
 
     where
-

docs/generated/sparse.save_npz.rst

@@ -0,0 +1,6 @@
+save\_npz
+=========
+
+.. currentmodule:: sparse
+
+.. autofunction:: save_npz

sparse/__init__.py

@@ -2,6 +2,7 @@
 from .dok import DOK
 from .sparse_array import SparseArray
 from .utils import random
+from .io import save_npz, load_npz
 
 from ._version import get_versions
 __version__ = get_versions()['version']

sparse/io.py

@@ -0,0 +1,104 @@
+import numpy as np
+
+from .coo.core import COO
+
+
+def save_npz(filename, matrix, compressed=True):
+    """ Save a sparse matrix to disk in numpy's ``.npz`` format.
+    Note: This is not binary compatible with scipy's ``save_npz()``.
+    Will save a file that can only be opend with this package's ``load_npz()``.
+
+    Parameters
+    ----------
+    filename : string or file
+        Either the file name (string) or an open file (file-like object)
+        where the data will be saved. If file is a string or a Path, the
+        ``.npz`` extension will be appended to the file name if it is not
+        already there
+    matrix : COO
+        The matrix to save to disk
+    compressed : bool
+        Whether to save in compressed or uncompressed mode
+
+    Example
+    --------
+    Store sparse matrix to disk, and load it again:
+
+    >>> import os
+    >>> import sparse
+    >>> import numpy as np
+    >>> dense_mat = np.array([[[0., 0.], [0., 0.70677779]], [[0., 0.], [0., 0.86522495]]])
+    >>> mat = sparse.COO(dense_mat)
+    >>> mat
+    <COO: shape=(2, 2, 2), dtype=float64, nnz=2>
+    >>> sparse.save_npz('mat.npz', mat)
+    >>> loaded_mat = sparse.load_npz('mat.npz')
+    >>> loaded_mat
+    <COO: shape=(2, 2, 2), dtype=float64, nnz=2>
+    >>> loaded_mat.todense()
+    array([[[0.        , 0.        ],
+            [0.        , 0.70677779]],
+    <BLANKLINE>
+           [[0.        , 0.        ],
+            [0.        , 0.86522495]]])
+    >>> os.remove('mat.npz')
+
+    See Also
+    --------
+    load_npz
+    scipy.sparse.save_npz
+    scipy.sparse.load_npz
+    numpy.savez
+    numpy.load
+
+    """
+
+    nodes = {'data':   matrix.data,
+             'coords': matrix.coords,
+             'shape':  matrix.shape}
+
+    if compressed:
+        np.savez_compressed(filename, **nodes)
+    else:
+        np.savez(filename, **nodes)
+
+
+def load_npz(filename):
+    """ Load a sparse matrix in numpy's ``.npz`` format from disk.
+    Note: This is not binary compatible with scipy's ``save_npz()``
+    output. Will only load files saved by this package.
+
+    Parameters
+    ----------
+    filename : file-like object, string, or pathlib.Path
+        The file to read. File-like objects must support the
+        ``seek()`` and ``read()`` methods.
+
+    Returns
+    -------
+    COO
+        The sparse matrix at path ``filename``
+
+    Example
+    --------
+    See :obj:`save_npz` for usage examples.
+
+    See Also
+    --------
+    save_npz
+    scipy.sparse.save_npz
+    scipy.sparse.load_npz
+    numpy.savez
+    numpy.load
+
+    """
+
+    with np.load(filename) as fp:
+        try:
+            coords = fp['coords']
+            data = fp['data']
+            shape = tuple(fp['shape'])
+        except KeyError:
+            raise RuntimeError('The file {} does not contain a valid sparse matrix'.format(filename))
+
+    return COO(coords=coords, data=data, shape=shape, sorted=True, has_duplicates=False)

sparse/tests/test_io.py

@@ -0,0 +1,39 @@
+import os
+import tempfile
+import shutil
+import pytest
+import numpy as np
+
+import sparse
+
+from sparse import save_npz, load_npz
+from sparse.utils import assert_eq
+
+
+@pytest.mark.parametrize('compression', [True, False])
+def test_save_load_npz_file(compression):
+    x = sparse.random((2, 3, 4, 5), density=.25)
+    y = x.todense()
+
+    dir_name = tempfile.mkdtemp()
+    filename = os.path.join(dir_name, 'mat.npz')
+
+    save_npz(filename, x, compressed=compression)
+    z = load_npz(filename)
+    assert_eq(x, z)
+    assert_eq(y, z.todense())
+
+    shutil.rmtree(dir_name)
+
+
+def test_load_wrong_format_exception():
+    x = np.array([1, 2, 3])
+
+    dir_name = tempfile.mkdtemp()
+    filename = os.path.join(dir_name, 'mat.npz')
+
+    np.savez(filename, x)
+    with pytest.raises(RuntimeError):
+        load_npz(filename)
+
+    shutil.rmtree(dir_name)