capytaine 2.2__cp312-cp312-macosx_11_0_arm64.whl

capytaine/matrices/block.py

@@ -0,0 +1,590 @@
+"""This module implements block matrices to be used in Hierarchical Toeplitz matrices.
+
+It takes inspiration from the following works:
+
+* `openHmx module from Gypsilab by Matthieu Aussal (GPL licensed) <https://github.com/matthieuaussal/gypsilab>`_
+* `HierarchicalMatrices by Markus Neumann (GPL licensed) <https://github.com/maekke97/HierarchicalMatrices>`_
+"""
+# Copyright (C) 2017-2019 Matthieu Ancellin
+# See LICENSE file at <https://github.com/mancellin/capytaine>
+
+import logging
+
+from numbers import Number
+from typing import Tuple, List, Callable, Union, Iterable
+from itertools import cycle, accumulate, chain, product
+from collections.abc import Iterator
+
+import numpy as np
+
+from capytaine.matrices.low_rank import LowRankMatrix
+from capytaine.tools.optional_imports import import_optional_dependency
+
+LOG = logging.getLogger(__name__)
+
+
+class BlockMatrix:
+    """A (2D) matrix, stored as a set of submatrices (or blocks).
+
+    Parameters
+    ----------
+    blocks: list of list of matrices
+        The blocks of the block matrix.
+    check: bool, optional
+        Should the code perform sanity checks on the inputs? (default: True)
+
+    Attributes
+    ----------
+    shape: pair of ints
+        shape of the full matrix
+    nb_blocks: pair of ints
+        number of blocks in each directions.
+        Example::
+
+            AAAABB
+            AAAABB  ->  nb_blocks = (1, 2)
+            AAAABB
+    """
+
+    ndim = 2  # Other dimensions have not been implemented.
+
+    def __init__(self, blocks, _stored_block_shapes=None, check=True):
+
+        self._stored_nb_blocks = (len(blocks), len(blocks[0]))
+
+        self._stored_blocks = np.empty(self._stored_nb_blocks, dtype=object)
+        for i in range(len(blocks)):
+            for j in range(len(blocks[i])):
+                self._stored_blocks[i, j] = blocks[i][j]
+
+        if _stored_block_shapes is None:
+            self._stored_block_shapes = ([block.shape[0] for block in self._stored_blocks[:, 0]],
+                                         [block.shape[1] for block in self._stored_blocks[0, :]])
+        else:
+            # To avoid going down the tree if it is already known.
+            self._stored_block_shapes = _stored_block_shapes
+
+        # Block shape of the full matrix
+        self.nb_blocks = self._compute_nb_blocks()
+
+        # Total shape of the full matrix
+        self.shape = self._compute_shape()
+
+        LOG.debug(f"New block matrix: %s", self)
+
+        if check:
+            assert self._check_dimensions_of_blocks()
+            assert self._check_dtype()
+
+    def _compute_shape(self):
+        # In a dedicated routine because it will be overloaded by subclasses.
+        return sum(self._stored_block_shapes[0]), sum(self._stored_block_shapes[1])
+
+    def _compute_nb_blocks(self):
+        return self._stored_nb_blocks
+
+    def _check_dimensions_of_blocks(self) -> bool:
+        """Check that the dimensions of the blocks are consistent."""
+        if not all(block.ndim == self.ndim for line in self._stored_blocks for block in line):
+            return False
+
+        for line in self.all_blocks:
+            block_height = line[0].shape[0]
+            for block in line[1:]:
+                if not block.shape[0] == block_height:  # Same height on a given line
+                    return False
+
+        for col in np.moveaxis(self.all_blocks, 1, 0):
+            block_width = col[0].shape[1]
+            for block in col[1:]:
+                if not block.shape[1] == block_width:  # Same width on a given column
+                    return False
+        return True
+
+    def _check_dtype(self) -> bool:
+        """Check that the type of the blocks are consistent."""
+        for line in self._stored_blocks:
+            for block in line:
+                if block.dtype != self.dtype:
+                    return False
+        return True
+
+    # ACCESSING DATA
+
+    @property
+    def dtype(self):
+        """The type of data of all of the subblocks."""
+        try:
+            return self._stored_blocks[0][0].dtype
+        except AttributeError:
+            return None
+
+    @property
+    def all_blocks(self) -> np.ndarray:
+        """The matrix of matrices. For a full block matrix, all the blocks are stored in memory."""
+        return self._stored_blocks
+
+    @property
+    def block_shapes(self) -> Tuple[List[int], List[int]]:
+        """The shapes of the blocks composing the BlockMatrix.
+
+        Example::
+
+            AAAABB
+            AAAABB  ->  block_shapes = ([3], [4, 2])
+            AAAABB
+        """
+        return self._stored_block_shapes
+
+    def _stored_block_positions(self, global_frame=(0, 0)) -> Iterable[List[Tuple[int, int]]]:
+        """The position of each blocks in the matrix as a generator.
+        The list is used by subclasses where the same block may appear several times in different positions.
+
+        Parameters
+        ----------
+        global_frame: Tuple[int], optional
+            the coordinate of the top right corner. Default: 0, 0.
+
+        Example::
+
+            AAAABB
+            AAAABB  ->  list(matrix._stored_block_positions) = [[(0,0)], [(0, 4)], [(2, 0)], [(2, 4)]]
+            CCCCDD
+        """
+        x_acc = accumulate([0] + self.block_shapes[0][:-1])
+        y_acc = accumulate([0] + self.block_shapes[1][:-1])
+        return ([(global_frame[0] + x, global_frame[1] + y)] for x, y in product(x_acc, y_acc))
+
+    def _put_in_full_matrix(self, full_matrix: np.ndarray, where=(0, 0)) -> None:
+        """Copy the content of the block matrix in a matrix, which is modified in place."""
+        all_blocks_in_flat_iterator = (block for line in self._stored_blocks for block in line)
+        positions_of_all_blocks = self._stored_block_positions(global_frame=where)
+        for block, positions_of_the_block in zip(all_blocks_in_flat_iterator, positions_of_all_blocks):
+            if isinstance(block, BlockMatrix):
+                position_of_first_appearance = positions_of_the_block[0]
+                frame_of_first_appearance = (slice(position_of_first_appearance[0], position_of_first_appearance[0]+block.shape[0]),
+                                             slice(position_of_first_appearance[1], position_of_first_appearance[1]+block.shape[1]))
+                block._put_in_full_matrix(full_matrix, where=position_of_first_appearance)
+
+                for position in positions_of_the_block[1:]:  # For the other appearances, only copy the first appearance
+                    block_frame = (slice(position[0], position[0]+block.shape[0]),
+                                   slice(position[1], position[1]+block.shape[1]))
+                    full_matrix[block_frame] = full_matrix[frame_of_first_appearance]
+
+            else:
+                full_block = block if isinstance(block, np.ndarray) else block.full_matrix()
+                for position in positions_of_the_block:
+                    block_frame = (slice(position[0], position[0]+block.shape[0]),
+                                   slice(position[1], position[1]+block.shape[1]))
+                    full_matrix[block_frame] = full_block
+
+    def full_matrix(self, dtype=None) -> np.ndarray:
+        """Flatten the block structure and return a full matrix."""
+        if dtype is None: dtype = self.dtype
+        full_matrix = np.empty(self.shape, dtype=dtype)
+        self._put_in_full_matrix(full_matrix)
+        return full_matrix
+
+    def __array__(self, dtype=None):
+        return self.full_matrix(dtype=dtype)
+
+    def no_toeplitz(self):
+        """Recursively replace the block toeplitz matrices by usual block matrices.
+        WARNING: the block matrices may still contain several references to the same block."""
+        blocks = [[block.no_toeplitz() if isinstance(block, BlockMatrix) else block for block in line] for line in self.all_blocks]
+        return BlockMatrix(blocks)
+
+    def __deepcopy__(self, memo):
+        from copy import deepcopy
+        blocks = [[deepcopy(block) for block in line] for line in self._stored_blocks]
+        # The call to deepcopy does not use the memo on purpose:
+        # the goal is to replace references to the same block by references to different copies of the block.
+        return self.__class__(blocks)
+
+    @property
+    def stored_data_size(self):
+        """Return the number of entries actually stored in memory."""
+        size = 0
+        for line in self._stored_blocks:
+            for block in line:
+                if isinstance(block, np.ndarray):
+                    size += np.prod(block.shape)
+                else:
+                    size += block.stored_data_size
+        return size
+
+    @property
+    def density(self):
+        return self.stored_data_size/np.prod(self.shape)
+
+    @property
+    def sparcity(self):
+        return 1 - self.density
+
+    def __hash__(self):
+        # Temporary
+        return id(self)
+
+    # TRANSFORMING DATA
+
+    def _apply_unary_op(self, op: Callable) -> 'BlockMatrix':
+        """Helper function applying a function recursively on all submatrices."""
+        LOG.debug(f"Apply op {op.__name__} to {self}")
+        result = [[op(block) for block in line] for line in self._stored_blocks]
+        return self.__class__(result, _stored_block_shapes=self._stored_block_shapes, check=False)
+
+    def _apply_binary_op(self, op: Callable, other: 'BlockMatrix') -> 'BlockMatrix':
+        """Helper function applying a binary operator recursively on all submatrices."""
+        if isinstance(other, self.__class__) and self.nb_blocks == other.nb_blocks:
+            LOG.debug(f"Apply op {op.__name__} to {self} and {other}")
+            result = [
+                [op(block, other_block) for block, other_block in zip(line, other_line)]
+                for line, other_line in zip(self._stored_blocks, other._stored_blocks)
+            ]
+            return self.__class__(result, _stored_block_shapes=self._stored_block_shapes, check=False)
+        else:
+            return NotImplemented
+
+    def __add__(self, other: 'BlockMatrix') -> 'BlockMatrix':
+        from operator import add
+        return self._apply_binary_op(add, other)
+
+    def __radd__(self, other: 'BlockMatrix') -> 'BlockMatrix':
+        return self + other
+
+    def __neg__(self) -> 'BlockMatrix':
+        from operator import neg
+        return self._apply_unary_op(neg)
+
+    def __sub__(self, other: 'BlockMatrix') -> 'BlockMatrix':
+        from operator import sub
+        return self._apply_binary_op(sub, other)
+
+    def __rsub__(self, other: 'BlockMatrix') -> 'BlockMatrix':
+        from operator import sub
+        return other._apply_binary_op(sub, self)
+
+    def __mul__(self, other: Union['BlockMatrix', Number]) -> 'BlockMatrix':
+        if isinstance(other, Number):
+            return self._apply_unary_op(lambda x: other*x)
+        else:
+            from operator import mul
+            return self._apply_binary_op(mul, other)
+
+    def __rmul__(self, other: Union['BlockMatrix', Number]) -> 'BlockMatrix':
+        return self * other
+
+    def __truediv__(self, other: Union['BlockMatrix', Number]) -> 'BlockMatrix':
+        from numbers import Number
+        if isinstance(other, Number):
+            return self._apply_unary_op(lambda x: x/other)
+        else:
+            from operator import truediv
+            return self._apply_binary_op(truediv, other)
+
+    def __rtruediv__(self, other: Union['BlockMatrix', Number]) -> 'BlockMatrix':
+        from numbers import Number
+        if isinstance(other, Number):
+            return self._apply_unary_op(lambda x: other/x)
+        else:
+            return self._apply_binary_op(lambda x, y: y/x, other)
+
+    def matvec(self, other):
+        """Matrix vector product.
+        Named as such to be used as scipy LinearOperator."""
+        LOG.debug(f"Multiplication of {self} with a full vector of size {other.shape}.")
+        result = np.zeros(self.shape[0], dtype=other.dtype)
+        line_heights = self.block_shapes[0]
+        line_positions = list(accumulate(chain([0], line_heights)))
+        col_widths = self.block_shapes[1]
+        col_positions = list(accumulate(chain([0], col_widths)))
+        for line, line_position, line_height in zip(self.all_blocks, line_positions, line_heights):
+            line_slice = slice(line_position, line_position+line_height)
+            for block, col_position, col_width in zip(line, col_positions, col_widths):
+                col_slice = slice(col_position, col_position+col_width)
+                result[line_slice] += block @ other[col_slice]
+        return result
+
+    def rmatvec(self, other):
+        """Vector matrix product.
+        Named as such to be used as scipy LinearOperator."""
+        LOG.debug(f"Multiplication of a full vector of size {other.shape} with {self}.")
+        result = np.zeros(self.shape[1], dtype=other.dtype)
+        line_heights = self.block_shapes[0]
+        line_positions = list(accumulate(chain([0], line_heights)))
+        col_widths = self.block_shapes[1]
+        col_positions = list(accumulate(chain([0], col_widths)))
+        for col, col_position, col_width in zip(self.all_blocks.T, col_positions, col_widths):
+            col_slice = slice(col_position, col_position+col_width)
+            for block, line_position, line_height in zip(col, line_positions, line_heights):
+                line_slice = slice(line_position, line_position+line_height)
+                if isinstance(block, BlockMatrix):
+                    result[col_slice] += block.rmatvec(other[line_slice])
+                else:
+                    result[col_slice] += other[line_slice] @ block
+        return result
+
+    def matmat(self, other):
+        """Matrix-matrix product."""
+        if isinstance(other, BlockMatrix) and self.block_shapes[1] == other.block_shapes[0]:
+            LOG.debug(f"Multiplication of %s with %s", self, other)
+            own_blocks = self.all_blocks
+            other_blocks = np.moveaxis(other.all_blocks, 1, 0)
+            new_matrix = []
+            for own_line in own_blocks:
+                new_line = []
+                for other_col in other_blocks:
+                    new_line.append(sum(own_block @ other_block for own_block, other_block in zip(own_line, other_col)))
+                new_matrix.append(new_line)
+            return BlockMatrix(new_matrix, check=False)
+
+        elif isinstance(other, np.ndarray) and self.shape[1] == other.shape[0]:
+            LOG.debug(f"Multiplication of {self} with a full matrix of shape {other.shape}.")
+            # Cut the matrix and recursively call itself to use the code above.
+            from capytaine.matrices.builders import cut_matrix
+            cut_other = cut_matrix(other, self.block_shapes[1], [other.shape[1]], check=False)
+            return (self @ cut_other).full_matrix()
+
+    def __matmul__(self, other: Union['BlockMatrix', np.ndarray]) -> Union['BlockMatrix', np.ndarray]:
+        if not (isinstance(other, BlockMatrix) or isinstance(other, np.ndarray)):
+            return NotImplemented
+        elif other.ndim == 2:  # Other is a matrix
+            if other.shape[1] == 1:  # Actually a column vector
+                return self.matvec(other.flatten())
+            else:
+                return self.matmat(other)
+        elif other.ndim == 1:  # Other is a vector
+            return self.matvec(other)
+        else:
+            return NotImplemented
+
+    def __rmatmul__(self, other: Union['BlockMatrix', np.ndarray]) -> Union['BlockMatrix', np.ndarray]:
+        if not (isinstance(other, BlockMatrix) or isinstance(other, np.ndarray)):
+            return NotImplemented
+        elif other.ndim == 2:  # Other is a matrix
+            if other.shape[1] == 1:  # Actually a column vector
+                return self.rmatvec(other.flatten())
+            else:
+                return NotImplemented
+        elif other.ndim == 1:  # Other is a vector
+            return self.rmatvec(other)
+        else:
+            return NotImplemented
+
+    def astype(self, dtype: np.dtype) -> 'BlockMatrix':
+        return self._apply_unary_op(lambda x: x.astype(dtype))
+
+    def fft_of_list(*block_matrices, check=True):
+        """Compute the fft of a list of block matrices of the same type and shape.
+        The output is a list of block matrices of the same shape as the input ones.
+        The fft is computed element-wise, so the block structure does not cause any mathematical difficulty.
+        Returns an array of BlockMatrices.
+        """
+        class_of_matrices = type(block_matrices[0])
+        nb_blocks = block_matrices[0]._stored_nb_blocks
+
+        LOG.debug(f"FFT of {len(block_matrices)} {class_of_matrices.__name__} (stored blocks = {nb_blocks})")
+
+        if check:
+            # Check the validity of the shapes of the matrices given as input
+            shape = block_matrices[0].shape
+            assert [nb_blocks == matrix._stored_nb_blocks for matrix in block_matrices[1:]]
+            assert [shape == matrix.shape for matrix in block_matrices[1:]]
+            assert [class_of_matrices == type(matrix) for matrix in block_matrices[1:]]
+
+        # Initialize a vector of block matrices without values in the blocks.
+        result = np.empty(len(block_matrices), dtype=object)
+        for i in range(len(block_matrices)):
+            result[i] = class_of_matrices(np.empty(nb_blocks, dtype=object),
+                                          _stored_block_shapes=block_matrices[0]._stored_block_shapes,
+                                          check=False)
+
+        for i_block, j_block in product(range(nb_blocks[0]), range(nb_blocks[1])):
+            list_of_i_j_blocks = [block_matrices[i_matrix]._stored_blocks[i_block, j_block]
+                                  for i_matrix in range(len(block_matrices))]
+
+            if any(isinstance(block, np.ndarray) or isinstance(block, LowRankMatrix) for block in list_of_i_j_blocks):
+                list_of_i_j_blocks = [block if isinstance(block, np.ndarray) else block.full_matrix() for block in list_of_i_j_blocks]
+                fft_of_blocks = np.fft.fft(list_of_i_j_blocks, axis=0)
+            else:
+                fft_of_blocks = BlockMatrix.fft_of_list(*list_of_i_j_blocks, check=False)
+
+            for matrix, computed_block in zip(result, fft_of_blocks):
+                matrix._stored_blocks[i_block, j_block] = computed_block
+
+        return result
+
+    # COMPARISON AND REDUCTION
+
+    def __eq__(self, other: 'BlockMatrix') -> 'BlockMatrix[bool]':
+        from operator import eq
+        return self._apply_binary_op(eq, other)
+
+    def __invert__(self) -> 'BlockMatrix':
+        """Boolean not (~)"""
+        from operator import invert
+        return self._apply_unary_op(invert)
+
+    def __ne__(self, other: 'BlockMatrix') -> 'BlockMatrix[bool]':
+        return ~(self == other)
+
+    def all(self) -> bool:
+        for line in self._stored_blocks:
+            for block in line:
+                if not block.all():
+                    return False
+        return True
+
+    def any(self) -> bool:
+        for line in self._stored_blocks:
+            for block in line:
+                if block.any():
+                    return True
+        return False
+
+    def min(self) -> Number:
+        return min(block.min() for line in self._stored_blocks for block in line)
+
+    def max(self) -> Number:
+        return max(block.max() for line in self._stored_blocks for block in line)
+
+    # DISPLAYING DATA
+
+    @property
+    def str_shape(self):
+        blocks_str = []
+        for line in self.all_blocks:
+            for block in line:
+                if isinstance(block, BlockMatrix):
+                    blocks_str.append(block.str_shape)
+                elif isinstance(block, np.ndarray) or isinstance(block, LowRankMatrix):
+                    blocks_str.append("{}×{}".format(*block.shape))
+                else:
+                    blocks_str.append("?×?")
+
+        if len(set(blocks_str)) == 1:
+            return "{}×{}×[".format(*self.nb_blocks) + blocks_str[0] + "]"
+        else:
+            blocks_str = np.array(blocks_str).reshape(self.nb_blocks).tolist()
+            return str(blocks_str).replace("'", "")
+
+    def __str__(self):
+        if not hasattr(self, '_str'):
+            args = [self.str_shape]
+            if self.dtype not in [np.float64, float]:
+                args.append(f"dtype={self.dtype}")
+            self._str = f"{self.__class__.__name__}(" + ", ".join(args) + ")"
+        return self._str
+
+    def _repr_pretty_(self, p, cycle):
+        p.text(self.__str__())
+
+    display_color = cycle([f'C{i}' for i in range(10)])
+
+    def _patches(self,
+                 global_frame: Union[Tuple[int, int], np.ndarray]
+                 ):
+        """Helper function for displaying the shape of the matrix.
+        Recursively returns a list of rectangles representing the sub-blocks of the matrix.
+
+        Uses BlockMatrix.display_color to assign color to the blocks.
+        By default, it cycles through matplotlib default colors.
+        But if display_color is redefined as a callable, it is called with the block as argument.
+
+        Parameters
+        ----------
+        global_frame: tuple of ints
+            coordinates of the origin in the top left corner.
+
+        Returns
+        -------
+        list of matplotlib.patches.Rectangle
+        """
+        matplotlib_patches = import_optional_dependency("matplotlib.patches", "matplotlib")
+        Rectangle = matplotlib_patches.Rectangle
+
+        all_blocks_in_flat_iterator = (block for line in self._stored_blocks for block in line)
+        positions_of_all_blocks = self._stored_block_positions(global_frame=global_frame)
+        patches = []
+        for block, positions_of_the_block in zip(all_blocks_in_flat_iterator, positions_of_all_blocks):
+            position_of_first_appearance = positions_of_the_block[0]
+            # Exchange coordinates: row index i -> y, column index j -> x
+            position_of_first_appearance = np.array((position_of_first_appearance[1], position_of_first_appearance[0]))
+
+            if isinstance(block, BlockMatrix):
+                patches_of_this_block = block._patches(np.array((position_of_first_appearance[1], position_of_first_appearance[0])))
+            elif isinstance(block, np.ndarray):
+
+                if isinstance(self.display_color, Iterator):
+                    color = next(self.display_color)
+                elif callable(self.display_color):
+                    color = self.display_color(block)
+                else:
+                    color = np.random.rand(3)
+
+                patches_of_this_block = [Rectangle(position_of_first_appearance,
+                                                   block.shape[1], block.shape[0],
+                                                   edgecolor='k', facecolor=color)]
+            elif isinstance(block, LowRankMatrix):
+
+                if isinstance(self.display_color, Iterator):
+                    color = next(self.display_color)
+                elif callable(self.display_color):
+                    color = self.display_color(block)
+                else:
+                    color = np.random.rand(3)
+
+                patches_of_this_block = [
+                    # Left block
+                    Rectangle(position_of_first_appearance,
+                              block.left_matrix.shape[1], block.left_matrix.shape[0],
+                              edgecolor='k', facecolor=color),
+                    # Top block
+                    Rectangle(position_of_first_appearance,
+                              block.right_matrix.shape[1], block.right_matrix.shape[0],
+                              edgecolor='k', facecolor=color),
+                    # Rest of the matrix
+                    Rectangle(position_of_first_appearance,
+                              block.right_matrix.shape[1], block.left_matrix.shape[0],
+                              facecolor=color, alpha=0.2),
+                ]
+            else:
+                raise NotImplementedError()
+
+            patches.extend(patches_of_this_block)
+
+            # For the other appearances, copy the patches of the first appearance
+            for block_position in positions_of_the_block[1:]:
+                block_position = np.array((block_position[1], block_position[0]))
+                for patch in patches_of_this_block:  # A block can be made of several patches.
+                    shift = block_position - position_of_first_appearance
+                    patch_position = np.array(patch.get_xy()) + shift
+                    patches.append(Rectangle(patch_position, patch.get_width(), patch.get_height(),
+                                             facecolor=patch.get_facecolor(), alpha=0.2))
+
+        return patches
+
+    def plot_shape(self):
+        """Plot the structure of the matrix using matplotlib."""
+        matplotlib = import_optional_dependency("matplotlib")
+        plt = matplotlib.pyplot
+
+        plt.figure()
+        for patch in self._patches((0, 0)):
+            plt.gca().add_patch(patch)
+        plt.axis('equal')
+        plt.xlim(0, self.shape[1])
+        plt.ylim(0, self.shape[0])
+        plt.gca().invert_yaxis()
+        # plt.show()
+
+
+    def access_block_by_path(self, path):
+        """
+        Access a diagonal block in a block matrix from the path of the
+        corresponding leaf
+        """
+        this_block = self
+        for index in path:
+            this_block = this_block.all_blocks[index, index]
+        return this_block