1  """
 2  SparseArray module
 3  """
 4  
 5  import numpy as np
 6  
 7  # Conditional import
 8  try:
 9      from scipy.sparse import csr_matrix
10  
11      SCIPY = True
12  except ImportError:
13      SCIPY = False
14  
15  
16  class SparseArray:
17      """
18      Methods to load and save sparse arrays to file.
19      """
20  
21      def __init__(self):
22          """
23          Creates a SparseArray instance.
24          """
25  
26          if not SCIPY:
27              raise ImportError("SciPy is not available - install scipy to enable")
28  
29      def load(self, f):
30          """
31          Loads a sparse array from file.
32  
33          Args:
34              f: input file handle
35  
36          Returns:
37              sparse array
38          """
39  
40          # Load raw data
41          data, indices, indptr, shape = (
42              np.load(f, allow_pickle=False),
43              np.load(f, allow_pickle=False),
44              np.load(f, allow_pickle=False),
45              np.load(f, allow_pickle=False),
46          )
47  
48          # Load data into sparse array
49          return csr_matrix((data, indices, indptr), shape=shape)
50  
51      def save(self, f, array):
52          """
53          Saves a sparse array to file.
54  
55          Args:
56              f: output file handle
57              array: sparse array
58          """
59  
60          # Save sparse array to file
61          for x in [array.data, array.indices, array.indptr, array.shape]:
62              np.save(f, x, allow_pickle=False)