diff --git a/graphblas/matrix.py b/graphblas/matrix.py index 7e93c9f8a..5c338cc65 100644 --- a/graphblas/matrix.py +++ b/graphblas/matrix.py @@ -895,11 +895,20 @@ def select(self, op, thunk=None): if isinstance(op, str): op = select.from_string(op) else: - op, opclass = find_opclass(op) + if isinstance(op, MatrixExpression): + # Try to rewrite e.g. `A.select(A == 7)` to `gb.select.value(A == 7)` + if thunk is not None: + raise TypeError( + "thunk argument not None when calling select with mask or boolean object" + ) + expr = select._match_expr(self, op) + if expr is not None: + return expr + opclass = UNKNOWN_OPCLASS + else: + op, opclass = find_opclass(op) if opclass == UNKNOWN_OPCLASS: - # Optimization opportunity: we could be smarter and change e.g. `A.select(A == 7)` - # to `gb.select.value(A == 7)` or `A.select(gb.select.valueeq, 7)`. - # Should we allow e.g., `A.select(B.T)`? + # e.g., `A.select(B.S)` or `A.select(B < 7)` mask = _check_mask(op) if thunk is not None: raise TypeError( diff --git a/graphblas/operator.py b/graphblas/operator.py index 6cea8be6e..9e04f95eb 100644 --- a/graphblas/operator.py +++ b/graphblas/operator.py @@ -3068,6 +3068,23 @@ def get_semiring(monoid, binaryop, name=None): semiring.register_new = Semiring.register_new semiring.register_anonymous = Semiring.register_anonymous +select._binary_to_select.update( + { + binary.eq: select.valueeq, + binary.ne: select.valuene, + binary.le: select.valuele, + binary.lt: select.valuelt, + binary.ge: select.valuege, + binary.gt: select.valuegt, + binary.iseq: select.valueeq, + binary.isne: select.valuene, + binary.isle: select.valuele, + binary.islt: select.valuelt, + binary.isge: select.valuege, + binary.isgt: select.valuegt, + } +) + _str_to_unary = { "-": unary.ainv, "~": unary.lnot, diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py index 1d925b14e..581b2b508 100644 --- a/graphblas/select/__init__.py +++ b/graphblas/select/__init__.py @@ -4,6 +4,7 @@ # boolean and they default to `select` when called, rather than # `apply`). _delayed = {} +_binary_to_select = {} def __dir__(): @@ -48,6 +49,25 @@ def _resolve_expr(expr, callname, opname): return globals()[method](tensor, thunk) +def _match_expr(parent, expr): + """Match expressions to rewrite `A.select(A < 5)` into select expression + + The argument must match the parent, so this _won't_ be rewritten: `A.select(B < 5)` + """ + args = expr.args + op = expr.op + if ( + expr.method_name == "apply" + and len(args) == 2 + and op.parent in _binary_to_select + and args[1]._is_scalar + and (args[0] is parent or getattr(args[0], "T", None) is parent) + and op.type in _binary_to_select[op.parent] + ): + selectop = _binary_to_select[op.parent][op.type] + return selectop(parent, args[1]) + + def value(expr): """ An advanced select method which allows for easily expressing diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 801932d0f..a88320d83 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1140,9 +1140,18 @@ def test_select(A): assert w7.isequal(Aupper) with pytest.raises(TypeError, match="thunk"): A.select(select.valueeq, object()) + + +@autocompute +def test_select_bools_and_masks(A): + A3 = Matrix.from_values([0, 3, 3, 6], [3, 0, 2, 4], [3, 3, 3, 3], nrows=7, ncols=7) # Select with boolean and masks w8 = A.select((A == 3).new()).new() assert w8.isequal(A3) + w8b = A.select(A == 3).new() # we rewrite! + assert w8b.isequal(A3) + w8c = A.select(~(A != 3)).new() + assert w8c.isequal(A3) w9 = A.select(w8.S).new() assert w9.isequal(A3) w8[0, 1] = 0 @@ -1150,6 +1159,8 @@ def test_select(A): assert w10.isequal(A3) with pytest.raises(TypeError, match="thunk"): A.select(w8.V, 777) + with pytest.raises(TypeError, match="thunk"): + A.select(A == 3, 777) with pytest.raises(TypeError): A.select(A[0, :].new().S) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index a838c090c..14097ec39 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -688,9 +688,18 @@ def test_select(v): assert w6.isequal(result) with pytest.raises(TypeError, match="thunk"): v.select(select.valueeq, object()) + + +@autocompute +def test_select_bools_and_masks(v): # Select with boolean and masks + result = Vector.from_values([1, 3], [1, 1], size=7) w7 = v.select((v == 1).new()).new() assert w7.isequal(result) + w7b = v.select(v == 1).new() # we rewrite! + assert w7b.isequal(result) + w7c = v.select(~(v != 1)).new() + assert w7c.isequal(result) w8 = v.select(w7.S).new() assert w8.isequal(result) w7[4] = 0 @@ -698,6 +707,8 @@ def test_select(v): assert w9.isequal(result) with pytest.raises(TypeError, match="thunk"): v.select(w7.V, 777) + with pytest.raises(TypeError, match="thunk"): + v.select(v == 1, 777) with pytest.raises(TypeError): v.select(v.outer(v).new().S) # Make sure we use `replace` diff --git a/graphblas/vector.py b/graphblas/vector.py index b1dd9e1d9..5d9db3c40 100644 --- a/graphblas/vector.py +++ b/graphblas/vector.py @@ -827,10 +827,20 @@ def select(self, op, thunk=None): if isinstance(op, str): op = select.from_string(op) else: - op, opclass = find_opclass(op) + if isinstance(op, VectorExpression): + # Try to rewrite e.g. `v.select(v == 7)` to `gb.select.value(v == 7)` + if thunk is not None: + raise TypeError( + "thunk argument not None when calling select with mask or boolean object" + ) + expr = select._match_expr(self, op) + if expr is not None: + return expr + opclass = UNKNOWN_OPCLASS + else: + op, opclass = find_opclass(op) if opclass == UNKNOWN_OPCLASS: - # Optimization opportunity: we could be smarter and change e.g. `v.select(v == 7)` - # to `gb.select.value(v == 7)` or `v.select(gb.select.valueeq, 7)`. + # e.g., `v.select(w.S)` or `v.select(w < 7)` mask = _check_mask(op) if thunk is not None: raise TypeError(