Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/scipy/sparse/_index.py: 13%

1"""Indexing mixin for sparse matrix classes. 

2""" 

3import numpy as np 

4from ._sputils import isintlike 

5 

6try: 

7 INT_TYPES = (int, long, np.integer) 

8except NameError: 

9 # long is not defined in Python3 

10 INT_TYPES = (int, np.integer) 

11 

12 

13def _broadcast_arrays(a, b): 

14 """ 

15 Same as np.broadcast_arrays(a, b) but old writeability rules. 

16 

17 NumPy >= 1.17.0 transitions broadcast_arrays to return 

18 read-only arrays. Set writeability explicitly to avoid warnings. 

19 Retain the old writeability rules, as our Cython code assumes 

20 the old behavior. 

21 """ 

22 x, y = np.broadcast_arrays(a, b) 

23 x.flags.writeable = a.flags.writeable 

24 y.flags.writeable = b.flags.writeable 

25 return x, y 

26 

27 

28class IndexMixin: 

29 """ 

30 This class provides common dispatching and validation logic for indexing. 

31 """ 

32 def _raise_on_1d_array_slice(self): 

33 """We do not currently support 1D sparse arrays. 

34 

35 This function is called each time that a 1D array would 

36 result, raising an error instead. 

37 

38 Once 1D sparse arrays are implemented, it should be removed. 

39 """ 

40 if self._is_array: 

41 raise NotImplementedError( 

42 'We have not yet implemented 1D sparse slices; ' 

43 'please index using explicit indices, e.g. `x[:, [0]]`' 

44 ) 

45 

46 def __getitem__(self, key): 

47 row, col = self._validate_indices(key) 

48 

49 # Dispatch to specialized methods. 

50 if isinstance(row, INT_TYPES): 

51 if isinstance(col, INT_TYPES): 

52 return self._get_intXint(row, col) 

53 elif isinstance(col, slice): 

54 self._raise_on_1d_array_slice() 

55 return self._get_intXslice(row, col) 

56 elif col.ndim == 1: 

57 self._raise_on_1d_array_slice() 

58 return self._get_intXarray(row, col) 

59 elif col.ndim == 2: 

60 return self._get_intXarray(row, col) 

61 raise IndexError('index results in >2 dimensions') 

62 elif isinstance(row, slice): 

63 if isinstance(col, INT_TYPES): 

64 self._raise_on_1d_array_slice() 

65 return self._get_sliceXint(row, col) 

66 elif isinstance(col, slice): 

67 if row == slice(None) and row == col: 

68 return self.copy() 

69 return self._get_sliceXslice(row, col) 

70 elif col.ndim == 1: 

71 return self._get_sliceXarray(row, col) 

72 raise IndexError('index results in >2 dimensions') 

73 elif row.ndim == 1: 

74 if isinstance(col, INT_TYPES): 

75 self._raise_on_1d_array_slice() 

76 return self._get_arrayXint(row, col) 

77 elif isinstance(col, slice): 

78 return self._get_arrayXslice(row, col) 

79 else: # row.ndim == 2 

80 if isinstance(col, INT_TYPES): 

81 return self._get_arrayXint(row, col) 

82 elif isinstance(col, slice): 

83 raise IndexError('index results in >2 dimensions') 

84 elif row.shape[1] == 1 and (col.ndim == 1 or col.shape[0] == 1): 

85 # special case for outer indexing 

86 return self._get_columnXarray(row[:,0], col.ravel()) 

87 

88 # The only remaining case is inner (fancy) indexing 

89 row, col = _broadcast_arrays(row, col) 

90 if row.shape != col.shape: 

91 raise IndexError('number of row and column indices differ') 

92 if row.size == 0: 

93 return self.__class__(np.atleast_2d(row).shape, dtype=self.dtype) 

94 return self._get_arrayXarray(row, col) 

95 

96 def __setitem__(self, key, x): 

97 row, col = self._validate_indices(key) 

98 

99 if isinstance(row, INT_TYPES) and isinstance(col, INT_TYPES): 

100 x = np.asarray(x, dtype=self.dtype) 

101 if x.size != 1: 

102 raise ValueError('Trying to assign a sequence to an item') 

103 self._set_intXint(row, col, x.flat[0]) 

104 return 

105 

106 if isinstance(row, slice): 

107 row = np.arange(*row.indices(self.shape[0]))[:, None] 

108 else: 

109 row = np.atleast_1d(row) 

110 

111 if isinstance(col, slice): 

112 col = np.arange(*col.indices(self.shape[1]))[None, :] 

113 if row.ndim == 1: 

114 row = row[:, None] 

115 else: 

116 col = np.atleast_1d(col) 

117 

118 i, j = _broadcast_arrays(row, col) 

119 if i.shape != j.shape: 

120 raise IndexError('number of row and column indices differ') 

121 

122 from ._base import isspmatrix 

123 if isspmatrix(x): 

124 if i.ndim == 1: 

125 # Inner indexing, so treat them like row vectors. 

126 i = i[None] 

127 j = j[None] 

128 broadcast_row = x.shape[0] == 1 and i.shape[0] != 1 

129 broadcast_col = x.shape[1] == 1 and i.shape[1] != 1 

130 if not ((broadcast_row or x.shape[0] == i.shape[0]) and 

131 (broadcast_col or x.shape[1] == i.shape[1])): 

132 raise ValueError('shape mismatch in assignment') 

133 if x.shape[0] == 0 or x.shape[1] == 0: 

134 return 

135 x = x.tocoo(copy=True) 

136 x.sum_duplicates() 

137 self._set_arrayXarray_sparse(i, j, x) 

138 else: 

139 # Make x and i into the same shape 

140 x = np.asarray(x, dtype=self.dtype) 

141 if x.squeeze().shape != i.squeeze().shape: 

142 x = np.broadcast_to(x, i.shape) 

143 if x.size == 0: 

144 return 

145 x = x.reshape(i.shape) 

146 self._set_arrayXarray(i, j, x) 

147 

148 def _validate_indices(self, key): 

149 M, N = self.shape 

150 row, col = _unpack_index(key) 

151 

152 if isintlike(row): 

153 row = int(row) 

154 if row < -M or row >= M: 

155 raise IndexError('row index (%d) out of range' % row) 

156 if row < 0: 

157 row += M 

158 elif not isinstance(row, slice): 

159 row = self._asindices(row, M) 

160 

161 if isintlike(col): 

162 col = int(col) 

163 if col < -N or col >= N: 

164 raise IndexError('column index (%d) out of range' % col) 

165 if col < 0: 

166 col += N 

167 elif not isinstance(col, slice): 

168 col = self._asindices(col, N) 

169 

170 return row, col 

171 

172 def _asindices(self, idx, length): 

173 """Convert `idx` to a valid index for an axis with a given length. 

174 

175 Subclasses that need special validation can override this method. 

176 """ 

177 try: 

178 x = np.asarray(idx) 

179 except (ValueError, TypeError, MemoryError) as e: 

180 raise IndexError('invalid index') from e 

181 

182 if x.ndim not in (1, 2): 

183 raise IndexError('Index dimension must be 1 or 2') 

184 

185 if x.size == 0: 

186 return x 

187 

188 # Check bounds 

189 max_indx = x.max() 

190 if max_indx >= length: 

191 raise IndexError('index (%d) out of range' % max_indx) 

192 

193 min_indx = x.min() 

194 if min_indx < 0: 

195 if min_indx < -length: 

196 raise IndexError('index (%d) out of range' % min_indx) 

197 if x is idx or not x.flags.owndata: 

198 x = x.copy() 

199 x[x < 0] += length 

200 return x 

201 

202 def getrow(self, i): 

203 """Return a copy of row i of the matrix, as a (1 x n) row vector. 

204 """ 

205 M, N = self.shape 

206 i = int(i) 

207 if i < -M or i >= M: 

208 raise IndexError('index (%d) out of range' % i) 

209 if i < 0: 

210 i += M 

211 return self._get_intXslice(i, slice(None)) 

212 

213 def getcol(self, i): 

214 """Return a copy of column i of the matrix, as a (m x 1) column vector. 

215 """ 

216 M, N = self.shape 

217 i = int(i) 

218 if i < -N or i >= N: 

219 raise IndexError('index (%d) out of range' % i) 

220 if i < 0: 

221 i += N 

222 return self._get_sliceXint(slice(None), i) 

223 

224 def _get_intXint(self, row, col): 

225 raise NotImplementedError() 

226 

227 def _get_intXarray(self, row, col): 

228 raise NotImplementedError() 

229 

230 def _get_intXslice(self, row, col): 

231 raise NotImplementedError() 

232 

233 def _get_sliceXint(self, row, col): 

234 raise NotImplementedError() 

235 

236 def _get_sliceXslice(self, row, col): 

237 raise NotImplementedError() 

238 

239 def _get_sliceXarray(self, row, col): 

240 raise NotImplementedError() 

241 

242 def _get_arrayXint(self, row, col): 

243 raise NotImplementedError() 

244 

245 def _get_arrayXslice(self, row, col): 

246 raise NotImplementedError() 

247 

248 def _get_columnXarray(self, row, col): 

249 raise NotImplementedError() 

250 

251 def _get_arrayXarray(self, row, col): 

252 raise NotImplementedError() 

253 

254 def _set_intXint(self, row, col, x): 

255 raise NotImplementedError() 

256 

257 def _set_arrayXarray(self, row, col, x): 

258 raise NotImplementedError() 

259 

260 def _set_arrayXarray_sparse(self, row, col, x): 

261 # Fall back to densifying x 

262 x = np.asarray(x.toarray(), dtype=self.dtype) 

263 x, _ = _broadcast_arrays(x, row) 

264 self._set_arrayXarray(row, col, x) 

265 

266 

267def _unpack_index(index): 

268 """ Parse index. Always return a tuple of the form (row, col). 

269 Valid type for row/col is integer, slice, or array of integers. 

270 """ 

271 # First, check if indexing with single boolean matrix. 

272 from ._base import spmatrix, isspmatrix 

273 if (isinstance(index, (spmatrix, np.ndarray)) and 

274 index.ndim == 2 and index.dtype.kind == 'b'): 

275 return index.nonzero() 

276 

277 # Parse any ellipses. 

278 index = _check_ellipsis(index) 

279 

280 # Next, parse the tuple or object 

281 if isinstance(index, tuple): 

282 if len(index) == 2: 

283 row, col = index 

284 elif len(index) == 1: 

285 row, col = index[0], slice(None) 

286 else: 

287 raise IndexError('invalid number of indices') 

288 else: 

289 idx = _compatible_boolean_index(index) 

290 if idx is None: 

291 row, col = index, slice(None) 

292 elif idx.ndim < 2: 

293 return _boolean_index_to_array(idx), slice(None) 

294 elif idx.ndim == 2: 

295 return idx.nonzero() 

296 # Next, check for validity and transform the index as needed. 

297 if isspmatrix(row) or isspmatrix(col): 

298 # Supporting sparse boolean indexing with both row and col does 

299 # not work because spmatrix.ndim is always 2. 

300 raise IndexError( 

301 'Indexing with sparse matrices is not supported ' 

302 'except boolean indexing where matrix and index ' 

303 'are equal shapes.') 

304 bool_row = _compatible_boolean_index(row) 

305 bool_col = _compatible_boolean_index(col) 

306 if bool_row is not None: 

307 row = _boolean_index_to_array(bool_row) 

308 if bool_col is not None: 

309 col = _boolean_index_to_array(bool_col) 

310 return row, col 

311 

312 

313def _check_ellipsis(index): 

314 """Process indices with Ellipsis. Returns modified index.""" 

315 if index is Ellipsis: 

316 return (slice(None), slice(None)) 

317 

318 if not isinstance(index, tuple): 

319 return index 

320 

321 # TODO: Deprecate this multiple-ellipsis handling, 

322 # as numpy no longer supports it. 

323 

324 # Find first ellipsis. 

325 for j, v in enumerate(index): 

326 if v is Ellipsis: 

327 first_ellipsis = j 

328 break 

329 else: 

330 return index 

331 

332 # Try to expand it using shortcuts for common cases 

333 if len(index) == 1: 

334 return (slice(None), slice(None)) 

335 if len(index) == 2: 

336 if first_ellipsis == 0: 

337 if index[1] is Ellipsis: 

338 return (slice(None), slice(None)) 

339 return (slice(None), index[1]) 

340 return (index[0], slice(None)) 

341 

342 # Expand it using a general-purpose algorithm 

343 tail = [] 

344 for v in index[first_ellipsis+1:]: 

345 if v is not Ellipsis: 

346 tail.append(v) 

347 nd = first_ellipsis + len(tail) 

348 nslice = max(0, 2 - nd) 

349 return index[:first_ellipsis] + (slice(None),)*nslice + tuple(tail) 

350 

351 

352def _maybe_bool_ndarray(idx): 

353 """Returns a compatible array if elements are boolean. 

354 """ 

355 idx = np.asanyarray(idx) 

356 if idx.dtype.kind == 'b': 

357 return idx 

358 return None 

359 

360 

361def _first_element_bool(idx, max_dim=2): 

362 """Returns True if first element of the incompatible 

363 array type is boolean. 

364 """ 

365 if max_dim < 1: 

366 return None 

367 try: 

368 first = next(iter(idx), None) 

369 except TypeError: 

370 return None 

371 if isinstance(first, bool): 

372 return True 

373 return _first_element_bool(first, max_dim-1) 

374 

375 

376def _compatible_boolean_index(idx): 

377 """Returns a boolean index array that can be converted to 

378 integer array. Returns None if no such array exists. 

379 """ 

380 # Presence of attribute `ndim` indicates a compatible array type. 

381 if hasattr(idx, 'ndim') or _first_element_bool(idx): 

382 return _maybe_bool_ndarray(idx) 

383 return None 

384 

385 

386def _boolean_index_to_array(idx): 

387 if idx.ndim > 1: 

388 raise IndexError('invalid index shape') 

389 return np.where(idx)[0]