From 995c02266c772fc87be7bf35aded8dbc914eac2d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 4 Jul 2022 14:12:18 -0500 Subject: [PATCH 1/3] Allow masks to be combined with `|` and`&` Closes #234 --- .pre-commit-config.yaml | 6 +- graphblas/base.py | 2 +- graphblas/infix.py | 7 +- graphblas/mask.py | 141 ++++++++++++++++++++++++++++++++++- graphblas/tests/test_mask.py | 71 ++++++++++++++++++ 5 files changed, 221 insertions(+), 6 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7acc79bdc..bc599f7eb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ # To run: `pre-commit run --all-files` repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.2.0 + rev: v4.3.0 hooks: - id: check-added-large-files - id: check-yaml @@ -23,12 +23,12 @@ repos: - id: isort language_version: python3 - repo: https://github.com/asottile/pyupgrade - rev: v2.32.0 + rev: v2.34.0 hooks: - id: pyupgrade args: [--py38-plus] - repo: https://github.com/psf/black - rev: 22.3.0 + rev: 22.6.0 hooks: - id: black language_version: python3 diff --git a/graphblas/base.py b/graphblas/base.py index 491c9b31f..1ebf820ad 100644 --- a/graphblas/base.py +++ b/graphblas/base.py @@ -282,7 +282,7 @@ def __rand__(self, other): return NotImplemented from .infix import _ewise_infix_expr - return _ewise_infix_expr(self, other, method="ewise_mult", within="__rand__") + return _ewise_infix_expr(other, self, method="ewise_mult", within="__rand__") def __matmul__(self, other): if self._is_scalar: diff --git a/graphblas/infix.py b/graphblas/infix.py index 746b25c41..d8b5a6f33 100644 --- a/graphblas/infix.py +++ b/graphblas/infix.py @@ -2,6 +2,7 @@ from .base import _expect_op, _expect_type from .dtypes import BOOL from .expr import InfixExprBase +from .mask import Mask from .matrix import Matrix, MatrixExpression, TransposedMatrix from .monoid import land, lor from .scalar import Scalar, ScalarExpression @@ -352,7 +353,11 @@ def _ewise_infix_expr(left, right, *, method, within): elif method == "ewise_mult": return MatrixEwiseMultExpr(left, right) return MatrixEwiseAddExpr(left, right) - if left_type in types: + if within == "__or__" and isinstance(right, Mask): + return right.__ror__(left) + elif within == "__and__" and isinstance(right, Mask): + return right.__rand__(left) + elif left_type in types: left._expect_type(right, tuple(types), within=within, argname="right") elif right_type in types: right._expect_type(left, tuple(types), within=within, argname="left") diff --git a/graphblas/mask.py b/graphblas/mask.py index 949c8d560..cb5b5356c 100644 --- a/graphblas/mask.py +++ b/graphblas/mask.py @@ -1,6 +1,6 @@ import warnings -from . import utils +from . import monoid, utils from .binary import land, lor, pair from .dtypes import BOOL from .select import valuene @@ -89,6 +89,54 @@ def new(self, dtype=None, *, complement=False, mask=None, name=None): func = d[type(self), type(mask)] return func(self, mask, dtype, name) + def __and__(self, other): + """Return the intersection of two masks as a new mask. + + `new_mask = mask1 & mask2` is equivalent to the following: + + >>> val = Matrix(bool, nrows, ncols) + >>> val(mask1) << True + >>> val(mask2, replace=True) << True + >>> new_mask = val.S + + This uses faster recipes than the above for all combinations of input mask types, + and aims to be memory efficient when operating on complemented masks. + """ + from .base import _check_mask + + other = _check_mask(other) + complement = self.complement or other.complement + d = _COMPLEMENT_MASKS if complement else _COMBINE_MASKS + func = d[type(self), type(other)] + val = func(self, other, bool, None) + if complement: + return ComplementedStructuralMask(val) + else: + return StructuralMask(val) + + __rand__ = __and__ + + def __or__(self, other): + """Return the union of two masks as a new mask. + + `new_mask = mask1 | mask2` is equivalent to the following: + + >>> val = Matrix(bool, nrows, ncols) + >>> val(mask1) << True + >>> val(mask2) << True + >>> new_mask = val.S + + This uses faster recipes than the above for all combinations of input mask types, + and aims to be memory efficient when operating on complemented masks. + """ + from .base import _check_mask + + other = _check_mask(other) + func = _MASK_OR[type(self), type(other)] + return func(self, other) + + __ror__ = __or__ + class StructuralMask(Mask): __slots__ = () @@ -371,6 +419,97 @@ def _complement_A_CV(m1, m2, dtype, name): # CV-CV (ComplementedValueMask, ComplementedValueMask): _complement_CV_CV, } + + +def _combine_S_S_mask_or(m1, m2): + """S-S""" + val = monoid.any(one(m1.parent).new(bool) | one(m2.parent).new(bool)).new() + return StructuralMask(val) + + +def _combine_S_SV_mask_or(m1, m2): + """S-V""" + val = monoid.any(one(m1.parent).new(bool) | one(m2.parent).new(bool, mask=m2)).new() + return StructuralMask(val) + + +def _combine_SV_S_mask_or(m1, m2): + """V-S""" + val = monoid.any(one(m1.parent).new(bool, mask=m1) | one(m2.parent).new(bool)).new() + return StructuralMask(val) + + +def _complement_A_CS_mask_or(m1, m2): + """~S-CS, ~V-CS, ~CV-CS""" + val = one(m2.parent).new(bool, mask=~m1) + return ComplementedStructuralMask(val) + + +def _complement_CS_A_mask_or(m1, m2): + """~CS-S, ~CS-V, ~CS-CV""" + val = one(m1.parent).new(bool, mask=~m2) + return ComplementedStructuralMask(val) + + +def _complement_A_CV_mask_or(m1, m2): + """~S-CV, ~V-CV""" + val = valuene(m2.parent).new(bool, mask=~m1) + return ComplementedStructuralMask(val) + + +def _complement_CV_A_mask_or(m1, m2): + """~CV-S, ~CV-V""" + val = valuene(m1.parent).new(bool, mask=~m2) + return ComplementedStructuralMask(val) + + +def _combine_V_V_mask_or(m1, m2): + """V-V""" + val = monoid.any(valuene(m1.parent).new() | valuene(m2.parent).new()).new(bool) + return StructuralMask(val) + + +def _complement_CS_CS_mask_or(m1, m2): + """~CS-CS""" + val = pair(m1.parent & m2.parent).new(bool) + return ComplementedStructuralMask(val) + + +def _complement_CV_CV_mask_or(m1, m2): + """~CV-CV""" + val = valuene(land(m1.parent & m2.parent).new(bool)).new() + return ComplementedStructuralMask(val) + + +_MASK_OR = { + # S-S + (StructuralMask, StructuralMask): _combine_S_S_mask_or, + # S-V, V-S + (StructuralMask, ValueMask): _combine_S_SV_mask_or, + (ValueMask, StructuralMask): _combine_SV_S_mask_or, + # S-CS, CS-S + (StructuralMask, ComplementedStructuralMask): _complement_A_CS_mask_or, + (ComplementedStructuralMask, StructuralMask): _complement_CS_A_mask_or, + # S-CV, CV-S + (StructuralMask, ComplementedValueMask): _complement_A_CV_mask_or, + (ComplementedValueMask, StructuralMask): _complement_CV_A_mask_or, + # V-V + (ValueMask, ValueMask): _combine_V_V_mask_or, + # V-CS, CS-V + (ValueMask, ComplementedStructuralMask): _complement_A_CS_mask_or, + (ComplementedStructuralMask, ValueMask): _complement_CS_A_mask_or, + # V-CV, CV-V + (ValueMask, ComplementedValueMask): _complement_A_CV_mask_or, + (ComplementedValueMask, ValueMask): _complement_CV_A_mask_or, + # CS-CS + (ComplementedStructuralMask, ComplementedStructuralMask): _complement_CS_CS_mask_or, + # CS-CV, CV-CS + (ComplementedStructuralMask, ComplementedValueMask): _complement_CS_A_mask_or, + (ComplementedValueMask, ComplementedStructuralMask): _complement_A_CS_mask_or, + # CV-CV + (ComplementedValueMask, ComplementedValueMask): _complement_CV_CV_mask_or, +} + utils._output_types[StructuralMask] = StructuralMask utils._output_types[ValueMask] = ValueMask utils._output_types[ComplementedStructuralMask] = ComplementedStructuralMask diff --git a/graphblas/tests/test_mask.py b/graphblas/tests/test_mask.py index 246303f1f..e976c5a87 100644 --- a/graphblas/tests/test_mask.py +++ b/graphblas/tests/test_mask.py @@ -3,6 +3,7 @@ import pytest from graphblas import Vector +from graphblas.mask import Mask @pytest.mark.parametrize("as_matrix", [False, True]) @@ -53,3 +54,73 @@ def test_mask_new(as_matrix): else: with pytest.raises(TypeError, match="Mask must be"): m.new(mask=v1) + + +@pytest.mark.parametrize("as_matrix", [False, True]) +def test_mask_or(as_matrix): + for mask_dtype in [bool, int]: + v1 = Vector(mask_dtype, size=10) + v1[3:6] = 0 + v1[:3] = 10 + v2 = Vector(mask_dtype, size=10) + v2[1::3] = 0 + v2[::3] = 10 + if as_matrix: + v1 = v1._as_matrix() + v2 = v2._as_matrix() + masks = [v1.S, v1.V, ~v1.S, ~v1.V, v2.S, v2.V, ~v2.S, ~v2.V] + for m1, m2 in itertools.product(masks, masks): + expected = Vector(bool, size=10) + if as_matrix: + expected = expected._as_matrix() + expected(m1) << True + expected(m2) << True + result = (m1 | m2).new() + assert result.isequal(expected, check_dtype=True) + with pytest.raises(TypeError, match="Invalid mask"): + m1 | object() + with pytest.raises(TypeError, match="Invalid mask"): + object() | m1 + if v1.dtype == bool: + assert isinstance(m1 | v1, Mask) + assert isinstance(v1 | m1, Mask) + else: + with pytest.raises(TypeError, match="Mask must be"): + m1 | v1 + with pytest.raises(TypeError, match="Mask must be"): + v1 | m1 + + +@pytest.mark.parametrize("as_matrix", [False, True]) +def test_mask_and(as_matrix): + for mask_dtype in [bool, int]: + v1 = Vector(mask_dtype, size=10) + v1[3:6] = 0 + v1[:3] = 10 + v2 = Vector(mask_dtype, size=10) + v2[1::3] = 0 + v2[::3] = 10 + if as_matrix: + v1 = v1._as_matrix() + v2 = v2._as_matrix() + masks = [v1.S, v1.V, ~v1.S, ~v1.V, v2.S, v2.V, ~v2.S, ~v2.V] + for m1, m2 in itertools.product(masks, masks): + expected = Vector(bool, size=10) + if as_matrix: + expected = expected._as_matrix() + expected[...] << True + expected = expected.dup(mask=m1).dup(mask=m2) + result = (m1 & m2).new() + assert result.isequal(expected, check_dtype=True) + with pytest.raises(TypeError, match="Invalid mask"): + m1 & object() + with pytest.raises(TypeError, match="Invalid mask"): + object() & m1 + if v1.dtype == bool: + assert isinstance(m1 & v1, Mask) + assert isinstance(v1 & m1, Mask) + else: + with pytest.raises(TypeError, match="Mask must be"): + m1 & v1 + with pytest.raises(TypeError, match="Mask must be"): + v1 & m1 From 1a94f52920f6879164fea78a3d430c7f2d6ec0e0 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 5 Jul 2022 10:20:28 -0500 Subject: [PATCH 2/3] flake8 --- graphblas/tests/test_matrix.py | 30 ++++++++++++++++++------------ graphblas/tests/test_vector.py | 29 ++++++++++++++++++----------- 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 7f55e8456..88767ca4f 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -2276,15 +2276,18 @@ def test_import_export_auto(A, do_iso, methods): else: d = A2.ss.unpack(format, sort=sort, raw=raw) if in_method == "import": - import_func = getattr(Matrix.ss, f"import_{import_name}") + + def import_func(x, import_name, **kwargs): + return getattr(Matrix.ss, f"import_{import_name}")(**kwargs) + else: - def import_func(**kwargs): - getattr(A2.ss, f"pack_{import_name}")(**kwargs) - return A2 + def import_func(x, import_name, **kwargs): + getattr(x.ss, f"pack_{import_name}")(**kwargs) + return x d["format"] = import_format - other = import_func(take_ownership=take_ownership, **d) + other = import_func(A2, import_name, take_ownership=take_ownership, **d) if format == "bitmapc" and raw and import_format is None and import_name == "any": # It's 1d, so we can't tell we're column-oriented w/o format keyword assert other.isequal(A_orig.T) @@ -2293,7 +2296,7 @@ def import_func(**kwargs): assert other.ss.is_iso is do_iso d["format"] = "bad_format" with pytest.raises(ValueError, match="Invalid format"): - import_func(**d) + import_func(A2, import_name, **d) assert A.isequal(A_orig) assert A.ss.is_iso is do_iso assert A_orig.ss.is_iso is do_iso @@ -2317,15 +2320,18 @@ def import_func(**kwargs): else: d = C2.ss.unpack(format, raw=raw) if in_method == "import": - import_func = getattr(Matrix.ss, f"import_{import_name}") + + def import_func(x, import_name, **kwargs): + return getattr(Matrix.ss, f"import_{import_name}")(**kwargs) + else: - def import_func(**kwargs): - getattr(C2.ss, f"pack_{import_name}")(**kwargs) - return C2 + def import_func(x, import_name, **kwargs): + getattr(x.ss, f"pack_{import_name}")(**kwargs) + return x d["format"] = import_format - other = import_func(take_ownership=take_ownership, **d) + other = import_func(C2, import_name, take_ownership=take_ownership, **d) if format == "fullc" and raw and import_format is None and import_name == "any": # It's 1d, so we can't tell we're column-oriented w/o format keyword if do_iso: @@ -2340,7 +2346,7 @@ def import_func(**kwargs): assert other.ss.is_iso is do_iso d["format"] = "bad_format" with pytest.raises(ValueError, match="Invalid format"): - import_func(**d) + import_func(C2, import_name, **d) assert C.isequal(C_orig) assert C.ss.is_iso is do_iso assert C_orig.ss.is_iso is do_iso diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index a850c9971..1db61581d 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -1251,19 +1251,23 @@ def test_import_export_auto(v, do_iso, methods): d = v2.ss.unpack(format, sort=sort, raw=raw) if in_method == "import": import_func = getattr(Vector.ss, f"import_{import_name}") + + def import_func(x, import_name, **kwargs): + return getattr(Vector.ss, f"import_{import_name}")(**kwargs) + else: - def import_func(**kwargs): - getattr(v2.ss, f"pack_{import_name}")(**kwargs) - return v2 + def import_func(x, import_name, **kwargs): + getattr(x.ss, f"pack_{import_name}")(**kwargs) + return x d["format"] = import_format - other = import_func(take_ownership=take_ownership, **d) + other = import_func(v2, import_name, take_ownership=take_ownership, **d) assert other.isequal(v_orig) assert other.ss.is_iso is do_iso d["format"] = "bad_format" with pytest.raises(ValueError, match="Invalid format"): - import_func(**d) + import_func(v2, import_name, **d) assert v.isequal(v_orig) assert v.ss.is_iso is do_iso assert v_orig.ss.is_iso is do_iso @@ -1286,20 +1290,23 @@ def import_func(**kwargs): else: d = w2.ss.unpack(format, raw=raw) if in_method == "import": - import_func = getattr(Vector.ss, f"import_{import_name}") + + def import_func(x, import_name, **kwargs): + return getattr(Vector.ss, f"import_{import_name}")(**kwargs) + else: - def import_func(**kwargs): - getattr(w2.ss, f"pack_{import_name}")(**kwargs) - return w2 + def import_func(x, import_name, **kwargs): + getattr(x.ss, f"pack_{import_name}")(**kwargs) + return x d["format"] = import_format - other = import_func(take_ownership=take_ownership, **d) + other = import_func(w2, import_name, take_ownership=take_ownership, **d) assert other.isequal(w_orig) assert other.ss.is_iso is do_iso d["format"] = "bad_format" with pytest.raises(ValueError, match="Invalid format"): - import_func(**d) + import_func(w2, import_name, **d) assert w.isequal(w_orig) assert w.ss.is_iso is do_iso assert w_orig.ss.is_iso is do_iso From 7cd2fdd5eb8062ff6c2cabb8731a45444fac775b Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 7 Jul 2022 12:00:58 -0500 Subject: [PATCH 3/3] Fix recipe for mask "and" --- graphblas/mask.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/graphblas/mask.py b/graphblas/mask.py index cb5b5356c..4b47e5d2e 100644 --- a/graphblas/mask.py +++ b/graphblas/mask.py @@ -62,7 +62,7 @@ def new(self, dtype=None, *, complement=False, mask=None, name=None): >>> val = Matrix(...) >>> val(self) << True - >>> val(mask, replace=True) << True + >>> val(mask, replace=True) << val If `complement=` argument is True, then the *complement* will be returned. This is equivalent to the following (but uses more efficient recipes): @@ -96,7 +96,7 @@ def __and__(self, other): >>> val = Matrix(bool, nrows, ncols) >>> val(mask1) << True - >>> val(mask2, replace=True) << True + >>> val(mask2, replace=True) << val >>> new_mask = val.S This uses faster recipes than the above for all combinations of input mask types,