Attempt to introduce masking

Author: gsvgit

cfpq_matrix/abstract_optimized_matrix_decorator.py

@@ -31,3 +31,6 @@ def dtype(self) -> DataType:
 
     def to_unoptimized(self) -> Matrix:
         return self.base.to_unoptimized()
+
+    def to_mask(self) -> Matrix:
+        return self.base.to_mask()

cfpq_matrix/block/block_matrix.py

@@ -23,15 +23,23 @@ def optimize_similarly(self, other: Matrix) -> "OptimizedMatrix":
             self.base.optimize_similarly(other)
         )
 
+    def to_mask(self):
+        if self.block_matrix_space.is_single_cell(self.shape):
+            print ("@@@@@@@")
+            print (self.shape)
+            return self.base.to_mask()
+        else:
+            return None
+
 
 class CellBlockMatrix(BlockMatrix):
     def __init__(self, base: OptimizedMatrix, block_matrix_space: BlockMatrixSpace):
         assert block_matrix_space.is_single_cell(base.shape)
         super().__init__(base, block_matrix_space)
 
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = False) -> Matrix:
         if self.block_matrix_space.is_single_cell(other.shape):
-            return self.base.mxm(other, op, swap_operands=swap_operands)
+            return self.base.mxm(other, op, mask, swap_operands=swap_operands)
         return self.base.mxm(
             self.block_matrix_space.hyper_rotate(
                 other,
@@ -40,6 +48,7 @@ def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matri
                 else BlockMatrixOrientation.HORIZONTAL
             ),
             op=op,
+            mask=mask,
             swap_operands=swap_operands,
         )
 
@@ -83,20 +92,21 @@ def _force_init_orientation(
         self.discard_base_on_reformat = False
         return self.matrices[desired_orientation]
 
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = False) -> Matrix:
         if self.block_matrix_space.is_single_cell(other.shape):
             return self._force_init_orientation(
                 BlockMatrixOrientation.HORIZONTAL
                 if swap_operands
                 else BlockMatrixOrientation.VERTICAL
-            ).mxm(other, op, swap_operands=swap_operands)
+            ).mxm(other, op, mask, swap_operands=swap_operands)
         return self._force_init_orientation(
             BlockMatrixOrientation.VERTICAL
             if swap_operands
             else BlockMatrixOrientation.HORIZONTAL
         ).mxm(
             self.block_matrix_space.to_block_diag_matrix(other),
             op=op,
+            mask=mask,
             swap_operands=swap_operands
         )
 

cfpq_matrix/empty_optimized_matrix.py

@@ -14,14 +14,14 @@ def __init__(self, base: OptimizedMatrix):
     def base(self) -> OptimizedMatrix:
         return self._base
 
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = False) -> Matrix:
         if self.nvals == 0 or other.nvals == 0:
             if swap_operands:
                 assert self.shape[0] == other.shape[1]
                 return Matrix(self.dtype, self.shape[1], other.shape[0])
             assert self.shape[1] == other.shape[0]
             return Matrix(self.dtype, self.shape[0], other.shape[1])
-        return self.base.mxm(other, op, swap_operands)
+        return self.base.mxm(other, op, mask, swap_operands)
 
     def rsub(self, other: Matrix, op: SubOp) -> Matrix:
         if self.nvals == 0 or other.nvals == 0:

cfpq_matrix/format_optimized_matrix.py

@@ -30,6 +30,11 @@ def __new__(
     def base(self) -> OptimizedMatrix:
         return self._base
 
+    def to_mask(self) -> Matrix:
+        print ("#####")
+        #print (self.matrices[0].to_mask())
+        return list(self.matrices.values())[0].to_mask()
+
     def _force_init_format(self, desired_format: str) -> OptimizedMatrix:
         if desired_format not in self.matrices:
             base_matrix = self.base.to_unoptimized().dup()
@@ -42,16 +47,16 @@ def _force_init_format(self, desired_format: str) -> OptimizedMatrix:
         res = self.matrices[desired_format]
         return res
 
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask: Matrix, swap_operands: bool = False) -> Matrix:
         left_nvals = other.nvals if swap_operands else self.nvals
         right_nvals = self.nvals if swap_operands else other.nvals
         desired_format = "by_row" if left_nvals < right_nvals else "by_col"
 
         if desired_format in self.matrices or other.nvals < self.nvals / self.reformat_threshold:
             other.ss.config["format"] = desired_format
             reformatted_self = self._force_init_format(desired_format)
-            return reformatted_self.mxm(other, op, swap_operands=swap_operands)
-        return self.base.mxm(other, op, swap_operands=swap_operands)
+            return reformatted_self.mxm(other, op, mask, swap_operands=swap_operands)
+        return self.base.mxm(other, op, mask, swap_operands=swap_operands)
 
     def rsub(self, other: Matrix, op: SubOp) -> Matrix:
         return self.matrices.get(other.ss.config["format"], self.base).rsub(other, op)

cfpq_matrix/lazy_add_optimized_matrix.py

@@ -23,6 +23,9 @@ def base(self) -> OptimizedMatrix:
     def nvals(self) -> int:
         return sum(m.nvals for m in self.matrices)
 
+    def to_mask(self) -> Matrix:
+        return self.matrices[0].to_mask()
+
     def _map_and_fold(
             self,
             mapper,
@@ -53,10 +56,10 @@ def to_unoptimized(self) -> Matrix:
         self.force_combine_small_matrices(nvals_combine_threshold=float("inf"))
         return self.base.to_unoptimized()
 
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = False) -> Matrix:
         self.update_monoid(op.monoid)
         return self._map_and_fold(
-            mapper=lambda m: m.mxm(other, op=op, swap_operands=swap_operands),
+            mapper=lambda m: m.mxm(other, op=op, mask=mask, swap_operands=swap_operands),
             combiner=lambda acc, cur: acc.ewise_add(cur, op=op.monoid).new(),
             nvals_combine_threshold=other.nvals
         )

cfpq_matrix/matrix_to_optimized_adapter.py

@@ -30,15 +30,39 @@ def format(self) -> MatrixFormat:
     def dtype(self) -> DataType:
         return self.base.dtype
 
+    def to_mask(self) -> Matrix:
+        return self.base
+        
     def to_unoptimized(self) -> Matrix:
         return self.base
 
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
-        return (
-            other.mxm(self.base, op)
-            if swap_operands
-            else self.base.mxm(other, op)
-        ).new(self.dtype)
+    def mxm(self, other: Matrix, op: Semiring, mask: Matrix, swap_operands: bool = False) -> Matrix:
+        #   return (
+        #    other.mxm(self.base, op)
+        #    if swap_operands
+        #    else self.base.mxm(other, op)
+        #).new(self.dtype)
+        if swap_operands:
+            if not mask is None:
+                #print("Mask applied, swap operands") 
+                #result = Matrix(self.dtype,nrows=other.shape[0],ncols=self.shape[1])
+                #mask_t = Matrix(mask.dtype, ncols=mask.ncols, nrows=mask.nrows)
+                #mask_t << mask.T 
+                #result(~mask) << other.mxm(self.base, op)
+                #result(~mask) << other.mxm(self.base, op).new(self.dtype)
+                #return result
+                return other.mxm(self.base, op).new(self.dtype)
+            else: return other.mxm(self.base, op).new(self.dtype)
+        else:
+            if not mask is None:
+                print("Mask applied")
+                result = Matrix(self.dtype,nrows=self.shape[0],ncols=other.shape[1])
+                #result(~mask) << self.base.mxm(other, op)
+                result(~mask) << self.base.mxm(other, op).new(self.dtype)
+                return result
+                #return self.base.mxm(other, op).new(self.dtype)
+            else: return self.base.mxm(other, op).new(self.dtype)
+        
 
     def rsub(self, other: Matrix, op: SubOp) -> Matrix:
         return op(other, self.base)

cfpq_matrix/optimized_matrix.py

@@ -47,7 +47,7 @@ def to_unoptimized(self) -> Matrix:
         pass
 
     @abstractmethod
-    def mxm(self, other: Matrix, op: Semiring, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask: Matrix, swap_operands: bool = False) -> Matrix:
         pass
 
     @abstractmethod

cfpq_model/label_decomposed_graph.py

@@ -262,9 +262,16 @@ def mxm(
             if swap_operands:
                 rhs1, rhs2 = rhs2, rhs1
             if rhs1 in self.matrices and rhs2 in other.matrices:
+                #if lhs in self.matrices: 
+                #    print ("!!! Mask: ")
+                #    print(self.matrices[lhs])
+                #    print("-------------------")
+                #    print(self.matrices[rhs1].shape)
+                #    print(other.matrices[rhs2].shape)
                 mxm = self.matrices[rhs1].mxm(
                     other.matrices[rhs2],
                     swap_operands=swap_operands,
+                    mask=(self.matrices[lhs].to_mask() if lhs in self.matrices  and other.matrices[rhs2].shape == self.matrices[rhs1].shape else None),   #Matrix(dtype=self.matrices[rhs1].dtype, nrows=self.matrices[rhs1].nrows, ncols = self.matrices[rhs2].ncols)),
                     op=op,
                 )
                 accum.iadd_by_symbol(lhs, mxm, op.monoid)