From a967244b7c16599e4c435cc5f7278e7771cc0799 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 3 Jun 2022 22:57:32 -0500 Subject: [PATCH 1/3] Apply with a dict! --- graphblas/matrix.py | 32 +++++++++++++++++++++++++------- graphblas/operator.py | 22 ++++++++++++++++++++++ graphblas/tests/test_matrix.py | 17 +++++++++++++++++ graphblas/tests/test_vector.py | 13 +++++++++++++ graphblas/vector.py | 31 +++++++++++++++++++++++++------ 5 files changed, 102 insertions(+), 13 deletions(-) diff --git a/graphblas/matrix.py b/graphblas/matrix.py index 9d6595219..7ec03fc16 100644 --- a/graphblas/matrix.py +++ b/graphblas/matrix.py @@ -1,5 +1,6 @@ import itertools import warnings +from collections.abc import Mapping import numpy as np @@ -10,7 +11,14 @@ from .exceptions import DimensionMismatch, NoValue, check_status from .expr import AmbiguousAssignOrExtract, IndexerResolver, Updater from .mask import Mask, StructuralMask, ValueMask -from .operator import UNKNOWN_OPCLASS, find_opclass, get_semiring, get_typed_op, op_from_string +from .operator import ( + UNKNOWN_OPCLASS, + _dict_to_func, + find_opclass, + get_semiring, + get_typed_op, + op_from_string, +) from .scalar import ( _MATERIALIZE, Scalar, @@ -758,12 +766,18 @@ def kronecker(self, other, op=binary.times): ) def apply(self, op, right=None, *, left=None): - """ - GrB_Matrix_apply - Apply UnaryOp to each element of the calling Matrix + """Apply an operator, function, or mapping to each element of the Matrix. + + Apply UnaryOp to each element of the calling Matrix. + A BinaryOp can also be applied if a scalar is passed in as `left` or `right`, - effectively converting a BinaryOp into a UnaryOp - An IndexUnaryOp can also be applied with the thunk passed in as `right` + effectively converting a BinaryOp into a UnaryOp. + + An IndexUnaryOp can also be applied with the thunk passed in as `right`. + + A dict or Mapping can also be applied to map input values to specific output values. + If an input value isn't in the mapping, the result is the default value passed in + as `right` with a default of 0. For example, `A.apply({1: 10, 2: 20}, 100)`. """ method_name = "apply" extra_message = ( @@ -779,7 +793,11 @@ def apply(self, op, right=None, *, left=None): right = False # most basic form of 0 when unifying dtypes if left is not None: raise TypeError("Do not pass `left` when applying IndexUnaryOp") - + elif opclass == UNKNOWN_OPCLASS and isinstance(op, Mapping): + op = _dict_to_func(op, right) + right = None + if left is not None: + raise TypeError("Do not pass `left` when applying a Mapping") if left is None and right is None: op = get_typed_op(op, self.dtype, kind="unary") self._expect_op( diff --git a/graphblas/operator.py b/graphblas/operator.py index 9e04f95eb..b50805f52 100644 --- a/graphblas/operator.py +++ b/graphblas/operator.py @@ -3260,3 +3260,25 @@ def aggregator_from_string(string): from . import agg # noqa isort:skip agg.from_string = aggregator_from_string + + +def _dict_to_func(d, default): + # This probably doesn't work on UDTs, and we could probably be smarter with dtypes + if default is None: + default = False + keys, vals = zip(*d.items()) + keys = np.array(keys) + lookup_dtype(keys.dtype) + vals = np.array(vals) + lookup_dtype(vals.dtype) + p = np.argsort(keys) + keys = keys[p] + vals = vals[p] + + def func(x): + i = np.searchsorted(keys, x) + if i < keys.size and keys[i] == x: + return vals[i] + return default + + return func diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index e01572c89..ba8c1275b 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1115,6 +1115,23 @@ def test_apply_indexunary(A): A.apply(select.valueeq, left=s3) +def test_apply_dict(): + rows = [0, 0, 0, 0] + cols = [1, 3, 4, 6] + vals = [1, 1, 2, 0] + V = Matrix.from_values(rows, cols, vals) + # Use right as default + W1 = V.apply({1: 10, 2: 20}, 100).new() + expected = Matrix.from_values(rows, cols, [10, 10, 20, 100]) + assert W1.isequal(expected) + # Default is 0 if unspecified + W2 = V.apply({0: 10, 2: 20}).new() + expected = Matrix.from_values(rows, cols, [0, 0, 20, 10]) + assert W2.isequal(expected) + with pytest.raises(TypeError, match="left"): + V.apply({0: 10, 2: 20}, left=999) + + 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() diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 4c216e765..0023b113b 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -695,6 +695,19 @@ def test_apply_indexunary(v): v.apply(indexunary.valueeq, left=s2) +def test_apply_dict(v): + # Use right as default + w1 = v.apply({1: 10, 2: 20}, 100).new() + expected = Vector.from_values([1, 3, 4, 6], [10, 10, 20, 100]) + assert w1.isequal(expected) + # Default is 0 if unspecified + w2 = v.apply({0: 10, 2: 20}).new() + expected = Vector.from_values([1, 3, 4, 6], [0, 0, 20, 10]) + assert w2.isequal(expected) + with pytest.raises(TypeError, match="left"): + v.apply({0: 10, 2: 20}, left=999) + + def test_select(v): result = Vector.from_values([1, 3], [1, 1], size=7) w1 = v.select(select.valueeq, 1).new() diff --git a/graphblas/vector.py b/graphblas/vector.py index 5630e83e4..6980f9fd7 100644 --- a/graphblas/vector.py +++ b/graphblas/vector.py @@ -1,5 +1,6 @@ import itertools import warnings +from collections.abc import Mapping import numpy as np @@ -10,7 +11,14 @@ from .exceptions import DimensionMismatch, NoValue, check_status from .expr import AmbiguousAssignOrExtract, IndexerResolver, Updater from .mask import Mask, StructuralMask, ValueMask -from .operator import UNKNOWN_OPCLASS, find_opclass, get_semiring, get_typed_op, op_from_string +from .operator import ( + UNKNOWN_OPCLASS, + _dict_to_func, + find_opclass, + get_semiring, + get_typed_op, + op_from_string, +) from .scalar import ( _MATERIALIZE, Scalar, @@ -694,11 +702,18 @@ def vxm(self, other, op=semiring.plus_times): return expr def apply(self, op, right=None, *, left=None): - """ - GrB_Vector_apply - Apply UnaryOp to each element of the calling Vector + """Apply an operator, function, or mapping to each element of the Vector + + Apply UnaryOp to each element of the calling Vector. + A BinaryOp can also be applied if a scalar is passed in as `left` or `right`, - effectively converting a BinaryOp into a UnaryOp + effectively converting a BinaryOp into a UnaryOp. + + An IndexUnaryOp can also be applied with the thunk passed in as `right` + + A dict or Mapping can also be applied to map input values to specific output values. + If an input value isn't in the mapping, the result is the default value passed in + as `right` with a default of 0. For example, `v.apply({1: 10, 2: 20}, 100)`. """ method_name = "apply" extra_message = ( @@ -714,7 +729,11 @@ def apply(self, op, right=None, *, left=None): right = False # most basic form of 0 when unifying dtypes if left is not None: raise TypeError("Do not pass `left` when applying IndexUnaryOp") - + elif opclass == UNKNOWN_OPCLASS and isinstance(op, Mapping): + if left is not None: + raise TypeError("Do not pass `left` when applying a Mapping") + op = _dict_to_func(op, right) + right = None if left is None and right is None: op = get_typed_op(op, self.dtype, kind="unary") self._expect_op( From 89699a4392ace5824e54708f7af6bd53f142f469 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 3 Jun 2022 23:00:43 -0500 Subject: [PATCH 2/3] Better --- graphblas/matrix.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/graphblas/matrix.py b/graphblas/matrix.py index 7ec03fc16..ff0a2e3b8 100644 --- a/graphblas/matrix.py +++ b/graphblas/matrix.py @@ -794,10 +794,10 @@ def apply(self, op, right=None, *, left=None): if left is not None: raise TypeError("Do not pass `left` when applying IndexUnaryOp") elif opclass == UNKNOWN_OPCLASS and isinstance(op, Mapping): - op = _dict_to_func(op, right) - right = None if left is not None: raise TypeError("Do not pass `left` when applying a Mapping") + op = _dict_to_func(op, right) + right = None if left is None and right is None: op = get_typed_op(op, self.dtype, kind="unary") self._expect_op( From c2b31048e1addfa29db41281b6cf46bd1cc04646 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 3 Jun 2022 23:17:38 -0500 Subject: [PATCH 3/3] More tests --- graphblas/tests/test_vector.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 0023b113b..466e05e2d 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -704,8 +704,16 @@ def test_apply_dict(v): w2 = v.apply({0: 10, 2: 20}).new() expected = Vector.from_values([1, 3, 4, 6], [0, 0, 20, 10]) assert w2.isequal(expected) + # Scalar default can up-cast dtype + w3 = v.apply({1: 10, 2: 20}, 0.5).new() + expected = Vector.from_values([1, 3, 4, 6], [10, 10, 20, 0.5]) + assert w3.isequal(expected) with pytest.raises(TypeError, match="left"): v.apply({0: 10, 2: 20}, left=999) + with pytest.raises(ValueError, match="Unknown dtype"): + v.apply({0: 10, 2: object()}) + with pytest.raises(Exception): # This error and message should be better + v.apply({0: 10, 2: 20}, object()) def test_select(v):