From 24aeeca6f48c0833038cf1ac86349e5e9ffa4c44 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 30 May 2022 22:31:19 -0500 Subject: [PATCH 1/2] Recipe to perform select with a mask (or bool) argument. This allows e.g. `A.select(B.S)`. --- graphblas/matrix.py | 32 +++++++++++++++++++++----- graphblas/tests/test_matrix.py | 12 ++++++++++ graphblas/tests/test_vector.py | 27 ++++++++++++++++++++++ graphblas/vector.py | 41 ++++++++++++++++++++++++++++++---- 4 files changed, 103 insertions(+), 9 deletions(-) diff --git a/graphblas/matrix.py b/graphblas/matrix.py index 83b44c827..7e93c9f8a 100644 --- a/graphblas/matrix.py +++ b/graphblas/matrix.py @@ -5,12 +5,12 @@ from . import _automethods, backend, binary, ffi, lib, monoid, select, semiring, utils from ._ss.matrix import ss -from .base import BaseExpression, BaseType, call +from .base import BaseExpression, BaseType, _check_mask, call from .dtypes import _INDEX, FP64, lookup_dtype, unify from .exceptions import DimensionMismatch, NoValue, check_status from .expr import AmbiguousAssignOrExtract, IndexerResolver, Updater -from .mask import StructuralMask, ValueMask -from .operator import find_opclass, get_semiring, get_typed_op, op_from_string +from .mask import Mask, StructuralMask, ValueMask +from .operator import UNKNOWN_OPCLASS, find_opclass, get_semiring, get_typed_op, op_from_string from .scalar import ( _MATERIALIZE, Scalar, @@ -28,7 +28,7 @@ values_to_numpy_buffer, wrapdoc, ) -from .vector import Vector, VectorExpression, VectorIndexExpr +from .vector import Vector, VectorExpression, VectorIndexExpr, _select_mask ffi_new = ffi.new @@ -891,9 +891,31 @@ def select(self, op, thunk=None): Compute SelectOp at each element of the calling Matrix, keeping elements which return True. """ + method_name = "select" if isinstance(op, str): op = select.from_string(op) - method_name = "select" + 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)`? + mask = _check_mask(op) + if thunk is not None: + raise TypeError( + "thunk argument not None when calling select with mask or boolean object" + ) + self._expect_type(mask.parent, (Matrix, Mask), within=method_name, argname="op") + return MatrixExpression( + "select", + None, + [self, mask, _select_mask, (self, mask)], # [*expr_args, func, args] + expr_repr="{0.name}.select({1.name})", + nrows=self.nrows, + ncols=self.ncols, + dtype=self.dtype, + ) + if thunk is None: thunk = False # most basic form of 0 when unifying dtypes if type(thunk) is not Scalar: diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 94fe7b754..801932d0f 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1140,6 +1140,18 @@ def test_select(A): assert w7.isequal(Aupper) with pytest.raises(TypeError, match="thunk"): A.select(select.valueeq, object()) + # Select with boolean and masks + w8 = A.select((A == 3).new()).new() + assert w8.isequal(A3) + w9 = A.select(w8.S).new() + assert w9.isequal(A3) + w8[0, 1] = 0 + w10 = A.select(w8.V).new() + assert w10.isequal(A3) + with pytest.raises(TypeError, match="thunk"): + A.select(w8.V, 777) + with pytest.raises(TypeError): + A.select(A[0, :].new().S) @pytest.mark.slow diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 5f928621f..526580555 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -688,6 +688,33 @@ def test_select(v): assert w6.isequal(result) with pytest.raises(TypeError, match="thunk"): v.select(select.valueeq, object()) + # Select with boolean and masks + w7 = v.select((v == 1).new()).new() + assert w7.isequal(result) + w8 = v.select(w7.S).new() + assert w8.isequal(result) + w7[4] = 0 + w9 = v.select(w7.V).new() + assert w9.isequal(result) + with pytest.raises(TypeError, match="thunk"): + v.select(w7.V, 777) + with pytest.raises(TypeError): + v.select(v.outer(v).new().S) + # Make sure we use `replace` + w9 << 1 + w9 << v.select(w7.V) + assert w9.isequal(result) + # and with masks + w9 << 1 + w9[1] = 0 + w9(w9.V) << v.select(w7.V) + result2 = Vector.from_values([1, 3], [0, 1], size=7) + assert w9.isequal(result2) + w9 << 1 + w9[1] = 0 + w9(w9.V, replace=True) << v.select(w7.V) + del result2[1] + assert w9.isequal(result2) @pytest.mark.slow diff --git a/graphblas/vector.py b/graphblas/vector.py index 21d4591dd..58ecc9901 100644 --- a/graphblas/vector.py +++ b/graphblas/vector.py @@ -5,12 +5,12 @@ from . import _automethods, backend, binary, ffi, lib, monoid, select, semiring, utils from ._ss.vector import ss -from .base import BaseExpression, BaseType, call +from .base import BaseExpression, BaseType, _check_mask, call from .dtypes import _INDEX, FP64, INT64, lookup_dtype, unify from .exceptions import DimensionMismatch, NoValue, check_status from .expr import AmbiguousAssignOrExtract, IndexerResolver, Updater -from .mask import StructuralMask, ValueMask -from .operator import find_opclass, get_semiring, get_typed_op, op_from_string +from .mask import Mask, StructuralMask, ValueMask +from .operator import UNKNOWN_OPCLASS, find_opclass, get_semiring, get_typed_op, op_from_string from .scalar import ( _MATERIALIZE, Scalar, @@ -55,6 +55,19 @@ def _reposition(updater, indices, chunk): updater[indices] = chunk +def _select_mask(updater, obj, mask): + if updater.kwargs.get("mask") is None: + orig_kwargs = updater.kwargs + try: + updater.kwargs = dict(orig_kwargs, mask=mask, replace=True) + updater << obj + finally: + updater.kwargs = orig_kwargs + else: + # Can we do any better depending on accum, replace, and type of masks? + updater << obj.dup(mask=mask) + + class Vector(BaseType): """ GraphBLAS Sparse Vector @@ -807,9 +820,29 @@ def select(self, op, thunk=None): Compute SelectOp at each element of the calling Vector, keeping elements which return True. """ + method_name = "select" if isinstance(op, str): op = select.from_string(op) - method_name = "select" + 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)`. + mask = _check_mask(op) + if thunk is not None: + raise TypeError( + "thunk argument not None when calling select with mask or boolean object" + ) + self._expect_type(mask.parent, (Vector, Mask), within=method_name, argname="op") + return VectorExpression( + "select", + None, + [self, mask, _select_mask, (self, mask)], # [*expr_args, func, args] + expr_repr="{0.name}.select({1.name})", + size=self.size, + dtype=self.dtype, + ) + if thunk is None: thunk = False # most basic form of 0 when unifying dtypes if type(thunk) is not Scalar: From 27019d900654b1cfb3ff737eb1a2ee494573dbbb Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 31 May 2022 10:07:13 -0500 Subject: [PATCH 2/2] Fix select-mask with accum with no mask. --- graphblas/tests/test_vector.py | 12 ++++++++++++ graphblas/vector.py | 5 ++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 526580555..a838c090c 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -715,6 +715,18 @@ def test_select(v): w9(w9.V, replace=True) << v.select(w7.V) del result2[1] assert w9.isequal(result2) + w9 << 1 + w9[1] = 0 + w9(w9.V, binary.plus) << v.select(w7.V) + w8 << 1 + w8[1] = 0 + w8(w8.V, binary.plus) << v.dup(mask=w7.V) + assert w8.isequal(w9) + w9 << 1 + w9(binary.plus) << v.select(w7.V) + w8 << 1 + w8(binary.plus) << v.dup(mask=w7.V) + assert w8.isequal(w9) @pytest.mark.slow diff --git a/graphblas/vector.py b/graphblas/vector.py index 58ecc9901..b1dd9e1d9 100644 --- a/graphblas/vector.py +++ b/graphblas/vector.py @@ -59,7 +59,10 @@ def _select_mask(updater, obj, mask): if updater.kwargs.get("mask") is None: orig_kwargs = updater.kwargs try: - updater.kwargs = dict(orig_kwargs, mask=mask, replace=True) + if updater.kwargs.get("accum") is None: + updater.kwargs = dict(orig_kwargs, mask=mask, replace=True) + else: + updater.kwargs = dict(orig_kwargs, mask=mask) updater << obj finally: updater.kwargs = orig_kwargs