diff options
| -rw-r--r-- | polymatrix/polymatrix/impl.py | 11 | ||||
| -rw-r--r-- | polymatrix/polymatrix/mixins.py | 323 |
2 files changed, 296 insertions, 38 deletions
diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py index 152d5fa..37837df 100644 --- a/polymatrix/polymatrix/impl.py +++ b/polymatrix/polymatrix/impl.py @@ -1,8 +1,8 @@ import dataclassabc from polymatrix.polymatrix.abc import PolyMatrix -from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin -from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict +from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin, PolyMatrixAsAffineExpressionMixin +from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MonomialIndex @dataclassabc.dataclassabc(frozen=True) @@ -15,3 +15,10 @@ class PolyMatrixImpl(PolyMatrix): class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin): data: PolyDict shape: tuple[int, int] + + +@dataclassabc.dataclassabc(frozen=True) +class PolyMatrixAsAffineExpressionImpl(PolyMatrixAsAffineExpressionMixin): + data: PolyMatrixAsAffineExpressionMixin.MatrixType + shape: tuple[int, int] + slices: dict[MonomialIndex, PolyMatrixAsAffineExpressionMixin.MatrixType] diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py index ddd00fc..cbd9749 100644 --- a/polymatrix/polymatrix/mixins.py +++ b/polymatrix/polymatrix/mixins.py @@ -1,42 +1,40 @@ +from __future__ import annotations + import abc import typing +import itertools +import functools +import numpy as np -from typing import Iterable, cast +from numpy.typing import NDArray +from typing import Iterable, Sequence, Callable, cast, TYPE_CHECKING from typing_extensions import override -from polymatrix.polymatrix.typing import PolyDict, PolyMatrixDict, MatrixIndex +from polymatrix.polymatrix.typing import PolyDict, PolyMatrixDict, MatrixIndex, MonomialIndex, VariableIndex from polymatrix.utils.deprecation import deprecated +if TYPE_CHECKING: + from polymatrix.expression.mixins.variablemixin import VariableMixin + from polymatrix.expressionstate.mixins import ExpressionStateMixin + class PolyMatrixMixin(abc.ABC): + """ + Matrix with polynomial entries. + """ + @property @abc.abstractmethod def shape(self) -> tuple[int, int]: ... - @deprecated("Replaced by PolyMatrixMixin.entries()") - def gen_data(self) -> typing.Generator[tuple[tuple[int, int], PolyDict], None, None]: - yield from self.entries() - - @abc.abstractmethod - @deprecated("Replaced by PolyMatrixMixin.at(row, col).") - def get_poly(self, row: int, col: int) -> PolyDict | None: - # NOTE: in the new API when there is nothing at an entry the at() - # function returns an empty dictionary. This is because emtpy - # dicionaries are falsy so you can do - # - # `if p.at(row, col):` - # - # like it was possible with get_poly. However this is better because - # id does not make use of a null (None) type. - return self.at(row, col) or None - - # --- New API --- - @abc.abstractmethod def at(self, row: int, col: int) -> PolyDict: - """ Return the polynomial at the entry (row, col). + """ Return the polynomial at the entry (row, col). + If the entry is zero it returns an empty `PolyDict`, i.e. an empty + dictionary `{}`. .. code:: py + p.at(row, col) # if you have row and col p.at(*entry) # if you have a MatrixIndex """ @@ -45,9 +43,10 @@ class PolyMatrixMixin(abc.ABC): """ Iterate over the non-zero entries of the polynomial matrix. - TODO: docstrings, use like this (`p` is a polymatrix) + Typical usage looks like this (`p` is a polymatrix): .. code:: py + for entry, poly in p.entries(): for monomial, coeff in poly.terms(): # do something with coefficients @@ -58,20 +57,43 @@ class PolyMatrixMixin(abc.ABC): if poly := self.at(row, col): yield MatrixIndex(row, col), poly + # -- Old API --- + + @deprecated("Replaced by PolyMatrixMixin.entries()") + def gen_data(self) -> typing.Generator[tuple[tuple[int, int], PolyDict], None, None]: + yield from self.entries() + + @abc.abstractmethod + @deprecated("Replaced by PolyMatrixMixin.at(row, col).") + def get_poly(self, row: int, col: int) -> PolyDict | None: + # NOTE: in the new API when there is nothing at an entry the at() + # function returns an empty dictionary. This is because emtpy + # dicionaries are falsy so you can do + # + # `if p.at(row, col):` + # + # like it was possible with get_poly. However this is better because + # id does not make use of a null (None) type. + return self.at(row, col) or None + class PolyMatrixAsDictMixin( PolyMatrixMixin, abc.ABC, ): - """ - Polynomial matrix, stored as a dictionary. - """ + """ Matrix with polynomial entries, stored as a dictionary. """ @property @abc.abstractmethod def data(self) -> PolyMatrixDict: """" Get the dictionary. """ + @override + def at(self, row: int, col: int) -> PolyDict: + """ See :py:meth:`PolyMatrixMixin.at`. """ + return self.data.get(MatrixIndex(row, col)) or PolyDict.empty() + + # -- Old API --- @override def get_poly(self, row: int, col: int) -> PolyDict | None: @@ -79,13 +101,6 @@ class PolyMatrixAsDictMixin( return self.data[cast(MatrixIndex, (row, col))] return None - # --- New API --- - - @override - def at(self, row: int, col: int) -> PolyDict: - """ See :py:meth:`PolyMatrixMixin.at`. """ - return self.data.get(MatrixIndex(row, col)) or PolyDict.empty() - class BroadcastPolyMatrixMixin( PolyMatrixMixin, @@ -99,13 +114,249 @@ class BroadcastPolyMatrixMixin( @abc.abstractmethod def data(self) -> PolyDict: ... + @override + def at(self, row: int, col: int) -> PolyDict: + """ See :py:meth:`PolyMatrixMixin.at`. """ + return self.data + + # --- Old API --- + # overwrites the abstract method of `PolyMatrixMixin` def get_poly(self, col: int, row: int) -> PolyDict | None: return self.data or None - # --- New API --- + +class PolyMatrixAsAffineExpressionMixin( + PolyMatrixMixin, + abc.ABC +): + r""" + Matrix with polynomial entries, stored as an expression that is affine in + the monomials. + + For clarity we will show this through an example. Suppose :math:`x,y` are + polynomial variables, and we have a matrix of polynomials + + .. math:: + P(x, y) = \begin{bmatrix} + x + y^2 & x + 2\\ + 2x^2 + xy + 1 & y + \end{bmatrix} + \in \mathbf{R}_2(x, y)^{2 \times 2} + + We can rewrite :math:`P(x,y)` as an affine expression in terms of + the monomial basis generated by :math:`x` and :math:`y` and get + + .. math:: + + P(x,y) + = \begin{bmatrix} 0 & 2 \\ 1 & 0 \end{bmatrix} + + \begin{bmatrix} 1 & 1 \\ 0 & 0 \end{bmatrix} x + + \begin{bmatrix} 0 & 0 \\ 0 & 1 \end{bmatrix} y + + \begin{bmatrix} 0 & 0 \\ 2 & 0 \end{bmatrix} x^2 + + \begin{bmatrix} 0 & 0 \\ 1 & 0 \end{bmatrix} xy + + \begin{bmatrix} 1 & 0 \\ 0 & 0 \end{bmatrix} y^2 + + More generally for a :math:`p \in \mathbf{R}_d(x)^{n\times m}, x \in + \mathbf{R}^q` and the set of monomials of degree :math:`d` or lower + :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\}` by defining the set multi-indices of + :math:`\langle x \rangle_d` to be + + .. math:: + \mathcal{E} \langle x \rangle_d = \{ + (0, \ldots, 0), (1,0,\ldots,0), (0,1,\ldots,0), \ldots, + (2,0,\ldots,0), \ldots, (0,\ldots,d)\} + \subset \mathbf{Z}^d + + we can write + + .. math:: + p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha + + where each matrix :math:`A_\alpha \in \mathbf{R}^{n\times m}`. In the + context of this class we will call the :math:`A_\alpha` matrices + `affine coefficients`. + """ + + # Type of matrix that stores affifne coefficients + # TODO: Consider also scipy.sparse.csc_matrix instead of NDArray + # since slices are row slices + MatrixType = NDArray[np.float64] + + @property + @abc.abstractmethod + def data(self) -> MatrixType: + r""" + The big matrix that store all :math:`A_\alpha` matrices: + + .. math:: + A = \begin{bmatrix} + A_{\alpha_1} & A_{\alpha_2} & \cdots & A_{\alpha_N} + \end{bmatrix} + \in + \mathbf{R}^{n \times mN}, + \quad + N = \sum_{k = 0}^d {q + k \choose k} + """ + ... + + @property + @abc.abstractmethod + def slices(self) -> dict[MonomialIndex, tuple[int, int]]: + # column slices of the big array stored in self.data to get the individual A_alphas + ... @override def at(self, row: int, col: int) -> PolyDict: - """ See :py:meth:`PolyMatrixMixin.at`. """ - return self.data + # TODO: docstring + p = PolyDict.empty() + for monomial in self.slices.keys(): + p[monomial] = self.affine_coefficient(monomial)[row, col] + + return p + + # -- Special methods for affine expression + + def monomials_eval(self, x: tuple) -> dict[MonomialIndex, float]: + r""" + Evaluate all monomials up to degree :math:`d`. + + Given an :math:`x \in \mathbf{R}^q` this function computes the + 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\}`. + """ + 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 + # x = (x_1, x_2, ... x_q) to compute all quadratic terms (d = 2) we need to + # evaluate all points in the following upper triangular part of the grid: + # + # x1 x2 x3 ... xq + # ┌───────────────────────────────────────┐ + # x1 │ x1*x1 x1*x2 x1*x3 ... x1*xq │ + # │ │ + # x2 │ x2*x2 x2*x3 ... x2*xq │ + # │ │ + # x3 │ x3*x3 ... x3*xq │ + # │ │ + # ... │ ... ... │ + # │ │ + # xq │ xq*xq │ + # └───────────────────────────────────────┘ + # + # For d = 2 we compute the products by iterating k from 1 to q and then + # inside of that iterating j from k to q and multiply x[j] * x[k]. + # + # For d = 3 the shape is an irregular tetrahedron, and so when we iterate k + # from 1 to q then we must iterate over "triangular slices" of the + # tetrahedron. Like for d = 2 where we take only values that have j > k, + # for d = 3 we take all tetrahedron slices that are "higher" than level k. + # + # This pattern continues for d > 3 so we define a recursion over d by + # defining the triangle function below, which returns all monomial + # combinations of exactly degree d that are "higher" than j. + # + # To keep track of the values during the recursion we define a dictionary + # where the key is a tuple with length smaller or equal to q that contains + # the index of the variables that make up the product. For example: + # + # x_1 <-> (0,) x_1 * x_1 <-> (0, 0) + # 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) + + # TODO: rewrite recursive code using dynamic programming table to + # replace functools.cache + @functools.cache + def triangle(d, j): + # Base case + if d == 1: + return { + (i,): x[i] + for i in range(j, q) + } + + current = {} + for k in range(j, q): + # Recursion! Iterate over the slices which are given by the + # d-1 case and are higher than k + for i, v in triangle(d-1, k).items(): + current[i + (k,)] = v * x[k] + + return current + + # The monomials from degree 1 to d is the union of triangles with j = 0, + # i.e. the union of all terms of exactly degree d + terms = {} + for i in range(1, d+1): + terms |= triangle(i, 0) + + return terms + + # Get the highest degree we need to compute from the monomial indices + # TODO: this value could be cached + d = max(m.degree for m in self.slices.keys()) + + # Add the constant term + monomial_values = { + MonomialIndex.constant(): 1. + } + + # 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 + new_key: dict[int, int] = {} + for e in key: + if e in new_key: + new_key[e] += 1 + else: + new_key[e] = 1 + + index = MonomialIndex(VariableIndex(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: + r""" Get the affine coefficient :math:`A_\alpha` associated to :math:`x^\alpha`. """ + columns = range(*self.slices[monomial]) + return self.data[:, columns] + + def affine_coefficients_by_degrees(self) -> Iterable[tuple[int, MatrixType]]: + r""" + Iterate over the coefficients grouped by degree. + """ + monomials = self.slices.keys() + groups = itertools.groupby(monomials, lambda m: m.degree) + for degree, monomials in groups: + yield degree, np.vstack(list(self.data[m] for m in monomials)) + + def affine_eval(self, x: MatrixType) -> MatrixType: + r""" + Evaluate the affine expression + :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())) + + # Evaluate the affine expression + nrows, ncols = self.shape + return self.data @ (np.kron(monomials, np.eye(ncols))) + + def affine_eval_fn(self) -> Callable[[MatrixType], MatrixType]: + r""" + Get a function that when called evaluates the expression + :math:`p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha` + at :math:`x`. + """ + # TODO: docstring + # TODO: If slow consider replacing with toolz.functoolz.curry from ctoolz + return functools.partial(PolyMatrixAsAffineExpression.affine_eval, self) |