diff --git a/graphblas/__init__.py b/graphblas/__init__.py index b0e46a995..352c80857 100644 --- 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 index b0c4f8a83..fc3c95ccd 100644 --- 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 index 01fbf927e..06ae6c1cf 100644 --- 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 index 4406e01c0..d24ae62dc 100644 --- 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 index 215b45ea6..3520a2f11 100644 --- 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 diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py new file mode 100644 index 000000000..126ccf366 --- /dev/null +++ b/graphblas/select/__init__.py @@ -0,0 +1,93 @@ +# All items are dynamically added by classes in operator.py +# This module acts as a container of SelectOp instances (which are +# identical to IndexUnaryOp under the hood, except they must return +# boolean and they default to `select` when called, rather than +# `apply`). +_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}") + + +def _resolve_expr(expr, callname, opname): + from graphblas.matrix import MatrixExpression + from graphblas.vector import VectorExpression + + if not isinstance(expr, (VectorExpression, MatrixExpression)): + raise TypeError( + f"Expected VectorExpression or MatrixExpression; found {type(expr)}\n" + f"Typical usage: select.{callname}(x <= 5)" + ) + tensor = expr.args[0] + thunk = expr.args[1] + method = f"{opname}{expr.op.name}" + if method not in globals(): + # TODO: remove this once rowlt/rowge/collt/colge exist + # Convert thunk to Python int to avoid possible subtraction with uints + thunk = thunk.value + # Attempt to convert < into <= (rowlt is not part of official spec, but rowle is) + if expr.op.name == "lt": + method = f"{opname}le" + thunk -= 1 + # Attempt to convert >= into > (rowge is not part of official spec, but rowgt is) + elif expr.op.name == "ge": + method = f"{opname}gt" + thunk -= 1 + if method not in globals(): + raise ValueError(f"Unknown or unregistered select method: {method}") + return globals()[method](tensor, thunk) + + +def value(expr): + """ + An advanced select method which allows for easily expressing + value comparison logic. + + Example usage: + >>> gb.select.value(A > 0) + + The example will dispatch to `gb.select.valuegt(A, 0)` + while being nicer to read. + """ + return _resolve_expr(expr, "value", "value") + + +def row(expr): + """ + An advanced select method which allows for easily expressing + row index comparison logic. + + Example usage: + >>> gb.select.row(A <= 5) + + The example will dispatch to `gb.select.rowle(A, 5)` + while being potentially nicer to read. + """ + return _resolve_expr(expr, "row", "row") + + +def column(expr): + """ + An advanced select method which allows for easily expressing + column index comparison logic. + + Example usage: + >>> gb.select.column(A <= 5) + + The example will dispatch to `gb.select.colle(A, 5)` + while being potentially nicer to read. + """ + return _resolve_expr(expr, "column", "col") diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index e0592e3c0..ab12bd429 100644 --- 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, semiring, unary +from graphblas import agg, binary, dtypes, monoid, select, semiring, unary from graphblas.exceptions import ( DimensionMismatch, EmptyObject, @@ -1082,6 +1082,30 @@ def test_apply_binary(A): assert w1.isequal(w3) +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() + w2 = A.select("==", 3).new() + w3 = select.value(A == 3).new() + assert w1.isequal(A3) + assert w2.isequal(A3) + assert w3.isequal(A3) + + A2cols = Matrix.from_values([3, 0, 3, 5, 6], [0, 1, 2, 2, 2], [3, 2, 3, 1, 5], nrows=7, ncols=7) + w4 = select.colle(A, 2).new() + w5 = A.select("col<=", 2).new() + w6 = select.column(A < 3).new() + assert w4.isequal(A2cols) + assert w5.isequal(A2cols) + assert w6.isequal(A2cols) + + Aupper = Matrix.from_values( + [0, 0, 1, 2, 4, 1], [1, 3, 4, 5, 5, 6], [2, 3, 8, 1, 7, 4], nrows=7, ncols=7 + ) + w7 = A.select("TRIU").new() + assert w7.isequal(Aupper) + + def test_reduce_row(A): result = Vector.from_values([0, 1, 2, 3, 4, 5, 6], [5, 12, 1, 6, 7, 1, 15]) w = A.reduce_rowwise(monoid.plus).new() @@ -3230,6 +3254,9 @@ def test_udt(): result = unary.positioni(A).new() expected = Matrix.from_values([0, 0, 1, 1], [0, 1, 0, 1], [0, 0, 1, 1]) assert result.isequal(expected) + AB = unary.one(select.tril(A).new()).new() + BA = select.tril(unary.one(A).new()).new() + assert AB.isequal(BA) # Just make sure these work for aggop in [agg.any_value, agg.first, agg.last, agg.count]: diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index dfde54cab..347919e79 100644 --- 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, semiring, unary +from graphblas import agg, binary, dtypes, monoid, select, semiring, unary from graphblas.exceptions import ( DimensionMismatch, EmptyObject, @@ -647,6 +647,22 @@ def test_apply_empty(v): v.apply(binary.plus, s).new() +def test_select(v): + result = Vector.from_values([1, 3], [1, 1], size=7) + w1 = v.select(select.valueeq, 1).new() + w2 = v.select("==", 1).new() + w3 = select.rowle(v, 3).new() + w4 = v.select("row<=", 3).new() + w5 = select.value(v == 1).new() + w6 = select.row(v < 4).new() + assert w1.isequal(result) + assert w2.isequal(result) + assert w3.isequal(result) + assert w4.isequal(result) + assert w5.isequal(result) + assert w6.isequal(result) + + def test_reduce(v): s = v.reduce(monoid.plus).new() assert s == 4 diff --git a/graphblas/vector.py b/graphblas/vector.py index 69eda3cae..7ecbe92cf 100644 --- a/graphblas/vector.py +++ b/graphblas/vector.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.vector import ss from .base import BaseExpression, BaseType, call from .dtypes import _INDEX, FP64, INT64, lookup_dtype, unify @@ -708,6 +708,51 @@ def apply(self, op, right=None, *, left=None): size=self._size, ) + def select(self, op, thunk=None): + """ + GrB_Vector_select + Compute SelectOp at each element of the calling Vector, 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_Vector_select_{dtype_name}" + else: + cfunc_name = "GrB_Vector_select_Scalar" + return VectorExpression( + method_name, + cfunc_name, + [self, thunk], + op=op, + expr_repr="{0.name}.select({op}, thunk={1._expr_name})", + size=self._size, + dtype=self.dtype, + ) + def reduce(self, op=monoid.plus, *, allow_empty=True): """ GrB_Vector_reduce @@ -1067,6 +1112,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)) @@ -1137,6 +1183,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))