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""" 

3import numpy as np 

4from ._sputils import isintlike 

5 

6INT_TYPES = (int, np.integer) 

7 

8 

9def _broadcast_arrays(a, b): 

10 """ 

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

12 

13 NumPy >= 1.17.0 transitions broadcast_arrays to return 

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

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

16 the old behavior. 

17 """ 

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

19 x.flags.writeable = a.flags.writeable 

20 y.flags.writeable = b.flags.writeable 

21 return x, y 

22 

23 

24class IndexMixin: 

25 """ 

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

27 """ 

28 def _raise_on_1d_array_slice(self): 

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

30 

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

32 result, raising an error instead. 

33 

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

35 """ 

36 from scipy.sparse import sparray 

37 

38 if isinstance(self, sparray): 

39 raise NotImplementedError( 

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

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

42 ) 

43 

44 def __getitem__(self, key): 

45 row, col = self._validate_indices(key) 

46 

47 # Dispatch to specialized methods. 

48 if isinstance(row, INT_TYPES): 

49 if isinstance(col, INT_TYPES): 

50 return self._get_intXint(row, col) 

51 elif isinstance(col, slice): 

52 self._raise_on_1d_array_slice() 

53 return self._get_intXslice(row, col) 

54 elif col.ndim == 1: 

55 self._raise_on_1d_array_slice() 

56 return self._get_intXarray(row, col) 

57 elif col.ndim == 2: 

58 return self._get_intXarray(row, col) 

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

60 elif isinstance(row, slice): 

61 if isinstance(col, INT_TYPES): 

62 self._raise_on_1d_array_slice() 

63 return self._get_sliceXint(row, col) 

64 elif isinstance(col, slice): 

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

66 return self.copy() 

67 return self._get_sliceXslice(row, col) 

68 elif col.ndim == 1: 

69 return self._get_sliceXarray(row, col) 

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

71 elif row.ndim == 1: 

72 if isinstance(col, INT_TYPES): 

73 self._raise_on_1d_array_slice() 

74 return self._get_arrayXint(row, col) 

75 elif isinstance(col, slice): 

76 return self._get_arrayXslice(row, col) 

77 else: # row.ndim == 2 

78 if isinstance(col, INT_TYPES): 

79 return self._get_arrayXint(row, col) 

80 elif isinstance(col, slice): 

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

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

83 # special case for outer indexing 

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

85 

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

87 row, col = _broadcast_arrays(row, col) 

88 if row.shape != col.shape: 

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

90 if row.size == 0: 

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

92 return self._get_arrayXarray(row, col) 

93 

94 def __setitem__(self, key, x): 

95 row, col = self._validate_indices(key) 

96 

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

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

99 if x.size != 1: 

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

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

102 return 

103 

104 if isinstance(row, slice): 

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

106 else: 

107 row = np.atleast_1d(row) 

108 

109 if isinstance(col, slice): 

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

111 if row.ndim == 1: 

112 row = row[:, None] 

113 else: 

114 col = np.atleast_1d(col) 

115 

116 i, j = _broadcast_arrays(row, col) 

117 if i.shape != j.shape: 

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

119 

120 from ._base import issparse 

121 if issparse(x): 

122 if i.ndim == 1: 

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

124 i = i[None] 

125 j = j[None] 

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

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

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

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

130 raise ValueError('shape mismatch in assignment') 

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

132 return 

133 x = x.tocoo(copy=True) 

134 x.sum_duplicates() 

135 self._set_arrayXarray_sparse(i, j, x) 

136 else: 

137 # Make x and i into the same shape 

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

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

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

141 if x.size == 0: 

142 return 

143 x = x.reshape(i.shape) 

144 self._set_arrayXarray(i, j, x) 

145 

146 def _validate_indices(self, key): 

147 M, N = self.shape 

148 row, col = _unpack_index(key) 

149 

150 if isintlike(row): 

151 row = int(row) 

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

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

154 if row < 0: 

155 row += M 

156 elif not isinstance(row, slice): 

157 row = self._asindices(row, M) 

158 

159 if isintlike(col): 

160 col = int(col) 

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

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

163 if col < 0: 

164 col += N 

165 elif not isinstance(col, slice): 

166 col = self._asindices(col, N) 

167 

168 return row, col 

169 

170 def _asindices(self, idx, length): 

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

172 

173 Subclasses that need special validation can override this method. 

174 """ 

175 try: 

176 x = np.asarray(idx) 

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

178 raise IndexError('invalid index') from e 

179 

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

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

182 

183 if x.size == 0: 

184 return x 

185 

186 # Check bounds 

187 max_indx = x.max() 

188 if max_indx >= length: 

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

190 

191 min_indx = x.min() 

192 if min_indx < 0: 

193 if min_indx < -length: 

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

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

196 x = x.copy() 

197 x[x < 0] += length 

198 return x 

199 

200 def _getrow(self, i): 

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

202 """ 

203 M, N = self.shape 

204 i = int(i) 

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

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

207 if i < 0: 

208 i += M 

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

210 

211 def _getcol(self, i): 

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

213 """ 

214 M, N = self.shape 

215 i = int(i) 

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

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

218 if i < 0: 

219 i += N 

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

221 

222 def _get_intXint(self, row, col): 

223 raise NotImplementedError() 

224 

225 def _get_intXarray(self, row, col): 

226 raise NotImplementedError() 

227 

228 def _get_intXslice(self, row, col): 

229 raise NotImplementedError() 

230 

231 def _get_sliceXint(self, row, col): 

232 raise NotImplementedError() 

233 

234 def _get_sliceXslice(self, row, col): 

235 raise NotImplementedError() 

236 

237 def _get_sliceXarray(self, row, col): 

238 raise NotImplementedError() 

239 

240 def _get_arrayXint(self, row, col): 

241 raise NotImplementedError() 

242 

243 def _get_arrayXslice(self, row, col): 

244 raise NotImplementedError() 

245 

246 def _get_columnXarray(self, row, col): 

247 raise NotImplementedError() 

248 

249 def _get_arrayXarray(self, row, col): 

250 raise NotImplementedError() 

251 

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

253 raise NotImplementedError() 

254 

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

256 raise NotImplementedError() 

257 

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

259 # Fall back to densifying x 

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

261 x, _ = _broadcast_arrays(x, row) 

262 self._set_arrayXarray(row, col, x) 

263 

264 

265def _unpack_index(index): 

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

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

268 """ 

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

270 from ._base import _spbase, issparse 

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

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

273 return index.nonzero() 

274 

275 # Parse any ellipses. 

276 index = _check_ellipsis(index) 

277 

278 # Next, parse the tuple or object 

279 if isinstance(index, tuple): 

280 if len(index) == 2: 

281 row, col = index 

282 elif len(index) == 1: 

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

284 else: 

285 raise IndexError('invalid number of indices') 

286 else: 

287 idx = _compatible_boolean_index(index) 

288 if idx is None: 

289 row, col = index, slice(None) 

290 elif idx.ndim < 2: 

291 return _boolean_index_to_array(idx), slice(None) 

292 elif idx.ndim == 2: 

293 return idx.nonzero() 

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

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

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

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

298 raise IndexError( 

299 'Indexing with sparse matrices is not supported ' 

300 'except boolean indexing where matrix and index ' 

301 'are equal shapes.') 

302 bool_row = _compatible_boolean_index(row) 

303 bool_col = _compatible_boolean_index(col) 

304 if bool_row is not None: 

305 row = _boolean_index_to_array(bool_row) 

306 if bool_col is not None: 

307 col = _boolean_index_to_array(bool_col) 

308 return row, col 

309 

310 

311def _check_ellipsis(index): 

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

313 if index is Ellipsis: 

314 return (slice(None), slice(None)) 

315 

316 if not isinstance(index, tuple): 

317 return index 

318 

319 # Find any Ellipsis objects. 

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

321 if not ellipsis_indices: 

322 return index 

323 if len(ellipsis_indices) > 1: 

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

325 

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

327 i, = ellipsis_indices 

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

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

330 

331 

332def _maybe_bool_ndarray(idx): 

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

334 """ 

335 idx = np.asanyarray(idx) 

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

337 return idx 

338 return None 

339 

340 

341def _first_element_bool(idx, max_dim=2): 

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

343 array type is boolean. 

344 """ 

345 if max_dim < 1: 

346 return None 

347 try: 

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

349 except TypeError: 

350 return None 

351 if isinstance(first, bool): 

352 return True 

353 return _first_element_bool(first, max_dim-1) 

354 

355 

356def _compatible_boolean_index(idx): 

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

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

359 """ 

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

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

362 return _maybe_bool_ndarray(idx) 

363 return None 

364 

365 

366def _boolean_index_to_array(idx): 

367 if idx.ndim > 1: 

368 raise IndexError('invalid index shape') 

369 return np.where(idx)[0]