diff options
-rw-r--r-- | polymatrix/polymatrix/mixins.py | 19 |
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 |