From 89b39576696b17e1dd3c8469282b5e583ae3c5b4 Mon Sep 17 00:00:00 2001 From: Nao Pross Date: Fri, 10 May 2024 22:13:10 +0200 Subject: Clean up PolyMatrixAsAffineExpression --- polymatrix/polymatrix/impl.py | 1 + polymatrix/polymatrix/index.py | 20 ++++++++------ polymatrix/polymatrix/init.py | 21 ++++++++++---- polymatrix/polymatrix/mixins.py | 61 ++++++++++++++++++++--------------------- 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? -- cgit v1.2.1