Create SlicePolyMatrix, add bounds check for PolyMatrices

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]