diff options
| -rw-r--r-- | polymatrix/expression/mixins/diagexprmixin.py | 6 | ||||
| -rw-r--r-- | polymatrix/expression/mixins/symmetricexprmixin.py | 3 | ||||
| -rw-r--r-- | polymatrix/expression/mixins/transposeexprmixin.py | 3 | ||||
| -rw-r--r-- | polymatrix/polymatrix/impl.py | 37 |
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] |