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

1"""Indexing mixin for sparse array/matrix classes. 

2""" 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7import numpy as np 

8from ._sputils import isintlike 

9 

10if TYPE_CHECKING: 

11 import numpy.typing as npt 

12 

13INT_TYPES = (int, np.integer) 

14 

15 

16def _broadcast_arrays(a, b): 

17 """ 

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

19 

20 NumPy >= 1.17.0 transitions broadcast_arrays to return 

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

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

23 the old behavior. 

24 """ 

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

26 x.flags.writeable = a.flags.writeable 

27 y.flags.writeable = b.flags.writeable 

28 return x, y 

29 

30 

31class IndexMixin: 

32 """ 

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

34 """ 

35 def _raise_on_1d_array_slice(self): 

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

37 

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

39 result, raising an error instead. 

40 

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

42 """ 

43 from scipy.sparse import sparray 

44 

45 if isinstance(self, sparray): 

46 raise NotImplementedError( 

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

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

49 ) 

50 

51 def __getitem__(self, key): 

52 row, col = self._validate_indices(key) 

53 

54 # Dispatch to specialized methods. 

55 if isinstance(row, INT_TYPES): 

56 if isinstance(col, INT_TYPES): 

57 return self._get_intXint(row, col) 

58 elif isinstance(col, slice): 

59 self._raise_on_1d_array_slice() 

60 return self._get_intXslice(row, col) 

61 elif col.ndim == 1: 

62 self._raise_on_1d_array_slice() 

63 return self._get_intXarray(row, col) 

64 elif col.ndim == 2: 

65 return self._get_intXarray(row, col) 

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

67 elif isinstance(row, slice): 

68 if isinstance(col, INT_TYPES): 

69 self._raise_on_1d_array_slice() 

70 return self._get_sliceXint(row, col) 

71 elif isinstance(col, slice): 

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

73 return self.copy() 

74 return self._get_sliceXslice(row, col) 

75 elif col.ndim == 1: 

76 return self._get_sliceXarray(row, col) 

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

78 elif row.ndim == 1: 

79 if isinstance(col, INT_TYPES): 

80 self._raise_on_1d_array_slice() 

81 return self._get_arrayXint(row, col) 

82 elif isinstance(col, slice): 

83 return self._get_arrayXslice(row, col) 

84 else: # row.ndim == 2 

85 if isinstance(col, INT_TYPES): 

86 return self._get_arrayXint(row, col) 

87 elif isinstance(col, slice): 

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

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

90 # special case for outer indexing 

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

92 

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

94 row, col = _broadcast_arrays(row, col) 

95 if row.shape != col.shape: 

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

97 if row.size == 0: 

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

99 return self._get_arrayXarray(row, col) 

100 

101 def __setitem__(self, key, x): 

102 row, col = self._validate_indices(key) 

103 

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

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

106 if x.size != 1: 

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

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

109 return 

110 

111 if isinstance(row, slice): 

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

113 else: 

114 row = np.atleast_1d(row) 

115 

116 if isinstance(col, slice): 

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

118 if row.ndim == 1: 

119 row = row[:, None] 

120 else: 

121 col = np.atleast_1d(col) 

122 

123 i, j = _broadcast_arrays(row, col) 

124 if i.shape != j.shape: 

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

126 

127 from ._base import issparse 

128 if issparse(x): 

129 if i.ndim == 1: 

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

131 i = i[None] 

132 j = j[None] 

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

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

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

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

137 raise ValueError('shape mismatch in assignment') 

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

139 return 

140 x = x.tocoo(copy=True) 

141 x.sum_duplicates() 

142 self._set_arrayXarray_sparse(i, j, x) 

143 else: 

144 # Make x and i into the same shape 

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

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

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

148 if x.size == 0: 

149 return 

150 x = x.reshape(i.shape) 

151 self._set_arrayXarray(i, j, x) 

152 

153 def _validate_indices(self, key): 

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

155 from ._base import _spbase 

156 if (isinstance(key, (_spbase, np.ndarray)) and 

157 key.ndim == 2 and key.dtype.kind == 'b'): 

158 if key.shape != self.shape: 

159 raise IndexError('boolean index shape does not match array shape') 

160 row, col = key.nonzero() 

161 else: 

162 row, col = _unpack_index(key) 

163 M, N = self.shape 

164 

165 def _validate_bool_idx( 

166 idx: npt.NDArray[np.bool_], 

167 axis_size: int, 

168 axis_name: str 

169 ) -> npt.NDArray[np.int_]: 

170 if len(idx) != axis_size: 

171 raise IndexError( 

172 f"boolean {axis_name} index has incorrect length: {len(idx)} " 

173 f"instead of {axis_size}" 

174 ) 

175 return _boolean_index_to_array(idx) 

176 

177 if isintlike(row): 

178 row = int(row) 

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

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

181 if row < 0: 

182 row += M 

183 elif (bool_row := _compatible_boolean_index(row)) is not None: 

184 row = _validate_bool_idx(bool_row, M, "row") 

185 elif not isinstance(row, slice): 

186 row = self._asindices(row, M) 

187 

188 if isintlike(col): 

189 col = int(col) 

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

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

192 if col < 0: 

193 col += N 

194 elif (bool_col := _compatible_boolean_index(col)) is not None: 

195 col = _validate_bool_idx(bool_col, N, "column") 

196 elif not isinstance(col, slice): 

197 col = self._asindices(col, N) 

198 

199 return row, col 

200 

201 def _asindices(self, idx, length): 

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

203 

204 Subclasses that need special validation can override this method. 

205 """ 

206 try: 

207 x = np.asarray(idx) 

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

209 raise IndexError('invalid index') from e 

210 

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

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

213 

214 if x.size == 0: 

215 return x 

216 

217 # Check bounds 

218 max_indx = x.max() 

219 if max_indx >= length: 

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

221 

222 min_indx = x.min() 

223 if min_indx < 0: 

224 if min_indx < -length: 

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

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

227 x = x.copy() 

228 x[x < 0] += length 

229 return x 

230 

231 def _getrow(self, i): 

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

233 """ 

234 M, N = self.shape 

235 i = int(i) 

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

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

238 if i < 0: 

239 i += M 

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

241 

242 def _getcol(self, i): 

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

244 """ 

245 M, N = self.shape 

246 i = int(i) 

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

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

249 if i < 0: 

250 i += N 

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

252 

253 def _get_intXint(self, row, col): 

254 raise NotImplementedError() 

255 

256 def _get_intXarray(self, row, col): 

257 raise NotImplementedError() 

258 

259 def _get_intXslice(self, row, col): 

260 raise NotImplementedError() 

261 

262 def _get_sliceXint(self, row, col): 

263 raise NotImplementedError() 

264 

265 def _get_sliceXslice(self, row, col): 

266 raise NotImplementedError() 

267 

268 def _get_sliceXarray(self, row, col): 

269 raise NotImplementedError() 

270 

271 def _get_arrayXint(self, row, col): 

272 raise NotImplementedError() 

273 

274 def _get_arrayXslice(self, row, col): 

275 raise NotImplementedError() 

276 

277 def _get_columnXarray(self, row, col): 

278 raise NotImplementedError() 

279 

280 def _get_arrayXarray(self, row, col): 

281 raise NotImplementedError() 

282 

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

284 raise NotImplementedError() 

285 

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

287 raise NotImplementedError() 

288 

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

290 # Fall back to densifying x 

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

292 x, _ = _broadcast_arrays(x, row) 

293 self._set_arrayXarray(row, col, x) 

294 

295 

296def _unpack_index(index) -> tuple[ 

297 int | slice | npt.NDArray[np.bool_ | np.int_], 

298 int | slice | npt.NDArray[np.bool_ | np.int_] 

299]: 

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

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

302 """ 

303 # Parse any ellipses. 

304 index = _check_ellipsis(index) 

305 

306 # Next, parse the tuple or object 

307 if isinstance(index, tuple): 

308 if len(index) == 2: 

309 row, col = index 

310 elif len(index) == 1: 

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

312 else: 

313 raise IndexError('invalid number of indices') 

314 else: 

315 idx = _compatible_boolean_index(index) 

316 if idx is None: 

317 row, col = index, slice(None) 

318 elif idx.ndim < 2: 

319 return idx, slice(None) 

320 elif idx.ndim == 2: 

321 return idx.nonzero() 

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

323 from ._base import issparse 

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

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

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

327 raise IndexError( 

328 'Indexing with sparse matrices is not supported ' 

329 'except boolean indexing where matrix and index ' 

330 'are equal shapes.') 

331 return row, col 

332 

333 

334def _check_ellipsis(index): 

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

336 if index is Ellipsis: 

337 return (slice(None), slice(None)) 

338 

339 if not isinstance(index, tuple): 

340 return index 

341 

342 # Find any Ellipsis objects. 

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

344 if not ellipsis_indices: 

345 return index 

346 if len(ellipsis_indices) > 1: 

347 raise IndexError("an index can only have a single ellipsis ('...')") 

348 

349 # Replace the Ellipsis object with 0, 1, or 2 null-slices as needed. 

350 i, = ellipsis_indices 

351 num_slices = max(0, 3 - len(index)) 

352 return index[:i] + (slice(None),) * num_slices + index[i + 1:] 

353 

354 

355def _maybe_bool_ndarray(idx): 

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

357 """ 

358 idx = np.asanyarray(idx) 

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

360 return idx 

361 return None 

362 

363 

364def _first_element_bool(idx, max_dim=2): 

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

366 array type is boolean. 

367 """ 

368 if max_dim < 1: 

369 return None 

370 try: 

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

372 except TypeError: 

373 return None 

374 if isinstance(first, bool): 

375 return True 

376 return _first_element_bool(first, max_dim-1) 

377 

378 

379def _compatible_boolean_index(idx): 

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

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

382 """ 

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

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

385 return _maybe_bool_ndarray(idx) 

386 return None 

387 

388 

389def _boolean_index_to_array(idx): 

390 if idx.ndim > 1: 

391 raise IndexError('invalid index shape') 

392 return np.where(idx)[0]