Add grouping of affine coefficients by variables

This is meant to be used in the sumofsquares package

1 files changed, 18 insertions, 1 deletions

diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index 3ca4f66..e13b868 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -360,13 +360,30 @@ class PolyMatrixAsAffineExpressionMixin(
         return self.data[:, columns]
 
     def affine_coefficients_by_degrees(self) -> Iterable[tuple[int, MatrixType]]:
-        r"""
+        """
         Iterate over the coefficients grouped by degree.
         """
         groups = itertools.groupby(self.slices.keys(), lambda m: m.degree)
         for degree, monomials in groups:
             yield degree, np.hstack(list(self.data[:, self.slices[m]] for m in monomials))
 
+    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))
+
+
     def affine_eval(self, x: MatrixType) -> MatrixType:
         r"""
         Evaluate the affine expression