Introduce index types and new PolyMatrix API

7 files changed, 170 insertions, 60 deletions

diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py
index 27b6260..152d5fa 100644
--- a/polymatrix/polymatrix/impl.py
+++ b/polymatrix/polymatrix/impl.py
@@ -2,16 +2,16 @@ import dataclassabc
 
 from polymatrix.polymatrix.abc import PolyMatrix
 from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin
-from polymatrix.polymatrix.typing import PolynomialData
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict
 
 
 @dataclassabc.dataclassabc(frozen=True)
 class PolyMatrixImpl(PolyMatrix):
-    data: dict[tuple[int, int], PolynomialData] # FIXME: use polymatrix.typing
+    data: PolyMatrixDict
     shape: tuple[int, int]
 
 
 @dataclassabc.dataclassabc(frozen=True)
 class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin):
-    data: tuple[tuple[int], float] # FIXME: use polymatrix.typing
+    data: PolyDict
     shape: tuple[int, int]
diff --git a/polymatrix/polymatrix/init.py b/polymatrix/polymatrix/init.py
index c4912c7..704ebe5 100644
--- a/polymatrix/polymatrix/init.py
+++ b/polymatrix/polymatrix/init.py
@@ -1,10 +1,10 @@
 from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl
-from polymatrix.polymatrix.typing import PolynomialData
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict
 
 
 # FIXME: use polymatrix.typing
 def init_poly_matrix(
-    data: dict[tuple[int, int], PolynomialData],
+    data: PolyMatrixDict,
     shape: tuple[int, int],
 ):
     return PolyMatrixImpl(
@@ -16,7 +16,7 @@ def init_poly_matrix(
 # NP: cosider renaming to scalar cf. comment on broadcast_poly_matrix and BroadcastPolymatrix
 # FIXME: use polymatrix.typing
 def init_broadcast_poly_matrix(
-    data: PolynomialData,
+    data: PolyDict,
     shape: tuple[int, int],
 ):
     return BroadcastPolyMatrixImpl(
diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index d4381e7..ddd00fc 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -1,68 +1,111 @@
 import abc
 import typing
 
-from polymatrix.polymatrix.typing import PolynomialData
+from typing import Iterable, cast
+from typing_extensions import override
+
+from polymatrix.polymatrix.typing import PolyDict, PolyMatrixDict, MatrixIndex
+from polymatrix.utils.deprecation import deprecated
 
 
 class PolyMatrixMixin(abc.ABC):
-    # NP: is this still in use? cf. Filter operation, that breaks the
-    # NP: assumption that shape can be computed in advance
-    # FIXME: classmethod as properties are no longer allowed, change to abstractmethod
     @property
-    @abc.abstractclassmethod
+    @abc.abstractmethod
     def shape(self) -> tuple[int, int]: ...
 
-    # NP: the "data" is an iteration over the entries of the matrix
-    # FIXME: rename to something meaningful like entries()
-    # NP: the fact that it returns the inner dictionary directly leaks the abstraction,
-    # NP: as discussed in in meeting #6 IMHO it is better to return indices and query like in mdpoly
-    def gen_data(self) -> typing.Generator[tuple[tuple[int, int], PolynomialData], None, None]:
-        for row in range(self.shape[0]):
-            for col in range(self.shape[1]):
-                polynomial = self.get_poly(row, col)
-                if polynomial is None:
-                    continue
-
-                yield (row, col), polynomial
-
-    
-    # NP: get a polynomial at a given entry
-    # NP: why is this an abstract class method if it uses self?
-    # FIXME: change to abstractmethod
-    @abc.abstractclassmethod
-    def get_poly(self, row: int, col: int) -> PolynomialData | None: ...
-
-
-# NP: As far as I can tell the dict is the only way you ever store a polymatrix
-# NP: (cf. PolynomialData), why separate interface for PolyMatrixAsDictMixin?
-# NP: Also applies to BroadcastPolyMatrixMixin class below.
+    @deprecated("Replaced by PolyMatrixMixin.entries()")
+    def gen_data(self) -> typing.Generator[tuple[tuple[int, int], PolyDict], None, None]:
+        yield from self.entries()
+ 
+    @abc.abstractmethod
+    @deprecated("Replaced by PolyMatrixMixin.at(row, col).")
+    def get_poly(self, row: int, col: int) -> PolyDict | None:
+        # NOTE: in the new API when there is nothing at an entry the at()
+        # function returns an empty dictionary. This is because emtpy
+        # dicionaries are falsy so you can do
+        #
+        # `if p.at(row, col):`
+        #
+        # like it was possible with get_poly. However this is better because
+        # id does not make use of a null (None) type.
+        return self.at(row, col) or None
+
+    # --- New API ---
+
+    @abc.abstractmethod
+    def at(self, row: int, col: int) -> PolyDict:
+        """ Return the polynomial at the entry (row, col). 
+
+        .. code:: py
+            p.at(row, col) # if you have row and col
+            p.at(*entry) # if you have a MatrixIndex
+        """
+
+    def entries(self) -> Iterable[tuple[MatrixIndex, PolyDict]]:
+        """
+        Iterate over the non-zero entries of the polynomial matrix.
+
+        TODO: docstrings, use like this (`p` is a polymatrix)
+
+        .. code:: py
+            for entry, poly in p.entries():
+                for monomial, coeff in poly.terms():
+                    # do something with coefficients
+        """
+        nrows, ncols = self.shape
+        for row in range(nrows):
+            for col in range(ncols):
+                if poly := self.at(row, col):
+                    yield MatrixIndex(row, col), poly
+
+
 class PolyMatrixAsDictMixin(
     PolyMatrixMixin,
     abc.ABC,
 ):
+    """
+    Polynomial matrix, stored as a dictionary.
+    """
+
     @property
     @abc.abstractmethod
-    def data(self) -> dict[tuple[int, int], PolynomialData]: ...
+    def data(self) -> PolyMatrixDict:
+        """" Get the dictionary. """
 
-    # overwrites the abstract method of `PolyMatrixMixin`
-    # NP: Returning none is not great design, from math perspective it would
-    # NP: make more sense to return zero
-    def get_poly(self, row: int, col: int) -> PolynomialData | None:
+
+    @override
+    def get_poly(self, row: int, col: int) -> PolyDict | None:
         if (row, col) in self.data:
-            return self.data[row, col]
+            return self.data[cast(MatrixIndex, (row, col))]
+        return None
+
+    # --- New API ---
+
+    @override
+    def at(self, row: int, col: int) -> PolyDict:
+        """ See :py:meth:`PolyMatrixMixin.at`. """
+        return self.data.get(MatrixIndex(row, col)) or PolyDict.empty()
 
 
-# NP: What does broadcast mean here? Is this a scalar multivariate polynomial
-# NP: that does not require the (row, col) indexing?
-# NP: If yes, consider new name ScalarPolynomial, ScalarPolyMatrix?
 class BroadcastPolyMatrixMixin(
     PolyMatrixMixin,
     abc.ABC,
 ):
+    """
+    TODO: docstring, similar to numpy broadcasting.
+    https://numpy.org/doc/stable/user/basics.broadcasting.html
+    """
     @property
     @abc.abstractmethod
-    def data(self) -> PolynomialData: ...
+    def data(self) -> PolyDict: ...
 
     # overwrites the abstract method of `PolyMatrixMixin`
-    def get_poly(self, col: int, row: int) -> PolynomialData | None:
+    def get_poly(self, col: int, row: int) -> PolyDict | None:
+        return self.data or None
+
+    # --- New API ---
+
+    @override
+    def at(self, row: int, col: int) -> PolyDict:
+        """ See :py:meth:`PolyMatrixMixin.at`. """
         return self.data
diff --git a/polymatrix/polymatrix/typing.py b/polymatrix/polymatrix/typing.py
index e304e5c..43d7210 100644
--- a/polymatrix/polymatrix/typing.py
+++ b/polymatrix/polymatrix/typing.py
@@ -1,9 +1,73 @@
-# monomial x1**2 x2
-# with indices {x1: 0, x2: 1}
-# is represented as ((0, 2), (1, 1))
+from __future__ import annotations
+from typing import NamedTuple, Iterable
 
-# NP: consider renaming Data to Index (data does not mean anything to me)
-# NP: introduce index classes?
+# TODO: remove these types, they are here for backward compatiblity
 MonomialData = tuple[tuple[int, int], ...]
 PolynomialData = dict[MonomialData, int | float]
 PolynomialMatrixData = dict[tuple[int, int], PolynomialData]
+
+
+class VariableIndex(NamedTuple):
+    """
+    Index for a variable raised to an integral power.
+    """
+    variable: int # index in ExpressionState object
+    power: int
+
+
+    def __lt__(self, other):
+        """ Variables indices can be sorted with respect to variable index. """
+        return self.variable < other.variable
+
+
+class MonomialIndex(tuple[VariableIndex]):
+    """
+    Index for a monomial, i.e. a product of `VariableIndex`.
+    """
+
+    @staticmethod
+    def empty() -> MonomialIndex:
+        return MonomialIndex(tuple())
+
+
+
+class PolyDict(dict[MonomialIndex, int | float]):
+    """ Polynomial, stored as a dictionary. """
+
+    # NOTE: this class directly inherits from dict instead of
+    # collections.UserDict, hence it is not possible to override any of the
+    # __dunder__ methods. If you like to change behaviour of one of these
+    # methods change to inherit from UserDict, this could incur a slight
+    # performance cost, because dict is straight CPython and thus faster
+
+    @staticmethod
+    def empty() -> PolyDict:
+        return PolyDict({})
+
+
+    def terms(self) -> Iterable[tuple[MonomialIndex, int | float]]:
+        """ Iterate over terms with a non-zero coefficient. """
+        # This is an alias for readability
+        yield from self.items()
+
+    def monomials(self) -> Iterable[MonomialIndex]:
+        """ Get monomials that compose the polynomial. """
+        # This is an alias for readability
+        yield from self.keys()
+
+    def coefficients(self) -> Iterable[int | float]:
+        """ Get the coefficients of the polynomial. """
+        # This is an alias for readability
+        yield from self.values()
+
+
+class MatrixIndex(NamedTuple):
+    """ Index to represent an entry of a matrix. """
+    row: int
+    col: int
+
+
+class PolyMatrixDict(dict[MatrixIndex, PolyDict]):
+    """ Matrix whose entries are polynomials, stored as a dictionary. """
+
+    # NOTE: Same as NOTE in PolyDict.
diff --git a/polymatrix/polymatrix/utils/mergemonomialindices.py b/polymatrix/polymatrix/utils/mergemonomialindices.py
index ec8f309..9ff2a63 100644
--- a/polymatrix/polymatrix/utils/mergemonomialindices.py
+++ b/polymatrix/polymatrix/utils/mergemonomialindices.py
@@ -1,12 +1,12 @@
-from polymatrix.polymatrix.typing import MonomialData
+from polymatrix.polymatrix.typing import MonomialIndex
 from polymatrix.polymatrix.utils.sortmonomialindices import sort_monomial_indices
 
 
 # NP: this is a werid case of the product
 # NP: merge is not intutive IMHO, find better name
 def merge_monomial_indices(
-    monomials: tuple[MonomialData, ...],
-) -> MonomialData:
+    monomials: tuple[MonomialIndex, ...],
+) -> MonomialIndex:
     """
     (x1**2 x2, x2**2)  ->  x1**2 x2**3
 
@@ -17,9 +17,8 @@ def merge_monomial_indices(
             ((1, 2),)              # x2**2
         )  ->  ((0, 2), (1, 3))    # x1**1 x2**3
     """
-
     if len(monomials) == 0:
-        return tuple()
+        return MonomialIndex.empty()
 
     elif len(monomials) == 1:
         return monomials[0]
diff --git a/polymatrix/polymatrix/utils/multiplypolynomial.py b/polymatrix/polymatrix/utils/multiplypolynomial.py
index 4b54813..5b7a62c 100644
--- a/polymatrix/polymatrix/utils/multiplypolynomial.py
+++ b/polymatrix/polymatrix/utils/multiplypolynomial.py
@@ -2,14 +2,14 @@ import itertools
 import math
 
 from polymatrix.polymatrix.utils.mergemonomialindices import merge_monomial_indices
-from polymatrix.polymatrix.typing import PolynomialData
+from polymatrix.polymatrix.typing import PolyDict
 
 
 # NP: compute index of product of polynomials
 def multiply_polynomial(
-    left: PolynomialData,
-    right: PolynomialData,
-    result: PolynomialData,
+    left: PolyDict,
+    right: PolyDict,
+    result: PolyDict,
 ) -> None:
     """
     Multiplies two polynomials `left` and `right` and adds the result to the mutable polynomial `result`.
diff --git a/polymatrix/polymatrix/utils/sortmonomialindices.py b/polymatrix/polymatrix/utils/sortmonomialindices.py
index ff4ce59..69d121a 100644
--- a/polymatrix/polymatrix/utils/sortmonomialindices.py
+++ b/polymatrix/polymatrix/utils/sortmonomialindices.py
@@ -1,7 +1,11 @@
 from polymatrix.polymatrix.typing import MonomialData
 
 
+from polymatrix.utils.deprecation import deprecated
+
 # 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,
 ) -> MonomialData: