diff options
Diffstat (limited to '')
| -rw-r--r-- | polymatrix/__init__.py | 4 | ||||
| -rw-r--r-- | polymatrix/polymatrix/init.py | 29 | ||||
| -rw-r--r-- | polymatrix/polymatrix/mixins.py | 71 |
3 files changed, 93 insertions, 11 deletions
diff --git a/polymatrix/__init__.py b/polymatrix/__init__.py index af1d5fa..a808d92 100644 --- a/polymatrix/__init__.py +++ b/polymatrix/__init__.py @@ -10,11 +10,14 @@ from polymatrix.expression import product as internal_product from polymatrix.expression.to import to_constant as internal_to_constant from polymatrix.expression.to import to_sympy as internal_to_sympy from polymatrix.denserepr.from_ import from_polymatrix +from polymatrix.polymatrix.init import to_affine_expression Expression = internal_Expression ExpressionState = internal_ExpressionState init_expression_state = internal_init_expression_state +make_state = init_expression_state + from_ = internal_from v_stack = internal_v_stack h_stack = internal_h_stack @@ -25,4 +28,5 @@ to_sympy_repr = internal_to_sympy to_sympy = internal_to_sympy to_matrix_repr = from_polymatrix to_dense = from_polymatrix +to_affine = to_affine_expression from_names = internal_from_names diff --git a/polymatrix/polymatrix/init.py b/polymatrix/polymatrix/init.py index 9a85a0d..f9a3de1 100644 --- a/polymatrix/polymatrix/init.py +++ b/polymatrix/polymatrix/init.py @@ -5,8 +5,9 @@ import numpy as np from typing import TYPE_CHECKING -from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl +from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl, PolyMatrixAsAffineExpressionImpl from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MatrixIndex, MonomialIndex, VariableIndex +from polymatrix.polymatrix.mixins import PolyMatrixAsAffineExpressionMixin if TYPE_CHECKING: from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin, PolyMatrixMixin @@ -48,3 +49,29 @@ def init_broadcast_poly_matrix( data=data, shape=shape, ) + + +def to_affine_expression(p: PolyMatrixMixin) -> PolyMatrixAsAffineExpressionMixin: + + if isinstance(p, PolyMatrixAsAffineExpressionMixin): + 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())) + + nrows, ncols = p.shape + data = np.zeros((nrows, ncols * len(monomials))) + + # Generate slices to index sub-matrices + slices = { + m: tuple(range(i*ncols, (i+1)*ncols)) + for i, m in enumerate(monomials) + } + + for (row, col), poly in p.entries(): + for monomial, coeff in poly.terms(): + slice = slices[monomial] + data[row, slice[col]] += coeff + + return PolyMatrixAsAffineExpressionImpl(data, p.shape, slices) diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py index 25b968c..b33379b 100644 --- a/polymatrix/polymatrix/mixins.py +++ b/polymatrix/polymatrix/mixins.py @@ -4,6 +4,7 @@ import abc import typing import itertools import functools +import math import numpy as np from numpy.typing import NDArray @@ -220,7 +221,10 @@ class PolyMatrixAsAffineExpressionMixin( # -- Special methods for affine expression - def monomials_eval(self, x: tuple) -> dict[MonomialIndex, float]: + def monomials_eval(self, x: Sequence[float] | MatrixType, monomials: Iterable[MonomialIndex]) -> dict[MonomialIndex, float]: + raise NotImplementedError + + def monomials_eval_all(self, x: Sequence[float] | MatrixType) -> dict[MonomialIndex, float]: r""" Evaluate all monomials up to degree :math:`d`. @@ -228,6 +232,20 @@ class PolyMatrixAsAffineExpressionMixin( numerical values of :math:`\langle x \rangle_d = \{1, x_1, x_2, \ldots, 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(v for m in self.slices.keys() + for v in m)) + + # TODO: this value should be cached + q = len(variables) + + if len(x) != q: + # TODO: improve error message, retrieve variable names from state? + raise ValueError(f"The value must be a {q} dimensional, however " + f"{x} is {len(x)} dimensional.") + def compute_with_triangles(d, x): # To compute all monomials we need to compute all integer points contained # in a d-dimensional triangle-like shape. For example given a q-dimensional @@ -300,7 +318,7 @@ class PolyMatrixAsAffineExpressionMixin( return terms # Get the highest degree we need to compute from the monomial indices - # TODO: this value could be cached + # TODO: this value should be cached d = max(m.degree for m in self.slices.keys()) # Add the constant term @@ -309,7 +327,6 @@ class PolyMatrixAsAffineExpressionMixin( } # Compute all monomials and convert key to MonomialIndex - # TODO: modify compute_with_triangles to directly use MonomialIndex for key, val in compute_with_triangles(d, x).items(): # See comment above for structure of key # Count occurrences to convert to MonomialIndex @@ -320,7 +337,7 @@ class PolyMatrixAsAffineExpressionMixin( else: new_key[e] = 1 - index = MonomialIndex(VariableIndex(k, v) for k, v in new_key.items()) + index = MonomialIndex(VariableIndex(variables[k].index, v) for k, v in new_key.items()) monomial_values[index] = val # Sort the values by monomial index, which have a total order @@ -346,10 +363,9 @@ class PolyMatrixAsAffineExpressionMixin( r""" Iterate over the coefficients grouped by degree. """ - monomials = self.slices.keys() - groups = itertools.groupby(monomials, lambda m: m.degree) + groups = itertools.groupby(self.slices.keys(), lambda m: m.degree) for degree, monomials in groups: - yield degree, np.vstack(list(self.data[m] for m in monomials)) + yield degree, np.hstack(list(self.data[:, self.slices[m]] for m in monomials)) def affine_eval(self, x: MatrixType) -> MatrixType: r""" @@ -357,12 +373,47 @@ class PolyMatrixAsAffineExpressionMixin( :math:`p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha` at a given :math:`x`. """ - # Compute all powers of x - monomials = np.array(tuple(self.monomials_eval(tuple(x)).items())) + # 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)) + + if len(x) != q: + # TODO: improve error message, retrieve variable names from state? + raise ValueError(f"The value must be a {q} dimensional, however " + f"{x} is {len(x)} dimensional.") + + # Get the highest degree we need to compute from the monomial indices + # TODO: this value should be cached + d = max(m.degree for m in self.slices.keys()) + + # Compute number of monomials that are generated by that + # TODO: this value should be cached + n = sum(math.comb(k + q - 1, k) for k in range(0, d+1)) + + # Heuristic on efficiency, computing all monomials exploits the + # structure to reduce the number of computations necessary. But the + # efficiency gain is useless if most of the computed terms are not + # used. + if len(self.slices) > n / 2: + # Compute all powers of x up to degree d + all_monomials = self.monomials_eval_all(x) + monomials = { + m: all_monomials[m] + for m in self.slices.keys() + } + + else: + # Compute only the necessary terms + monomials = self.monomials_eval(x, self.slices.keys()) + + mvalues = np.array(tuple(monomials.values())) + # Evaluate the affine expression _, ncols = self.shape - return self.data @ (np.kron(monomials, np.eye(ncols))) + return self.data @ (np.kron(mvalues, np.eye(ncols)).T) def affine_eval_fn(self) -> Callable[[MatrixType], MatrixType]: r""" |