jim22k · May 6, 2022 · May 5, 2022 · May 6, 2022 · May 6, 2022 · eriknw
diff --git a/graphblas/__init__.py b/graphblas/__init__.py
@@ -59,6 +59,7 @@ def get_config():
     "expr",
     "ffi",
     "formatting",
+    "indexunary",
     "infix",
     "io",
     "lib",

diff --git a/graphblas/indexunary/__init__.py b/graphblas/indexunary/__init__.py
@@ -0,0 +1,19 @@
+# All items are dynamically added by classes in operator.py
+# This module acts as a container of IndexUnaryOp instances
+_delayed = {}
+from graphblas import operator  # noqa isort:skip
+
+del operator
+
+
+def __dir__():
+    return globals().keys() | _delayed.keys()
+
+
+def __getattr__(key):
+    if key in _delayed:
+        func, kwargs = _delayed.pop(key)
+        rv = func(**kwargs)
+        globals()[key] = rv
+        return rv
+    raise AttributeError(f"module {__name__!r} has no attribute {key!r}")
diff --git a/graphblas/matrix.py b/graphblas/matrix.py
@@ -10,7 +10,7 @@
 from .exceptions import DimensionMismatch, NoValue, check_status
 from .expr import AmbiguousAssignOrExtract, IndexerResolver, Updater
 from .mask import StructuralMask, ValueMask
-from .operator import get_semiring, get_typed_op
+from .operator import find_opclass, get_semiring, get_typed_op, op_from_string
 from .scalar import _MATERIALIZE, Scalar, ScalarExpression, _as_scalar, _scalar_index
 from .utils import (
     _CArray,
@@ -740,11 +740,23 @@ def apply(self, op, right=None, *, left=None):
         Apply UnaryOp to each element of the calling Matrix
         A BinaryOp can also be applied if a scalar is passed in as `left` or `right`,
             effectively converting a BinaryOp into a UnaryOp
+        An IndexUnaryOp can also be applied with the thunk passed in as `right`
         """
         method_name = "apply"
         extra_message = (
-            "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar."
+            "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar "
+            "or IndexUnaryOp with `right` thunk."
         )
+        if isinstance(op, str):
+            op = op_from_string(op)
+        op, opclass = find_opclass(op)
+        if opclass in {"IndexUnaryOp", "SelectOp"}:
+            # Provide default value for index unary
+            if right is None:
+                right = False  # most basic form of 0 when unifying dtypes
+            if left is not None:
+                raise TypeError("Do not pass `left` when applying IndexUnaryOp")
+
         if left is None and right is None:
             op = get_typed_op(op, self.dtype, kind="unary")
             self._expect_op(
@@ -772,7 +784,7 @@ def apply(self, op, right=None, *, left=None):
                         op=op,
                     )
             op = get_typed_op(op, left.dtype, self.dtype, is_left_scalar=True, kind="binary")
-            if op.opclass == "Monoid":
+            if opclass == "Monoid":
                 op = op.binaryop
             else:
                 self._expect_op(
@@ -807,26 +819,33 @@ def apply(self, op, right=None, *, left=None):
                         extra_message="Literal scalars also accepted.",
                         op=op,
                     )
-            op = get_typed_op(op, self.dtype, right.dtype, is_right_scalar=True, kind="binary")
-            if op.opclass == "Monoid":
-                op = op.binaryop
-            else:
-                self._expect_op(
-                    op,
-                    "BinaryOp",
-                    within=method_name,
-                    argname="op",
-                    extra_message=extra_message,
+            if opclass in {"IndexUnaryOp", "SelectOp"}:
+                op = get_typed_op(
+                    op, self.dtype, right.dtype, is_right_scalar=True, kind="indexunary"
                 )
+                cfunc_method = "IndexOp"
+            else:
+                op = get_typed_op(op, self.dtype, right.dtype, is_right_scalar=True, kind="binary")
+                cfunc_method = "BinaryOp2nd"
+                if opclass == "Monoid":
+                    op = op.binaryop
+                else:
+                    self._expect_op(
+                        op,
+                        "BinaryOp",
+                        within=method_name,
+                        argname="op",
+                        extra_message=extra_message,
+                    )
             if right._is_cscalar:
                 if right.dtype._is_udt:
                     dtype_name = "UDT"
                     right = _Pointer(right)
                 else:
                     dtype_name = right.dtype.name
-                cfunc_name = f"GrB_Matrix_apply_BinaryOp2nd_{dtype_name}"
+                cfunc_name = f"GrB_Matrix_apply_{cfunc_method}_{dtype_name}"
             else:
-                cfunc_name = "GrB_Matrix_apply_BinaryOp2nd_Scalar"
+                cfunc_name = f"GrB_Matrix_apply_{cfunc_method}_Scalar"
             args = [self, right]
             expr_repr = "{0.name}.apply({op}, right={1._expr_name})"
         else:

diff --git a/graphblas/operator.py b/graphblas/operator.py
@@ -36,7 +36,7 @@
 
 _STANDARD_OPERATOR_NAMES = set()
 
-from . import binary, monoid, op, semiring, unary, select  # noqa isort:skip
+from . import binary, indexunary, monoid, op, select, semiring, unary  # noqa isort:skip
 
 ffi_new = ffi.new
 UNKNOWN_OPCLASS = "UnknownOpClass"
@@ -189,6 +189,16 @@ def __call__(self, val):
         )
 
 
+class TypedBuiltinIndexUnaryOp(TypedOpBase):
+    __slots__ = ()
+    opclass = "IndexUnaryOp"
+
+    def __call__(self, val, thunk=None):
+        if thunk is None:
+            thunk = False  # most basic form of 0 when unifying dtypes
+        return _call_op(self, val, right=thunk)
+
+
 class TypedBuiltinSelectOp(TypedOpBase):
     __slots__ = ()
     opclass = "SelectOp"
@@ -1110,8 +1120,70 @@ def __reduce__(self):
     __call__ = TypedBuiltinUnaryOp.__call__
 
 
+class IndexUnaryOp(OpBase):
+    __slots__ = "is_positional"
+    _module = indexunary
+    _modname = "indexunary"
+    _is_udt = False
+    _custom_dtype = None
+    _typed_class = TypedBuiltinIndexUnaryOp
+    _parse_config = {
+        "trim_from_front": 4,
+        "num_underscores": 1,
+        "re_exprs": [
+            re.compile("^GrB_(ROWINDEX|COLINDEX)_(INT32|INT64)$"),
+        ],
+        "re_exprs_return_bool": [
+            re.compile("^GrB_(TRIL|TRIU|DIAG|OFFDIAG|COLLE|COLGT|ROWLE|ROWGT)$"),
+            re.compile(
+                "^GrB_(VALUEEQ|VALUENE|VALUEGT|VALUEGE|VALUELT|VALUELE)"
+                "_(BOOL|INT8|UINT8|INT16|UINT16|INT32|UINT32|INT64|UINT64|FP32|FP64)$"
+            ),
+            re.compile("^GxB_(VALUEEQ|VALUENE)_(FC32|FC64)$"),
+        ],
+    }
+    # fmt: off
+    _positional = {"tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt",
+                   "rowindex", "colindex"}
+    # fmt: on
+
+    @classmethod
+    def _initialize(cls):
+        super()._initialize(include_in_ops=False)
+        # Update type information to include UINT64 for positional ops
+        for name in ("tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt"):
+            op = getattr(indexunary, name)
+            typed_op = op._typed_ops[BOOL]
+            output_type = op.types[BOOL]
+            if UINT64 not in op.types:
+                op.types[UINT64] = output_type
+                op._typed_ops[UINT64] = typed_op
+                op.coercions[UINT64] = BOOL
+        for name in ("rowindex", "colindex"):
+            op = getattr(indexunary, name)
+            typed_op = op._typed_ops[INT64]
+            output_type = op.types[INT64]
+            if UINT64 not in op.types:
+                op.types[UINT64] = output_type
+                op._typed_ops[UINT64] = typed_op
+                op.coercions[UINT64] = INT64
+        cls._initialized = True
+
+    def __init__(
+        self,
+        name,
+        *,
+        anonymous=False,
+        is_positional=False,
+    ):
+        super().__init__(name, anonymous=anonymous)
+        self.is_positional = is_positional
+
+    __call__ = TypedBuiltinIndexUnaryOp.__call__
+
+
 class SelectOp(OpBase):
-    __slots__ = ("is_positional",)
+    __slots__ = "is_positional"
     _module = select
     _modname = "select"
     _is_udt = False
@@ -1134,20 +1206,15 @@ class SelectOp(OpBase):
     @classmethod
     def _initialize(cls):
         super()._initialize(include_in_ops=False)
-        # Update type information to handle more than bool
-        bool_names = ("tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt")
-        types = [INT8, UINT8, INT16, UINT16, INT32, UINT32, INT64, UINT64, FP32, FP64]
-        if _supports_complex:
-            types.extend([FC32, FC64])
-        for name in bool_names:
+        # Update type information to include UINT64 for positional ops
+        for name in ("tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt"):
             op = getattr(select, name)
             typed_op = op._typed_ops[BOOL]
             output_type = op.types[BOOL]
-            for dtype in types:
-                if dtype not in op.types:
-                    op.types[dtype] = output_type
-                    op._typed_ops[dtype] = typed_op
-                    op.coercions[dtype] = BOOL
+            if UINT64 not in op.types:
+                op.types[UINT64] = output_type
+                op._typed_ops[UINT64] = typed_op
+                op.coercions[UINT64] = BOOL
         cls._initialized = True
 
     def __init__(
@@ -2475,10 +2542,12 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala
                     raise ValueError(
                         f"Unknown binary or aggregator string: {op!r}.  Example usage: '+[int]'"
                     ) from None
+
         else:
             raise ValueError(
                 f"Unable to get op from string {op!r}.  `kind=` argument must be provided as "
-                '"unary", "binary", "monoid", "semiring", or "binary|aggregator".'
+                '"unary", "binary", "monoid", "semiring", "indexunary", "select", '
+                'or "binary|aggregator".'
             )
         return get_typed_op(
             op,
@@ -2610,6 +2679,7 @@ def get_semiring(monoid, binaryop, name=None):
 # Now initialize all the things!
 try:
     UnaryOp._initialize()
+    IndexUnaryOp._initialize()
     SelectOp._initialize()
     BinaryOp._initialize()
     Monoid._initialize()
@@ -2726,6 +2796,12 @@ def unary_from_string(string):
     return _from_string(string, unary, _str_to_unary, "abs[int]")
 
 
+def indexunary_from_string(string):
+    # "select" is a variant of IndexUnary, so the string abbreviations in
+    # _str_to_select are appropriate to reuse here
+    return _from_string(string, indexunary, _str_to_select, "row_index")
+
+
 def select_from_string(string):
     return _from_string(string, select, _str_to_select, "tril")
 
@@ -2764,6 +2840,7 @@ def op_from_string(string):
         semiring_from_string,
         # Give binary precedence over this
         select_from_string,
+        indexunary_from_string,
     ]:
         try:
             return func(string)
@@ -2773,6 +2850,7 @@ def op_from_string(string):
 
 
 unary.from_string = unary_from_string
+indexunary.from_string = indexunary_from_string
 select.from_string = select_from_string
 binary.from_string = binary_from_string
 monoid.from_string = monoid_from_string

diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py
@@ -9,7 +9,7 @@
 from numpy.testing import assert_array_equal
 
 import graphblas
-from graphblas import agg, binary, dtypes, monoid, select, semiring, unary
+from graphblas import agg, binary, dtypes, indexunary, monoid, select, semiring, unary
 from graphblas.exceptions import (
     DimensionMismatch,
     EmptyObject,
@@ -1082,12 +1082,45 @@ def test_apply_binary(A):
     assert w1.isequal(w3)
 
 
+def test_apply_indexunary(A):
+    ridx = [3, 0, 3, 5, 6, 0, 6, 1, 6, 2, 4, 1]
+    cidx = [0, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 6]
+    Ar = Matrix.from_values(ridx, cidx, ridx)
+    r1 = A.apply("rowindex").new()
+    r2 = A.apply(indexunary.rowindex).new()
+    r3 = indexunary.rowindex(A).new()
+    assert r1.isequal(Ar)
+    assert r2.isequal(Ar)
+    assert r3.isequal(Ar)
+
+    Ac = Matrix.from_values(ridx, cidx, [c + 2 for c in cidx])
+    c1 = A.apply("colindex", 2).new()
+    c2 = A.apply(indexunary.colindex, 2).new()
+    c3 = indexunary.colindex(A, thunk=2).new()
+    assert c1.isequal(Ac)
+    assert c2.isequal(Ac)
+    assert c3.isequal(Ac)
+
+    A3 = Matrix.from_values(ridx, cidx, [1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0], dtype=bool)
+    s3 = Scalar.from_value(3, dtypes.INT64)
+    w1 = A.apply(indexunary.valueeq, s3).new()
+    w2 = A.apply(select.valueeq, s3).new()
+    w3 = A.apply("==", s3).new()
+    w4 = indexunary.valueeq(A, s3).new()
+    assert w1.isequal(A3)
+    assert w2.isequal(A3)
+    assert w3.isequal(A3)
+    assert w4.isequal(A3)
+
+
 def test_select(A):
     A3 = Matrix.from_values([0, 3, 3, 6], [3, 0, 2, 4], [3, 3, 3, 3], nrows=7, ncols=7)
     w1 = A.select(select.valueeq, 3).new()
+    w1b = A.select(indexunary.valueeq, 3).new()
     w2 = A.select("==", 3).new()
     w3 = select.value(A == 3).new()
     assert w1.isequal(A3)
+    assert w1b.isequal(A3)
     assert w2.isequal(A3)
     assert w3.isequal(A3)
 

diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py
@@ -9,7 +9,7 @@
 from numpy.testing import assert_array_equal
 
 import graphblas
-from graphblas import agg, binary, dtypes, monoid, select, semiring, unary
+from graphblas import agg, binary, dtypes, indexunary, monoid, select, semiring, unary
 from graphblas.exceptions import (
     DimensionMismatch,
     EmptyObject,
@@ -647,6 +647,29 @@ def test_apply_empty(v):
         v.apply(binary.plus, s).new()
 
 
+def test_apply_indexunary(v):
+    idx = [1, 3, 4, 6]
+    vr = Vector.from_values(idx, idx)
+    r1 = v.apply(indexunary.rowindex).new()
+    r2 = v.apply("rowindex").new()
+    r3 = indexunary.rowindex(v).new()
+    assert r1.isequal(vr)
+    assert r2.isequal(vr)
+    assert r3.isequal(vr)
+
+    v2 = Vector.from_values(idx, [False, False, True, False])
+    s2 = Scalar(dtypes.INT64)
+    s2.value = 2
+    w1 = v.apply("==", s2).new()
+    w2 = v.apply(indexunary.valueeq, s2).new()
+    w3 = v.apply(select.valueeq, s2).new()
+    w4 = indexunary.valueeq(v, s2).new()
+    assert w1.isequal(v2)
+    assert w2.isequal(v2)
+    assert w3.isequal(v2)
+    assert w4.isequal(v2)
+
+
 def test_select(v):
     result = Vector.from_values([1, 3], [1, 1], size=7)
     w1 = v.select(select.valueeq, 1).new()
-Original file line number
+Diff line change
@@ Expand Up / @@ -59,6 +59,7 @@ def get_config(): @@
         "expr",
         "ffi",
         "formatting",
+        "indexunary",
         "infix",
         "io",
         "lib",
@@ Expand Down @@