python-graphblas · eriknw · Apr 18, 2022 · Apr 16, 2022 · Apr 16, 2022 · Apr 16, 2022
diff --git a/graphblas/__init__.py b/graphblas/__init__.py
@@ -186,7 +186,7 @@ def find_spec(self, fullname, path, target=None):
         if fullname in _NEEDS_OPERATOR and "operator" not in globals():
             _load("operator")
             name = fullname[7:]  # Trim "graphblas."
-            if name in globals():
+            if name in globals():  # pragma: no cover
                 # Make sure we execute the module only once
                 module = globals()[name]
                 spec = module.__spec__
@@ -195,14 +195,14 @@ def find_spec(self, fullname, path, target=None):
 
 
 class _SkipLoad(_Loader):
-    def __init__(self, module, orig_loader):
+    def __init__(self, module, orig_loader):  # pragma: no cover
         self.module = module
         self.orig_loader = orig_loader
 
-    def create_module(self, spec):
+    def create_module(self, spec):  # pragma: no cover
         return self.module
 
-    def exec_module(self, module):
+    def exec_module(self, module):  # pragma: no cover
         # Don't execute the module, but restore the original loader
         module.__spec__.loader = self.orig_loader
 

diff --git a/graphblas/_automethods.py b/graphblas/_automethods.py
@@ -230,6 +230,10 @@ def reduce_scalar(self):
     return self._get_value("reduce_scalar")
 
 
+def reposition(self):
+    return self._get_value("reposition")
+
+
 def select(self):
     return self._get_value("select")
 
@@ -353,6 +357,7 @@ def __ixor__(self, other):
         "ewise_add",
         "ewise_mult",
         "ewise_union",
+        "reposition",
         "select",
         "ss",
         "to_values",

diff --git a/graphblas/base.py b/graphblas/base.py
@@ -423,10 +423,14 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None):
 
         if input_mask is not None:
             raise TypeError("`input_mask` argument may only be used for extract")
-        if expr.op is not None and expr.op.opclass == "Aggregator":
+        elif expr.op is not None and expr.op.opclass == "Aggregator":
             updater = self(mask=mask, accum=accum, replace=replace)
             expr.op._new(updater, expr)
             return
+        elif expr.cfunc_name is None:  # Custom recipe
+            updater = self(mask=mask, accum=accum, replace=replace)
+            expr.args[-2](updater, *expr.args[-1])
+            return
 
         # Normalize mask and separate out complement and structural flags
         if mask is None:
@@ -540,9 +544,9 @@ def __init__(
         self.bt = bt
         self.op = op
         if expr_repr is None:
-            if len(args) == 1:
+            if len(args) == 1 or cfunc_name is None and len(args) == 3:
                 expr_repr = "{0.name}.{method_name}({op})"
-            elif len(args) == 2:
+            elif len(args) == 2 or cfunc_name is None and len(args) == 4:
                 expr_repr = "{0.name}.{method_name}({1.name}, op={op})"
             else:  # pragma: no cover
                 raise ValueError(f"No default expr_repr for len(args) == {len(args)}")
@@ -573,6 +577,8 @@ def _new(self, dtype, mask, name, **kwargs):
         if self.op is not None and self.op.opclass == "Aggregator":
             updater = output(mask=mask)
             self.op._new(updater, self)
+        elif self.cfunc_name is None:  # Custom recipe
+            self.args[-2](output(mask=mask), *self.args[-1])
         elif mask is None:
             output.update(self)
         else:
@@ -583,13 +589,15 @@ def _new(self, dtype, mask, name, **kwargs):
     dup = new
 
     def _format_expr(self):
-        return self.expr_repr.format(*self.args, method_name=self.method_name, op=self.op)
+        args = self.args[:-2] if self.cfunc_name is None else self.args
+        return self.expr_repr.format(*args, method_name=self.method_name, op=self.op)
 
     def _format_expr_html(self):
         expr_repr = self.expr_repr.replace(".name", "._name_html").replace(
             "._expr_name", "._expr_name_html"
         )
-        return expr_repr.format(*self.args, method_name=self.method_name, op=self.op)
+        args = self.args[:-2] if self.cfunc_name is None else self.args
+        return expr_repr.format(*args, method_name=self.method_name, op=self.op)
 
     _expect_type = _expect_type
     _expect_op = _expect_op

diff --git a/graphblas/formatting.py b/graphblas/formatting.py
@@ -109,7 +109,7 @@ def _update_matrix_dataframe(df, matrix, rows, row_offset, columns, column_offse
         if type(matrix) is TransposedMatrix:
             parent = matrix._matrix
             submatrix = Matrix(parent.dtype, parent._nrows, parent._ncols, name="")
-            if parent._nvals > 0:
+            if parent._nvals > 0:  # pragma: no branch
                 # Get val to support iso-valued matrices
                 val = parent.reduce_scalar(monoid.any, allow_empty=False).new(name="")
                 submatrix(parent.S)[columns, rows] = val
@@ -257,7 +257,7 @@ def _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, *, mask=None)
         if min(nonzero._nvals, num_rows) > 2 * df.count().sum():
             rows, cols, vals = nonzero.ss.head(num_rows, sort=True)
             if mask.complement:
-                if not vals.flags.writeable:
+                if not vals.flags.writeable:  # pragma: no branch
                     vals = vals.copy()
                 vals[:] = 0
             df = pd.DataFrame({"row": rows, "col": cols, "val": vals})
@@ -307,7 +307,7 @@ def _get_vector_dataframe(vector, max_rows, min_rows, max_columns, *, mask=None)
         if min(nonzero._nvals, num_rows) > 2 * df.count().sum():
             indices, vals = nonzero.ss.head(num_rows, sort=True)
             if mask.complement:
-                if not vals.flags.writeable:
+                if not vals.flags.writeable:  # pragma: no branch
                     vals = vals.copy()
                 vals[:] = 0
             df = pd.DataFrame({"index": indices, "val": vals})

diff --git a/graphblas/infix.py b/graphblas/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))
+    reposition = wrapdoc(Vector.reposition)(property(_automethods.reposition))
     select = wrapdoc(Vector.select)(property(_automethods.select))
     ss = wrapdoc(Vector.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Vector.to_pygraphblas)(property(_automethods.to_pygraphblas))
@@ -206,6 +207,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))
+    reposition = wrapdoc(Matrix.reposition)(property(_automethods.reposition))
     select = wrapdoc(Matrix.select)(property(_automethods.select))
     ss = wrapdoc(Matrix.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Matrix.to_pygraphblas)(property(_automethods.to_pygraphblas))

diff --git a/graphblas/matrix.py b/graphblas/matrix.py
@@ -26,6 +26,29 @@
 ffi_new = ffi.new
 
 
+# Custom recipes
+def _m_add_v(updater, left, right, op):
+    full = Vector(right.dtype, left._nrows, name="v_full")
+    full[:] = 0
+    temp = full.outer(right, binary.second).new(name="M_temp", mask=updater.kwargs.get("mask"))
+    updater << left.ewise_add(temp, op, require_monoid=False)
+
+
+def _m_mult_v(updater, left, right, op):
+    updater << left.mxm(right.diag(name="M_temp"), get_semiring(monoid.any, op))
+
+
+def _m_union_v(updater, left, right, left_default, right_default, op):
+    full = Vector(right.dtype, left._nrows, name="v_full")
+    full[:] = 0
+    temp = full.outer(right, binary.second).new(name="M_temp", mask=updater.kwargs.get("mask"))
+    updater << left.ewise_union(temp, op, left_default=left_default, right_default=right_default)
+
+
+def _reposition(updater, indices, chunk):
+    updater[indices] = chunk
+
+
 class Matrix(BaseType):
     """
     GraphBLAS Sparse Matrix
@@ -477,17 +500,20 @@ def ewise_add(self, other, op=monoid.plus, *, require_monoid=True):
             self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op")
         if other.ndim == 1:
             # Broadcast rowwise from the right
-            # Can we do `C(M.S) << plus(A | v)` -> `C(M.S) << plus(any_second(M @ v.diag()) | A)`?
             if self._ncols != other._size:
-                # Check this before we compute a possibly large matrix below
                 raise DimensionMismatch(
                     "Dimensions not compatible for broadcasting Vector from the right "
                     f"to rows of Matrix in {method_name}.  Matrix.ncols (={self._ncols}) "
                     f"must equal Vector.size (={other._size})."
                 )
-            full = Vector(other.dtype, self._nrows, name="v_full")
-            full[:] = 0
-            other = full.outer(other, binary.second).new(name="M_temp")
+            return MatrixExpression(
+                method_name,
+                None,
+                [self, other, _m_add_v, (self, other, op)],  # [*expr_args, func, args]
+                nrows=self._nrows,
+                ncols=self._ncols,
+                op=op,
+            )
         expr = MatrixExpression(
             method_name,
             f"GrB_Matrix_eWiseAdd_{op.opclass}",
@@ -515,7 +541,21 @@ def ewise_mult(self, other, op=binary.times):
         # Per the spec, op may be a semiring, but this is weird, so don't.
         self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op")
         if other.ndim == 1:
-            return self.mxm(other.diag(name="M_temp"), get_semiring(monoid.any, op))
+            # Broadcast rowwise from the right
+            if self._ncols != other._size:
+                raise DimensionMismatch(
+                    "Dimensions not compatible for broadcasting Vector from the right "
+                    f"to rows of Matrix in {method_name}.  Matrix.ncols (={self._ncols}) "
+                    f"must equal Vector.size (={other._size})."
+                )
+            return MatrixExpression(
+                method_name,
+                None,
+                [self, other, _m_mult_v, (self, other, op)],  # [*expr_args, func, args]
+                nrows=self._nrows,
+                ncols=self._ncols,
+                op=op,
+            )
         expr = MatrixExpression(
             method_name,
             f"GrB_Matrix_eWiseMult_{op.opclass}",
@@ -584,27 +624,32 @@ def ewise_union(self, other, op, left_default, right_default):
         self._expect_op(op, ("BinaryOp", "Monoid"), within=method_name, argname="op")
         if op.opclass == "Monoid":
             op = op.binaryop
+        expr_repr = "{0.name}.{method_name}({2.name}, {op}, {1._expr_name}, {3._expr_name})"
         if other.ndim == 1:
             # Broadcast rowwise from the right
-            # Can we do `C(M.S) << plus(A | v)` -> `C(M.S) << plus(any_second(M @ v.diag()) | A)`?
             if self._ncols != other._size:
-                # Check this before we compute a possibly large matrix below
                 raise DimensionMismatch(
                     "Dimensions not compatible for broadcasting Vector from the right "
                     f"to rows of Matrix in {method_name}.  Matrix.ncols (={self._ncols}) "
                     f"must equal Vector.size (={other._size})."
                 )
-            full = Vector(other.dtype, self._nrows, name="v_full")
-            full[:] = 0
-            other = full.outer(other, binary.second).new(name="M_temp")
+            return MatrixExpression(
+                method_name,
+                None,
+                [self, left, other, right, _m_union_v, (self, other, left, right, op)],
+                expr_repr=expr_repr,
+                nrows=self._nrows,
+                ncols=self._ncols,
+                op=op,
+            )
         expr = MatrixExpression(
             method_name,
             "GxB_Matrix_eWiseUnion",
             [self, left, other, right],
             op=op,
             at=self._is_transposed,
             bt=other._is_transposed,
-            expr_repr="{0.name}.{method_name}({2.name}, {op}, {1._expr_name}, {3._expr_name})",
+            expr_repr=expr_repr,
         )
         if self.shape != other.shape:
             expr.new(name="")  # incompatible shape; raise now
@@ -917,6 +962,67 @@ def reduce_scalar(self, op=monoid.plus, *, allow_empty=True):
             is_cscalar=not allow_empty,
         )
 
+    # Unofficial methods
+    def reposition(self, row_offset, column_offset, *, nrows=None, ncols=None):
+        """Reposition values by adding `row_offset` and `column_offset` to the indices.
+
+        Positive offset moves values to the right (or down), negative to the left (or up).
+        Values repositioned outside of the new Matrix are dropped (i.e., they don't wrap around).
+
+        This is not a standard GraphBLAS method.  This is implemented with an extract and assign.
+
+        Parameters
+        ----------
+        row_offset : int
+        column_offset : int
+        nrows : int, optional
+            The nrows of the new Matrix.  If not specified, same nrows as input Matrix.
+        ncols : int, optional
+            The ncols of the new Matrix.  If not specified, same ncols as input Matrix.
+
+        """
+        if nrows is None:
+            nrows = self._nrows
+        else:
+            nrows = int(nrows)
+        if ncols is None:
+            ncols = self._ncols
+        else:
+            ncols = int(ncols)
+        row_offset = int(row_offset)
+        if row_offset < 0:
+            row_start = -row_offset
+            row_stop = row_start + nrows
+        else:
+            row_start = 0
+            row_stop = max(0, nrows - row_offset)
+        col_offset = int(column_offset)
+        if col_offset < 0:
+            col_start = -col_offset
+            col_stop = col_start + ncols
+        else:
+            col_start = 0
+            col_stop = max(0, ncols - col_offset)
+        if self._is_transposed:
+            chunk = (
+                self._matrix[col_start:col_stop, row_start:row_stop].new(name="M_repositioning").T
+            )
+        else:
+            chunk = self[row_start:row_stop, col_start:col_stop].new(name="M_repositioning")
+        indices = (
+            slice(row_start + row_offset, row_start + row_offset + chunk._nrows),
+            slice(col_start + col_offset, col_start + col_offset + chunk._ncols),
+        )
+        return MatrixExpression(
+            "reposition",
+            None,
+            [self, _reposition, (indices, chunk)],  # [*expr_args, func, args]
+            expr_repr="{2.name}.reposition(%d, %d)" % (row_offset, column_offset),
+            nrows=nrows,
+            ncols=ncols,
+            dtype=self.dtype,
+        )
+
     ##################################
     # Extract and Assign index methods
     ##################################
@@ -1474,6 +1580,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))
+    reposition = wrapdoc(Matrix.reposition)(property(_automethods.reposition))
     select = wrapdoc(Matrix.select)(property(_automethods.select))
     ss = wrapdoc(Matrix.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Matrix.to_pygraphblas)(property(_automethods.to_pygraphblas))
@@ -1553,6 +1660,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))
+    reposition = wrapdoc(Matrix.reposition)(property(_automethods.reposition))
     select = wrapdoc(Matrix.select)(property(_automethods.select))
     ss = wrapdoc(Matrix.ss)(property(_automethods.ss))
     to_pygraphblas = wrapdoc(Matrix.to_pygraphblas)(property(_automethods.to_pygraphblas))
@@ -1660,6 +1768,7 @@ def _name_html(self):
     reduce_rowwise = Matrix.reduce_rowwise
     reduce_columnwise = Matrix.reduce_columnwise
     reduce_scalar = Matrix.reduce_scalar
+    reposition = Matrix.reposition
 
     # Operator sugar
     __or__ = Matrix.__or__