Repository: scipy/scipy
Pull request: BUG: sparse: fix ellipsis and bool array indexing combination
Fixes #21747 

I added tests for the cases with ellipsis and boolean arrays combined in an index. Line 2819 in `tests/test_base.py` tests the issue here and failed until the fix.

`A[..., [True, False, True]]` leads to cyclic logic: the ellipsis cant figure out how big to expand without knowing the size of the bool array, and the bool array can't check its shape until it knows which axis it applies to (which requires knowing how much ellipsis is expanded). I have switched to a two pass iteration over the indexing `key`. The first pass processes any ellipsis and preprocesses any arrays in the index `key`. Then after ellipsis has been expanded, and arrays have been converted to `numpy.array`, we can go through the more involved pass over the index `key`. That involves checking the indices range and building a canonical index as well as finding the desired shape to return to `__getitem__`.

We were going to need a two pass system for nD indexing eventually anyway. I guess we needed it already. I also noticed that we will need to update the code for nD indexing in order to handle the distinction between contiguous fancy index dimensions `[:, arr1, arr2]` and separated fancy index dimensions.  `[arr1, :, arr2]` so I added a "Todo" comment.

Problem statement:
BUG: boolean indexing of sparse arrays broken on `main`
### Describe your issue.

Boolean indexing of sparse arrays is broken if an Ellipsis is provided.

### Reproducing Code Example

```python
import numpy as np

dense = np.zeros((1, 3), dtype=bool)
dense[0, 0] = True
array = csc_array(dense)
mask = np.array([True, True, False])
array[:, mask] # works!
array[..., mask] # error below
```


### Error message

```shell
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/deepak/miniforge3/envs/xarray-upstream/lib/python3.12/site-packages/scipy/sparse/_index.py", line 30, in __getitem__
    index, new_shape = self._validate_indices(key)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/deepak/miniforge3/envs/xarray-upstream/lib/python3.12/site-packages/scipy/sparse/_index.py", line 237, in _validate_indices
    raise IndexError(
IndexError: bool index 1 has shape (1,) instead of (3,)
```


^ Note that it's trying to index along axis 0 instead of axis 1 as requested.


### SciPy/NumPy/Python version and system information

```shell
1.15.0.dev0+git20241023.59fd061 2.2.0.dev0+git20241022.020e78c sys.version_info(major=3, minor=12, micro=2, releaselevel='final', serial=0)
Build Dependencies:
  blas:
    detection method: extraframeworks
    found: true
    include directory: unknown
    lib directory: unknown
    name: Accelerate
    openblas configuration: unknown
    pc file directory: unknown
    version: unknown
  lapack:
    detection method: extraframeworks
    found: true
    include directory: unknown
    lib directory: unknown
    name: Accelerate
    openblas configuration: unknown
    pc file directory: unknown
    version: unknown
  pybind11:
    detection method: config-tool
    include directory: unknown
    name: pybind11
    version: 2.13.6
Compilers:
  c:
    commands: cc
    linker: ld64
    name: clang
    version: 15.0.0
  c++:
    commands: c++
    linker: ld64
    name: clang
    version: 15.0.0
  cython:
    commands: cython
    linker: cython
    name: cython
    version: 3.0.11
  fortran:
    commands: gfortran
    linker: ld64
    name: gcc
    version: 13.3.0
  pythran:
    include directory: ../../../../../../private/var/folders/q5/6nls1_g145s2fl1lbvx2zby40000gn/T/pip-build-env-jfqgu1eg/overlay/lib/python3.12/site-packages/pythran
    version: 0.16.1
Machine Information:
  build:
    cpu: aarch64
    endian: little
    family: aarch64
    system: darwin
  cross-compiled: false
  host:
    cpu: aarch64
    endian: little
    family: aarch64
    system: darwin
Python Information:
  path: /private/var/folders/q5/6nls1_g145s2fl1lbvx2zby40000gn/T/cibw-run-8jvo626d/cp312-macosx_arm64/build/venv/bin/python
  version: '3.12'
```

Hints from issue discussion:
This might be related to gh-21615, though I haven't thought too deeply about it. @dschult FYI
Yes -- I can verify that this is a problem, and one related to gh-21615, so I initially thought it would have been fixed by [gh-21616](https://github.com/scipy/scipy/pull/21616) but it isn't fixed. I will take a look.

Thanks @dcherian !!

Commit message: fix indexing with two pass check for ellipsis and preprocess array indexing

Patch to review:
diff --git a/scipy/sparse/_index.py b/scipy/sparse/_index.py
index ad488f5de851..97373ecdf627 100644
--- a/scipy/sparse/_index.py
+++ b/scipy/sparse/_index.py
@@ -206,31 +206,60 @@ def _validate_indices(self, key):
             key = [key]
 
         ellps_pos = None
-        idx_shape = []
-        index = []
-        index_ndim = 0
-        array_indices = []
+        index_1st = []
+        prelim_ndim = 0
         for i, idx in enumerate(key):
             if idx is Ellipsis:
                 if ellps_pos is not None:
                     raise IndexError('an index can only have a single ellipsis')
                 ellps_pos = i
             elif idx is None:
+                index_1st.append(idx)
+            elif isinstance(idx, slice) or isintlike(idx):
+                index_1st.append(idx)
+                prelim_ndim += 1
+            elif (ix := _compatible_boolean_index(idx, self.ndim)) is not None:
+                index_1st.append(ix)
+                prelim_ndim += ix.ndim
+            elif issparse(idx):
+                # TODO: make sparse matrix indexing work for sparray
+                raise IndexError(
+                    'Indexing with sparse matrices is not supported '
+                    'except boolean indexing where matrix and index '
+                    'are equal shapes.')
+            else:  # dense array
+                index_1st.append(np.asarray(idx))
+                prelim_ndim += 1
+        ellip_slices = (self.ndim - prelim_ndim) * [slice(None)]
+        if ellip_slices:
+            if ellps_pos is None:
+                ellps_pos = prelim_ndim
+            index_1st = index_1st[:ellps_pos] + ellip_slices + index_1st[ellps_pos:]
+
+        # second pass (have processed ellipsis and preprocessed arrays)
+        idx_shape = []
+        index_ndim = 0
+        index = []
+        array_indices = []
+        for i, idx in enumerate(index_1st):
+            if idx is None:
                 idx_shape.append(1)
             elif isinstance(idx, slice):
                 index.append(idx)
-                Ms = self._shape[index_ndim] if ellps_pos is None else self._shape[-1]
+                Ms = self._shape[index_ndim]
                 len_slice = len(range(*idx.indices(Ms)))
                 idx_shape.append(len_slice)
                 index_ndim += 1
             elif isintlike(idx):
-                N = self._shape[index_ndim] if ellps_pos is None else self._shape[-1]
+                N = self._shape[index_ndim]
                 if not (-N <= idx < N):
                     raise IndexError(f'index ({idx}) out of range')
                 idx = int(idx + N if idx < 0 else idx)
                 index.append(idx)
                 index_ndim += 1
-            elif (ix := _compatible_boolean_index(idx, self.ndim)) is not None:
+            # bool array (checked in first pass)
+            elif idx.dtype.kind == 'b':
+                ix = idx
                 tmp_ndim = index_ndim + ix.ndim
                 mid_shape = self._shape[index_ndim:tmp_ndim]
                 if ix.shape != mid_shape:
@@ -240,15 +269,8 @@ def _validate_indices(self, key):
                 index.extend(ix.nonzero())
                 array_indices.extend(range(index_ndim, tmp_ndim))
                 index_ndim = tmp_ndim
-            elif issparse(idx):
-                # TODO: make sparse matrix indexing work for sparray
-                raise IndexError(
-                    'Indexing with sparse matrices is not supported '
-                    'except boolean indexing where matrix and index '
-                    'are equal shapes.')
             else:  # dense array
                 N = self._shape[index_ndim]
-                idx = np.array(idx)
                 idx = self._asindices(idx, N)
                 index.append(idx)
                 array_indices.append(index_ndim)
@@ -260,22 +282,16 @@ def _validate_indices(self, key):
         if len(array_indices) > 1:
             idx_arrays = _broadcast_arrays(*(index[i] for i in array_indices))
             if any(idx_arrays[0].shape != ix.shape for ix in idx_arrays[1:]):
-                raise IndexError('array indices after broadcast differ in shape')
+                shapes = [str(ix.shape) for ix in idx_arrays]
+                msg = (f'shape mismatch: indexing arrays could not be broadcast '
+                       f'together with shapes {" ".join(shapes)}')
+                raise IndexError(msg)
+            # Todo: for nD (adjacenct arrays stay, separated move to start)
             idx_shape = list(idx_arrays[0].shape) + idx_shape
         elif len(array_indices) == 1:
             arr_index = array_indices[0]
             arr_shape = list(index[arr_index].shape)
             idx_shape = idx_shape[:arr_index] + arr_shape + idx_shape[arr_index:]
-
-        # add slice(None) (which is colon) to fill out full index
-        nslice = self.ndim - index_ndim
-        if nslice > 0:
-            if ellps_pos is None:
-                ellps_pos = index_ndim
-            index = index[:ellps_pos] + [slice(None)] * nslice + index[ellps_pos:]
-            mid_shape = list(self.shape[ellps_pos : ellps_pos + nslice])
-            idx_shape = idx_shape[:ellps_pos] + mid_shape + idx_shape[ellps_pos:]
-
         if (ndim := len(idx_shape)) > 2:
             raise IndexError(f'Only 1D or 2D arrays allowed. Index makes {ndim}D')
         return tuple(index), tuple(idx_shape)
diff --git a/scipy/sparse/tests/test_base.py b/scipy/sparse/tests/test_base.py
index a4c3e35e5659..b049c7b4388d 100644
--- a/scipy/sparse/tests/test_base.py
+++ b/scipy/sparse/tests/test_base.py
@@ -2803,6 +2803,28 @@ def test_ellipsis_slicing(self):
         assert_equal(a[1, 1, ...], b[1, 1, ...])
         assert_equal(a[1, ..., 1], b[1, ..., 1])
 
+    def test_ellipsis_fancy_bool(self):
+        numpy_a = self.asdense(arange(50).reshape(5, 10))
+        a = self.spcreator(numpy_a)
+
+        ix5 = [True, False, True, False, True]
+        ix10 = [False] * 5 + ix5  # same number of True values as ix5
+        ix10_6True = ix5 + ix5  # not same number of True values as ix5
+        full_ix = [ix10] * 5
+
+        assert_array_equal(toarray(a[full_ix, ...]), numpy_a[full_ix, ...])
+        assert_array_equal(toarray(a[..., full_ix]), numpy_a[..., full_ix])
+
+        assert_array_equal(toarray(a[ix5, ...]), numpy_a[ix5, ...])
+        assert_array_equal(toarray(a[..., ix10]), numpy_a[..., ix10])
+
+        assert_array_equal(toarray(a[ix5, ..., ix10]), numpy_a[ix5, ..., ix10])
+        assert_array_equal(toarray(a[..., ix5, ix10]), numpy_a[..., ix5, ix10])
+        assert_array_equal(toarray(a[ix5, ix10, ...]), numpy_a[ix5, ix10, ...])
+
+        with assert_raises(ValueError, match="shape mismatch"):
+            a[ix5, ix10_6True]
+
     def test_ellipsis_fancy_slicing(self):
         b = self.asdense(arange(50).reshape(5, 10))
         a = self.spcreator(b)


Code review comments:
diff_hunk: @@ -260,22 +282,16 @@ def _validate_indices(self, key):
         if len(array_indices) > 1:
             idx_arrays = _broadcast_arrays(*(index[i] for i in array_indices))
             if any(idx_arrays[0].shape != ix.shape for ix in idx_arrays[1:]):
-                raise IndexError('array indices after broadcast differ in shape')
+                shapes = [str(ix.shape) for ix in idx_arrays]
+                msg = (f'shape mismatch: indexing arrays could not be broadcast '
+                       f'together with shapes {" ".join(shapes)}')
path: scipy/sparse/_index.py
text: @user1:
```suggestion
                shapes = " ".join(str(ix.shape) for ix in idx_arrays)
                msg = (f'shape mismatch: indexing arrays could not be broadcast '
                       f'together with shapes {shapes}')
```

diff_hunk: @@ -260,22 +282,16 @@ def _validate_indices(self, key):
         if len(array_indices) > 1:
             idx_arrays = _broadcast_arrays(*(index[i] for i in array_indices))
             if any(idx_arrays[0].shape != ix.shape for ix in idx_arrays[1:]):
-                raise IndexError('array indices after broadcast differ in shape')
+                shapes = [str(ix.shape) for ix in idx_arrays]
+                msg = (f'shape mismatch: indexing arrays could not be broadcast '
+                       f'together with shapes {" ".join(shapes)}')
+                raise IndexError(msg)
+            # Todo: for nD (adjacenct arrays stay, separated move to start)
path: scipy/sparse/_index.py
text: @user1:
```suggestion
            # TODO: handle this for nD (adjacent arrays stay, separated move to start)
```