eriknw · Apr 14, 2022 · Apr 11, 2022 · Apr 11, 2022 · Apr 11, 2022 · Apr 12, 2022
diff --git a/graphblas/__init__.py b/graphblas/__init__.py
@@ -68,6 +68,7 @@ def get_config():
     "operator",
     "recorder",
     "scalar",
+    "select",
     "semiring",
     "ss",
     "tests",

diff --git a/graphblas/_automethods.py b/graphblas/_automethods.py
@@ -230,6 +230,10 @@ def reduce_scalar(self):
     return self._get_value("reduce_scalar")
 
 
+def select(self):
+    return self._get_value("select")
+
+
 def ss(self):
     return self._get_value("ss")
 
@@ -349,6 +353,7 @@ def __ixor__(self, other):
         "ewise_add",
         "ewise_mult",
         "ewise_union",
+        "select",
         "ss",
         "to_values",
     }

diff --git a/graphblas/infix.py b/graphblas/infix.py
@@ -89,6 +89,7 @@ def shape(self):
     nvals = wrapdoc(Vector.nvals)(property(_automethods.nvals))
     outer = wrapdoc(Vector.outer)(property(_automethods.outer))
     reduce = wrapdoc(Vector.reduce)(property(_automethods.reduce))
+    select = wrapdoc(Vector.select)(property(_automethods.select))
     ss = wrapdoc(Vector.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Vector.to_pygraphblas)(property(_automethods.to_pygraphblas))
     to_values = wrapdoc(Vector.to_values)(property(_automethods.to_values))
@@ -205,6 +206,7 @@ def shape(self):
     reduce_columnwise = wrapdoc(Matrix.reduce_columnwise)(property(_automethods.reduce_columnwise))
     reduce_rowwise = wrapdoc(Matrix.reduce_rowwise)(property(_automethods.reduce_rowwise))
     reduce_scalar = wrapdoc(Matrix.reduce_scalar)(property(_automethods.reduce_scalar))
+    select = wrapdoc(Matrix.select)(property(_automethods.select))
     ss = wrapdoc(Matrix.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Matrix.to_pygraphblas)(property(_automethods.to_pygraphblas))
     to_values = wrapdoc(Matrix.to_values)(property(_automethods.to_values))

diff --git a/graphblas/matrix.py b/graphblas/matrix.py
@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from . import _automethods, backend, binary, ffi, lib, monoid, semiring, utils
+from . import _automethods, backend, binary, ffi, lib, monoid, select, semiring, utils
 from ._ss.matrix import ss
 from .base import BaseExpression, BaseType, call
 from .dtypes import _INDEX, FP64, lookup_dtype, unify
@@ -791,6 +791,53 @@ def apply(self, op, right=None, *, left=None):
             bt=self._is_transposed,
         )
 
+    def select(self, op, thunk=None):
+        """
+        GrB_Matrix_select
+        Compute SelectOp at each element of the calling Matrix, keeping
+        elements which return True.
+        """
+        if isinstance(op, str):
+            op = select.from_string(op)
+        method_name = "select"
+        if thunk is None:
+            thunk = False  # most basic form of 0 when unifying dtypes
+        if type(thunk) is not Scalar:
+            dtype = self.dtype if (self.dtype._is_udt and not op.is_positional) else None
+            try:
+                thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="")
+            except TypeError:
+                thunk = self._expect_type(
+                    thunk,
+                    Scalar,
+                    within=method_name,
+                    keyword_name="thunk",
+                    extra_message="Literal scalars also accepted.",
+                    op=op,
+                )
+        op = get_typed_op(op, self.dtype, thunk.dtype, is_right_scalar=True, kind="select")
+        self._expect_op(op, ("SelectOp", "IndexUnaryOp"), within=method_name, argname="op")
+        if thunk._is_cscalar:
+            if thunk.dtype._is_udt:
+                dtype_name = "UDT"
+                thunk = _Pointer(thunk)
+            else:
+                dtype_name = thunk.dtype.name
+            cfunc_name = f"GrB_Matrix_select_{dtype_name}"
+        else:
+            cfunc_name = "GrB_Matrix_select_Scalar"
+        return MatrixExpression(
+            method_name,
+            cfunc_name,
+            [self, thunk],
+            op=op,
+            expr_repr="{0.name}.select({op}, thunk={1._expr_name})",
+            nrows=self._nrows,
+            ncols=self._ncols,
+            dtype=self.dtype,
+            at=self._is_transposed,
+        )
+
     def reduce_rowwise(self, op=monoid.plus):
         """
         GrB_Matrix_reduce
@@ -1427,6 +1474,7 @@ def shape(self):
     reduce_columnwise = wrapdoc(Matrix.reduce_columnwise)(property(_automethods.reduce_columnwise))
     reduce_rowwise = wrapdoc(Matrix.reduce_rowwise)(property(_automethods.reduce_rowwise))
     reduce_scalar = wrapdoc(Matrix.reduce_scalar)(property(_automethods.reduce_scalar))
+    select = wrapdoc(Matrix.select)(property(_automethods.select))
     ss = wrapdoc(Matrix.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Matrix.to_pygraphblas)(property(_automethods.to_pygraphblas))
     to_values = wrapdoc(Matrix.to_values)(property(_automethods.to_values))
@@ -1505,6 +1553,7 @@ def shape(self):
     reduce_columnwise = wrapdoc(Matrix.reduce_columnwise)(property(_automethods.reduce_columnwise))
     reduce_rowwise = wrapdoc(Matrix.reduce_rowwise)(property(_automethods.reduce_rowwise))
     reduce_scalar = wrapdoc(Matrix.reduce_scalar)(property(_automethods.reduce_scalar))
+    select = wrapdoc(Matrix.select)(property(_automethods.select))
     ss = wrapdoc(Matrix.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Matrix.to_pygraphblas)(property(_automethods.to_pygraphblas))
     to_values = wrapdoc(Matrix.to_values)(property(_automethods.to_values))
@@ -1607,6 +1656,7 @@ def _name_html(self):
     mxm = Matrix.mxm
     kronecker = Matrix.kronecker
     apply = Matrix.apply
+    select = Matrix.select
     reduce_rowwise = Matrix.reduce_rowwise
     reduce_columnwise = Matrix.reduce_columnwise
     reduce_scalar = Matrix.reduce_scalar

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  # noqa isort:skip
+from . import binary, monoid, op, semiring, unary, select  # noqa isort:skip
 
 ffi_new = ffi.new
 UNKNOWN_OPCLASS = "UnknownOpClass"
@@ -60,8 +60,8 @@ def __getattr__(self, key):
         self.__getattribute__(key)  # raises
 
 
-def _call_op(op, left, right=None, **kwargs):
-    if right is None:
+def _call_op(op, left, right=None, thunk=None, **kwargs):
+    if right is None and thunk is None:
         if isinstance(left, InfixExprBase):
             # op(A & B), op(A | B), op(A @ B)
             return getattr(left.left, left.method_name)(left.right, op, **kwargs)
@@ -78,12 +78,15 @@ def _call_op(op, left, right=None, **kwargs):
             f"    - {op!r}(1, A)"
         )
 
-    # op(A, 1) -> apply (or select once available)
+    # op(A, 1) -> apply (or select if thunk provided)
     from .matrix import Matrix, TransposedMatrix
     from .vector import Vector
 
     if output_type(left) in {Vector, Matrix, TransposedMatrix}:
-        return left.apply(op, right=right, **kwargs)
+        if thunk is not None:
+            return left.select(op, thunk=thunk, **kwargs)
+        else:
+            return left.apply(op, right=right, **kwargs)
     elif output_type(right) in {Vector, Matrix, TransposedMatrix}:
         return right.apply(op, left=left, **kwargs)
     raise TypeError(
@@ -186,6 +189,16 @@ def __call__(self, val):
         )
 
 
+class TypedBuiltinSelectOp(TypedOpBase):
+    __slots__ = ()
+    opclass = "SelectOp"
+
+    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, thunk=thunk)
+
+
 class TypedBuiltinBinaryOp(TypedOpBase):
     __slots__ = ()
     opclass = "BinaryOp"
@@ -742,7 +755,11 @@ def _find(cls, funcname):
         return rv
 
     @classmethod
-    def _initialize(cls):
+    def _initialize(cls, include_in_ops=True):
+        """
+        include_in_ops determines whether the operators are included in the
+        `gb.ops` namespace in addition to the defined module.
+        """
         if cls._initialized:
             return
         # Read in the parse configs
@@ -782,7 +799,7 @@ def _initialize(cls):
                                 obj = cls(name, is_positional=name in cls._positional)
                             setattr(cls._module, name, obj)
                             _STANDARD_OPERATOR_NAMES.add(f"{cls._modname}.{name}")
-                            if not hasattr(op, name):
+                            if include_in_ops and not hasattr(op, name):
                                 setattr(op, name, obj)
                         else:
                             obj = getattr(cls._module, name)
@@ -1092,6 +1109,58 @@ def __reduce__(self):
     __call__ = TypedBuiltinUnaryOp.__call__
 
 
+class SelectOp(OpBase):
+    __slots__ = ("is_positional",)
+    _module = select
+    _modname = "select"
+    _is_udt = False
+    _typed_class = TypedBuiltinSelectOp
+    _parse_config = {
+        "trim_from_front": 4,
+        "num_underscores": 1,
+        "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)$"),
+        ],
+    }
+    _positional = {"tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt"}
+
+    @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:
+            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
+        cls._initialized = True
+
+    def __init__(
+        self,
+        name,
+        *,
+        anonymous=False,
+        is_positional=False,
+    ):
+        super().__init__(name, anonymous=anonymous)
+        self.is_positional = is_positional
+
+    __call__ = TypedBuiltinSelectOp.__call__
+
+
 def _floordiv(x, y):
     return x // y  # pragma: no cover
 
@@ -2315,6 +2384,8 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala
     elif isinstance(op, str):
         if kind == "unary":
             op = unary_from_string(op)
+        elif kind == "select":
+            op = select_from_string(op)
         elif kind == "binary":
             op = binary_from_string(op)
         elif kind == "monoid":
@@ -2460,6 +2531,7 @@ def get_semiring(monoid, binaryop, name=None):
 # Now initialize all the things!
 try:
     UnaryOp._initialize()
+    SelectOp._initialize()
     BinaryOp._initialize()
     Monoid._initialize()
     Semiring._initialize()
@@ -2474,6 +2546,18 @@ def get_semiring(monoid, binaryop, name=None):
     "-": unary.ainv,
     "~": unary.lnot,
 }
+_str_to_select = {
+    "<": select.valuelt,
+    ">": select.valuegt,
+    "<=": select.valuele,
+    ">=": select.valuege,
+    "!=": select.valuene,
+    "==": select.valueeq,
+    "col<=": select.colle,
+    "col>": select.colgt,
+    "row<=": select.rowle,
+    "row>": select.rowgt,
+}
 _str_to_binary = {
     "<": binary.lt,
     ">": binary.gt,
@@ -2554,6 +2638,10 @@ def unary_from_string(string):
     return _from_string(string, unary, _str_to_unary, "abs[int]")
 
 
+def select_from_string(string):
+    return _from_string(string, select, _str_to_select, "tril")
+
+
 def binary_from_string(string):
     return _from_string(string, binary, _str_to_binary, "+[int]")
 
@@ -2586,6 +2674,8 @@ def op_from_string(string):
         binary_from_string,
         monoid_from_string,
         semiring_from_string,
+        # Give binary precedence over this
+        select_from_string,
     ]:
         try:
             return func(string)
@@ -2595,6 +2685,7 @@ def op_from_string(string):
 
 
 unary.from_string = unary_from_string
+select.from_string = select_from_string
 binary.from_string = binary_from_string
 monoid.from_string = monoid_from_string
 semiring.from_string = semiring_from_string