diff options
-rw-r--r-- | polymatrix/polymatrix/impl.py | 25 | ||||
-rw-r--r-- | polymatrix/polymatrix/index.py | 10 | ||||
-rw-r--r-- | polymatrix/polymatrix/utils.py | 55 |
3 files changed, 87 insertions, 3 deletions
diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py index 07b9ffa..bb42313 100644 --- a/polymatrix/polymatrix/impl.py +++ b/polymatrix/polymatrix/impl.py @@ -10,6 +10,31 @@ class PolyMatrixImpl(PolyMatrix): data: PolyMatrixDict shape: tuple[int, int] + def __str__(self) -> str: + nrows, ncols = self.shape + columns = [] + for c in range(ncols): + rows, row_width = [], 0 + for r in range(nrows): + if (r, c) in self.data: + entry = str(self.data[r, c]) + rows.append(entry) + + if len(entry) > row_width: + row_width = len(entry) + else: + rows.append("0") + columns.append(tuple(row.rjust(row_width) for row in rows)) + + string = "" + for r in range(nrows): + column = ", ".join(columns[c][r] for c in range(ncols)) + string += f"[ {column} ]\n" + + return string + + + @dataclassabc.dataclassabc(frozen=True) class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin): diff --git a/polymatrix/polymatrix/index.py b/polymatrix/polymatrix/index.py index 7fe3620..83eccf5 100644 --- a/polymatrix/polymatrix/index.py +++ b/polymatrix/polymatrix/index.py @@ -28,6 +28,11 @@ class VariableIndex(NamedTuple): """ Variables indices can be sorted with respect to variable index. """ return self.index < other.index + def __str__(self) -> str: + if self.power == 1: + return f"V({self.index:d})" + return f"V({self.index:d})^{self.power:d}" + class MonomialIndex(tuple[VariableIndex]): r""" @@ -48,7 +53,7 @@ class MonomialIndex(tuple[VariableIndex]): return "MonomialIndex(" + ",".join(repr(v) for v in self) + ")" def __str__(self): - return "MonomialIndex(" + ",".join(f"{v.index}^{v.power}" for v in self) + ")" + return " * ".join(str(v) for v in sorted(self)) @property def degree(self) -> int: @@ -158,6 +163,9 @@ class PolyDict(UserDict[MonomialIndex, int | float]): def __repr__(self): return f"{self.__class__.__qualname__}({super().__repr__()})" + def __str__(self) -> str: + return " + ".join(f"{str(c)} * {str(m)}" for m, c in self.items()) + def __setitem__(self, key: Iterable[VariableIndex] | MonomialIndex, value: int | float): if not isinstance(key, MonomialIndex): key = MonomialIndex(key) diff --git a/polymatrix/polymatrix/utils.py b/polymatrix/polymatrix/utils.py index 289302a..d1f9d51 100644 --- a/polymatrix/polymatrix/utils.py +++ b/polymatrix/polymatrix/utils.py @@ -1,10 +1,62 @@ import itertools import math -from polymatrix.polymatrix.index import MonomialData, MonomialIndex, PolyDict +from polymatrix.expressionstate import ExpressionState +from polymatrix.polymatrix.index import MonomialData, MonomialIndex, PolyDict, PolyMatrixDict +from polymatrix.polymatrix.mixins import PolyMatrixMixin, PolyMatrixAsDictMixin from polymatrix.utils.deprecation import deprecated +def pretty_format_polydict(p: PolyDict, state: ExpressionState) -> str: + """ Pretty print a polynomial """ + strings = [] + for monomial, coeff in p.terms(): + s = f"{coeff}" + for var in monomial: + if var.power == 1: + s += f" * {state.get_name(var.index)}" + else: + s += f" * {state.get_name(var.index)}^{var.power:d}" + + strings.append(s) + return " + ".join(strings) + + +def pretty_format_polymatrixdict(p: PolyMatrixDict, shape: tuple[int, int], state: ExpressionState) -> str: + """ Pretty print a polymatrix """ + nrows, ncols = shape + columns = [] + for c in range(ncols): + rows, row_width = [], 0 + for r in range(nrows): + if (r, c) in p: + entry = pretty_format_polydict(p[r, c], state) + rows.append(entry) + + if len(entry) > row_width: + row_width = len(entry) + else: + rows.append("0") + + columns.append(tuple(row.rjust(row_width) for row in rows)) + + string = "" + for r in range(nrows): + column = ", ".join(columns[c][r] for c in range(ncols)) + string += f"[ {column} ]\n" + + return string + + +def pretty_format_polymatrix(p: PolyMatrixMixin, state: ExpressionState) -> str: + """ Pretty print a polymatrix """ + if isinstance(p, PolyMatrixAsDictMixin): + return pretty_format_polymatrixdict(p.data, (p.shape[0], p.shape[1]), state) + + else: + # TODO: add pretty print for other functions + raise NotImplementedError + # NP: this is a werid case of the product # NP: merge is not intutive IMHO, find better name @@ -74,7 +126,6 @@ def multiply_polynomial( # NP: sort mononial indices with respect to variable index in state object - @deprecated("With new index types you can use sorted() on MonomialIndex") def sort_monomial_indices( monomial: MonomialData, |