capytaine 2.3.1__cp314-cp314-macosx_14_0_arm64.whl

capytaine/matrices/block_toeplitz.py

@@ -0,0 +1,325 @@
+"""This modules implements block Toeplitz matrices to be used in Hierarchical Toeplitz matrices.
+
+The module also contains several special cases such as block symmetric Toeplitz matrices and block circulant matrices.
+All classes inherits from the BlockMatrix class.
+"""
+# Copyright (C) 2017-2019 Matthieu Ancellin
+# See LICENSE file at <https://github.com/mancellin/capytaine>
+
+import logging
+from typing import Tuple, List, Set, Iterable
+
+import numpy as np
+
+from capytaine.matrices.block import BlockMatrix
+
+LOG = logging.getLogger(__name__)
+
+
+################################################################################
+#                            Block Toeplitz matrix                             #
+################################################################################
+
+class BlockToeplitzMatrix(BlockMatrix):
+    """A (2D) block Toeplitz matrix, stored as a list of blocks.
+    All blocks should have the same shape.
+
+    Stored in the backend as a 1×(2N-1) array of arrays."""
+
+    # INITIALIZATION
+
+    def _compute_shape(self) -> Tuple[int, int]:
+        # The full shape is found by multiplying the shape of the blocks. All of them have the same shape.
+        return (self._stored_block_shapes[0][0]*self.nb_blocks[0],
+                self._stored_block_shapes[1][0]*self.nb_blocks[1])
+
+    def _compute_nb_blocks(self) -> Tuple[int, int]:
+        """Will be overridden by subclasses."""
+        assert self._stored_nb_blocks[1] % 2 == 1, "Expecting an odd number of blocks to build a Toeplitz matrix"
+        n = (self._stored_nb_blocks[1]+1)//2
+        return n, n
+
+    def _check_dimensions_of_blocks(self) -> bool:
+        for block in self._stored_blocks[0, :]:
+            if not block.shape == self.block_shape:  # All blocks have same shape
+                return False
+        return True
+
+    # ACCESSING DATA
+
+    @property
+    def block_shapes(self) -> Tuple[List[int], List[int]]:
+        """The shapes of the blocks composing the block matrix.
+        Actually, they should be all the same."""
+        return ([self._stored_block_shapes[0][0]]*self.nb_blocks[0],
+                [self._stored_block_shapes[1][0]]*self.nb_blocks[1])
+
+    @property
+    def block_shape(self) -> Tuple[int, int]:
+        """The shape of any of the blocks."""
+        return self._stored_block_shapes[0][0], self._stored_block_shapes[1][0]
+
+    def _block_indices_of(self, k: int) -> Set[Tuple[int, int]]:
+        """The block indices at which the stored block k can be found in the full matrix of size n.
+        Will be overridden by subclasses."""
+        n = self.nb_blocks[0]
+
+        if k < n:
+            i, j = 0, k  # Upper triangle
+        elif n <= k < 2*n:
+            i, j = 2*n-1-k, 0  # Lower triangle
+        else:
+            raise AttributeError
+
+        indices = set()
+        while i < n and j < n:
+            indices.add((i, j))
+            i, j = i+1, j+1  # Going along the diagonal
+
+        return indices
+
+    @property
+    def all_blocks(self):
+        all_blocks = np.empty(self.nb_blocks, dtype=object)
+        for k in range(self._stored_nb_blocks[1]):
+            for i, j in self._block_indices_of(k):
+                all_blocks[i, j] = self._stored_blocks[0, k]
+        return all_blocks
+
+    def _positions_of(self, k: int, global_frame=(0, 0)) -> List[Tuple[int, int]]:
+        """The positions in the full matrix at which the block k from the first line can also be found."""
+        shape = self.block_shape
+        return sorted([(global_frame[0] + i*shape[0], global_frame[1] + j*shape[1])
+                       for i, j in self._block_indices_of(k)])
+
+    def _stored_block_positions(self, global_frame=(0, 0)) -> Iterable[List[Tuple[int, int]]]:
+        """The position of each blocks in the matrix.
+
+        Example::
+
+            AABB
+            AABB  ->  list(matrix._stored_block_positions) = [[(0,0), (2, 2)], [(0, 2), (2, 0)]]
+            BBAA
+            BBAA
+        """
+        return (self._positions_of(k, global_frame=global_frame) for k in range(self._stored_nb_blocks[1]))
+
+    # # TRANSFORMING DATA
+
+    @property
+    def circulant_super_matrix(self):
+        if not hasattr(self, '_circulant_super_matrix'):
+            self._circulant_super_matrix = BlockCirculantMatrix(
+                self._stored_blocks,
+                _stored_block_shapes=self._stored_block_shapes,
+                check=False)
+        return self._circulant_super_matrix
+
+    def matvec(self, other):
+        """Matrix vector product.
+        Named as such to be used as scipy LinearOperator."""
+        LOG.debug(f"Product of {self} with vector of shape {other.shape}")
+        A = self.circulant_super_matrix
+        b = np.concatenate([other, np.zeros(A.shape[1] - self.shape[1])])
+        return (A @ b)[:self.shape[0]]
+
+    def rmatvec(self, other):
+        """Matrix vector product.
+        Named as such to be used as scipy LinearOperator."""
+        LOG.debug(f"Product of vector of shape {other.shape} with {self}")
+        if other.ndim == 2 and other.shape[0] == 1:  # Actually a 1×N matrix
+            other = other[0, :]
+        A = self.circulant_super_matrix
+        b = np.concatenate([other, np.zeros(A.shape[0] - self.shape[0])])
+        return (A.rmatvec(b))[:self.shape[1]]
+
+
+################################################################################
+#                       Block symmetric Toeplitz matrix                        #
+################################################################################
+
+class BlockSymmetricToeplitzMatrix(BlockToeplitzMatrix):
+    """A (2D) block symmetric Toeplitz matrix, stored as a list of blocks.
+    All blocks should have the same shape.
+
+    Stored in the backend as a 1×N array of arrays."""
+
+    def _compute_nb_blocks(self) -> Tuple[int, int]:
+        n = self._stored_nb_blocks[1]
+        return n, n
+
+    def _block_indices_of(self, k: int) -> Set[Tuple[int, int]]:
+        """The block indices at which the stored block k can be found in the full matrix."""
+        n = self.nb_blocks[0]
+        assert k < n
+
+        if k == 0:
+            return BlockToeplitzMatrix._block_indices_of(self, 0)
+        else:
+            return BlockToeplitzMatrix._block_indices_of(self, k).union(
+                BlockToeplitzMatrix._block_indices_of(self, 2*n-k-1))
+
+    @property
+    def circulant_super_matrix(self):
+        if not hasattr(self, '_circulant_super_matrix'):
+            self._circulant_super_matrix = EvenBlockSymmetricCirculantMatrix(
+                self._stored_blocks,
+                _stored_block_shapes=self._stored_block_shapes,
+                check=False)
+        return self._circulant_super_matrix
+
+
+################################################################################
+#                            Block circulant matrix                            #
+################################################################################
+
+class BlockCirculantMatrix(BlockToeplitzMatrix):
+    """A (2D) block circulant matrix, stored as a list of blocks.
+    All blocks should have the same shape.
+
+    Stored in the backend as a 1×N array of arrays."""
+
+    def _compute_nb_blocks(self) -> Tuple[int, int]:
+        n = self._stored_nb_blocks[1]
+        return n, n
+
+    def _block_indices_of(self, k: int) -> Set[Tuple[int, int]]:
+        """The block indices at which the stored block k can be found in the full matrix."""
+        n = self.nb_blocks[0]
+        assert k < n
+
+        if k == 0:
+            return BlockToeplitzMatrix._block_indices_of(self, 0)
+        else:
+            return BlockToeplitzMatrix._block_indices_of(self, k).union(
+                BlockToeplitzMatrix._block_indices_of(self, n+k-1))
+
+    # LINEAR SYSTEMS
+
+    def block_diagonalize(self):
+        """Returns a vector of matrices"""
+        if not hasattr(self, 'block_diagonalization'):
+            if all(isinstance(matrix, BlockMatrix) for matrix in self._stored_blocks[0, :]):
+                self.block_diagonalization = BlockMatrix.fft_of_list(*self.all_blocks[:, 0])
+            else:
+                stacked_blocks = np.empty((self.nb_blocks[1],) + self.block_shape, dtype=self.dtype)
+                for i, block in enumerate(self.all_blocks[:, 0]):
+                    stacked_blocks[i] = block.full_matrix() if not isinstance(block, np.ndarray) else block
+                self.block_diagonalization =  np.fft.fft(stacked_blocks, axis=0)
+        return self.block_diagonalization
+
+    def matvec(self, other):
+        """Matrix vector product.
+        Named as such to be used as scipy LinearOperator."""
+        LOG.debug(f"Product of {self} with vector of shape {other.shape}")
+        fft_of_vector = np.fft.fft(np.reshape(other, (self.nb_blocks[0], self.block_shape[1], 1)), axis=0)
+        blocks_of_diagonalization = self.block_diagonalize()
+        try:  # Try to run it as vectorized numpy arrays.
+            fft_of_result = blocks_of_diagonalization @ fft_of_vector
+        # When the above fails, numpy 1.15 returns a TypeError, whereas numpy 1.16 returns a ValueError.
+        except (TypeError, ValueError):  # Or do the same thing with list comprehension.
+            fft_of_result = np.array([block @ vec for block, vec in zip(blocks_of_diagonalization, fft_of_vector)])
+        result = np.fft.ifft(fft_of_result, axis=0).reshape(self.shape[0])
+        if self.dtype == complex or other.dtype == complex:
+            return np.asarray(result)
+        else:
+            return np.asarray(np.real(result))
+
+    def rmatvec(self, other):
+        """Matrix vector product.
+        Named as such to be used as scipy LinearOperator."""
+        other = np.conjugate(other)
+        fft_of_vector = np.fft.ifft(np.reshape(other, (self.nb_blocks[0], 1, self.block_shape[0])), axis=0)
+        blocks_of_diagonalization = self.block_diagonalize()
+        try:  # Try to run it as vectorized numpy arrays.
+            fft_of_result = fft_of_vector @ blocks_of_diagonalization
+        # When the above fails, numpy 1.15 returns a TypeError, whereas numpy 1.16 returns a ValueError.
+        except (TypeError, ValueError):
+            # Instead we do the same thing with list comprehension.
+            fft_of_result = np.array(
+                [block.rmatvec(vec.flatten()) for block, vec in zip(blocks_of_diagonalization, fft_of_vector)]
+            )
+        result = np.fft.fft(fft_of_result, axis=0).reshape(self.shape[1])
+        if self.dtype == complex or other.dtype == complex:
+            return np.asarray(result)
+        else:
+            return np.asarray(np.real(result))
+
+
+###########################################################################
+#                    Block symmetric circulant matrix                     #
+###########################################################################
+
+class EvenBlockSymmetricCirculantMatrix(BlockCirculantMatrix, BlockSymmetricToeplitzMatrix):
+    """A block symmetric circulant matrix, with an even number of blocks.
+
+    Examples::
+
+        ABCB
+        BABC
+        CBAB
+        BCBA
+
+        ABCDCB
+        BABCDB
+        CBABCD
+        DCBABC
+        CDCBAB
+        BCDCBA
+
+    Stored in the backend as a 1×(N/2+1) array of arrays."""
+
+    def _compute_nb_blocks(self) -> Tuple[int, int]:
+        """The number of blocks in the full matrix."""
+        n = (self._stored_nb_blocks[1] - 1)*2
+        return n, n
+
+    def _block_indices_of(self, k: int) -> List[Tuple[int, int]]:
+        n = self.nb_blocks[0]
+        assert k < n/2 + 1
+        if k == 0:
+            return BlockToeplitzMatrix._block_indices_of(self, 0)
+        else:
+            return (BlockToeplitzMatrix._block_indices_of(self, k) |
+                BlockToeplitzMatrix._block_indices_of(self, n+k-1) |
+                BlockToeplitzMatrix._block_indices_of(self, n-k)   |
+                BlockToeplitzMatrix._block_indices_of(self, 2*n-k-1)
+                    )
+
+
+class OddBlockSymmetricCirculantMatrix(BlockCirculantMatrix, BlockSymmetricToeplitzMatrix):
+    """A block symmetric circulant matrix, with an odd number of blocks.
+
+    Examples::
+
+        ABCCB
+        BABCC
+        CBABC
+        CCBAB
+        BCCBA
+
+        ABCDDCB
+        BABCDDB
+        CBABCDD
+        DCBABCD
+        DDCBABC
+        CDDCBAB
+        BCDDCBA
+
+    Stored in the backend as a 1×(N+1)/2 array of arrays."""
+
+    def _compute_nb_blocks(self) -> Tuple[int, int]:
+        n = self._stored_nb_blocks[1]*2 - 1
+        return n, n
+
+    def _block_indices_of(self, k: int) -> List[Tuple[int, int]]:
+        n = self.nb_blocks[0]
+        assert k < (n+1)/2
+        if k == 0:
+            return BlockToeplitzMatrix._block_indices_of(self, 0)
+        else:
+            return (BlockToeplitzMatrix._block_indices_of(self, k) |
+                BlockToeplitzMatrix._block_indices_of(self, n+k-1) |
+                BlockToeplitzMatrix._block_indices_of(self, n-k)   |
+                BlockToeplitzMatrix._block_indices_of(self, 2*n-k-1)
+                    )

capytaine/matrices/builders.py

@@ -0,0 +1,89 @@
+"""This module contains some helpful functions to create block matrices."""
+# Copyright (C) 2017-2019 Matthieu Ancellin
+# See LICENSE file at <https://github.com/mancellin/capytaine>
+
+import logging
+from itertools import accumulate
+
+import numpy as np
+
+from capytaine.matrices.block import BlockMatrix
+from capytaine.matrices.low_rank import LowRankMatrix
+
+LOG = logging.getLogger(__name__)
+
+
+def cut_matrix(full_matrix, x_shapes, y_shapes, check=False):
+    """Transform a numpy array into a block matrix of numpy arrays.
+
+    Parameters
+    ----------
+    full_matrix: numpy array
+        The matrix to split into blocks.
+    x_shapes: sequence of int
+        The columns at which to split the blocks.
+    y_shapes: sequence of int
+        The lines at which to split the blocks.
+    check: bool, optional
+        Check to dimensions and type of the matrix after creation (default: False).
+
+    Return
+    ------
+    BlockMatrix
+        The same matrix as the input one but in block form.
+    """
+    new_block_matrix = []
+    for i, di in zip(accumulate([0] + x_shapes[:-1]), x_shapes):
+        line = []
+        for j, dj in zip(accumulate([0] + x_shapes[:-1]), y_shapes):
+            line.append(full_matrix[i:i+di, j:j+dj])
+        new_block_matrix.append(line)
+    return BlockMatrix(new_block_matrix, check=check)
+
+
+def random_block_matrix(x_shapes, y_shapes, rng=np.random.default_rng()):
+    """A random block matrix."""
+    return cut_matrix(rng.uniform(size=(sum(x_shapes), sum(y_shapes))), x_shapes, y_shapes)
+
+
+def full_like(A, value, dtype=np.float64):
+    """A matrix of the same kind and shape as A but filled with a single value."""
+    if isinstance(A, BlockMatrix):
+        new_matrix = []
+        for i in range(A._stored_nb_blocks[0]):
+            line = []
+            for j in range(A._stored_nb_blocks[1]):
+                line.append(full_like(A._stored_blocks[i, j], value, dtype=dtype))
+            new_matrix.append(line)
+        return A.__class__(new_matrix)
+    elif isinstance(A, LowRankMatrix):
+        return LowRankMatrix(np.ones((A.shape[0], 1)), np.full((1, A.shape[1]), value))
+    elif isinstance(A, np.ndarray):
+        return np.full_like(A, value, dtype=dtype)
+
+
+def zeros_like(A, dtype=np.float64):
+    """A matrix of the same kind and shape as A but filled with zeros."""
+    return full_like(A, 0.0, dtype=dtype)
+
+
+def ones_like(A, dtype=np.float64):
+    """A matrix of the same kind and shape as A but filled with ones."""
+    return full_like(A, 1.0, dtype=dtype)
+
+
+def identity_like(A, dtype=np.float64):
+    """A identity matrix of the same kind and shape as A."""
+    if isinstance(A, BlockMatrix):
+        I = []
+        for i in range(A._stored_nb_blocks[0]):
+            line = []
+            for j in range(A._stored_nb_blocks[1]):
+                if i == j:
+                    line.append(identity_like(A._stored_blocks[i, j], dtype=dtype))
+                else:
+                    line.append(zeros_like(A._stored_blocks[i, j], dtype=dtype))
+            I.append(line)
+        return A.__class__(I)
+    elif isinstance(A, np.ndarray):
+        return np.eye(A.shape[0], A.shape[1], dtype=dtype)

capytaine/matrices/linear_solvers.py

@@ -0,0 +1,232 @@
+"""The linear solvers used in Capytaine.
+
+They are based on numpy solvers with a thin layer for the handling of Hierarchical Toeplitz matrices.
+"""
+# Copyright (C) 2017-2019 Matthieu Ancellin
+# See LICENSE file at <https://github.com/capytaine/capytaine>
+
+import logging
+
+import numpy as np
+from scipy import linalg as sl
+from scipy.sparse import linalg as ssl
+
+from capytaine.matrices.block import BlockMatrix
+from capytaine.matrices.block_toeplitz import BlockSymmetricToeplitzMatrix, BlockCirculantMatrix
+
+LOG = logging.getLogger(__name__)
+
+
+# DIRECT SOLVER
+
+def solve_directly(A, b):
+    assert isinstance(b, np.ndarray) and A.ndim == b.ndim+1 and A.shape[-2] == b.shape[-1]
+    if isinstance(A, BlockCirculantMatrix):
+        LOG.debug("\tSolve linear system %s", A)
+        blocks_of_diagonalization = A.block_diagonalize()
+        fft_of_rhs = np.fft.fft(np.reshape(b, (A.nb_blocks[0], A.block_shape[0])), axis=0)
+        try:  # Try to run it as vectorized numpy arrays.
+            fft_of_result = np.linalg.solve(blocks_of_diagonalization, fft_of_rhs[..., np.newaxis])[..., 0]
+        except np.linalg.LinAlgError:  # Or do the same thing with list comprehension.
+            fft_of_result = np.array([solve_directly(block, vec) for block, vec in zip(blocks_of_diagonalization, fft_of_rhs)])
+        result = np.fft.ifft(fft_of_result, axis=0).reshape((A.shape[1],))
+        return result
+
+    elif isinstance(A, BlockSymmetricToeplitzMatrix):
+        if A.nb_blocks == (2, 2):
+            LOG.debug("\tSolve linear system %s", A)
+            A1, A2 = A._stored_blocks[0, :]
+            b1, b2 = b[:len(b)//2], b[len(b)//2:]
+            x_plus = solve_directly(A1 + A2, b1 + b2)
+            x_minus = solve_directly(A1 - A2, b1 - b2)
+            return np.concatenate([x_plus + x_minus, x_plus - x_minus])/2
+        else:
+            # Not implemented
+            LOG.debug("\tSolve linear system %s", A)
+            return solve_directly(A.full_matrix(), b)
+
+    elif isinstance(A, BlockMatrix):
+        LOG.debug("\tSolve linear system %s", A)
+        return solve_directly(A.full_matrix(), b)
+
+    elif isinstance(A, np.ndarray):
+        LOG.debug(f"\tSolve linear system (size: {A.shape}) with numpy direct solver.")
+        return np.linalg.solve(A, b)
+
+    else:
+        raise ValueError(f"Unrecognized type of matrix to solve: {A}")
+
+
+# CACHED LU DECOMPOSITION
+class LUSolverWithCache:
+    """Solve linear system with the LU decomposition.
+
+    The latest LU decomposition is kept in memory, if a system with the same matrix needs to be solved again, then the decomposition is reused.
+
+    Most of the complexity of this class comes from:
+    1. @lru_cache does not work because numpy arrays are not hashable. So a basic cache system has been recoded from scratch.
+    2. To be the default solver for the BasicMatrixEngine, the solver needs to support matrices for problems with one or two reflection symmetries.
+    Hence, a custom way to cache the LU decomposition of the matrices involved in the direct linear resolution of the symmetric problem.
+    """
+    def __init__(self):
+        self.cached_matrix = None
+        self.cached_decomp = None
+
+    def solve(self, A, b):
+        return self.solve_with_decomp(self.cached_lu_decomp(A), b)
+
+    def lu_decomp(self, A):
+        """Return the LU decomposition of A.
+        If A is BlockSymmetricToeplitzMatrix, then return a list of LU decompositions for each block of the block diagonalisation of the matrix.
+        """
+        if isinstance(A, BlockSymmetricToeplitzMatrix) and A.nb_blocks == (2, 2):
+            A1, A2 = A._stored_blocks[0, :]
+            return [self.lu_decomp(A1 + A2), self.lu_decomp(A1 - A2)]
+        elif isinstance(A, np.ndarray):
+            return sl.lu_factor(A)
+        else:
+            raise NotImplementedError("Cached LU solver is only implemented for dense matrices and 2×2 BlockSymmetricToeplitzMatrix.")
+
+    def cached_lu_decomp(self, A):
+        if not(A is self.cached_matrix):
+            self.cached_matrix = A
+            LOG.debug(f"Computing and caching LU decomposition")
+            self.cached_decomp = self.lu_decomp(A)
+        else:
+            LOG.debug(f"Using cached LU decomposition")
+        return self.cached_decomp
+
+    def solve_with_decomp(self, decomp, b):
+        """Solve the system using the precomputed LU decomposition.
+        TODO: find a better way to differentiate a LU decomposition (returned as tuple by sl.lu_factor)
+        and a set of LU decomposition (stored as a list in self.lu_decomp).
+        """
+        if isinstance(decomp, list):  # The matrix was a BlockSymmetricToeplitzMatrix
+            b1, b2 = b[:len(b)//2], b[len(b)//2:]
+            x_plus = self.solve_with_decomp(decomp[0], b1 + b2)
+            x_minus = self.solve_with_decomp(decomp[1], b1 - b2)
+            return np.concatenate([x_plus + x_minus, x_plus - x_minus])/2
+        elif isinstance(decomp, tuple):  # The matrix was a np.ndarray
+            return sl.lu_solve(decomp, b)
+        else:
+            raise NotImplementedError("Cached LU solver is only implemented for dense matrices and 2×2 BlockSymmetricToeplitzMatrix.")
+
+
+# ITERATIVE SOLVER
+
+class Counter:
+    def __init__(self):
+        self.nb_iter = 0
+
+    def __call__(self, *args, **kwargs):
+        self.nb_iter += 1
+
+
+def solve_gmres(A, b):
+    LOG.debug(f"Solve with GMRES for {A}.")
+
+    if LOG.isEnabledFor(logging.INFO):
+        counter = Counter()
+        x, info = ssl.gmres(A, b, atol=1e-6, callback=counter, callback_type="pr_norm")
+        LOG.info(f"End of GMRES after {counter.nb_iter} iterations.")
+
+    else:
+        x, info = ssl.gmres(A, b, atol=1e-6)
+
+    if info > 0:
+        raise RuntimeError(f"No convergence of the GMRES after {info} iterations.\n"
+                            "This can be due to overlapping panels or irregular frequencies.\n"
+                            "In the latter case, using a direct solver can help (https://github.com/mancellin/capytaine/issues/30).")
+
+    return x
+
+def gmres_no_fft(A, b):
+    LOG.debug(f"Solve with GMRES for {A} without using FFT.")
+
+    x, info = ssl.gmres(A.no_toeplitz() if isinstance(A, BlockMatrix) else A, b, atol=1e-6)
+
+    if info != 0:
+        LOG.warning(f"No convergence of the GMRES. Error code: {info}")
+
+    return x
+
+
+# PRECONDITIONED SOLVER
+
+def _block_Jacobi_coarse_corr(A, b, x0, R, RA, AcLU, DLU, diag_shapes, n):
+    """
+    Performs a single step of the block-Jacobi method with coarse correction.
+    Can be used as a preconditioner for matrix A.
+
+    Parameters
+    ----------
+    A: BlockMatrix
+        System matrix
+    b: array
+        System right hand side vector
+    x0: array
+        Initial guess of the solution
+    R: array
+        Coarse space restriction matrix
+    RA: array
+        Precomputed product of R and A
+    AcLU: list
+        LU decomposition data of the coarse system matrix Ac (output of
+        scipy.linalg.lu_factor)
+    DLU: list
+        List of LU decomposition data for the diagonal blocks of A
+    diag_shapes: list
+        List of shapes of diagonal blocks
+    n: integer
+        Size of the coarse problem (e.g. number of bodies simulated)
+
+    Returns
+    -------
+    array
+        Action of a step of the method on vector x0
+    """
+    # x_ps = x after the pre-smoothing step (block-Jacobi)
+    x_ps = np.zeros(A.shape[0], dtype=complex)
+
+    # the diagonal blocks of A have already been put to zero in build_matrices
+    # they are not needed anymore
+    q = b - A@x0
+    # loop over diagonal blocks
+    for kk in range(n):
+        local_slice = slice(sum(diag_shapes[:kk]), sum(diag_shapes[:kk+1]))
+        local_rhs = q[local_slice]
+        local_sol = sl.lu_solve(DLU[kk], local_rhs, check_finite=False)
+
+        x_ps[local_slice] = local_sol
+
+    r_c = R@b - RA@x_ps #restricted residual
+    e_c = sl.lu_solve(AcLU, r_c, check_finite=False)
+    # update
+    return x_ps + R.T@e_c
+
+def solve_precond_gmres(A_and_precond_data, b):
+    """
+    Implementation of the preconditioner presented in
+    `<https://doi.org/10.1007/978-3-031-50769-4_14>`.
+    """
+    A, R, RA, AcLU, DLU, diag_shapes, n, PinvA = A_and_precond_data
+    N = A.shape[0]
+
+    Pinvb = _block_Jacobi_coarse_corr(A, b, np.zeros(N, dtype=complex), R, RA, AcLU, DLU, diag_shapes, n)
+
+    LOG.debug(f"Solve with GMRES for {A}.")
+
+    if LOG.isEnabledFor(logging.INFO):
+        counter = Counter()
+        x, info = ssl.gmres(PinvA, Pinvb, atol=1e-6, callback=counter)
+        LOG.info(f"End of GMRES after {counter.nb_iter} iterations.")
+
+    else:
+        x, info = ssl.gmres(PinvA, Pinvb, atol=1e-6)
+
+    if info > 0:
+        raise RuntimeError(f"No convergence of the GMRES after {info} iterations.\n"
+                            "This can be due to overlapping panels or irregular frequencies.\n"
+                            "In the latter case, using a direct solver can help (https://github.com/mancellin/capytaine/issues/30).")
+
+    return x