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

1"""Indexing mixin for sparse array/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 from scipy.sparse import sparray 

38 

39 if isinstance(self, sparray): 

40 raise NotImplementedError( 

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

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

43 ) 

44 

45 def __getitem__(self, key): 

46 row, col = self._validate_indices(key) 

47 

48 # Dispatch to specialized methods. 

49 if isinstance(row, INT_TYPES): 

50 if isinstance(col, INT_TYPES): 

51 return self._get_intXint(row, col) 

52 elif isinstance(col, slice): 

53 self._raise_on_1d_array_slice() 

54 return self._get_intXslice(row, col) 

55 elif col.ndim == 1: 

56 self._raise_on_1d_array_slice() 

57 return self._get_intXarray(row, col) 

58 elif col.ndim == 2: 

59 return self._get_intXarray(row, col) 

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

61 elif isinstance(row, slice): 

62 if isinstance(col, INT_TYPES): 

63 self._raise_on_1d_array_slice() 

64 return self._get_sliceXint(row, col) 

65 elif isinstance(col, slice): 

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

67 return self.copy() 

68 return self._get_sliceXslice(row, col) 

69 elif col.ndim == 1: 

70 return self._get_sliceXarray(row, col) 

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

72 elif row.ndim == 1: 

73 if isinstance(col, INT_TYPES): 

74 self._raise_on_1d_array_slice() 

75 return self._get_arrayXint(row, col) 

76 elif isinstance(col, slice): 

77 return self._get_arrayXslice(row, col) 

78 else: # row.ndim == 2 

79 if isinstance(col, INT_TYPES): 

80 return self._get_arrayXint(row, col) 

81 elif isinstance(col, slice): 

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

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

84 # special case for outer indexing 

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

86 

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

88 row, col = _broadcast_arrays(row, col) 

89 if row.shape != col.shape: 

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

91 if row.size == 0: 

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

93 return self._get_arrayXarray(row, col) 

94 

95 def __setitem__(self, key, x): 

96 row, col = self._validate_indices(key) 

97 

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

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

100 if x.size != 1: 

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

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

103 return 

104 

105 if isinstance(row, slice): 

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

107 else: 

108 row = np.atleast_1d(row) 

109 

110 if isinstance(col, slice): 

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

112 if row.ndim == 1: 

113 row = row[:, None] 

114 else: 

115 col = np.atleast_1d(col) 

116 

117 i, j = _broadcast_arrays(row, col) 

118 if i.shape != j.shape: 

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

120 

121 from ._base import issparse 

122 if issparse(x): 

123 if i.ndim == 1: 

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

125 i = i[None] 

126 j = j[None] 

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

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

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

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

131 raise ValueError('shape mismatch in assignment') 

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

133 return 

134 x = x.tocoo(copy=True) 

135 x.sum_duplicates() 

136 self._set_arrayXarray_sparse(i, j, x) 

137 else: 

138 # Make x and i into the same shape 

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

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

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

142 if x.size == 0: 

143 return 

144 x = x.reshape(i.shape) 

145 self._set_arrayXarray(i, j, x) 

146 

147 def _validate_indices(self, key): 

148 M, N = self.shape 

149 row, col = _unpack_index(key) 

150 

151 if isintlike(row): 

152 row = int(row) 

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

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

155 if row < 0: 

156 row += M 

157 elif not isinstance(row, slice): 

158 row = self._asindices(row, M) 

159 

160 if isintlike(col): 

161 col = int(col) 

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

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

164 if col < 0: 

165 col += N 

166 elif not isinstance(col, slice): 

167 col = self._asindices(col, N) 

168 

169 return row, col 

170 

171 def _asindices(self, idx, length): 

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

173 

174 Subclasses that need special validation can override this method. 

175 """ 

176 try: 

177 x = np.asarray(idx) 

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

179 raise IndexError('invalid index') from e 

180 

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

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

183 

184 if x.size == 0: 

185 return x 

186 

187 # Check bounds 

188 max_indx = x.max() 

189 if max_indx >= length: 

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

191 

192 min_indx = x.min() 

193 if min_indx < 0: 

194 if min_indx < -length: 

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

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

197 x = x.copy() 

198 x[x < 0] += length 

199 return x 

200 

201 def _getrow(self, i): 

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

203 """ 

204 M, N = self.shape 

205 i = int(i) 

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

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

208 if i < 0: 

209 i += M 

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

211 

212 def _getcol(self, i): 

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

214 """ 

215 M, N = self.shape 

216 i = int(i) 

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

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

219 if i < 0: 

220 i += N 

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

222 

223 def _get_intXint(self, row, col): 

224 raise NotImplementedError() 

225 

226 def _get_intXarray(self, row, col): 

227 raise NotImplementedError() 

228 

229 def _get_intXslice(self, row, col): 

230 raise NotImplementedError() 

231 

232 def _get_sliceXint(self, row, col): 

233 raise NotImplementedError() 

234 

235 def _get_sliceXslice(self, row, col): 

236 raise NotImplementedError() 

237 

238 def _get_sliceXarray(self, row, col): 

239 raise NotImplementedError() 

240 

241 def _get_arrayXint(self, row, col): 

242 raise NotImplementedError() 

243 

244 def _get_arrayXslice(self, row, col): 

245 raise NotImplementedError() 

246 

247 def _get_columnXarray(self, row, col): 

248 raise NotImplementedError() 

249 

250 def _get_arrayXarray(self, row, col): 

251 raise NotImplementedError() 

252 

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

254 raise NotImplementedError() 

255 

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

257 raise NotImplementedError() 

258 

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

260 # Fall back to densifying x 

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

262 x, _ = _broadcast_arrays(x, row) 

263 self._set_arrayXarray(row, col, x) 

264 

265 

266def _unpack_index(index): 

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

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

269 """ 

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

271 from ._base import _spbase, issparse 

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

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

274 return index.nonzero() 

275 

276 # Parse any ellipses. 

277 index = _check_ellipsis(index) 

278 

279 # Next, parse the tuple or object 

280 if isinstance(index, tuple): 

281 if len(index) == 2: 

282 row, col = index 

283 elif len(index) == 1: 

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

285 else: 

286 raise IndexError('invalid number of indices') 

287 else: 

288 idx = _compatible_boolean_index(index) 

289 if idx is None: 

290 row, col = index, slice(None) 

291 elif idx.ndim < 2: 

292 return _boolean_index_to_array(idx), slice(None) 

293 elif idx.ndim == 2: 

294 return idx.nonzero() 

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

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

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

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

299 raise IndexError( 

300 'Indexing with sparse matrices is not supported ' 

301 'except boolean indexing where matrix and index ' 

302 'are equal shapes.') 

303 bool_row = _compatible_boolean_index(row) 

304 bool_col = _compatible_boolean_index(col) 

305 if bool_row is not None: 

306 row = _boolean_index_to_array(bool_row) 

307 if bool_col is not None: 

308 col = _boolean_index_to_array(bool_col) 

309 return row, col 

310 

311 

312def _check_ellipsis(index): 

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

314 if index is Ellipsis: 

315 return (slice(None), slice(None)) 

316 

317 if not isinstance(index, tuple): 

318 return index 

319 

320 # Find any Ellipsis objects. 

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

322 if not ellipsis_indices: 

323 return index 

324 if len(ellipsis_indices) > 1: 

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

326 DeprecationWarning, stacklevel=2) 

327 first_ellipsis = ellipsis_indices[0] 

328 

329 # Try to expand it using shortcuts for common cases 

330 if len(index) == 1: 

331 return (slice(None), slice(None)) 

332 if len(index) == 2: 

333 if first_ellipsis == 0: 

334 if index[1] is Ellipsis: 

335 return (slice(None), slice(None)) 

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

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

338 

339 # Expand it using a general-purpose algorithm 

340 tail = [] 

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

342 if v is not Ellipsis: 

343 tail.append(v) 

344 nd = first_ellipsis + len(tail) 

345 nslice = max(0, 2 - nd) 

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

347 

348 

349def _maybe_bool_ndarray(idx): 

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

351 """ 

352 idx = np.asanyarray(idx) 

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

354 return idx 

355 return None 

356 

357 

358def _first_element_bool(idx, max_dim=2): 

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

360 array type is boolean. 

361 """ 

362 if max_dim < 1: 

363 return None 

364 try: 

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

366 except TypeError: 

367 return None 

368 if isinstance(first, bool): 

369 return True 

370 return _first_element_bool(first, max_dim-1) 

371 

372 

373def _compatible_boolean_index(idx): 

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

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

376 """ 

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

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

379 return _maybe_bool_ndarray(idx) 

380 return None 

381 

382 

383def _boolean_index_to_array(idx): 

384 if idx.ndim > 1: 

385 raise IndexError('invalid index shape') 

386 return np.where(idx)[0]