From 103d53d0925abc82fc6f15a603aa42a7b6638f48 Mon Sep 17 00:00:00 2001
From: Nao Pross <np@0hm.ch>
Date: Tue, 7 May 2024 01:53:32 +0200
Subject: Add MonomialIndex.combinations_of_degree

To improve PolyMatrixAsAffineExpression.affine_coefficients_by_degrees
---
 polymatrix/polymatrix/index.py  | 26 +++++++++++++++++++++++++-
 polymatrix/polymatrix/mixins.py | 34 ++++++++++++----------------------
 2 files changed, 37 insertions(+), 23 deletions(-)

diff --git a/polymatrix/polymatrix/index.py b/polymatrix/polymatrix/index.py
index 384feb8..d31e987 100644
--- a/polymatrix/polymatrix/index.py
+++ b/polymatrix/polymatrix/index.py
@@ -2,7 +2,8 @@ from __future__ import annotations
 
 from typing import NamedTuple, Iterable, cast
 from collections import UserDict
-from itertools import dropwhile
+from itertools import dropwhile, combinations_with_replacement
+from functools import reduce
 
 # TODO: remove these types, they are here for backward compatiblity
 MonomialData = tuple[tuple[int, int], ...]
@@ -41,6 +42,12 @@ class MonomialIndex(tuple[VariableIndex]):
     Constant monomial indices are considered equal.
     """
 
+    def __repr__(self):
+        return "MonomialIndex(" + ",".join(repr(v) for v in self) + ")"
+
+    def __str__(self):
+        return "MonomialIndex(" + ",".join(f"{v.index}^{v.power}" for v in self) + ")"
+
     @property
     def degree(self) -> int:
         """ Degree of the monomial """
@@ -85,6 +92,21 @@ class MonomialIndex(tuple[VariableIndex]):
         """ Sort the variable indices inside the monomial index. """
         return MonomialIndex(tuple(sorted(index)))
 
+    @staticmethod
+    def combinations_of_degree(variables: Iterable[VariableIndex], degree: int) -> Iterable[MonomialIndex]:
+        """
+        Compute the indices of all monomial combinations `variables` of exactly
+        degree `degree`.
+        """
+        combinations = []
+        for comb in combinations_with_replacement(variables, degree):
+            # Convert variables into monomials and mutliply them
+            monomials = map(lambda v: MonomialIndex((v,)), comb)
+            r = reduce(MonomialIndex.product, monomials, MonomialIndex.constant())
+            combinations.append(r)
+
+        return sorted(combinations)
+
     @staticmethod
     def product(left: MonomialIndex, right: MonomialIndex) -> MonomialIndex:
         """
@@ -113,6 +135,8 @@ class MonomialIndex(tuple[VariableIndex]):
 
     @staticmethod
     def differentiate(index: MonomialIndex, wrt: int) -> MonomialIndex | None:
+        # TODO: implement this and then PolyDict.differentiate to replace
+        # expression.utils.gederivativemonomials.differentiate_polynomial
         raise NotImplementedError
         
 
diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index 81aa2dd..ce6d09b 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -239,7 +239,8 @@ class PolyMatrixAsAffineExpressionMixin(
 
         # Get involved variables
         # TODO: this value should be cached
-        variables = sorted(set(v for m in self.slices.keys()
+        variables = sorted(set(VariableIndex(v.index, 1)
+                               for m in self.slices.keys()
                                for v in m))
 
         # TODO: this value should be cached
@@ -359,31 +360,20 @@ class PolyMatrixAsAffineExpressionMixin(
         columns = range(*self.slices[monomial])
         return self.data[:, columns]
 
-    def affine_coefficients_by_degrees(self) -> Iterable[tuple[int, tuple[MonomialIndex], MatrixType]]:
+    def affine_coefficients_by_degrees(self) -> Iterable[tuple[int, MatrixType]]:
         """
         Iterate over the coefficients grouped by degree.
         """
-        groups = itertools.groupby(self.slices.keys(), lambda m: m.degree)
-        for degree, monomials in groups:
-            monomials = tuple(monomials)
-            columns = np.hstack(list(self.data[:, self.slices[m]] for m in monomials))
-            yield degree, monomials, columns
+        # Get involved variables
+        # TODO: this value should be cached
+        variables = sorted(set(VariableIndex(v.index, power=1)
+                               for m in self.slices.keys()
+                               for v in m))
 
-    def affine_coefficients_by_variable(self) -> Iterable[tuple[int, MatrixType]]:
-        r"""
-        If the affine expression is linear, that is, :math:`|\alpha| \leq 1`
-        for all :math:`\alpha`. Iterate over the coefficients grouped by
-        variable.
-        """
-        if any(m.degree > 1 for m in self.slices.keys()):
-            raise RuntimeError("Affine expression is not linear! "
-                               "This operation is only allowed for linear affine expresions.")
-
-        # if len(m) is zero it means that it is a constant
-        # a more explicit way would be MonomialIndex.is_constant(m)
-        groups = itertools.groupby(self.slices.keys(), lambda m: m[0] if len(m) > 0 else m)
-        for v, monomials in groups:
-            yield v, np.hstack(list(self.data[:, self.slices[m]] for m in monomials))
+        for degree in range(self.degree +1):
+            monomials = MonomialIndex.combinations_of_degree(variables, degree)
+            columns = tuple(self.affine_coefficient(m) for m in monomials)
+            yield degree, np.hstack(columns)
 
     def affine_eval(self, x: MatrixType) -> MatrixType:
         r"""
-- 
cgit v1.2.1