Add apply_IndexOp_Scalar to matrix and vector functional API

Add matrix_apply_indexop and vector_apply_indexop wrapping the
GrB_Matrix_apply_IndexOp_Scalar and GrB_Vector_apply_IndexOp_Scalar
C functions, with doctests. Update BFS example to use the new
vector_apply_indexop instead of raw ffi calls.

Author: michelp

suitesparse_graphblas/api/examples/bfs.py

@@ -29,7 +29,7 @@
     [0, 0, 1, 2]
 """
 
-from suitesparse_graphblas import check_status, ffi, lib
+from suitesparse_graphblas import lib
 from suitesparse_graphblas.api import matrix, scalar, vector
 
 
@@ -69,6 +69,8 @@ def bfs(A, src, compute_level=True, compute_parent=True):
         semiring = lib.GrB_MIN_FIRST_SEMIRING_INT64
         frontier = vector.vector_new(lib.GrB_INT64, n)
         vector.set_int64(frontier, src, src)
+        index_thunk = scalar.scalar_new(lib.GrB_INT64)
+        scalar.set_int64(index_thunk, 0)
     else:
         # Level-only mode: frontier is BOOL.
         # Semiring: ANY_PAIR just checks reachability.
@@ -106,12 +108,8 @@ def bfs(A, src, compute_level=True, compute_parent=True):
             # Convert frontier values to their own indices (ROWINDEX).
             # After this, frontier(i) == i for every stored entry,
             # so the next vxm propagates node i as the parent ID.
-            check_status(
-                frontier,
-                lib.GrB_Vector_apply_IndexOp_INT64(
-                    frontier[0], ffi.NULL, ffi.NULL,
-                    lib.GrB_ROWINDEX_INT64, frontier[0], 0, ffi.NULL,
-                ),
+            vector.vector_apply_indexop(
+                frontier, lib.GrB_ROWINDEX_INT64, frontier, index_thunk,
             )
 
         current_level += 1

suitesparse_graphblas/api/matrix.py

@@ -505,6 +505,34 @@ def matrix_apply_second(C, op, A, y, mask=None, accum=None, desc=None):
     ))
 
 
+def matrix_apply_indexop(C, op, A, thunk, mask=None, accum=None, desc=None):
+    """Apply an index unary operator to each entry of a matrix with a scalar thunk.
+
+    >>> from suitesparse_graphblas import scalar
+    >>> C = matrix_new(lib.GrB_INT64, 3, 3)
+    >>> A = matrix_new(lib.GrB_INT64, 3, 3)
+    >>> set_int64(A, 10, 0, 0)
+    >>> set_int64(A, 20, 1, 1)
+    >>> s = scalar.scalar_new(lib.GrB_INT64)
+    >>> scalar.set_int64(s, 0)
+    >>> matrix_apply_indexop(C, lib.GrB_ROWINDEX_INT64, A, s)
+    >>> get_int64(C, 0, 0) == 0
+    True
+    >>> get_int64(C, 1, 1) == 1
+    True
+
+    """
+    check_status(C, lib.GrB_Matrix_apply_IndexOp_Scalar(
+        C[0],
+        ffi.NULL if mask is None else mask[0],
+        ffi.NULL if accum is None else accum,
+        op,
+        A[0],
+        thunk[0],
+        ffi.NULL if desc is None else desc,
+    ))
+
+
 def matrix_select(C, op, A, thunk, mask=None, accum=None, desc=None):
     """Select entries from a matrix using an index unary operator and scalar thunk.
 

suitesparse_graphblas/api/vector.py

@@ -432,6 +432,34 @@ def vector_apply_second(w, op, u, y, mask=None, accum=None, desc=None):
     ))
 
 
+def vector_apply_indexop(w, op, u, thunk, mask=None, accum=None, desc=None):
+    """Apply an index unary operator to each entry of a vector with a scalar thunk.
+
+    >>> from suitesparse_graphblas import scalar
+    >>> w = vector_new(lib.GrB_INT64, 3)
+    >>> u = vector_new(lib.GrB_INT64, 3)
+    >>> set_int64(u, 10, 0)
+    >>> set_int64(u, 20, 2)
+    >>> s = scalar.scalar_new(lib.GrB_INT64)
+    >>> scalar.set_int64(s, 0)
+    >>> vector_apply_indexop(w, lib.GrB_ROWINDEX_INT64, u, s)
+    >>> get_int64(w, 0) == 0
+    True
+    >>> get_int64(w, 2) == 2
+    True
+
+    """
+    check_status(w, lib.GrB_Vector_apply_IndexOp_Scalar(
+        w[0],
+        ffi.NULL if mask is None else mask[0],
+        ffi.NULL if accum is None else accum,
+        op,
+        u[0],
+        thunk[0],
+        ffi.NULL if desc is None else desc,
+    ))
+
+
 def vector_select(w, op, u, thunk, mask=None, accum=None, desc=None):
     """Select entries from a vector using an index unary operator and scalar thunk.