Start refactoring model slicer

src/inversion_ideas/operators.py

@@ -204,3 +204,251 @@ def diagonal(self):
             raise TypeError()
         diagonal = self.block.diagonal()
         return self.slice_matrix.T @ diagonal
+
+
+class ExpandedSquareMatrix(LinearOperator):
+    r"""
+    Represents a square matrix expanded with zeros.
+
+    Use this class to represent hessian matrices that are built from smaller matrices
+    that operate only on a subset of the entire model. Only a portion of that hessian
+    will be non-zero (the one related to the relevant model elements for that objective
+    function), while the rest of the matrix will be filled of zeros.
+
+    Parameters
+    ----------
+    matrix : array, sparse array or LinearOperator
+        Square matrix to be expanded
+    shape : tuple of int
+        Shape of the expanded matrix.
+    slice_ : slice or list of slices
+        Slice(s) used to slice portions of the square matrix. The start and stop should
+        be positive or zero, start should be lower than stop, and step should be None.
+
+    Notes
+    -----
+    Given a square matrix :math:`\mathbf{A}`, this
+    ``LinearOperator`` represents an *expanded* version of :math:`\mathbf{A}` like:
+
+    .. math::
+
+        \mathbf{H} = \begin{bmatrix}
+            0 & 0          & 0\\
+            0 & \mathbf{A} & 0\\
+            0 & 0          & 0
+        \end{bmatrix},
+
+    where :math:`\mathbf{A}` sits in the diagonal of :math:`\mathbf{H}`.
+    The ``slice_`` argument is used to specify the location of the matrix
+    :math:`\mathbf{A}` within the expanded :math:`\mathbf{H}` matrix.
+
+    When passing multiple ``slice_``s, it can represent more complex expanded
+    matrix like:
+
+    .. math::
+
+        \mathbf{H} = \begin{bmatrix}
+            0 & 0                & 0 & 0                 & 0 & 0      & 0 & 0                   & 0 & 0                  & 0 \\
+            0 & \mathbf{A}_{1,1} & 0 & \mathbf{A}_{2,1}  & 0 & \dots  & 0 & \mathbf{A}_{n-1,1}  & 0 & \mathbf{A}_{n,1}   & 0 \\
+            0 & 0                & 0 & 0                 & 0 & 0      & 0 & 0                   & 0 & 0                  & 0 \\
+            0 & \vdots           & 0 & \vdots            & 0 & \ddots & 0 & \vdots              & 0 & \vdots             & 0 \\
+            0 & 0                & 0 & 0                 & 0 & 0      & 0 & 0                   & 0 & 0                  & 0 \\
+            0 & \mathbf{A}_{1,n} & 0 & \mathbf{A}_{2,n}  & 0 & \dots  & 0 & \mathbf{A}_{n-1,n}  & 0 & \mathbf{A}_{n,n}   & 0 \\
+            0 & 0                & 0 & 0                 & 0 & 0      & 0 & 0                   & 0 & 0                  & 0
+        \end{bmatrix},
+
+    where each :math:`\mathbf{A}_{i,j}` is a square submatrix of :math:`\mathbf{A}`.
+    """
+
+    def __init__(
+        self,
+        matrix: npt.NDArray | LinearOperator | SparseArray,
+        shape: tuple[int],
+        slice_: slice | list[slice],
+    ):
+        if len(shape) != 2:
+            # TODO: add msg
+            raise ValueError()
+
+        if shape[0] != shape[1]:
+            # TODO: add msg
+            raise ValueError()
+
+        if matrix.ndim != 2:
+            # TODO: add msg
+            raise ValueError()
+
+        if matrix.shape[0] != matrix.shape[1]:
+            # TODO: add msg
+            raise ValueError()
+
+        # TODO: add sanity checks for slices
+        #   - start >= 0 and stop > 0
+        #   - step is None
+        #   - start < stop
+
+        super().__init__(shape=shape, dtype=matrix.dtype)
+        self._matrix = matrix
+        self._slice = slice_
+
+        if self.reduced_size != matrix.shape[0]:
+            # TODO: add msg
+            raise ValueError()
+
+    @property
+    def matrix(self):
+        return self._matrix
+
+    @property
+    def slice_(self):
+        return self._slice
+
+    @property
+    def slices(self):
+        if isinstance(self.slice_, slice):
+            return [self.slice_]
+        return self._slice
+
+    @property
+    def reduced_size(self):
+        # TODO: choose better name
+        return sum(s.stop - s.start for s in self.slices)
+
+    @property
+    def slicer_matrix(self):
+        if not hasattr(self, "_slicer_matrix"):
+            shape = (self.reduced_size, self.shape[0])
+            slicer_matrix = Zero(shape=shape, dtype=self.dtype)
+            offset = 0
+            for slice_ in self.slices:
+                slicer_matrix += _Slicer(
+                    start=slice_.start, stop=slice_.stop, shape=shape, offset=offset
+                )
+                offset += slice_.stop - slice_.start
+            if isinstance(slicer_matrix, Zero):
+                msg = "No slices were found!"
+                raise ValueError(msg)
+            self._slicer_matrix = slicer_matrix
+        return self._slicer_matrix
+
+    def _matvec(self, x):
+        """
+        Dot product between the matrix and a vector.
+        """
+        return self.slicer_matrix.T @ (self.matrix @ (self.slicer_matrix @ x))
+
+    def _rmatvec(self, x):
+        """
+        Dot product between the transposed matrix and a vector.
+        """
+        return self.slicer_matrix @ (self.matrix.T @ (self.slicer_matrix.T @ x))
+
+    def _matmat(self, x):
+        """
+        Dot product between the matrix and a vector.
+        """
+        raise NotImplementedError()
+
+    def diagonal(self):
+        """
+        Diagonal of the matrix.
+        """
+        raise NotImplementedError()
+
+    def toarray(self):
+        """
+        Return expanded matrix as a dense array.
+
+        .. warning::
+
+            This method could create a very large 2D array that might not fit in memory.
+        """
+        raise NotImplementedError()
+
+
+class _Slicer(LinearOperator):
+    """ """
+
+    def __init__(
+        self, start: int, stop: int, shape: tuple[int, int], offset: int | None = None
+    ):
+        if shape[0] < start - stop:
+            # TODO: add msg
+            # The size of the slice cannot be larger than the output array.
+            raise ValueError()
+
+        if shape[0] > shape[1]:
+            # TODO: add msg.
+            # The sliced array cannot be larger than the array that will be sliced.
+            raise ValueError()
+
+        # TODO: add sanity checks for offset
+
+        super().__init__(shape=shape, dtype=np.float64)
+        self._start = start
+        self._stop = stop
+        self._offset = offset
+
+    @property
+    def start(self):
+        return self._start
+
+    @property
+    def stop(self):
+        return self._stop
+
+    @property
+    def slice(self):
+        return slice(self._start, self._stop)
+
+    @property
+    def offset(self) -> int:
+        if self._offset is None:
+            return 0
+        return self._offset
+
+    def _matvec(self, x):
+        """
+        Slices the array.
+        """
+        result = np.zeros(self.shape[0], dtype=x.dtype)
+        out_slice = slice(self.offset, self.offset + self.slice.stop - self.slice.start)
+        result[out_slice] = x[self.slice]
+        return result
+
+    def _rmatvec(self, x):
+        """
+        Expand the array.
+        """
+        result = np.zeros(self.shape[1], dtype=x.dtype)
+        out_slice = slice(self.offset, self.offset + self.slice.stop - self.slice.start)
+        result[self.slice] = x[out_slice]
+        return result
+
+    def _matmat(self, x):
+        """
+        Dot product between the matrix and a vector.
+        """
+        pass
+        # raise NotImplementedError()
+
+    def toarray(self):
+        # Idea: represent this as a sparse diagonal array and use the toarray method to
+        # return a dense representation.
+        raise NotImplementedError()
+
+
+class Zero(LinearOperator):
+    """
+    Null linear operator.
+    """
+
+    def _matvec(self, x):
+        return np.zeros(shape=self.shape[0], dtype=self.dtype)
+
+    def _adjoint(self):
+        new_shape = (self.shape[1], self.shape[0])
+        return Zero(shape=new_shape, dtype=self.dtype)
+
+    def _matmat(self, x):
+        return np.zeros(shape=(self.shape[0], x.shape[1]), dtype=self.dtype)

src/inversion_ideas/wires.py

@@ -11,7 +11,7 @@
 from scipy.sparse import dia_array, diags_array
 from scipy.sparse.linalg import LinearOperator
 
-from .operators import BlockSquareMatrix
+from .operators import BlockSquareMatrix, ExpandedSquareMatrix
 from .typing import Model, SparseArray
 
 
@@ -244,6 +244,10 @@ def full_size(self) -> int: ...
     @abstractmethod
     def wires(self) -> Wires: ...
 
+    @property
+    @abstractmethod
+    def slice(self) -> slice | list[slice]: ...
+
     @property
     @abstractmethod
     def slice_matrix(self) -> dia_array[np.float64]: ...
@@ -267,7 +271,8 @@ def expand_array(self, array: npt.NDArray) -> npt.NDArray:
 
     def expand_matrix(
         self, matrix: npt.NDArray | LinearOperator | SparseArray
-    ) -> "BlockSquareMatrix":
+    ) -> "ExpandedSquareMatrix":
+        # ) -> "BlockSquareMatrix":
         """
         Expand a square matrix into a ``BlockSquareMatrix`` filling blocks with zeros.
 
@@ -282,7 +287,9 @@ def expand_matrix(
             LinearOperator that represents the matrix filled with zeros outside of the
             block.
         """
-        return BlockSquareMatrix(block=matrix, slice_matrix=self.slice_matrix)
+        shape = (self.full_size, self.full_size)
+        return ExpandedSquareMatrix(matrix=matrix, shape=shape, slice_=self.slice)
+        # return BlockSquareMatrix(block=matrix, slice_matrix=self.slice_matrix)
 
 
 class ModelSlice(_BaseModelSlice):
@@ -369,6 +376,10 @@ def names(self) -> list[str]:
     def slices(self) -> dict[str, slice]:
         return self._slices
 
+    @property
+    def slice(self) -> list[slice]:
+        return list(self._slices.values())
+
     @property
     def wires(self) -> Wires:
         return self._wires