Add more __str__ methods to PolyMatrix for debugging

4 files changed, 41 insertions, 8 deletions

diff --git a/polymatrix/expression/mixins/diagexprmixin.py b/polymatrix/expression/mixins/diagexprmixin.py
index b7666f0..fb4aad9 100644
--- a/polymatrix/expression/mixins/diagexprmixin.py
+++ b/polymatrix/expression/mixins/diagexprmixin.py
@@ -37,6 +37,9 @@ class DiagExprMixin(ExpressionBaseMixin):
                 underlying: PolyMatrixMixin
                 shape: tuple[int, int]
 
+                def __str__(self):
+                    return f"diag(\n{self.underlying})"
+
                 def at(self, row: int, col: int) -> PolyDict:
                     if row != col:
                         return PolyDict.empty()
@@ -59,6 +62,9 @@ class DiagExprMixin(ExpressionBaseMixin):
                 underlying: PolyMatrixMixin
                 shape: tuple[int, int]
 
+                def __str__(self):
+                    return f"diag(\n{self.underlying})"
+
                 def at(self, row: int, _col: int) -> PolyDict:
                     return self.underlying.at(row, row)
 
diff --git a/polymatrix/expression/mixins/symmetricexprmixin.py b/polymatrix/expression/mixins/symmetricexprmixin.py
index dd77bef..da168da 100644
--- a/polymatrix/expression/mixins/symmetricexprmixin.py
+++ b/polymatrix/expression/mixins/symmetricexprmixin.py
@@ -35,6 +35,9 @@ class SymmetricExprMixin(ExpressionBaseMixin):
         class SymmetricPolyMatrix(PolyMatrixMixin):
             underlying: PolyMatrixMixin
 
+            def __str__(self):
+                return f"symmetric(\n{self.underlying})"
+
             @property
             def shape(self) -> tuple[int, int]:
                 return self.underlying.shape
diff --git a/polymatrix/expression/mixins/transposeexprmixin.py b/polymatrix/expression/mixins/transposeexprmixin.py
index 16139ba..c4e3f4e 100644
--- a/polymatrix/expression/mixins/transposeexprmixin.py
+++ b/polymatrix/expression/mixins/transposeexprmixin.py
@@ -32,6 +32,9 @@ class TransposeExprMixin(ExpressionBaseMixin):
             underlying: PolyMatrixMixin
             shape: tuple[int, int]
 
+            def __str__(self):
+                return f"transpose(\n{self.underlying})"
+
             def at(self, row: int, col: int) -> PolyDict:
                 return self.underlying.at(col, row)
 
diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py
index a6a1acf..ac426b7 100644
--- a/polymatrix/polymatrix/impl.py
+++ b/polymatrix/polymatrix/impl.py
@@ -1,11 +1,12 @@
-import dataclassabc
+from dataclassabc import dataclassabc
+from itertools import groupby
 
 from polymatrix.polymatrix.abc import PolyMatrix
 from polymatrix.polymatrix.mixins import PolyMatrixMixin, BroadcastPolyMatrixMixin, BlockPolyMatrixMixin, SlicePolyMatrixMixin, PolyMatrixAsAffineExpressionMixin
 from polymatrix.polymatrix.index import PolyMatrixDict, PolyDict, MonomialIndex
 
 
-@dataclassabc.dataclassabc(frozen=True)
+@dataclassabc(frozen=True)
 class PolyMatrixImpl(PolyMatrix):
     data: PolyMatrixDict
     shape: tuple[int, int]
@@ -34,28 +35,48 @@ class PolyMatrixImpl(PolyMatrix):
         return string
 
 
-
-
-@dataclassabc.dataclassabc(frozen=True)
+@dataclassabc(frozen=True)
 class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin):
     polynomial: PolyDict
     shape: tuple[int, int]
 
+    def __str__(self):
+        return f"broadcast({self.polynomial}, {self.shape})"
 
-@dataclassabc.dataclassabc(frozen=True)
+
+@dataclassabc(frozen=True)
 class BlockPolyMatrixImpl(BlockPolyMatrixMixin):
     blocks: dict[tuple[range, range], PolyMatrixMixin]
     shape: tuple[int, int]
 
+    # FIXME: added for debugging, but ugly af
+    def __str__(self):
+        string = "block(\n"
+
+        def by_row(ranges):
+            return ranges[0].start
+        
+        def by_col(ranges):
+            return ranges[1].start
+
+        rows_strings = []
+        for r, row in groupby(sorted(self.blocks.keys(), key=by_row), key=by_row):
+            cols_strings = []
+            for blk in sorted(row, key=by_col):
+                cols_strings.append(str(self.blocks[blk]))
+            rows_strings.append(",\n".join(cols_strings))
+
+        return string + "\n".join(rows_strings) + ")"
+
 
-@dataclassabc.dataclassabc(frozen=True)
+@dataclassabc(frozen=True)
 class SlicePolyMatrixImpl(SlicePolyMatrixMixin):
     reference: PolyMatrixMixin
     shape: tuple[int, int]
     slice: tuple[tuple[int, ...], tuple[int, ...]]
 
 
-@dataclassabc.dataclassabc(frozen=True)
+@dataclassabc(frozen=True)
 class PolyMatrixAsAffineExpressionImpl(PolyMatrixAsAffineExpressionMixin):
     affine_coefficients: PolyMatrixAsAffineExpressionMixin.MatrixType
     shape: tuple[int, int]