ENH: sparse: make two sputils public for easier index array casting (scipy#22043)

* make index tools public

* update release notes

* Add examples to docstring

Author: dschult

doc/source/reference/sparse.migration_to_sparray.rst

@@ -449,19 +449,23 @@ The function signatures are::
     def get_index_dtype(arrays=(), maxval=None, check_contents=False):
     def safely_cast_index_arrays(A, idx_dtype=np.int32, msg=""):
 
-Example idioms include the following::
+Example idioms include the following for ``get_index_dtype``::
 
    .. code-block:: python
 
        # select index dtype before construction based on shape
        shape = (3, 3)
-       idx_dtype = scipy.sparse._sputils.get_index_dtype(maxval=max(shape))
+       idx_dtype = scipy.sparse.get_index_dtype(maxval=max(shape))
        indices = np.array([0, 1, 0], dtype=idx_dtype)
        indptr = np.arange(3, dtype=idx_dtype)
        A = csr_array((data, indices, indptr), shape=shape)
 
-       # rescast after construction, raising exception before overflow
-       indices, indptr = scipy.sparse._sputils.safely_cast_index_arrays(B, np.int32)
+and for ``safely_cast_index_arrays``::
+
+   .. code-block:: python
+
+       # rescast after construction, raising exception if shape too big
+       indices, indptr = scipy.sparse.safely_cast_index_arrays(B, np.int32)
        B.indices, B.indptr = indices, indptr
 
 Other

doc/source/release/1.15.0-notes.rst

@@ -194,8 +194,8 @@ and support several Array API compatible array libraries in addition to NumPy
   incompatible data types, such as ``float16``.
 - ``min``, ``max``, ``argmin``, and ``argmax`` now support computation
   over nonzero elements only via the new ``explicit`` argument.
-- New function ``safely_cast_index_arrays`` has been added
-  to facilitate casting challenges in ``sparse``.
+- New functions ``get_index_dtype`` and ``safely_cast_index_arrays`` are
+  available to facilitate index array casting in ``sparse``.
 
 
 `scipy.spatial` improvements

scipy/sparse/__init__.py

@@ -102,6 +102,8 @@
    save_npz - Save a sparse array to a file using ``.npz`` format.
    load_npz - Load a sparse array from a file using ``.npz`` format.
    find - Return the indices and values of the nonzero elements
+   get_index_dtype - determine a good dtype for index arrays.
+   safely_cast_index_arrays - cast index array dtype or raise if shape too big
 
 Identifying sparse arrays
 -------------------------
@@ -307,6 +309,7 @@
 from ._extract import *
 from ._matrix import spmatrix
 from ._matrix_io import *
+from ._sputils import get_index_dtype, safely_cast_index_arrays
 
 # For backward compatibility with v0.19.
 from . import csgraph

scipy/sparse/_sputils.py

@@ -32,13 +32,13 @@ def upcast(*args):
     --------
     >>> from scipy.sparse._sputils import upcast
     >>> upcast('int32')
-    <type 'numpy.int32'>
+    <class 'numpy.int32'>
     >>> upcast('bool')
-    <type 'numpy.bool_'>
+    <class 'numpy.bool'>
     >>> upcast('int32','float32')
-    <type 'numpy.float64'>
+    <class 'numpy.float64'>
     >>> upcast('bool',complex,float)
-    <type 'numpy.complex128'>
+    <class 'numpy.complex128'>
 
     """
 
@@ -172,6 +172,7 @@ def safely_cast_index_arrays(A, idx_dtype=np.int32, msg=""):
         The array for which index arrays should be downcast.
     idx_dtype : dtype
         Desired dtype. Should be an integer dtype (default: ``np.int32``).
+        Most of scipy.sparse uses either int64 or int32.
     msg : string, optional
         A string to be added to the end of the ValueError message
         if the array shape is too big to fit in `idx_dtype`.
@@ -193,6 +194,24 @@ def safely_cast_index_arrays(A, idx_dtype=np.int32, msg=""):
     ValueError
         If the array has shape that would not fit in the new dtype, or if
         the sparse format does not use index arrays.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from scipy import sparse
+    >>> data = [3]
+    >>> coords = (np.array([3]), np.array([1]))  # Note: int64 arrays
+    >>> A = sparse.coo_array((data, coords))
+    >>> A.coords[0].dtype
+    dtype('int64')
+
+    >>> # rescast after construction, raising exception if shape too big
+    >>> coords = sparse.safely_cast_index_arrays(A, np.int32)
+    >>> A.coords[0] is coords[0]  # False if casting is needed
+    False
+    >>> A.coords = coords  # set the index dtype of A
+    >>> A.coords[0].dtype
+    dtype('int32')
     """
     if not msg:
         msg = f"dtype {idx_dtype}"
@@ -262,6 +281,25 @@ def get_index_dtype(arrays=(), maxval=None, check_contents=False):
     dtype : dtype
         Suitable index data type (int32 or int64)
 
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from scipy import sparse
+    >>> # select index dtype based on shape
+    >>> shape = (3, 3)
+    >>> idx_dtype = sparse.get_index_dtype(maxval=max(shape))
+    >>> data = [1.1, 3.0, 1.5]
+    >>> indices = np.array([0, 1, 0], dtype=idx_dtype)
+    >>> indptr = np.array([0, 2, 3, 3], dtype=idx_dtype)
+    >>> A = sparse.csr_array((data, indices, indptr), shape=shape)
+    >>> A.indptr.dtype
+    dtype('int32')
+
+    >>> # select based on larger of existing arrays and shape
+    >>> shape = (3, 3)
+    >>> idx_dtype = sparse.get_index_dtype(A.indptr, maxval=max(shape))
+    >>> idx_dtype
+    <class 'numpy.int32'>
     """
     # not using intc directly due to misinteractions with pythran
     if np.intc().itemsize != 4: