diff options
| -rw-r--r-- | polymatrix/polymatrix/index.py | 26 | ||||
| -rw-r--r-- | 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 """ @@ -86,6 +93,21 @@ class MonomialIndex(tuple[VariableIndex]): 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: """ Compute the index of the product of two monomials. @@ -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""" |