import numpy as np
from ._coo.core import COO
from ._compressed import GCXS
[docs]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()``.
This binary format is not currently stable. 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 : SparseArray
The matrix to save to disk
compressed : bool
Whether to save in compressed or uncompressed mode
Examples
--------
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, fill_value=0.0>
>>> sparse.save_npz('mat.npz', mat)
>>> loaded_mat = sparse.load_npz('mat.npz')
>>> loaded_mat
<COO: shape=(2, 2, 2), dtype=float64, nnz=2, fill_value=0.0>
>>> os.remove('mat.npz')
See Also
--------
load_npz
scipy.sparse.save_npz
scipy.sparse.load_npz
numpy.savez
numpy.load
"""
nodes = {
"data": matrix.data,
"shape": matrix.shape,
"fill_value": matrix.fill_value,
}
if type(matrix) == COO:
nodes["coords"] = matrix.coords
elif type(matrix) == GCXS:
nodes["indices"] = matrix.indices
nodes["indptr"] = matrix.indptr
nodes["compressed_axes"] = matrix.compressed_axes
if compressed:
np.savez_compressed(filename, **nodes)
else:
np.savez(filename, **nodes)
[docs]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. This binary format is not currently stable.
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
-------
SparseArray
The sparse matrix at path ``filename``.
Examples
--------
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"])
fill_value = fp["fill_value"][()]
return COO(
coords=coords,
data=data,
shape=shape,
sorted=True,
has_duplicates=False,
fill_value=fill_value,
)
except KeyError:
pass
try:
data = fp["data"]
indices = fp["indices"]
indptr = fp["indptr"]
comp_axes = fp["compressed_axes"]
shape = tuple(fp["shape"])
fill_value = fp["fill_value"][()]
return GCXS(
(data, indices, indptr),
shape=shape,
fill_value=fill_value,
compressed_axes=comp_axes,
)
except KeyError:
raise RuntimeError(
"The file {!s} does not contain a valid sparse matrix".format(filename)
)