Repository: scipy/scipy
Pull request: BUG: sparse.csgraph: Support int64 indices in traversal.pyx
#### Reference issue
Fixes gh-18716.

#### What does this implement/fix?
I added fused type signatures for breadth-first + depth-first traversals, along with connected components. This allows the existing API to support sparse arrays/matrices with int64 index dtypes. After gh-18509, this situation is more likely to occur in user code.

#### Additional information
You may notice that the body of each modified function is unchanged by this PR. This means that matrices large enough to actually need int64 index arrays will still not work properly, but those were unsupported before anyway.

The introduction of fused dtypes increases the shared library size (not surprisingly). On my machine, the `_traversal.cpython-311-x86_64-linux-gnu.so` is currently 554K on disk, and with this PR it grows to 2.3M. A nontrivial portion of this comes from the way we treat the undirected code paths, as the `.transpose().tocsr()` operation doesn't guarantee that the index dtype of the result will be the same as the original matrix.

*Cython note*: I needed to add extra `def` functions for each of the `cdef` functions that I modified, otherwise the Cython compiler would complain to me about "no suitable method found" and "Invalid use of fused types, type cannot be specialized". I don't fully understand why adding the indirection solves the problem, but it seems okay for now.

Problem statement:
Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'
I am trying to get connected components in an adjacency matrix which is quite big but sparse. However when I run the code, I get this error: 
ValueError: Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'
Exception ignored in: 'scipy.sparse.csgraph._traversal._connected_components_undirected'
ValueError: Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'

When I run the code for a smaller matrix everything works fine.


### Reproducing Code Example

```python
from scipy import sparse
adj_matrix = sparse.load_npz(adj_matrix_address).astype(int)

# adj_matrix  only has zeros and ones
# <3556401x3556401 sparse matrix of type '<class 'numpy.int64'>'
#	with 6917339542 stored elements in Compressed Sparse Row format>

n_cluster, cluster_labels = sparse.csgraph.connected_components(csgraph=adj_matrix, directed=False, return_labels=True)
```


### Error message

```shell
ValueError: Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'
Exception ignored in: 'scipy.sparse.csgraph._traversal._connected_components_undirected'
ValueError: Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'
```


### SciPy/NumPy/Python version and system information

```shell
lapack_mkl_info:
  NOT AVAILABLE
openblas_lapack_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_mkl_info:
  NOT AVAILABLE
blis_info:
  NOT AVAILABLE
openblas_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    runtime_library_dirs = ['/usr/local/lib']
```

Hints from issue discussion:
Could you upload your matrix as a .npz file? Otherwise this will not be possible to debug.
I don't think uploading the matrix will be necessary here, actually. Instead, could you print `adj_matrix.indices.dtype` and also `adj_matrix.indptr.dtype`?
I suspect this is a fallout of #18509 (or something related); a simple repro is
```
import numpy as np; from scipy import sparse as sp
print(sp.csgraph.connected_components(sp.coo_array((np.array([1]), np.array([[0], [1]])), (2, 2))))
```
which prints `(1, array([0, 0], dtype=int32))` with scipy 1.10.1 and errors with
```
ValueError: Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'
Exception ignored in: 'scipy.sparse.csgraph._traversal._connected_components_undirected'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
ValueError: Buffer dtype mismatch, expected 'ITYPE_t' but got 'long'
(0, array([-9999, -9999], dtype=int32))
```
(as above) with scipy 1.11.1.
In particular, note the "Exception ignored", which is pretty awkward: even though the exception gets printed to stdout, the error is then discarded and an incorrect value is actually returned.
Thanks for the reproducer, @anntzer ! This definitely looks like fallout from #18509, similar to the recently-fixed #18603.

Looking at the Cython code in question, we may be able to fix this by using a fused index type in _traversal.pyx, as in https://github.com/scipy/scipy/blob/ddee711eb0ab3f7985f2af738eab230898702e2f/scipy/sparse/csgraph/_traversal.pyx#L647-L650

There's a precedent for this in other csgraph Cython code: https://github.com/scipy/scipy/blob/3c89445b6439f3ce7bffc4cf11c6407c39faedc5/scipy/sparse/csgraph/_reordering.pyx#L102-L103

The other option is to do what we did in #18644 and downcast the index arrays before calling the Cython code, but I'd prefer to try the fused type first.

Commit message: Fix tests on 32bit arch

Patch to review:
diff --git a/scipy/sparse/csgraph/_traversal.pyx b/scipy/sparse/csgraph/_traversal.pyx
index 06dc2e25f0f7..2c7ea86dd1ba 100644
--- a/scipy/sparse/csgraph/_traversal.pyx
+++ b/scipy/sparse/csgraph/_traversal.pyx
@@ -352,12 +352,22 @@ cpdef breadth_first_order(csgraph, i_start,
         return node_list[:length]
 
 
-cdef unsigned int _breadth_first_directed(
-                           unsigned int head_node,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors) noexcept:
+def _breadth_first_directed(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors):
+    return _breadth_first_directed2(head_node, indices, indptr,
+                                    node_list, predecessors)
+
+
+cdef unsigned int _breadth_first_directed2(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors) noexcept:
     # Inputs:
     #  head_node: (input) index of the node from which traversal starts
     #  indices: (input) CSR indices of graph
@@ -390,15 +400,25 @@ cdef unsigned int _breadth_first_directed(
 
     return i_nl
 
-
-cdef unsigned int _breadth_first_undirected(
-                           unsigned int head_node,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices1,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr1,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices2,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr2,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors) noexcept:
+def _breadth_first_undirected(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices1,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr1,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indices2,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indptr2,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors):
+    return _breadth_first_undirected2(head_node, indices1, indptr1, indices2,
+                                      indptr2, node_list, predecessors)
+
+cdef unsigned int _breadth_first_undirected2(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices1,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr1,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indices2,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indptr2,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors) noexcept:
     # Inputs:
     #  head_node: (input) index of the node from which traversal starts
     #  indices1: (input) CSR indices of graph
@@ -538,14 +558,27 @@ cpdef depth_first_order(csgraph, i_start,
         return node_list[:length]
 
 
-cdef unsigned int _depth_first_directed(
-                           unsigned int head_node,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] root_list,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] flag) noexcept:
+def _depth_first_directed(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] root_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] flag):
+    return _depth_first_directed2(head_node, indices, indptr,
+                                  node_list, predecessors,
+                                  root_list, flag)
+
+
+cdef unsigned int _depth_first_directed2(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] root_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] flag) noexcept:
     cdef unsigned int i, i_nl_end, cnode, pnode
     cdef unsigned int N = node_list.shape[0]
     cdef int no_children, i_root
@@ -582,16 +615,32 @@ cdef unsigned int _depth_first_directed(
     return i_nl_end
 
 
-cdef unsigned int _depth_first_undirected(
-                           unsigned int head_node,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices1,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr1,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices2,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr2,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] root_list,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] flag) noexcept:
+def _depth_first_undirected(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices1,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr1,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indices2,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indptr2,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] root_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] flag):
+    return _depth_first_undirected2(head_node, indices1, indptr1,
+                                    indices2, indptr2,
+                                    node_list, predecessors,
+                                    root_list, flag)
+
+
+cdef unsigned int _depth_first_undirected2(
+        unsigned int head_node,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices1,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr1,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indices2,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indptr2,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] node_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] predecessors,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] root_list,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] flag) noexcept:
     cdef unsigned int i, i_nl_end, cnode, pnode
     cdef unsigned int N = node_list.shape[0]
     cdef int no_children, i_root
@@ -644,10 +693,17 @@ cdef unsigned int _depth_first_undirected(
     return i_nl_end
 
 
-cdef int _connected_components_directed(
-                                 np.ndarray[ITYPE_t, ndim=1, mode='c'] indices,
-                                 np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr,
-                                 np.ndarray[ITYPE_t, ndim=1, mode='c'] labels) noexcept:
+def _connected_components_directed(
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] labels):
+    return _connected_components_directed2(indices, indptr, labels)
+
+
+cdef int _connected_components_directed2(
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] labels) noexcept:
     """
     Uses an iterative version of Tarjan's algorithm to find the
     strongly connected components of a directed graph represented as a
@@ -766,12 +822,24 @@ cdef int _connected_components_directed(
     labels += (N - 1)
     return (N - 1) - label
 
-cdef int _connected_components_undirected(
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices1,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr1,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indices2,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] indptr2,
-                           np.ndarray[ITYPE_t, ndim=1, mode='c'] labels) noexcept:
+
+def _connected_components_undirected(
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices1,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr1,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indices2,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indptr2,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] labels):
+    return _connected_components_undirected2(indices1, indptr1,
+                                             indices2, indptr2,
+                                             labels)
+
+
+cdef int _connected_components_undirected2(
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indices1,
+        np.ndarray[int32_or_int64, ndim=1, mode='c'] indptr1,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indices2,
+        np.ndarray[int32_or_int64_b, ndim=1, mode='c'] indptr2,
+        np.ndarray[ITYPE_t, ndim=1, mode='c'] labels) noexcept:
 
     cdef int v, w, j, label, SS_head
     cdef int N = labels.shape[0]
diff --git a/scipy/sparse/csgraph/parameters.pxi b/scipy/sparse/csgraph/parameters.pxi
index 91cbdd640441..5391bedf4d4f 100644
--- a/scipy/sparse/csgraph/parameters.pxi
+++ b/scipy/sparse/csgraph/parameters.pxi
@@ -10,6 +10,11 @@ ctypedef fused int32_or_int64:
     np.int32_t
     np.int64_t
 
+# Another copy of the same fused type, for working with mixed-type functions.
+ctypedef fused int32_or_int64_b:
+    np.int32_t
+    np.int64_t
+
 # EPS is the precision of DTYPE
 DEF DTYPE_EPS = 1E-15
 
diff --git a/scipy/sparse/csgraph/tests/test_connected_components.py b/scipy/sparse/csgraph/tests/test_connected_components.py
index 681f29fd48a3..fec8db4a26a8 100644
--- a/scipy/sparse/csgraph/tests/test_connected_components.py
+++ b/scipy/sparse/csgraph/tests/test_connected_components.py
@@ -1,6 +1,6 @@
 import numpy as np
 from numpy.testing import assert_equal, assert_array_almost_equal
-from scipy.sparse import csgraph
+from scipy.sparse import csgraph, csr_array
 
 
 def test_weak_connections():
@@ -97,3 +97,23 @@ def test_fully_connected_graph():
     g = np.ones((4, 4))
     n_components, labels = csgraph.connected_components(g)
     assert_equal(n_components, 1)
+
+
+def test_int64_indices_undirected():
+    # See https://github.com/scipy/scipy/issues/18716
+    g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2))
+    assert g.indices.dtype == np.int64
+    n, labels = csgraph.connected_components(g)
+    assert n == 1
+    assert_array_almost_equal(labels, [0, 0])
+
+
+def test_int64_indices_directed():
+    # See https://github.com/scipy/scipy/issues/18716
+    g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2))
+    assert g.indices.dtype == np.int64
+    n, labels = csgraph.connected_components(g, directed=True,
+                                             connection='strong')
+    assert n == 2
+    assert_array_almost_equal(labels, [1, 0])
+
diff --git a/scipy/sparse/csgraph/tests/test_traversal.py b/scipy/sparse/csgraph/tests/test_traversal.py
index 026fbe275484..6494c93558b6 100644
--- a/scipy/sparse/csgraph/tests/test_traversal.py
+++ b/scipy/sparse/csgraph/tests/test_traversal.py
@@ -1,5 +1,6 @@
 import numpy as np
 from numpy.testing import assert_array_almost_equal
+from scipy.sparse import csr_array
 from scipy.sparse.csgraph import (breadth_first_tree, depth_first_tree,
     csgraph_to_dense, csgraph_from_dense)
 
@@ -66,3 +67,21 @@ def test_graph_depth_first_trivial_graph():
         bfirst_test = depth_first_tree(csgraph, 0, directed)
         assert_array_almost_equal(csgraph_to_dense(bfirst_test),
                                   bfirst)
+
+
+def test_breadth_first_int64_indices():
+    # See https://github.com/scipy/scipy/issues/18716
+    g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2))
+    assert g.indices.dtype == np.int64
+    for directed in [True, False]:
+        bf = breadth_first_tree(g, 0, directed=directed)
+        assert_array_almost_equal(csgraph_to_dense(bf), [[0, 1], [0, 0]])
+
+
+def test_depth_first_int64_indices():
+    # See https://github.com/scipy/scipy/issues/18716
+    g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2))
+    assert g.indices.dtype == np.int64
+    for directed in [True, False]:
+        df = depth_first_tree(g, 0, directed=directed)
+        assert_array_almost_equal(csgraph_to_dense(df), [[0, 1], [0, 0]])


Code review comments:
diff_hunk: @@ -97,3 +97,23 @@ def test_fully_connected_graph():
     g = np.ones((4, 4))
     n_components, labels = csgraph.connected_components(g)
     assert_equal(n_components, 1)
+
+
+def test_int64_indices_undirected():
+    # See https://github.com/scipy/scipy/issues/18716
+    g = csr_array(([1], np.array([[0], [1]], dtype=np.int64)), shape=(2, 2))
+    assert g.indices.dtype == np.int64
+    n, labels = csgraph.connected_components(g)
path: scipy/sparse/csgraph/tests/test_connected_components.py
text: @user1:
Nit: for explicitness.
```suggestion
    n, labels = csgraph.connected_components(g, directed=False)
```