diff options
| -rw-r--r-- | polymatrix/polymatrix/impl.py | 8 | ||||
| -rw-r--r-- | polymatrix/polymatrix/init.py | 48 | ||||
| -rw-r--r-- | polymatrix/polymatrix/mixins.py | 87 |
3 files changed, 116 insertions, 27 deletions
diff --git a/polymatrix/polymatrix/impl.py b/polymatrix/polymatrix/impl.py index 3ed65dc..1d3e1b9 100644 --- a/polymatrix/polymatrix/impl.py +++ b/polymatrix/polymatrix/impl.py @@ -1,7 +1,7 @@ import dataclassabc from polymatrix.polymatrix.abc import PolyMatrix -from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin, PolyMatrixAsAffineExpressionMixin +from polymatrix.polymatrix.mixins import PolyMatrixMixin, BroadcastPolyMatrixMixin, SlicePolyMatrixMixin, PolyMatrixAsAffineExpressionMixin from polymatrix.polymatrix.index import PolyMatrixDict, PolyDict, MonomialIndex @@ -18,6 +18,12 @@ class BroadcastPolyMatrixImpl(BroadcastPolyMatrixMixin): @dataclassabc.dataclassabc(frozen=True) +class SlicePolyMatrixImpl(SlicePolyMatrixMixin): + reference: PolyMatrixMixin + shape: tuple[int, int] + + +@dataclassabc.dataclassabc(frozen=True) class PolyMatrixAsAffineExpressionImpl(PolyMatrixAsAffineExpressionMixin): affine_coefficients: PolyMatrixAsAffineExpressionMixin.MatrixType shape: tuple[int, int] diff --git a/polymatrix/polymatrix/init.py b/polymatrix/polymatrix/init.py index 77c544c..0778a50 100644 --- a/polymatrix/polymatrix/init.py +++ b/polymatrix/polymatrix/init.py @@ -5,7 +5,7 @@ import numpy as np from typing import TYPE_CHECKING -from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl, PolyMatrixAsAffineExpressionImpl +from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl, SlicePolyMatrixImpl, PolyMatrixAsAffineExpressionImpl from polymatrix.polymatrix.index import PolyMatrixDict, PolyDict, MatrixIndex, MonomialIndex, VariableIndex from polymatrix.polymatrix.mixins import PolyMatrixAsAffineExpressionMixin @@ -53,6 +53,52 @@ def init_broadcast_poly_matrix( ) +def init_slice_poly_matrix( + reference: PolyMatrixMixin, + slices: tuple[int | Iterable[int], int | Iterable[int]] +) -> SlicePolyMatrixMixin: + + formatted_slices: list[tuple] = [(), ()] + shape = [0, 0] + + for i, (what, el, numel) in enumerate(zip(("Row", "Column"), slices, reference.shape)): + if isinstance(el, int): + if not (0 <= el < numel): + raise IndexError(f"{what} {el} is out of range in shape {reference.shape}.") + + # convert to tuple, and we are done + formatted_slices[i] = (el,) + shape[i] = 1 + + elif isinstance(el, slice): + # convert to range + el = range(el.start or 0, el.stop or numel, el.step or 1) + + if not (0 <= el.start < numel): + raise IndexError(f"{what} start {el} is out of range in shape {reference.shape}.") + + # range does not include stop + if not (0 <= el.stop <= numel): + raise IndexError(f"{what} stop {el} is out of range in shape {reference.shape}.") + + formatted_slices[i] = tuple(el) + shape[i] = len(formatted_slices[i]) + + elif isinstance(el, tuple): + # FIXME: this does not handle all edge cases, what if the rows and column + # are not ranges / slices but tuples and result in a non-rectangular slice? + # e.g. it has a hole in the middle. Need to check + formatted_slices[i] = el + shape[i] = len(el) + + else: + raise TypeError("{what} {el} of type {type(el)} is not a valid slice type.") + + + return SlicePolyMatrixImpl(reference=reference, shape=shape, slice=tuple(formatted_slices)) + + +# FIXME: rename to init_affine_expression def to_affine_expression(p: PolyMatrixMixin) -> PolyMatrixAsAffineExpressionMixin: """ Convert a polymatrix into a PolyMatrixAsAffineExpressionMixin. """ if isinstance(p, PolyMatrixAsAffineExpressionMixin): diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py index 559a72b..63eeb18 100644 --- a/polymatrix/polymatrix/mixins.py +++ b/polymatrix/polymatrix/mixins.py @@ -1,12 +1,11 @@ from __future__ import annotations -import abc import typing -import itertools import functools import math import numpy as np +from abc import ABC, abstractmethod from numpy.typing import NDArray from typing import Iterable, Sequence, Callable, cast, TYPE_CHECKING from typing_extensions import override @@ -19,16 +18,16 @@ if TYPE_CHECKING: from polymatrix.expressionstate.mixins import ExpressionStateMixin -class PolyMatrixMixin(abc.ABC): +class PolyMatrixMixin(ABC): """ Matrix with polynomial entries. """ @property - @abc.abstractmethod + @abstractmethod def shape(self) -> tuple[int, int]: ... - @abc.abstractmethod + @abstractmethod def at(self, row: int, col: int) -> PolyDict: """ Return the polynomial at the entry (row, col). If the entry is zero it returns an empty `PolyDict`, i.e. an empty @@ -77,20 +76,24 @@ class PolyMatrixMixin(abc.ABC): return self.at(row, col) or None -class PolyMatrixAsDictMixin( - PolyMatrixMixin, - abc.ABC, -): +class PolyMatrixAsDictMixin(PolyMatrixMixin, ABC): """ Matrix with polynomial entries, stored as a dictionary. """ @property - @abc.abstractmethod + @abstractmethod def data(self) -> PolyMatrixDict: """" Get the dictionary. """ @override def at(self, row: int, col: int) -> PolyDict: """ See :py:meth:`PolyMatrixMixin.at`. """ + # Check bounds + if not (0 <= row < self.shape[0]): + raise IndexError(f"Row {row} out of range, shape is {self.shape}") + + if not (0 <= col < self.shape[1]): + raise IndexError(f"Column {col} out of range, shape is {self.shape}") + return self.data.get(MatrixIndex(row, col)) or PolyDict.empty() # -- Old API --- @@ -102,16 +105,13 @@ class PolyMatrixAsDictMixin( return None -class BroadcastPolyMatrixMixin( - PolyMatrixMixin, - abc.ABC, -): +class BroadcastPolyMatrixMixin(PolyMatrixMixin, ABC): """ TODO: docstring, similar to numpy broadcasting. https://numpy.org/doc/stable/user/basics.broadcasting.html """ @property - @abc.abstractmethod + @abstractmethod def data(self) -> PolyDict: ... @override @@ -121,15 +121,46 @@ class BroadcastPolyMatrixMixin( # --- Old API --- - # overwrites the abstract method of `PolyMatrixMixin` + @override def get_poly(self, col: int, row: int) -> PolyDict | None: return self.data or None -class PolyMatrixAsAffineExpressionMixin( - PolyMatrixMixin, - abc.ABC -): +class SlicePolyMatrixMixin(PolyMatrixMixin, ABC): + """ Slice of a poly matrix. """ + @property + @abstractmethod + def reference(self) -> PolyMatrixMixin: + """ Polymatrix which the slice was taken from. """ + + @property + @abstractmethod + def slice(self) -> tuple[tuple[int, ...], tuple[int, ...]]: + """ Row and column indices to take from reference. """ + # TODO: consider changing this member to be of type + # tuple[slice | tuple[int, ...], slice | tuple[int, ...]] + # so that one can optimize for storage, especially if you have huge + # matrices and access 1:300. In the current implementation there would + # be a tuple with 300 elements (bad for storage, though free fast to + # access in at() method) + + @override + def at(self, row: int, col: int) -> PolyDict: + """ See :py:meth:`PolyMatrixMixin.at`. """ + # Check bounds + if not (0 <= row < self.shape[0]): + raise IndexError(f"Row {row} out of range, shape is {self.shape}") + + if not (0 <= col < self.shape[1]): + raise IndexError(f"Column {col} out of range, shape is {self.shape}") + + ref_row = self.slice[0][row] + ref_col = self.slice[1][col] + + return self.reference.at(ref_row, ref_col) + + +class PolyMatrixAsAffineExpressionMixin(PolyMatrixMixin, ABC): r""" Matrix with polynomial entries, stored as an expression that is affine in the monomials. @@ -190,7 +221,7 @@ class PolyMatrixAsAffineExpressionMixin( MatrixType = NDArray[np.float64] @property - @abc.abstractmethod + @abstractmethod def affine_coefficients(self) -> MatrixType: r""" The big matrix that store all :math:`A_\alpha` matrices: @@ -206,7 +237,7 @@ class PolyMatrixAsAffineExpressionMixin( """ @property - @abc.abstractmethod + @abstractmethod def slices(self) -> dict[MonomialIndex, tuple[int, int]]: r""" Map from monomial indices to column slices of the big matrix that @@ -214,7 +245,7 @@ class PolyMatrixAsAffineExpressionMixin( """ @property - @abc.abstractmethod + @abstractmethod def degree(self) -> int: """ Maximal degree of the affine expressions """ # Value of max(m.degree for m in self.slices.keys()) @@ -222,13 +253,19 @@ class PolyMatrixAsAffineExpressionMixin( @property - @abc.abstractmethod + @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: - # TODO: docstring + # Check bounds + if not (0 <= row < self.shape[0]): + raise IndexError(f"Row {row} out of range, shape is {self.shape}") + + if not (0 <= col < self.shape[1]): + raise IndexError(f"Column {col} out of range, shape is {self.shape}") + p = PolyDict.empty() for monomial in self.slices.keys(): p[monomial] = self.affine_coefficient(monomial)[row, col] |