Clean up PolyMatrixAsAffineExpression

4 files changed, 58 insertions, 45 deletions

diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py
index 2ac4dc0..3ed65dc 100644
--- a/polymatrix/polymatrix/impl.py
+++ b/polymatrix/polymatrix/impl.py
@@ -23,3 +23,4 @@ class PolyMatrixAsAffineExpressionImpl(PolyMatrixAsAffineExpressionMixin):
     shape: tuple[int, int]
     slices: dict[MonomialIndex, PolyMatrixAsAffineExpressionMixin.MatrixType]
     degree: int
+    variable_indices: tuple[int]
diff --git a/polymatrix/polymatrix/index.py b/polymatrix/polymatrix/index.py
index e50df43..2a46625 100644
--- a/polymatrix/polymatrix/index.py
+++ b/polymatrix/polymatrix/index.py
@@ -98,14 +98,18 @@ class MonomialIndex(tuple[VariableIndex]):
         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)
+        if degree == 0:
+            yield MonomialIndex.constant()
+
+        else:
+            combinations = []
+            for comb in combinations_with_replacement(variables, degree):
+                # Convert variables into monomials and mutliply them
+                monomials = (MonomialIndex((v,)) for v in comb)
+                r = reduce(MonomialIndex.product, monomials, MonomialIndex.constant())
+                combinations.append(r)
+
+            yield from sorted(combinations)
 
     @staticmethod
     def product(left: MonomialIndex, right: MonomialIndex) -> MonomialIndex:
diff --git a/polymatrix/polymatrix/init.py b/polymatrix/polymatrix/init.py
index d45b02e..77c544c 100644
--- a/polymatrix/polymatrix/init.py
+++ b/polymatrix/polymatrix/init.py
@@ -59,9 +59,17 @@ def to_affine_expression(p: PolyMatrixMixin) -> PolyMatrixAsAffineExpressionMixi
         return p
 
     # get all monomials that come up in p, and sort them
-    monomials = sorted(set(m for _, entry in p.entries() for m
-                           in entry.monomials()))
+    monomials = tuple(sorted(set(m for _, entry in p.entries() for m
+                                 in entry.monomials())))
 
+    # get all variables that come up in expression, and sort them
+    variables = tuple(sorted(set(v.index
+                                 for m in monomials
+                                 for v in m)))
+
+    max_degree = max(m.degree for m in monomials) 
+
+    # Prepare big array
     nrows, ncols = p.shape
     data = np.zeros((nrows, ncols * len(monomials)))
     
@@ -71,12 +79,15 @@ def to_affine_expression(p: PolyMatrixMixin) -> PolyMatrixAsAffineExpressionMixi
         for i, m in enumerate(monomials)
     }
 
+    # Copy coefficients
     for (row, col), poly in p.entries():
         for monomial, coeff in poly.terms():
             slice = slices[monomial]
             data[row, slice[col]] += coeff
 
-    max_degree = max(m.degree for m in slices.keys())
-
     return PolyMatrixAsAffineExpressionImpl(
-            affine_coefficients=data, shape=p.shape, slices=slices, degree=max_degree)
+            affine_coefficients=data,
+            shape=p.shape,
+            slices=slices,
+            degree=max_degree,
+            variable_indices=variables)
diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index 95c875a..559a72b 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -217,6 +217,14 @@ class PolyMatrixAsAffineExpressionMixin(
     @abc.abstractmethod
     def degree(self) -> int:
         """ Maximal degree of the affine expressions """
+        # Value of max(m.degree for m in self.slices.keys())
+        # cached at construction
+
+
+    @property
+    @abc.abstractmethod
+    def variable_indices(self) -> tuple[int]:
+        """ Indices of the variables involved in the expression, sorted. """
 
     @override
     def at(self, row: int, col: int) -> PolyDict:
@@ -241,14 +249,8 @@ class PolyMatrixAsAffineExpressionMixin(
         x_q, x_1^2, x_1 x_2, \ldots, x_q^2, x_1^3, \ldots, x_q^d\}`.
         """
 
-        # Get involved variables
-        # TODO: this value should be cached
-        variables = sorted(set(VariableIndex(v.index, 1)
-                               for m in self.slices.keys()
-                               for v in m))
-
-        # TODO: this value should be cached
-        q = len(variables)
+        # Number of variables in the expression
+        q = len(self.variable_indices)
 
         if len(x) != q:
             # TODO: improve error message, retrieve variable names from state?
@@ -294,7 +296,7 @@ class PolyMatrixAsAffineExpressionMixin(
             #    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)
 
@@ -342,14 +344,22 @@ class PolyMatrixAsAffineExpressionMixin(
                 else:
                     new_key[e] = 1
 
-            index = MonomialIndex(VariableIndex(variables[k].index, v) for k, v in new_key.items())
+            index = MonomialIndex(VariableIndex(self.variable_indices[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:
+    def affine_coefficient(self, monomial: MonomialIndex | VariableIndex | int) -> MatrixType:
         r""" Get the affine coefficient :math:`A_\alpha` associated to :math:`x^\alpha`. """
+        if isinstance(monomial, int):
+            # interpret as linear
+            monomial = VariableIndex(monomial, power=1)
+
+        if isinstance(monomial, VariableIndex):
+            monomial = MonomialIndex((monomial,))
+
         if monomial not in self.slices.keys():
             nrows, _ = self.shape
             return np.zeros(self.shape)
@@ -357,29 +367,19 @@ class PolyMatrixAsAffineExpressionMixin(
         columns = self.slices[monomial]
         return self.affine_coefficients[:, columns]
 
-    def affine_coefficients_by_degrees(
-            self,
-            variables: Iterable[VariableIndex] | Iterable[MonomialIndex] | None = None
-        ) -> Iterable[tuple[int, MatrixType]]:
+    def affine_coefficients_by_degrees(self, variables: Sequence[VariableIndex] | None = None) -> Iterable[tuple[int, MatrixType]]:
         """
         Iterate over the coefficients grouped by degree.
         """
         if not variables:
-            # Get involved variables
-            # TODO: this value should be cached
-            variables = tuple(set(VariableIndex(v.index, power=1)
-                                   for m in self.slices.keys()
-                                   for v in m))
-
-        # FIXME: allow passing linear monomial indices
-        # elif isinstance(variables[0], MonomialIndex):
-        #     for v in variables:
-        #         if v.degree != 1:
-        #             raise ValueError("Variables ...")
-        #         variables = map(lambda m: m[0], variables)
+            # cast to tuple because it is used multiple times in the loop
+            variables = tuple(VariableIndex(i, power=1) for i in self.variable_indices)
 
+        else:
+            # TODO: document behaviour on why it is allowed to pass varirables
+            # it has to do with sumofsquares/solver/cvxopt.py
+            variables = tuple(sorted(variables))
 
-        variables = sorted(variables)
         for degree in range(self.degree +1):
             monomials = MonomialIndex.combinations_of_degree(variables, degree)
             columns = tuple(self.affine_coefficient(m) for m in monomials)
@@ -405,10 +405,7 @@ class PolyMatrixAsAffineExpressionMixin(
         """
 
         # Get the number of variables involved
-        # TODO: this value should be cached
-        q = len(set(v.index
-                    for m in self.slices.keys()
-                    for v in m))
+        q = len(self.variable_indices)
 
         if len(x) != q:
             # TODO: improve error message, retrieve variable names from state?