From c48095b611525e00b8d8992ac8729fd7b1d4bc70 Mon Sep 17 00:00:00 2001 From: Jim Kitchen Date: Mon, 11 Apr 2022 15:14:31 -0500 Subject: [PATCH 1/6] select operation Two ways of calling: 1. A.select(op, thunk) 2. gb.select.op(A, thunk) When calling `A.select(op)`, the op may be passed as a string. An additional set of "special" select functions take a comparator expression. An example: `gb.select.row(A <= 2)`. Row, column, and value are included, indicating what part of A is being compared in the select statement. --- grblas/__init__.py | 1 + grblas/_automethods.py | 4 ++ grblas/infix.py | 2 + grblas/matrix.py | 36 +++++++++++++ grblas/operator.py | 100 +++++++++++++++++++++++++++++++++--- grblas/select/__init__.py | 56 ++++++++++++++++++++ grblas/tests/test_matrix.py | 26 +++++++++- grblas/tests/test_vector.py | 18 ++++++- grblas/vector.py | 34 ++++++++++++ 9 files changed, 267 insertions(+), 10 deletions(-) create mode 100644 grblas/select/__init__.py diff --git a/grblas/__init__.py b/grblas/__init__.py index 16da3c332..f45590cbe 100644 --- a/grblas/__init__.py +++ b/grblas/__init__.py @@ -68,6 +68,7 @@ def get_config(): "operator", "recorder", "scalar", + "select", "semiring", "ss", "tests", diff --git a/grblas/_automethods.py b/grblas/_automethods.py index d8974c41c..f743a04cb 100644 --- a/grblas/_automethods.py +++ b/grblas/_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") diff --git a/grblas/infix.py b/grblas/infix.py index 0f16489ba..292b7740d 100644 --- a/grblas/infix.py +++ b/grblas/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)) @@ -201,6 +202,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/grblas/matrix.py b/grblas/matrix.py index a25b3855f..f798520be 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -731,6 +731,39 @@ def apply(self, op, right=None, *, left=None): bt=self._is_transposed, ) + def select(self, op, thunk=0): + """ + GrB_Matrix_select + Compute IndexUnaryOp at each element of the calling Matrix, keeping + elements which return True. + """ + method_name = "select" + op = get_typed_op(op, self.dtype, kind="indexunary") + self._expect_op(op, "IndexUnaryOp", within=method_name, argname="op") + if type(thunk) is not Scalar: + try: + thunk = Scalar.from_value(thunk, 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, + ) + cfunc_name = f"GrB_Matrix_select_{self.dtype}" + 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, + ) + def reduce_rowwise(self, op=monoid.plus): """ GrB_Matrix_reduce @@ -1347,6 +1380,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)) @@ -1425,6 +1459,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)) @@ -1527,6 +1562,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/grblas/operator.py b/grblas/operator.py index dc23b6bb9..57e73773e 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -17,7 +17,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" @@ -45,8 +45,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) @@ -63,12 +63,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( @@ -128,6 +131,18 @@ def __call__(self, val): ) +class TypedBuiltinIndexUnaryOp(TypedOpBase): + __slots__ = () + opclass = "IndexUnaryOp" + + def __call__(self, val, thunk=None): + if isinstance(val, InfixExprBase): + raise NotImplementedError("select with infix coming soon...") + if thunk is None: + thunk = 0 + return _call_op(self, val, thunk=thunk) + + class TypedBuiltinBinaryOp(TypedOpBase): __slots__ = () opclass = "BinaryOp" @@ -604,7 +619,7 @@ def _find(cls, funcname): return rv @classmethod - def _initialize(cls): + def _initialize(cls, merge_namespace=True): if cls._initialized: return # Read in the parse configs @@ -644,7 +659,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 merge_namespace and not hasattr(op, name): setattr(op, name, obj) else: obj = getattr(cls._module, name) @@ -894,6 +909,53 @@ def __reduce__(self): __call__ = TypedBuiltinUnaryOp.__call__ +class IndexUnaryOp(OpBase): + _module = select + _modname = "select" + _typed_class = TypedBuiltinIndexUnaryOp + _parse_config = { + "trim_from_front": 4, + "num_underscores": 1, + # "re_exprs": [ + # re.compile( + # "^GrB_(ROWINDEX|COLINDEX|DIAGINDEX)" + # "_(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)$"), + ], + } + + @classmethod + def _initialize(cls): + super()._initialize(merge_namespace=False) + # Update type information to handle more than bool + # fmt: off + bool_names = ("tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt") + types = ("INT8", "UINT8", "INT16", "UINT16", "INT32", "UINT32", "INT64", "UINT64", + "FP32", "FP64", "FC32", "FC64") + # fmt: on + 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 and (_supports_complex or dtype not in {"FC32", "FC64"}): + op.types[dtype] = output_type + op._typed_ops[dtype] = typed_op + op.coercions[dtype] = "BOOL" + # TODO: add "special" ops: row, column, index, value + cls._initialized = True + + __call__ = TypedBuiltinIndexUnaryOp.__call__ + + def _floordiv(x, y): return x // y # cast to integer @@ -1421,7 +1483,7 @@ def identities(self): @classmethod def _initialize(cls): - super()._initialize() + super()._initialize(merge_namespace=False) lor = monoid.lor._typed_ops["BOOL"] land = monoid.land._typed_ops["BOOL"] for cur_op, typed_op in [ @@ -1776,6 +1838,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 == "indexunary": + op = indexunary_from_string(op) elif kind == "binary": op = binary_from_string(op) elif kind == "monoid": @@ -1921,6 +1985,7 @@ def get_semiring(monoid, binaryop, name=None): # Now initialize all the things! try: UnaryOp._initialize() + IndexUnaryOp._initialize() BinaryOp._initialize() Monoid._initialize() Semiring._initialize() @@ -1935,6 +2000,18 @@ def get_semiring(monoid, binaryop, name=None): "-": unary.ainv, "~": unary.lnot, } +_str_to_indexunary = { + "<": 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, @@ -2015,6 +2092,10 @@ def unary_from_string(string): return _from_string(string, unary, _str_to_unary, "abs[int]") +def indexunary_from_string(string): + return _from_string(string, select, _str_to_indexunary, "tril") + + def binary_from_string(string): return _from_string(string, binary, _str_to_binary, "+[int]") @@ -2047,6 +2128,8 @@ def op_from_string(string): binary_from_string, monoid_from_string, semiring_from_string, + # Give binary precedence over this + indexunary_from_string, ]: try: return func(string) @@ -2056,6 +2139,7 @@ def op_from_string(string): unary.from_string = unary_from_string +select.from_string = indexunary_from_string binary.from_string = binary_from_string monoid.from_string = monoid_from_string semiring.from_string = semiring_from_string diff --git a/grblas/select/__init__.py b/grblas/select/__init__.py new file mode 100644 index 000000000..95d7c2316 --- /dev/null +++ b/grblas/select/__init__.py @@ -0,0 +1,56 @@ +# All items are dynamically added by classes in operator.py +# This module acts as a container of IndexUnaryOp instances +_delayed = {} +from grblas 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 grblas.base import BaseExpression + + if not isinstance(expr, BaseExpression) and hasattr(expr, "select"): + 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(): + # Attempt to convert < into <= (rowlt is not part of official spec, but rowle is) + if expr.op.name == "lt": + method = f"{opname}le" + thunk.value -= 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.value -= 1 + if method not in globals(): + raise ValueError(f"Unknown or unregistered select method: {method}") + return globals()[method](tensor, thunk) + + +def value(expr): + return _resolve_expr(expr, "value", "value") + + +def row(expr): + return _resolve_expr(expr, "row", "row") + + +def column(expr): + return _resolve_expr(expr, "column", "col") diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index 9149b2f5a..73f9175ab 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -9,7 +9,7 @@ from numpy.testing import assert_array_equal import grblas -from grblas import agg, binary, dtypes, monoid, semiring, unary +from grblas import agg, binary, dtypes, monoid, semiring, unary, select from grblas.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() diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 8cc512e3b..79e25b6b2 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -9,7 +9,7 @@ from numpy.testing import assert_array_equal import grblas -from grblas import agg, binary, dtypes, monoid, semiring, unary +from grblas import agg, binary, dtypes, monoid, semiring, unary, select from grblas.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/grblas/vector.py b/grblas/vector.py index 83cf1872a..0d95f59bd 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -636,6 +636,38 @@ def apply(self, op, right=None, *, left=None): size=self._size, ) + def select(self, op, thunk=0): + """ + GrB_Vector_select + Compute IndexUnaryOp at each element of the calling Vector, keeping + elements which return True. + """ + method_name = "select" + op = get_typed_op(op, self.dtype, kind="indexunary") + self._expect_op(op, "IndexUnaryOp", within=method_name, argname="op") + if type(thunk) is not Scalar: + try: + thunk = Scalar.from_value(thunk, 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, + ) + cfunc_name = f"GrB_Vector_select_{self.dtype}" + 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 @@ -975,6 +1007,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)) @@ -1045,6 +1078,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)) From 9f4f89171d3e5d18a145b08620024a4dfb903248 Mon Sep 17 00:00:00 2001 From: Jim Kitchen Date: Mon, 11 Apr 2022 16:24:38 -0500 Subject: [PATCH 2/6] Fix import sorting --- grblas/tests/test_matrix.py | 2 +- grblas/tests/test_vector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/grblas/tests/test_matrix.py b/grblas/tests/test_matrix.py index 821f9babf..53c6f2ce9 100644 --- a/grblas/tests/test_matrix.py +++ b/grblas/tests/test_matrix.py @@ -9,7 +9,7 @@ from numpy.testing import assert_array_equal import grblas -from grblas import agg, binary, dtypes, monoid, semiring, unary, select +from grblas import agg, binary, dtypes, monoid, select, semiring, unary from grblas.exceptions import ( DimensionMismatch, EmptyObject, diff --git a/grblas/tests/test_vector.py b/grblas/tests/test_vector.py index 4d3085e4b..f917ca83c 100644 --- a/grblas/tests/test_vector.py +++ b/grblas/tests/test_vector.py @@ -9,7 +9,7 @@ from numpy.testing import assert_array_equal import grblas -from grblas import agg, binary, dtypes, monoid, semiring, unary, select +from grblas import agg, binary, dtypes, monoid, select, semiring, unary from grblas.exceptions import ( DimensionMismatch, EmptyObject, From c1c8db8884ade2a1b14e450ef7d54572d39df975 Mon Sep 17 00:00:00 2001 From: Jim Kitchen Date: Tue, 12 Apr 2022 17:19:08 -0500 Subject: [PATCH 3/6] Fixes based on feedback - Use correct dtype based on thunk - Handle Scalar thunk properly - Rename IndexUnaryOp to SelectOp --- grblas/_automethods.py | 1 + grblas/matrix.py | 23 +++++++++++---- grblas/operator.py | 59 ++++++++++++++++++++++----------------- grblas/select/__init__.py | 11 ++++++-- grblas/vector.py | 23 +++++++++++---- 5 files changed, 76 insertions(+), 41 deletions(-) diff --git a/grblas/_automethods.py b/grblas/_automethods.py index f743a04cb..7dad668d3 100644 --- a/grblas/_automethods.py +++ b/grblas/_automethods.py @@ -353,6 +353,7 @@ def __ixor__(self, other): "ewise_add", "ewise_mult", "ewise_union", + "select", "ss", "to_values", } diff --git a/grblas/matrix.py b/grblas/matrix.py index 6d3d42b10..dcd9a1a96 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -785,18 +785,19 @@ def apply(self, op, right=None, *, left=None): bt=self._is_transposed, ) - def select(self, op, thunk=0): + def select(self, op, thunk=None): """ GrB_Matrix_select - Compute IndexUnaryOp at each element of the calling Matrix, keeping + Compute SelectOp at each element of the calling Matrix, keeping elements which return True. """ method_name = "select" - op = get_typed_op(op, self.dtype, kind="indexunary") - self._expect_op(op, "IndexUnaryOp", within=method_name, argname="op") + 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 else None try: - thunk = Scalar.from_value(thunk, is_cscalar=None, name="") + thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: thunk = self._expect_type( thunk, @@ -806,7 +807,17 @@ def select(self, op, thunk=0): extra_message="Literal scalars also accepted.", op=op, ) - cfunc_name = f"GrB_Matrix_select_{self.dtype}" + 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, diff --git a/grblas/operator.py b/grblas/operator.py index 316cdbdc0..5c772e647 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -189,15 +189,13 @@ def __call__(self, val): ) -class TypedBuiltinIndexUnaryOp(TypedOpBase): +class TypedBuiltinSelectOp(TypedOpBase): __slots__ = () - opclass = "IndexUnaryOp" + opclass = "SelectOp" def __call__(self, val, thunk=None): - if isinstance(val, InfixExprBase): - raise NotImplementedError("select with infix coming soon...") if thunk is None: - thunk = 0 + thunk = False # most basic form of 0 when unifying dtypes return _call_op(self, val, thunk=thunk) @@ -760,7 +758,11 @@ def _find(cls, funcname): return rv @classmethod - def _initialize(cls, merge_namespace=True): + 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 @@ -800,7 +802,7 @@ def _initialize(cls, merge_namespace=True): obj = cls(name, is_positional=name in cls._positional) setattr(cls._module, name, obj) _STANDARD_OPERATOR_NAMES.add(f"{cls._modname}.{name}") - if merge_namespace and not hasattr(op, name): + if include_in_ops and not hasattr(op, name): setattr(op, name, obj) else: obj = getattr(cls._module, name) @@ -1110,20 +1112,14 @@ def __reduce__(self): __call__ = TypedBuiltinUnaryOp.__call__ -class IndexUnaryOp(OpBase): +class SelectOp(OpBase): _module = select _modname = "select" _is_udt = False - _typed_class = TypedBuiltinIndexUnaryOp + _typed_class = TypedBuiltinSelectOp _parse_config = { "trim_from_front": 4, "num_underscores": 1, - # "re_exprs": [ - # re.compile( - # "^GrB_(ROWINDEX|COLINDEX|DIAGINDEX)" - # "_(INT32|INT64)$" - # ), - # ], "re_exprs_return_bool": [ re.compile("^GrB_(TRIL|TRIU|DIAG|OFFDIAG|COLLE|COLGT|ROWLE|ROWGT)$"), re.compile( @@ -1133,10 +1129,11 @@ class IndexUnaryOp(OpBase): re.compile("^GxB_(VALUEEQ|VALUENE)_(FC32|FC64)$"), ], } + _positional = {"tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt"} @classmethod def _initialize(cls): - super()._initialize(merge_namespace=False) + super()._initialize(include_in_ops=False) # Update type information to handle more than bool # fmt: off bool_names = ("tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt") @@ -1153,7 +1150,17 @@ def _initialize(cls): op.coercions[dtype] = BOOL cls._initialized = True - __call__ = TypedBuiltinIndexUnaryOp.__call__ + 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): @@ -1955,7 +1962,7 @@ def _is_udt(self): @classmethod def _initialize(cls): - super()._initialize(merge_namespace=False) + super()._initialize() lor = monoid.lor._typed_ops[BOOL] land = monoid.land._typed_ops[BOOL] for cur_op, typed_op in [ @@ -2380,8 +2387,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 == "indexunary": - op = indexunary_from_string(op) + elif kind == "select": + op = select_from_string(op) elif kind == "binary": op = binary_from_string(op) elif kind == "monoid": @@ -2527,7 +2534,7 @@ def get_semiring(monoid, binaryop, name=None): # Now initialize all the things! try: UnaryOp._initialize() - IndexUnaryOp._initialize() + SelectOp._initialize() BinaryOp._initialize() Monoid._initialize() Semiring._initialize() @@ -2542,7 +2549,7 @@ def get_semiring(monoid, binaryop, name=None): "-": unary.ainv, "~": unary.lnot, } -_str_to_indexunary = { +_str_to_select = { "<": select.valuelt, ">": select.valuegt, "<=": select.valuele, @@ -2634,8 +2641,8 @@ def unary_from_string(string): return _from_string(string, unary, _str_to_unary, "abs[int]") -def indexunary_from_string(string): - return _from_string(string, select, _str_to_indexunary, "tril") +def select_from_string(string): + return _from_string(string, select, _str_to_select, "tril") def binary_from_string(string): @@ -2671,7 +2678,7 @@ def op_from_string(string): monoid_from_string, semiring_from_string, # Give binary precedence over this - indexunary_from_string, + select_from_string, ]: try: return func(string) @@ -2681,7 +2688,7 @@ def op_from_string(string): unary.from_string = unary_from_string -select.from_string = indexunary_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/grblas/select/__init__.py b/grblas/select/__init__.py index 95d7c2316..99c688872 100644 --- a/grblas/select/__init__.py +++ b/grblas/select/__init__.py @@ -1,5 +1,8 @@ # All items are dynamically added by classes in operator.py -# This module acts as a container of IndexUnaryOp instances +# 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 grblas import operator # noqa isort:skip @@ -31,14 +34,16 @@ def _resolve_expr(expr, callname, opname): thunk = expr.args[1] method = f"{opname}{expr.op.name}" if method not in globals(): + # Convert thunk to Python int to avoid possible subtraction with uints + thunk = int(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.value -= 1 + 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.value -= 1 + thunk -= 1 if method not in globals(): raise ValueError(f"Unknown or unregistered select method: {method}") return globals()[method](tensor, thunk) diff --git a/grblas/vector.py b/grblas/vector.py index 14da60bec..259267d88 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -698,18 +698,19 @@ def apply(self, op, right=None, *, left=None): size=self._size, ) - def select(self, op, thunk=0): + def select(self, op, thunk=None): """ GrB_Vector_select - Compute IndexUnaryOp at each element of the calling Vector, keeping + Compute SelectOp at each element of the calling Vector, keeping elements which return True. """ method_name = "select" - op = get_typed_op(op, self.dtype, kind="indexunary") - self._expect_op(op, "IndexUnaryOp", within=method_name, argname="op") + 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 else None try: - thunk = Scalar.from_value(thunk, is_cscalar=None, name="") + thunk = Scalar.from_value(thunk, dtype, is_cscalar=None, name="") except TypeError: thunk = self._expect_type( thunk, @@ -719,7 +720,17 @@ def select(self, op, thunk=0): extra_message="Literal scalars also accepted.", op=op, ) - cfunc_name = f"GrB_Vector_select_{self.dtype}" + 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, From e3a4cfe55d68ff2112e8e16b44eac92e95b19a4f Mon Sep 17 00:00:00 2001 From: Jim Kitchen Date: Tue, 12 Apr 2022 18:02:40 -0500 Subject: [PATCH 4/6] Fix for Windows --- graphblas/operator.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/graphblas/operator.py b/graphblas/operator.py index 881775d88..01e7186ef 100644 --- a/graphblas/operator.py +++ b/graphblas/operator.py @@ -1132,16 +1132,16 @@ class SelectOp(OpBase): def _initialize(cls): super()._initialize(include_in_ops=False) # Update type information to handle more than bool - # fmt: off bool_names = ("tril", "triu", "diag", "offdiag", "colle", "colgt", "rowle", "rowgt") - types = (INT8, UINT8, INT16, UINT16, INT32, UINT32, INT64, UINT64, FP32, FP64, FC32, FC64) - # fmt: on + 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 and (_supports_complex or dtype not in {FC32, FC64}): + if dtype not in op.types: op.types[dtype] = output_type op._typed_ops[dtype] = typed_op op.coercions[dtype] = BOOL From 2439109e3febef3bad0645ccf564e8c4c6046b81 Mon Sep 17 00:00:00 2001 From: Jim Kitchen Date: Wed, 13 Apr 2022 14:38:17 -0500 Subject: [PATCH 5/6] Better handling of op.is_positional --- graphblas/matrix.py | 6 ++++-- graphblas/vector.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/graphblas/matrix.py b/graphblas/matrix.py index 87f2d94d9..a79d3e86d 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 @@ -797,11 +797,13 @@ def select(self, op, thunk=None): 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 else None + 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: diff --git a/graphblas/vector.py b/graphblas/vector.py index 3bd7fb43c..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 @@ -714,11 +714,13 @@ def select(self, op, thunk=None): 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 else None + 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: From 2db16967a24ba3388cecb6d00ab69be0bf5f0b9e Mon Sep 17 00:00:00 2001 From: Jim Kitchen Date: Thu, 14 Apr 2022 15:21:22 -0500 Subject: [PATCH 6/6] Updates based on feedback --- graphblas/matrix.py | 1 + graphblas/operator.py | 1 + graphblas/select/__init__.py | 38 +++++++++++++++++++++++++++++++--- graphblas/tests/test_matrix.py | 3 +++ 4 files changed, 40 insertions(+), 3 deletions(-) diff --git a/graphblas/matrix.py b/graphblas/matrix.py index a79d3e86d..d24ae62dc 100644 --- a/graphblas/matrix.py +++ b/graphblas/matrix.py @@ -835,6 +835,7 @@ def select(self, op, thunk=None): nrows=self._nrows, ncols=self._ncols, dtype=self.dtype, + at=self._is_transposed, ) def reduce_rowwise(self, op=monoid.plus): diff --git a/graphblas/operator.py b/graphblas/operator.py index a837a06c4..3520a2f11 100644 --- a/graphblas/operator.py +++ b/graphblas/operator.py @@ -1110,6 +1110,7 @@ def __reduce__(self): class SelectOp(OpBase): + __slots__ = ("is_positional",) _module = select _modname = "select" _is_udt = False diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py index 78380fb6a..126ccf366 100644 --- a/graphblas/select/__init__.py +++ b/graphblas/select/__init__.py @@ -23,9 +23,10 @@ def __getattr__(key): def _resolve_expr(expr, callname, opname): - from graphblas.base import BaseExpression + from graphblas.matrix import MatrixExpression + from graphblas.vector import VectorExpression - if not isinstance(expr, BaseExpression) and hasattr(expr, "select"): + if not isinstance(expr, (VectorExpression, MatrixExpression)): raise TypeError( f"Expected VectorExpression or MatrixExpression; found {type(expr)}\n" f"Typical usage: select.{callname}(x <= 5)" @@ -34,8 +35,9 @@ def _resolve_expr(expr, callname, opname): 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 = int(thunk.value) + 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" @@ -50,12 +52,42 @@ def _resolve_expr(expr, callname, opname): 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 bdb737194..ab12bd429 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -3254,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]: