Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import numpy as np 

2import scipy.sparse 

3 

4__all__ = ['save_npz', 'load_npz'] 

5 

6 

7# Make loading safe vs. malicious input 

8PICKLE_KWARGS = dict(allow_pickle=False) 

9 

10 

11def save_npz(file, matrix, compressed=True): 

12 """ Save a sparse matrix to a file using ``.npz`` format. 

13 

14 Parameters 

15 ---------- 

16 file : str or file-like object 

17 Either the file name (string) or an open file (file-like object) 

18 where the data will be saved. If file is a string, the ``.npz`` 

19 extension will be appended to the file name if it is not already 

20 there. 

21 matrix: spmatrix (format: ``csc``, ``csr``, ``bsr``, ``dia`` or coo``) 

22 The sparse matrix to save. 

23 compressed : bool, optional 

24 Allow compressing the file. Default: True 

25 

26 See Also 

27 -------- 

28 scipy.sparse.load_npz: Load a sparse matrix from a file using ``.npz`` format. 

29 numpy.savez: Save several arrays into a ``.npz`` archive. 

30 numpy.savez_compressed : Save several arrays into a compressed ``.npz`` archive. 

31 

32 Examples 

33 -------- 

34 Store sparse matrix to disk, and load it again: 

35 

36 >>> import scipy.sparse 

37 >>> sparse_matrix = scipy.sparse.csc_matrix(np.array([[0, 0, 3], [4, 0, 0]])) 

38 >>> sparse_matrix 

39 <2x3 sparse matrix of type '<class 'numpy.int64'>' 

40 with 2 stored elements in Compressed Sparse Column format> 

41 >>> sparse_matrix.todense() 

42 matrix([[0, 0, 3], 

43 [4, 0, 0]], dtype=int64) 

44 

45 >>> scipy.sparse.save_npz('/tmp/sparse_matrix.npz', sparse_matrix) 

46 >>> sparse_matrix = scipy.sparse.load_npz('/tmp/sparse_matrix.npz') 

47 

48 >>> sparse_matrix 

49 <2x3 sparse matrix of type '<class 'numpy.int64'>' 

50 with 2 stored elements in Compressed Sparse Column format> 

51 >>> sparse_matrix.todense() 

52 matrix([[0, 0, 3], 

53 [4, 0, 0]], dtype=int64) 

54 """ 

55 arrays_dict = {} 

56 if matrix.format in ('csc', 'csr', 'bsr'): 

57 arrays_dict.update(indices=matrix.indices, indptr=matrix.indptr) 

58 elif matrix.format == 'dia': 

59 arrays_dict.update(offsets=matrix.offsets) 

60 elif matrix.format == 'coo': 

61 arrays_dict.update(row=matrix.row, col=matrix.col) 

62 else: 

63 raise NotImplementedError('Save is not implemented for sparse matrix of format {}.'.format(matrix.format)) 

64 arrays_dict.update( 

65 format=matrix.format.encode('ascii'), 

66 shape=matrix.shape, 

67 data=matrix.data 

68 ) 

69 if compressed: 

70 np.savez_compressed(file, **arrays_dict) 

71 else: 

72 np.savez(file, **arrays_dict) 

73 

74 

75def load_npz(file): 

76 """ Load a sparse matrix from a file using ``.npz`` format. 

77 

78 Parameters 

79 ---------- 

80 file : str or file-like object 

81 Either the file name (string) or an open file (file-like object) 

82 where the data will be loaded. 

83 

84 Returns 

85 ------- 

86 result : csc_matrix, csr_matrix, bsr_matrix, dia_matrix or coo_matrix 

87 A sparse matrix containing the loaded data. 

88 

89 Raises 

90 ------ 

91 IOError 

92 If the input file does not exist or cannot be read. 

93 

94 See Also 

95 -------- 

96 scipy.sparse.save_npz: Save a sparse matrix to a file using ``.npz`` format. 

97 numpy.load: Load several arrays from a ``.npz`` archive. 

98 

99 Examples 

100 -------- 

101 Store sparse matrix to disk, and load it again: 

102 

103 >>> import scipy.sparse 

104 >>> sparse_matrix = scipy.sparse.csc_matrix(np.array([[0, 0, 3], [4, 0, 0]])) 

105 >>> sparse_matrix 

106 <2x3 sparse matrix of type '<class 'numpy.int64'>' 

107 with 2 stored elements in Compressed Sparse Column format> 

108 >>> sparse_matrix.todense() 

109 matrix([[0, 0, 3], 

110 [4, 0, 0]], dtype=int64) 

111 

112 >>> scipy.sparse.save_npz('/tmp/sparse_matrix.npz', sparse_matrix) 

113 >>> sparse_matrix = scipy.sparse.load_npz('/tmp/sparse_matrix.npz') 

114 

115 >>> sparse_matrix 

116 <2x3 sparse matrix of type '<class 'numpy.int64'>' 

117 with 2 stored elements in Compressed Sparse Column format> 

118 >>> sparse_matrix.todense() 

119 matrix([[0, 0, 3], 

120 [4, 0, 0]], dtype=int64) 

121 """ 

122 

123 with np.load(file, **PICKLE_KWARGS) as loaded: 

124 try: 

125 matrix_format = loaded['format'] 

126 except KeyError: 

127 raise ValueError('The file {} does not contain a sparse matrix.'.format(file)) 

128 

129 matrix_format = matrix_format.item() 

130 

131 if not isinstance(matrix_format, str): 

132 # Play safe with Python 2 vs 3 backward compatibility; 

133 # files saved with SciPy < 1.0.0 may contain unicode or bytes. 

134 matrix_format = matrix_format.decode('ascii') 

135 

136 try: 

137 cls = getattr(scipy.sparse, '{}_matrix'.format(matrix_format)) 

138 except AttributeError: 

139 raise ValueError('Unknown matrix format "{}"'.format(matrix_format)) 

140 

141 if matrix_format in ('csc', 'csr', 'bsr'): 

142 return cls((loaded['data'], loaded['indices'], loaded['indptr']), shape=loaded['shape']) 

143 elif matrix_format == 'dia': 

144 return cls((loaded['data'], loaded['offsets']), shape=loaded['shape']) 

145 elif matrix_format == 'coo': 

146 return cls((loaded['data'], (loaded['row'], loaded['col'])), shape=loaded['shape']) 

147 else: 

148 raise NotImplementedError('Load is not implemented for ' 

149 'sparse matrix of format {}.'.format(matrix_format))