Create PolyMatrixAsAffineExpression as replacement for DenseExpr

2 files changed, 296 insertions, 38 deletions

diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py
index 152d5fa..37837df 100644
--- a/polymatrix/polymatrix/impl.py
+++ b/polymatrix/polymatrix/impl.py
@@ -1,8 +1,8 @@
 import dataclassabc
 
 from polymatrix.polymatrix.abc import PolyMatrix
-from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin
-from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict
+from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin, PolyMatrixAsAffineExpressionMixin
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MonomialIndex
 
 
 @dataclassabc.dataclassabc(frozen=True)
@@ -15,3 +15,10 @@ class PolyMatrixImpl(PolyMatrix):
 class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin):
     data: PolyDict
     shape: tuple[int, int]
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class PolyMatrixAsAffineExpressionImpl(PolyMatrixAsAffineExpressionMixin):
+    data: PolyMatrixAsAffineExpressionMixin.MatrixType
+    shape: tuple[int, int]
+    slices: dict[MonomialIndex, PolyMatrixAsAffineExpressionMixin.MatrixType]
diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index ddd00fc..cbd9749 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -1,42 +1,40 @@
+from __future__ import annotations
+
 import abc
 import typing
+import itertools
+import functools
+import numpy as np
 
-from typing import Iterable, cast
+from numpy.typing import NDArray
+from typing import Iterable, Sequence, Callable, cast, TYPE_CHECKING
 from typing_extensions import override
 
-from polymatrix.polymatrix.typing import PolyDict, PolyMatrixDict, MatrixIndex
+from polymatrix.polymatrix.typing import PolyDict, PolyMatrixDict, MatrixIndex, MonomialIndex, VariableIndex
 from polymatrix.utils.deprecation import deprecated
 
+if TYPE_CHECKING:
+    from polymatrix.expression.mixins.variablemixin import VariableMixin
+    from polymatrix.expressionstate.mixins import ExpressionStateMixin
+
 
 class PolyMatrixMixin(abc.ABC):
+    """
+    Matrix with polynomial entries.
+    """
+
     @property
     @abc.abstractmethod
     def shape(self) -> tuple[int, int]: ...
 
-    @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). 
+        """ Return the polynomial at the entry (row, col).
+        If the entry is zero it returns an empty `PolyDict`, i.e. an empty
+        dictionary `{}`.
 
         .. code:: py
+
             p.at(row, col) # if you have row and col
             p.at(*entry) # if you have a MatrixIndex
         """
@@ -45,9 +43,10 @@ class PolyMatrixMixin(abc.ABC):
         """
         Iterate over the non-zero entries of the polynomial matrix.
 
-        TODO: docstrings, use like this (`p` is a polymatrix)
+        Typical usage looks like this (`p` is a polymatrix):
 
         .. code:: py
+
             for entry, poly in p.entries():
                 for monomial, coeff in poly.terms():
                     # do something with coefficients
@@ -58,20 +57,43 @@ class PolyMatrixMixin(abc.ABC):
                 if poly := self.at(row, col):
                     yield MatrixIndex(row, col), poly
 
+    # -- Old API ---
+
+    @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
+
 
 class PolyMatrixAsDictMixin(
     PolyMatrixMixin,
     abc.ABC,
 ):
-    """
-    Polynomial matrix, stored as a dictionary.
-    """
+    """ Matrix with polynomial entries, stored as a dictionary. """
 
     @property
     @abc.abstractmethod
     def data(self) -> PolyMatrixDict:
         """" Get the dictionary. """
 
+    @override
+    def at(self, row: int, col: int) -> PolyDict:
+        """ See :py:meth:`PolyMatrixMixin.at`. """
+        return self.data.get(MatrixIndex(row, col)) or PolyDict.empty()
+
+    # -- Old API ---
 
     @override
     def get_poly(self, row: int, col: int) -> PolyDict | None:
@@ -79,13 +101,6 @@ class PolyMatrixAsDictMixin(
             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()
-
 
 class BroadcastPolyMatrixMixin(
     PolyMatrixMixin,
@@ -99,13 +114,249 @@ class BroadcastPolyMatrixMixin(
     @abc.abstractmethod
     def data(self) -> PolyDict: ...
 
+    @override
+    def at(self, row: int, col: int) -> PolyDict:
+        """ See :py:meth:`PolyMatrixMixin.at`. """
+        return self.data
+
+    # --- Old API ---
+
     # overwrites the abstract method of `PolyMatrixMixin`
     def get_poly(self, col: int, row: int) -> PolyDict | None:
         return self.data or None
 
-    # --- New API ---
+
+class PolyMatrixAsAffineExpressionMixin(
+        PolyMatrixMixin,
+        abc.ABC
+):
+    r"""
+    Matrix with polynomial entries, stored as an expression that is affine in
+    the monomials.
+
+    For clarity we will show this through an example. Suppose :math:`x,y` are
+    polynomial variables, and we have a matrix of polynomials
+
+    .. math::
+        P(x, y) = \begin{bmatrix}
+            x + y^2  & x + 2\\
+            2x^2 + xy + 1 & y
+        \end{bmatrix}
+        \in \mathbf{R}_2(x, y)^{2 \times 2}
+
+    We can rewrite :math:`P(x,y)` as an affine expression in terms of
+    the monomial basis generated by :math:`x` and :math:`y` and get
+
+    .. math::
+
+        P(x,y)
+        = \begin{bmatrix} 0 & 2 \\ 1 & 0 \end{bmatrix}
+        + \begin{bmatrix} 1 & 1 \\ 0 & 0 \end{bmatrix} x
+        + \begin{bmatrix} 0 & 0 \\ 0 & 1 \end{bmatrix} y
+        + \begin{bmatrix} 0 & 0 \\ 2 & 0 \end{bmatrix} x^2
+        + \begin{bmatrix} 0 & 0 \\ 1 & 0 \end{bmatrix} xy
+        + \begin{bmatrix} 1 & 0 \\ 0 & 0 \end{bmatrix} y^2
+
+    More generally for a :math:`p \in \mathbf{R}_d(x)^{n\times m}, x \in
+    \mathbf{R}^q` and the set of monomials of degree :math:`d` or lower
+    :math:`\langle x \rangle_d = \{1, x_1, x_2, \ldots, x_q, x_1^2, x_1 x_2,
+    \ldots, x_q^2, x_1^3, \ldots, x_q^d\}` by defining the set multi-indices of
+    :math:`\langle x \rangle_d` to be 
+
+    .. math::
+        \mathcal{E} \langle x \rangle_d = \{
+            (0, \ldots, 0), (1,0,\ldots,0), (0,1,\ldots,0), \ldots,
+            (2,0,\ldots,0), \ldots, (0,\ldots,d)\}
+            \subset \mathbf{Z}^d
+
+    we can write
+    
+    .. math::
+        p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha
+
+    where each matrix :math:`A_\alpha \in \mathbf{R}^{n\times m}`. In the
+    context of this class we will call the :math:`A_\alpha` matrices
+    `affine coefficients`.
+    """
+
+    # Type of matrix that stores affifne coefficients
+    # TODO: Consider also scipy.sparse.csc_matrix instead of NDArray
+    # since slices are row slices
+    MatrixType = NDArray[np.float64]
+
+    @property
+    @abc.abstractmethod
+    def data(self) -> MatrixType:
+        r"""
+        The big matrix that store all :math:`A_\alpha` matrices:
+
+        .. math::
+            A = \begin{bmatrix}
+                A_{\alpha_1} & A_{\alpha_2} & \cdots & A_{\alpha_N}
+            \end{bmatrix}
+            \in
+            \mathbf{R}^{n \times mN},
+            \quad
+            N = \sum_{k = 0}^d {q + k \choose k}
+        """
+        ...
+
+    @property
+    @abc.abstractmethod
+    def slices(self) -> dict[MonomialIndex, tuple[int, int]]:
+        # column slices of the big array stored in self.data to get the individual A_alphas
+        ...
 
     @override
     def at(self, row: int, col: int) -> PolyDict:
-        """ See :py:meth:`PolyMatrixMixin.at`. """
-        return self.data
+        # TODO: docstring
+        p = PolyDict.empty()
+        for monomial in self.slices.keys():
+            p[monomial] = self.affine_coefficient(monomial)[row, col]
+
+        return p
+
+    # -- Special methods for affine expression
+
+    def monomials_eval(self, x: tuple) -> dict[MonomialIndex, float]:
+        r"""
+        Evaluate all monomials up to degree :math:`d`.
+
+        Given an :math:`x \in \mathbf{R}^q` this function computes the
+        numerical values of :math:`\langle x \rangle_d = \{1, x_1, x_2, \ldots,
+        x_q, x_1^2, x_1 x_2, \ldots, x_q^2, x_1^3, \ldots, x_q^d\}`.
+        """
+        def compute_with_triangles(d, x):
+            # To compute all monomials we need to compute all integer points contained
+            # in a d-dimensional triangle-like shape. For example given a q-dimensional
+            # x = (x_1, x_2, ... x_q) to compute all quadratic terms (d = 2) we need to
+            # evaluate all points in the following upper triangular part of the grid:
+            #
+            #              x1       x2      x3      ...     xq
+            #          ┌───────────────────────────────────────┐
+            #      x1  │ x1*x1    x1*x2    x1*x3    ...  x1*xq │
+            #          │                                       │
+            #      x2  │          x2*x2    x2*x3    ...  x2*xq │
+            #          │                                       │
+            #      x3  │                   x3*x3    ...  x3*xq │
+            #          │                                       │
+            #     ...  │                            ...   ...  │
+            #          │                                       │
+            #      xq  │                                 xq*xq │
+            #          └───────────────────────────────────────┘
+            #
+            # For d = 2 we compute the products by iterating k from 1 to q and then
+            # inside of that iterating j from k to q and multiply x[j] * x[k].
+            #
+            # For d = 3 the shape is an irregular tetrahedron, and so when we iterate k
+            # from 1 to q then we must iterate over "triangular slices" of the
+            # tetrahedron. Like for d = 2 where we take only values that have j > k,
+            # for d = 3 we take all tetrahedron slices that are "higher" than level k.
+            #
+            # This pattern continues for d > 3 so we define a recursion over d by
+            # defining the triangle function below, which returns all monomial
+            # combinations of exactly degree d that are "higher" than j.
+            #
+            # To keep track of the values during the recursion we define a dictionary
+            # where the key is a tuple with length smaller or equal to q that contains
+            # the index of the variables that make up the product. For example:
+            #
+            #    x_1 <-> (0,)    x_1 * x_1       <-> (0, 0)
+            #    x_2 <-> (1,)    x_2 * x_3       <-> (1, 2)
+            #    x_3 <-> (2,)    x_1 * x_2 * x_2 <-> (0, 1, 1)
+            #
+            
+            # Number variables
+            q = len(x)
+
+            # TODO: rewrite recursive code using dynamic programming table to
+            # replace functools.cache
+            @functools.cache
+            def triangle(d, j):
+                # Base case
+                if d == 1:
+                    return {
+                        (i,): x[i]
+                        for i in range(j, q)
+                    }
+
+                current = {}
+                for k in range(j, q):
+                    # Recursion! Iterate over the slices which are given by the
+                    # d-1 case and are higher than k
+                     for i, v in triangle(d-1, k).items():
+                        current[i + (k,)] = v * x[k]
+
+                return current
+
+            # The monomials from degree 1 to d is the union of triangles with j = 0,
+            # i.e. the union of all terms of exactly degree d
+            terms = {}
+            for i in range(1, d+1):
+                terms |= triangle(i, 0)
+
+            return terms
+
+        # Get the highest degree we need to compute from the monomial indices
+        # TODO: this value could be cached
+        d = max(m.degree for m in self.slices.keys())
+
+        # Add the constant term
+        monomial_values = {
+            MonomialIndex.constant(): 1.
+        }
+
+        # Compute all monomials and convert key to MonomialIndex
+        # TODO: modify compute_with_triangles to directly use MonomialIndex
+        for key, val in compute_with_triangles(d, x).items():
+            # See comment above for structure of key
+            # Count occurrences to convert to MonomialIndex
+            new_key: dict[int, int] = {}
+            for e in key:
+                if e in new_key:
+                    new_key[e] += 1
+                else:
+                    new_key[e] = 1
+
+            index = MonomialIndex(VariableIndex(k, v) for k, v in new_key.items())
+            monomial_values[index] = val
+        
+        # Sort the values by monomial index, which have a total order
+        return dict(sorted(monomial_values.items()))
+
+    def affine_coefficient(self, monomial: MonomialIndex) -> MatrixType:
+        r""" Get the affine coefficient :math:`A_\alpha` associated to :math:`x^\alpha`. """
+        columns = range(*self.slices[monomial])
+        return self.data[:, columns]
+
+    def affine_coefficients_by_degrees(self) -> Iterable[tuple[int, MatrixType]]:
+        r"""
+        Iterate over the coefficients grouped by degree.
+        """
+        monomials = self.slices.keys()
+        groups = itertools.groupby(monomials, lambda m: m.degree)
+        for degree, monomials in groups:
+            yield degree, np.vstack(list(self.data[m] for m in monomials))
+
+    def affine_eval(self, x: MatrixType) -> MatrixType:
+        r"""
+        Evaluate the affine expression
+        :math:`p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha`
+        at a given :math:`x`.
+        """
+        # Compute all powers of x
+        monomials = np.array(tuple(self.monomials_eval(tuple(x)).items()))
+
+        # Evaluate the affine expression
+        nrows, ncols = self.shape
+        return self.data @ (np.kron(monomials, np.eye(ncols)))
+
+    def affine_eval_fn(self) -> Callable[[MatrixType], MatrixType]:
+        r"""
+        Get a function that when called evaluates the expression
+        :math:`p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha`
+        at :math:`x`.
+        """
+        # TODO: docstring
+        # TODO: If slow consider replacing with toolz.functoolz.curry from ctoolz
+        return functools.partial(PolyMatrixAsAffineExpression.affine_eval, self)