Coverage for /usr/lib/python3/dist-packages/scipy/sparse/_index.py: 12%

1"""Indexing mixin for sparse matrix classes. 

2""" 

3import numpy as np 

4from warnings import warn 

5from ._sputils import isintlike 

6 

7INT_TYPES = (int, np.integer) 

8 

9 

10def _broadcast_arrays(a, b): 

11 """ 

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

13 

14 NumPy >= 1.17.0 transitions broadcast_arrays to return 

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

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

17 the old behavior. 

18 """ 

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

20 x.flags.writeable = a.flags.writeable 

21 y.flags.writeable = b.flags.writeable 

22 return x, y 

23 

24 

25class IndexMixin: 

26 """ 

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

28 """ 

29 def _raise_on_1d_array_slice(self): 

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

31 

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

33 result, raising an error instead. 

34 

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

36 """ 

37 if self._is_array: 

38 raise NotImplementedError( 

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

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

41 ) 

42 

43 def __getitem__(self, key): 

44 row, col = self._validate_indices(key) 

45 

46 # Dispatch to specialized methods. 

47 if isinstance(row, INT_TYPES): 

48 if isinstance(col, INT_TYPES): 

49 return self._get_intXint(row, col) 

50 elif isinstance(col, slice): 

51 self._raise_on_1d_array_slice() 

52 return self._get_intXslice(row, col) 

53 elif col.ndim == 1: 

54 self._raise_on_1d_array_slice() 

55 return self._get_intXarray(row, col) 

56 elif col.ndim == 2: 

57 return self._get_intXarray(row, col) 

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

59 elif isinstance(row, slice): 

60 if isinstance(col, INT_TYPES): 

61 self._raise_on_1d_array_slice() 

62 return self._get_sliceXint(row, col) 

63 elif isinstance(col, slice): 

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

65 return self.copy() 

66 return self._get_sliceXslice(row, col) 

67 elif col.ndim == 1: 

68 return self._get_sliceXarray(row, col) 

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

70 elif row.ndim == 1: 

71 if isinstance(col, INT_TYPES): 

72 self._raise_on_1d_array_slice() 

73 return self._get_arrayXint(row, col) 

74 elif isinstance(col, slice): 

75 return self._get_arrayXslice(row, col) 

76 else: # row.ndim == 2 

77 if isinstance(col, INT_TYPES): 

78 return self._get_arrayXint(row, col) 

79 elif isinstance(col, slice): 

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

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

82 # special case for outer indexing 

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

84 

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

86 row, col = _broadcast_arrays(row, col) 

87 if row.shape != col.shape: 

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

89 if row.size == 0: 

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

91 return self._get_arrayXarray(row, col) 

92 

93 def __setitem__(self, key, x): 

94 row, col = self._validate_indices(key) 

95 

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

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

98 if x.size != 1: 

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

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

101 return 

102 

103 if isinstance(row, slice): 

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

105 else: 

106 row = np.atleast_1d(row) 

107 

108 if isinstance(col, slice): 

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

110 if row.ndim == 1: 

111 row = row[:, None] 

112 else: 

113 col = np.atleast_1d(col) 

114 

115 i, j = _broadcast_arrays(row, col) 

116 if i.shape != j.shape: 

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

118 

119 from ._base import issparse 

120 if issparse(x): 

121 if i.ndim == 1: 

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

123 i = i[None] 

124 j = j[None] 

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

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

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

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

129 raise ValueError('shape mismatch in assignment') 

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

131 return 

132 x = x.tocoo(copy=True) 

133 x.sum_duplicates() 

134 self._set_arrayXarray_sparse(i, j, x) 

135 else: 

136 # Make x and i into the same shape 

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

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

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

140 if x.size == 0: 

141 return 

142 x = x.reshape(i.shape) 

143 self._set_arrayXarray(i, j, x) 

144 

145 def _validate_indices(self, key): 

146 M, N = self.shape 

147 row, col = _unpack_index(key) 

148 

149 if isintlike(row): 

150 row = int(row) 

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

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

153 if row < 0: 

154 row += M 

155 elif not isinstance(row, slice): 

156 row = self._asindices(row, M) 

157 

158 if isintlike(col): 

159 col = int(col) 

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

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

162 if col < 0: 

163 col += N 

164 elif not isinstance(col, slice): 

165 col = self._asindices(col, N) 

166 

167 return row, col 

168 

169 def _asindices(self, idx, length): 

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

171 

172 Subclasses that need special validation can override this method. 

173 """ 

174 try: 

175 x = np.asarray(idx) 

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

177 raise IndexError('invalid index') from e 

178 

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

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

181 

182 if x.size == 0: 

183 return x 

184 

185 # Check bounds 

186 max_indx = x.max() 

187 if max_indx >= length: 

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

189 

190 min_indx = x.min() 

191 if min_indx < 0: 

192 if min_indx < -length: 

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

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

195 x = x.copy() 

196 x[x < 0] += length 

197 return x 

198 

199 def _getrow(self, i): 

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

201 """ 

202 M, N = self.shape 

203 i = int(i) 

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

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

206 if i < 0: 

207 i += M 

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

209 

210 def _getcol(self, i): 

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

212 """ 

213 M, N = self.shape 

214 i = int(i) 

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

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

217 if i < 0: 

218 i += N 

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

220 

221 def _get_intXint(self, row, col): 

222 raise NotImplementedError() 

223 

224 def _get_intXarray(self, row, col): 

225 raise NotImplementedError() 

226 

227 def _get_intXslice(self, row, col): 

228 raise NotImplementedError() 

229 

230 def _get_sliceXint(self, row, col): 

231 raise NotImplementedError() 

232 

233 def _get_sliceXslice(self, row, col): 

234 raise NotImplementedError() 

235 

236 def _get_sliceXarray(self, row, col): 

237 raise NotImplementedError() 

238 

239 def _get_arrayXint(self, row, col): 

240 raise NotImplementedError() 

241 

242 def _get_arrayXslice(self, row, col): 

243 raise NotImplementedError() 

244 

245 def _get_columnXarray(self, row, col): 

246 raise NotImplementedError() 

247 

248 def _get_arrayXarray(self, row, col): 

249 raise NotImplementedError() 

250 

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

252 raise NotImplementedError() 

253 

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

255 raise NotImplementedError() 

256 

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

258 # Fall back to densifying x 

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

260 x, _ = _broadcast_arrays(x, row) 

261 self._set_arrayXarray(row, col, x) 

262 

263 

264def _unpack_index(index): 

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

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

267 """ 

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

269 from ._base import _spbase, issparse 

270 if (isinstance(index, (_spbase, np.ndarray)) and 

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

272 return index.nonzero() 

273 

274 # Parse any ellipses. 

275 index = _check_ellipsis(index) 

276 

277 # Next, parse the tuple or object 

278 if isinstance(index, tuple): 

279 if len(index) == 2: 

280 row, col = index 

281 elif len(index) == 1: 

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

283 else: 

284 raise IndexError('invalid number of indices') 

285 else: 

286 idx = _compatible_boolean_index(index) 

287 if idx is None: 

288 row, col = index, slice(None) 

289 elif idx.ndim < 2: 

290 return _boolean_index_to_array(idx), slice(None) 

291 elif idx.ndim == 2: 

292 return idx.nonzero() 

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

294 if issparse(row) or issparse(col): 

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

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

297 raise IndexError( 

298 'Indexing with sparse matrices is not supported ' 

299 'except boolean indexing where matrix and index ' 

300 'are equal shapes.') 

301 bool_row = _compatible_boolean_index(row) 

302 bool_col = _compatible_boolean_index(col) 

303 if bool_row is not None: 

304 row = _boolean_index_to_array(bool_row) 

305 if bool_col is not None: 

306 col = _boolean_index_to_array(bool_col) 

307 return row, col 

308 

309 

310def _check_ellipsis(index): 

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

312 if index is Ellipsis: 

313 return (slice(None), slice(None)) 

314 

315 if not isinstance(index, tuple): 

316 return index 

317 

318 # Find any Ellipsis objects. 

319 ellipsis_indices = [i for i, v in enumerate(index) if v is Ellipsis] 

320 if not ellipsis_indices: 

321 return index 

322 if len(ellipsis_indices) > 1: 

323 warn('multi-Ellipsis indexing is deprecated will be removed in v1.13.', 

324 DeprecationWarning, stacklevel=2) 

325 first_ellipsis = ellipsis_indices[0] 

326 

327 # Try to expand it using shortcuts for common cases 

328 if len(index) == 1: 

329 return (slice(None), slice(None)) 

330 if len(index) == 2: 

331 if first_ellipsis == 0: 

332 if index[1] is Ellipsis: 

333 return (slice(None), slice(None)) 

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

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

336 

337 # Expand it using a general-purpose algorithm 

338 tail = [] 

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

340 if v is not Ellipsis: 

341 tail.append(v) 

342 nd = first_ellipsis + len(tail) 

343 nslice = max(0, 2 - nd) 

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

345 

346 

347def _maybe_bool_ndarray(idx): 

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

349 """ 

350 idx = np.asanyarray(idx) 

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

352 return idx 

353 return None 

354 

355 

356def _first_element_bool(idx, max_dim=2): 

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

358 array type is boolean. 

359 """ 

360 if max_dim < 1: 

361 return None 

362 try: 

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

364 except TypeError: 

365 return None 

366 if isinstance(first, bool): 

367 return True 

368 return _first_element_bool(first, max_dim-1) 

369 

370 

371def _compatible_boolean_index(idx): 

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

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

374 """ 

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

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

377 return _maybe_bool_ndarray(idx) 

378 return None 

379 

380 

381def _boolean_index_to_array(idx): 

382 if idx.ndim > 1: 

383 raise IndexError('invalid index shape') 

384 return np.where(idx)[0]