Try to rotate mask. Select lhs matrix with maximal nnz as a mask.

Author: gsvgit

cfpq_matrix/block/block_matrix.py

@@ -32,10 +32,14 @@ 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, mask:Matrix, swap_operands: bool = False) -> Matrix:
+    def mxm(self, other: Matrix, op: Semiring, mask: Matrix, swap_operands: bool = False) -> Matrix:
+        print ("Self shape")
+        print (self.shape)
+        print ("Other shape")
+        print (other.shape)
         if self.block_matrix_space.is_single_cell(other.shape):
-            return self.base.mxm(other, op, mask, swap_operands=swap_operands)
-        return self.base.mxm(
+            return self.base.mxm(other, op, mask, swap_operands=swap_operands)        
+        return self.base.mxm(            
             self.block_matrix_space.hyper_rotate(
                 other,
                 BlockMatrixOrientation.VERTICAL
@@ -94,14 +98,19 @@ def mxm(self, other: Matrix, op: Semiring, mask:Matrix, swap_operands: bool = Fa
                 if swap_operands
                 else BlockMatrixOrientation.VERTICAL
             ).mxm(other, op, None, swap_operands=swap_operands)
+        mask = (self.block_matrix_space.hyper_rotate(
+                mask,
+                BlockMatrixOrientation.VERTICAL
+                if swap_operands
+                else BlockMatrixOrientation.HORIZONTAL)) if not mask is None  else mask
         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=None,
+            mask=mask,
             swap_operands=swap_operands
         )
 

cfpq_model/label_decomposed_graph.py

@@ -262,14 +262,21 @@ def mxm(
             if swap_operands:
                 rhs1, rhs2 = rhs2, rhs1
             if rhs1 in self.matrices and rhs2 in other.matrices:
-                #if swap_operands:
-                mask = (accum.matrices[lhs].to_mask() 
-                        if lhs in accum.matrices and other.matrices[rhs2].shape == self.matrices[rhs1].shape 
-                        else None)
-                #else:
-                #    mask = (self.matrices[lhs].to_mask() 
-                #            if lhs in self.matrices and other.matrices[rhs2].shape == self.matrices[rhs1].shape 
-                #            else None)
+
+                if lhs in accum.matrices:
+                    print("Mask shape")
+                    print(accum.matrices[lhs].shape)                
+
+                matrix_for_mask = max(((accum.matrices[lhs]
+                                           if lhs in accum.matrices
+                                           else None),
+                                       (self.matrices[lhs]
+                                           if lhs in self.matrices
+                                           else None),
+                                       (other.matrices[lhs]
+                                           if lhs in other.matrices
+                                           else None)), key = lambda m: m.nvals if not m is None else -1)
+                mask = matrix_for_mask if isinstance(matrix_for_mask, Matrix) or matrix_for_mask is None else matrix_for_mask.to_mask()
 
                 mxm = self.matrices[rhs1].mxm(
                     other.matrices[rhs2],