pygeoinf 1.2.6__py3-none-any.whl → 1.2.8__py3-none-any.whl

pygeoinf/linear_operators.py

@@ -18,12 +18,18 @@ Key Classes
 """
 
 from __future__ import annotations
-from typing import Callable, List, Optional, Any, Union, Tuple, TYPE_CHECKING
+from typing import Callable, List, Optional, Any, Union, Tuple, TYPE_CHECKING, Dict
+
+from collections import defaultdict
 
 import numpy as np
+import scipy.sparse as sp
 from scipy.sparse.linalg import LinearOperator as ScipyLinOp
 from scipy.sparse import diags
 
+
+from joblib import Parallel, delayed
+
 # from .operators import Operator
 from .nonlinear_operators import NonLinearOperator
 
@@ -49,8 +55,8 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
 
     This class represents a linear map `L(x) = Ax` and provides rich
     functionality for linear algebraic operations. It specializes
-    `NonLinearOperator`, correctly defining its derivative as the operator
-    itself.
+    `NonLinearOperator`, with the derivative mapping taking the
+    required form (i.e., the derivative is just the operator itself).
 
     Key features include operator algebra (`@`, `+`, `*`), automatic
     derivation of adjoint (`.adjoint`) and dual (`.dual`) operators, and
@@ -67,7 +73,6 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         *,
         dual_mapping: Optional[Callable[[Any], Any]] = None,
         adjoint_mapping: Optional[Callable[[Any], Any]] = None,
-        thread_safe: bool = False,
         dual_base: Optional[LinearOperator] = None,
         adjoint_base: Optional[LinearOperator] = None,
     ) -> None:
@@ -80,9 +85,14 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
             mapping (callable): The function defining the linear mapping.
             dual_mapping (callable, optional): The action of the dual operator.
             adjoint_mapping (callable, optional): The action of the adjoint.
-            thread_safe (bool, optional): True if the mapping is thread-safe.
             dual_base (LinearOperator, optional): Internal use for duals.
             adjoint_base (LinearOperator, optional): Internal use for adjoints.
+
+        Notes:
+            If neither the dual or adjoint mappings are provided, an they are
+            deduced internally using a correction but very inefficient method.
+            In general this functionality should not be relied on other than
+            for operators between low-dimensional spaces.
         """
         super().__init__(
             domain, codomain, self._mapping_impl, derivative=self._derivative_impl
@@ -90,7 +100,6 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         self._mapping = mapping
         self._dual_base: Optional[LinearOperator] = dual_base
         self._adjoint_base: Optional[LinearOperator] = adjoint_base
-        self._thread_safe: bool = thread_safe
         self.__adjoint_mapping: Callable[[Any], Any]
         self.__dual_mapping: Callable[[Any], Any]
 
@@ -249,83 +258,107 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
     def from_matrix(
         domain: HilbertSpace,
         codomain: HilbertSpace,
-        matrix: Union[np.ndarray, ScipyLinOp],
+        matrix: Union[np.ndarray, sp.sparray, ScipyLinOp],
         /,
         *,
         galerkin: bool = False,
-    ) -> LinearOperator:
+    ) -> MatrixLinearOperator:
         """
-        Creates a LinearOperator from its matrix representation.
+        Creates the most appropriate LinearOperator from a matrix representation.
 
-        This factory defines a `LinearOperator` using a concrete matrix that
-        acts on the component vectors of the abstract Hilbert space vectors.
+        This factory method acts as a dispatcher, inspecting the type of the
+        input matrix and returning the most specialized and optimized operator
+        subclass (e.g., Dense, Sparse, or DiagonalSparse). It also handles
+        matrix-free `scipy.sparse.linalg.LinearOperator` objects.
 
         Args:
             domain: The operator's domain space.
             codomain: The operator's codomain space.
-            matrix: The matrix representation (NumPy array or SciPy
-                LinearOperator). Shape must be `(codomain.dim, domain.dim)`.
-            galerkin: If `True`, the matrix is interpreted in its "weak form"
-                or Galerkin representation (`M_ij = <basis_j, A(basis_i)>`),
-                which maps a vector's components to the components of its
-                *dual*. This is crucial as it ensures a self-adjoint
-                operator is represented by a symmetric matrix. If `False`
-                (default), it's a standard component-to-component map.
+            matrix: The matrix representation (NumPy ndarray, SciPy sparray,
+                    or SciPy LinearOperator).
+            galerkin: If `True`, the matrix is interpreted in Galerkin form.
 
         Returns:
-            A new `LinearOperator` defined by the matrix action.
+            An instance of the most appropriate MatrixLinearOperator subclass.
         """
+        # The order of these checks is important: from most specific to most general.
 
-        assert matrix.shape == (codomain.dim, domain.dim)
+        # 1. Check for the most specific diagonal-sparse format
+        if isinstance(matrix, sp.dia_array):
+            diagonals_tuple = (matrix.data, matrix.offsets)
+            return DiagonalSparseMatrixLinearOperator(
+                domain, codomain, diagonals_tuple, galerkin=galerkin
+            )
 
-        if galerkin:
+        # 2. Check for any other modern sparse format
+        elif isinstance(matrix, sp.sparray):
+            return SparseMatrixLinearOperator(
+                domain, codomain, matrix, galerkin=galerkin
+            )
 
-            def mapping(x: Any) -> Any:
-                cx = domain.to_components(x)
-                cyp = matrix @ cx
-                yp = codomain.dual.from_components(cyp)
-                return codomain.from_dual(yp)
+        # 3. Check for a dense NumPy array
+        elif isinstance(matrix, np.ndarray):
+            return DenseMatrixLinearOperator(
+                domain, codomain, matrix, galerkin=galerkin
+            )
 
-            def adjoint_mapping(y: Any) -> Any:
-                cy = codomain.to_components(y)
-                cxp = matrix.T @ cy
-                xp = domain.dual.from_components(cxp)
-                return domain.from_dual(xp)
+        # 4. Check for a matrix-free SciPy LinearOperator
+        elif isinstance(matrix, ScipyLinOp):
+            # This is matrix-free, so the general MatrixLinearOperator is the correct wrapper.
+            return MatrixLinearOperator(domain, codomain, matrix, galerkin=galerkin)
 
-            return LinearOperator(
-                domain,
-                codomain,
-                mapping,
-                adjoint_mapping=adjoint_mapping,
+        # 5. Handle legacy sparse matrix formats (optional but robust)
+        elif sp.issparse(matrix):
+            modern_array = sp.csr_array(matrix)
+            return SparseMatrixLinearOperator(
+                domain, codomain, modern_array, galerkin=galerkin
             )
 
+        # 6. Raise an error for unsupported types
         else:
-
-            def mapping(x: Any) -> Any:
-                cx = domain.to_components(x)
-                cy = matrix @ cx
-                return codomain.from_components(cy)
-
-            def dual_mapping(yp: Any) -> Any:
-                cyp = codomain.dual.to_components(yp)
-                cxp = matrix.T @ cyp
-                return domain.dual.from_components(cxp)
-
-            return LinearOperator(domain, codomain, mapping, dual_mapping=dual_mapping)
+            raise TypeError(f"Unsupported matrix type: {type(matrix)}")
 
     @staticmethod
     def self_adjoint_from_matrix(
-        domain: HilbertSpace, matrix: Union[np.ndarray, ScipyLinOp]
-    ) -> LinearOperator:
-        """Forms a self-adjoint operator from its Galerkin matrix."""
+        domain: HilbertSpace,
+        matrix: Union[np.ndarray, sp.sparray, ScipyLinOp],
+    ) -> MatrixLinearOperator:
+        """
+        Creates the most appropriate self-adjoint LinearOperator from a matrix.
 
-        def mapping(x: Any) -> Any:
-            cx = domain.to_components(x)
-            cyp = matrix @ cx
-            yp = domain.dual.from_components(cyp)
-            return domain.from_dual(yp)
+        This factory acts as a dispatcher, returning the most specialized
+        subclass for the given matrix type (e.g., Dense, Sparse).
 
-        return LinearOperator.self_adjoint(domain, mapping)
+        It ALWAYS assumes the provided matrix is the **Galerkin** representation
+        of the operator. The user is responsible for ensuring the input matrix
+        is symmetric (or self-adjoint for ScipyLinOp).
+
+        Args:
+            domain: The operator's domain and codomain space.
+            matrix: The symmetric matrix representation.
+
+        Returns:
+            An instance of the most appropriate MatrixLinearOperator subclass.
+        """
+        # Dispatch to the appropriate subclass, always with galerkin=True
+        if isinstance(matrix, sp.dia_array):
+            diagonals_tuple = (matrix.data, matrix.offsets)
+            return DiagonalSparseMatrixLinearOperator(
+                domain, domain, diagonals_tuple, galerkin=True
+            )
+        elif isinstance(matrix, sp.sparray):
+            return SparseMatrixLinearOperator(domain, domain, matrix, galerkin=True)
+        elif isinstance(matrix, np.ndarray):
+            return DenseMatrixLinearOperator(domain, domain, matrix, galerkin=True)
+        elif isinstance(matrix, ScipyLinOp):
+            return MatrixLinearOperator(domain, domain, matrix, galerkin=True)
+        elif sp.issparse(matrix):
+            modern_array = sp.csr_array(matrix)
+            return SparseMatrixLinearOperator(
+                domain, domain, modern_array, galerkin=True
+            )
+        else:
+            raise TypeError(f"Unsupported matrix type: {type(matrix)}")
 
     @staticmethod
     def from_tensor_product(
@@ -413,11 +446,6 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         else:
             return self._adjoint_base
 
-    @property
-    def thread_safe(self) -> bool:
-        """True if the operator's mapping is thread-safe."""
-        return self._thread_safe
-
     def matrix(
         self,
         /,
@@ -479,22 +507,34 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
                     return self.domain.dual.to_components(xp)
 
             def matmat(xmat: np.ndarray) -> np.ndarray:
-                n, k = xmat.shape
-                assert n == self.domain.dim
-                ymat = np.zeros((self.codomain.dim, k))
-                for j in range(k):
-                    cx = xmat[:, j]
-                    ymat[:, j] = matvec(cx)
-                return ymat
+                _n, k = xmat.shape
+                assert _n == self.domain.dim
+
+                if not parallel:
+                    ymat = np.zeros((self.codomain.dim, k))
+                    for j in range(k):
+                        ymat[:, j] = matvec(xmat[:, j])
+                    return ymat
+                else:
+                    result_cols = Parallel(n_jobs=n_jobs)(
+                        delayed(matvec)(xmat[:, j]) for j in range(k)
+                    )
+                    return np.column_stack(result_cols)
 
             def rmatmat(ymat: np.ndarray) -> np.ndarray:
-                m, k = ymat.shape
-                assert m == self.codomain.dim
-                xmat = np.zeros((self.domain.dim, k))
-                for j in range(k):
-                    cy = ymat[:, j]
-                    xmat[:, j] = rmatvec(cy)
-                return xmat
+                _m, k = ymat.shape
+                assert _m == self.codomain.dim
+
+                if not parallel:
+                    xmat = np.zeros((self.domain.dim, k))
+                    for j in range(k):
+                        xmat[:, j] = rmatvec(ymat[:, j])
+                    return xmat
+                else:
+                    result_cols = Parallel(n_jobs=n_jobs)(
+                        delayed(rmatvec)(ymat[:, j]) for j in range(k)
+                    )
+                    return np.column_stack(result_cols)
 
             return ScipyLinOp(
                 (self.codomain.dim, self.domain.dim),
@@ -504,6 +544,124 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
                 rmatmat=rmatmat,
             )
 
+    def extract_diagonal(
+        self,
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> np.ndarray:
+        """
+        Computes the main diagonal of the operator's matrix representation.
+
+        This method is highly parallelizable and memory-efficient, as it
+        avoids forming the full dense matrix.
+
+        Args:
+            galerkin: If True, computes the diagonal of the Galerkin matrix.
+            parallel: If True, computes the entries in parallel.
+            n_jobs: Number of parallel jobs to use.
+
+        Returns:
+            A NumPy array containing the diagonal entries.
+        """
+
+        dim = min(self.domain.dim, self.codomain.dim)
+        jobs = n_jobs if parallel else 1
+
+        def compute_entry(i: int) -> float:
+            """Worker function to compute a single diagonal entry."""
+            e_i = self.domain.basis_vector(i)
+            L_e_i = self(e_i)
+
+            if galerkin:
+                return self.domain.inner_product(e_i, L_e_i)
+            else:
+                return self.codomain.to_components(L_e_i)[i]
+
+        diagonal_entries = Parallel(n_jobs=jobs)(
+            delayed(compute_entry)(i) for i in range(dim)
+        )
+        return np.array(diagonal_entries)
+
+    def extract_diagonals(
+        self,
+        offsets: List[int],
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> Tuple[np.ndarray, List[int]]:
+        """
+        Computes specified diagonals of the operator's matrix representation.
+
+        This is a memory-efficient and parallelizable method that computes
+        the matrix one column at a time.
+
+        Args:
+            offsets: A list of diagonal offsets to extract (e.g., [0] for
+                the main diagonal, [-1, 0, 1] for a tridiagonal matrix).
+            galerkin: If True, computes the diagonals of the Galerkin matrix.
+            parallel: If True, computes columns in parallel.
+            n_jobs: Number of parallel jobs to use.
+
+        Returns:
+            A tuple containing:
+            - A NumPy array where each row is a diagonal.
+            - The list of offsets.
+            This format is compatible with scipy.sparse.spdiags.
+        """
+        dim = min(self.domain.dim, self.codomain.dim)
+        jobs = n_jobs if parallel else 1
+
+        # Prepare a thread-safe dictionary to store results
+
+        results: Dict[int, Dict[int, float]] = defaultdict(dict)
+
+        def compute_column_entries(j: int) -> Dict[int, Dict[int, float]]:
+            """
+            Worker function to compute all needed entries for column j.
+            """
+            e_j = self.domain.basis_vector(j)
+            L_e_j = self(e_j)
+
+            col_results = defaultdict(dict)
+
+            for k in offsets:
+                i = j - k
+                if 0 <= i < dim:
+                    if galerkin:
+                        e_i = self.domain.basis_vector(i)
+                        val = self.domain.inner_product(e_i, L_e_j)
+                    else:
+                        val = self.codomain.to_components(L_e_j)[i]
+                    col_results[k][i] = val
+            return col_results
+
+        # Run the computation in parallel
+        column_data = Parallel(n_jobs=jobs)(
+            delayed(compute_column_entries)(j) for j in range(dim)
+        )
+
+        # Aggregate results from the parallel computation
+        for col_dict in column_data:
+            for k, entries in col_dict.items():
+                results[k].update(entries)
+
+        # Format the results for spdiags
+        # The array must have padding for shorter off-diagonals.
+        diagonals_array = np.zeros((len(offsets), dim))
+        for idx, k in enumerate(offsets):
+            diag_entries = results[k]
+            for i, val in diag_entries.items():
+                j = i + k
+                if 0 <= j < dim:
+                    diagonals_array[idx, j] = val
+
+        return diagonals_array, offsets
+
     def random_svd(
         self,
         size_estimate: int,
@@ -517,7 +675,11 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         block_size: int = 10,
         parallel: bool = False,
         n_jobs: int = -1,
-    ) -> Tuple[LinearOperator, DiagonalLinearOperator, LinearOperator]:
+    ) -> Tuple[
+        DenseMatrixLinearOperator,
+        DiagonalSparseMatrixLinearOperator,
+        DenseMatrixLinearOperator,
+    ]:
         """
         Computes an approximate SVD using a randomized algorithm.
 
@@ -540,9 +702,9 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
             n_jobs: Number of jobs for parallelism.
 
         Returns:
-            left (LinearOperator): The left singular vector matrix.
-            singular_values (DiagonalLinearOperator): The singular values.
-            right (LinearOperator): The right singular vector matrix.
+            left (DenseMatrixLinearOperator): The left singular vector matrix.
+            singular_values (DiagonalSparseMatrixLinearOperator): The singular values.
+            right (DenseMatrixLinearOperator): The right singular vector matrix.
 
         Notes:
             The right factor is in transposed form. This means the original
@@ -572,7 +734,9 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         )
 
         euclidean = EuclideanSpace(qr_factor.shape[1])
-        diagonal = DiagonalLinearOperator(euclidean, euclidean, singular_values)
+        diagonal = DiagonalSparseMatrixLinearOperator.from_diagonal_values(
+            euclidean, euclidean, singular_values
+        )
 
         if galerkin:
             right = LinearOperator.from_matrix(
@@ -603,7 +767,7 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         block_size: int = 10,
         parallel: bool = False,
         n_jobs: int = -1,
-    ) -> Tuple[LinearOperator, DiagonalLinearOperator]:
+    ) -> Tuple[DenseMatrixLinearOperator, DiagonalSparseMatrixLinearOperator]:
         """
         Computes an approximate eigen-decomposition using a randomized algorithm.
 
@@ -625,8 +789,8 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
             n_jobs: Number of jobs for parallelism.
 
         Returns:
-            expansion (LinearOperator): Mapping from coefficients in eigen-basis to vectors.
-            eigenvaluevalues (DiagonalLinearOperator): The eigenvalues values.
+            expansion (DenseMatrixLinearOperator): Mapping from coefficients in eigen-basis to vectors.
+            eigenvaluevalues (DiagonalSparseMatrixLinearOperator): The eigenvalues values.
 
         """
         from .hilbert_space import EuclideanSpace
@@ -650,7 +814,9 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
 
         eigenvectors, eigenvalues = rm_eig(matrix, qr_factor)
         euclidean = EuclideanSpace(qr_factor.shape[1])
-        diagonal = DiagonalLinearOperator(euclidean, euclidean, eigenvalues)
+        diagonal = DiagonalSparseMatrixLinearOperator.from_diagonal_values(
+            euclidean, euclidean, eigenvalues
+        )
 
         expansion = LinearOperator.from_matrix(
             euclidean, self.domain, eigenvectors, galerkin=True
@@ -670,7 +836,7 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         block_size: int = 10,
         parallel: bool = False,
         n_jobs: int = -1,
-    ) -> LinearOperator:
+    ) -> DenseMatrixLinearOperator:
         """
         Computes an approximate Cholesky decomposition for a positive-definite
         self-adjoint operator using a randomized algorithm.
@@ -693,7 +859,7 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
             n_jobs: Number of jobs for parallelism.
 
         Returns:
-            factor (LinearOperator): A linear operator from a Euclidean space
+            factor (DenseMatrixLinearOperator): A linear operator from a Euclidean space
                 into the domain of the operator.
 
         Notes:
@@ -915,85 +1081,743 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         return self.matrix(dense=True).__str__()
 
 
-class DiagonalLinearOperator(LinearOperator):
-    """A LinearOperator that is diagonal in its component representation.
+class MatrixLinearOperator(LinearOperator):
+    """
+    A sub-class of LinearOperator for which the operator's action is
+    defined internally through its matrix representation.
+
+    This matrix can be either a dense numpy matrix or a
+    scipy LinearOperator.
+    """
+
+    def __init__(
+        self,
+        domain: HilbertSpace,
+        codomain: HilbertSpace,
+        matrix: Union[np.ndarray, ScipyLinOp],
+        /,
+        *,
+        galerkin=False,
+    ):
+        """
+        Args:
+            domain: The domain of the operator.
+            codomain: The codomain of the operator.
+            matrix: matrix representation of the linear operator in either standard
+                 or Galerkin form.
+            galerkin: If True, galerkin representation used. Default is false.
+        """
+        assert matrix.shape == (codomain.dim, domain.dim)
+
+        self._matrix = matrix
+        self._is_dense = isinstance(matrix, np.ndarray)
+        self._galerkin = galerkin
+
+        if galerkin:
+
+            def mapping(x: Any) -> Any:
+                cx = domain.to_components(x)
+                cyp = matrix @ cx
+                yp = codomain.dual.from_components(cyp)
+                return codomain.from_dual(yp)
+
+            def adjoint_mapping(y: Any) -> Any:
+                cy = codomain.to_components(y)
+                cxp = matrix.T @ cy
+                xp = domain.dual.from_components(cxp)
+                return domain.from_dual(xp)
+
+            super().__init__(domain, codomain, mapping, adjoint_mapping=adjoint_mapping)
 
-    This provides an efficient implementation for diagonal linear operators.
-    Its key feature is support for **functional calculus**, allowing for the
-    direct computation of operator functions like inverse (`.inverse`) or
+        else:
+
+            def mapping(x: Any) -> Any:
+                cx = domain.to_components(x)
+                cy = matrix @ cx
+                return codomain.from_components(cy)
+
+            def dual_mapping(yp: Any) -> Any:
+                cyp = codomain.dual.to_components(yp)
+                cxp = matrix.T @ cyp
+                return domain.dual.from_components(cxp)
+
+            super().__init__(domain, codomain, mapping, dual_mapping=dual_mapping)
+
+    @property
+    def is_dense(self) -> bool:
+        """
+        Returns True if the matrix representation is stored internally in dense form.
+        """
+        return self._is_dense
+
+    @property
+    def is_galerkin(self) -> bool:
+        """
+        Returns True if the matrix representation is stored in Galerkin form.
+        """
+        return self._galerkin
+
+    def _compute_dense_matrix(
+        self, galerkin: bool, parallel: bool, n_jobs: int
+    ) -> np.ndarray:
+        """
+        Overloaded method to efficiently compute the dense matrix.
+        """
+
+        if galerkin == self.is_galerkin and self.is_dense:
+            return self._matrix
+        else:
+            return super()._compute_dense_matrix(galerkin, parallel, n_jobs)
 
-    square root (`.sqrt`) by applying the function to the diagonal entries.
+    def extract_diagonal(
+        self,
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> np.ndarray:
+        """
+        Overload for efficiency.
+        """
+
+        if galerkin == self.is_galerkin and self.is_dense:
+            return self._matrix.diagonal()
+        else:
+            return super().extract_diagonal(
+                galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
+            )
+
+    def extract_diagonals(
+        self,
+        offsets: List[int],
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> Tuple[np.ndarray, List[int]]:
+        """
+        Overrides the base method for efficiency by extracting diagonals directly
+        from the stored dense matrix when possible.
+        """
+
+        if self.is_dense and galerkin == self.is_galerkin:
+            dim = self.domain.dim
+
+            diagonals_array = np.zeros((len(offsets), dim))
+
+            for i, k in enumerate(offsets):
+                diag_k = np.diag(self._matrix, k=k)
+
+                if k >= 0:
+                    diagonals_array[i, k : k + len(diag_k)] = diag_k
+                else:
+                    diagonals_array[i, : len(diag_k)] = diag_k
+
+            return diagonals_array, offsets
+
+        else:
+            return super().extract_diagonals(
+                offsets, galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
+            )
+
+
+class DenseMatrixLinearOperator(MatrixLinearOperator):
+    """
+    A specialisation of the MatrixLinearOperator class to instances where
+    the matrix representation is always provided as a numpy array.
+
+    This is a class provides some additional methods for component-wise access.
     """
 
     def __init__(
         self,
         domain: HilbertSpace,
         codomain: HilbertSpace,
-        diagonal_values: np.ndarray,
+        matrix: np.ndarray,
+        /,
+        *,
+        galerkin=False,
+    ):
+        """
+        domain: The domain of the operator.
+            codomain: The codomain of the operator.
+            matrix: matrix representation of the linear operator in either standard
+                 or Galerkin form.
+            galerkin: If True, galerkin representation used. Default is false.
+        """
+
+        if not isinstance(matrix, np.ndarray):
+            raise ValueError("Matrix must be input in dense form.")
+
+        super().__init__(domain, codomain, matrix, galerkin=galerkin)
+
+    @staticmethod
+    def from_linear_operator(
+        operator: LinearOperator,
         /,
         *,
         galerkin: bool = False,
-    ) -> None:
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> DenseMatrixLinearOperator:
         """
-        Initializes the DiagonalLinearOperator.
+        Converts a LinearOperator into a DenseMatrixLinearOperator by forming its dense matrix representation.
 
         Args:
-            domain (HilbertSpace): The domain of the operator.
-            codomain (HilbertSpace): The codomain of the operator.
-            diagonal_values (np.ndarray): The diagonal entries of the
-                operator's matrix representation.
-            galerkin (bool): If True, use the Galerkin representation.
+            operator: The operator to be converted.
+            galerkin: If True, the Galerkin representation is used. Default is False.
+            parallel: If True, dense matrix calculation is done in parallel. Default is False.
+            n_jobs: Number of jobs used for parallel calculations. Default is False.
         """
 
-        assert domain.dim == codomain.dim
-        assert domain.dim == len(diagonal_values)
-        self._diagonal_values: np.ndarray = diagonal_values
-        matrix = diags([diagonal_values], [0])
-        operator = LinearOperator.from_matrix(
-            domain, codomain, matrix, galerkin=galerkin
+        if isinstance(operator, DenseMatrixLinearOperator):
+            return operator
+
+        domain = operator.domain
+        codomain = operator.codomain
+
+        matrix = operator.matrix(
+            dense=True, galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
         )
-        super().__init__(
+
+        return DenseMatrixLinearOperator(domain, codomain, matrix, galerkin=galerkin)
+
+    def __getitem__(self, key: tuple[int, int] | int | slice) -> float | np.ndarray:
+        """
+        Provides direct, component-wise access to the underlying matrix.
+
+        This allows for intuitive slicing and indexing, like `op[i, j]` or `op[0, :]`.
+        Note: The access is on the stored matrix, which may be in either
+        standard or Galerkin form depending on how the operator was initialized.
+        """
+        return self._matrix[key]
+
+
+class SparseMatrixLinearOperator(MatrixLinearOperator):
+    """
+    A specialization for operators represented by a modern SciPy sparse array.
+
+    This class requires a `scipy.sparse.sparray` object (e.g., csr_array)
+    and provides optimized methods that delegate to efficient SciPy routines.
+
+    Upon initialization, the internal array is converted to the CSR
+    (Compressed Sparse Row) format to ensure consistently fast matrix-vector
+    products and row-slicing operations.
+    """
+
+    def __init__(
+        self,
+        domain: HilbertSpace,
+        codomain: HilbertSpace,
+        matrix: sp.sparray,
+        /,
+        *,
+        galerkin: bool = False,
+    ):
+        """
+        Args:
+            domain: The domain of the operator.
+            codomain: The codomain of the operator.
+            matrix: The sparse array representation of the linear operator.
+                    Must be a modern sparray object (e.g., csr_array).
+            galerkin: If True, the matrix is in Galerkin form. Defaults to False.
+        """
+        # Strict check for the modern sparse array type
+        if not isinstance(matrix, sp.sparray):
+            raise TypeError(
+                "Matrix must be a modern SciPy sparray object (e.g., csr_array)."
+            )
+
+        super().__init__(domain, codomain, matrix, galerkin=galerkin)
+        self._matrix = self._matrix.asformat("csr")
+
+    def __getitem__(self, key):
+        """Provides direct component access using SciPy's sparse indexing."""
+        return self._matrix[key]
+
+    def _compute_dense_matrix(
+        self, galerkin: bool, parallel: bool, n_jobs: int
+    ) -> np.ndarray:
+        """
+        Overrides the base method to efficiently compute the dense matrix.
+        """
+        # ⚡️ Fast path: Use the highly optimized .toarray() method.
+        if galerkin == self.is_galerkin:
+            return self._matrix.toarray()
+
+        # Fallback path for when a basis conversion is needed.
+        else:
+            return super()._compute_dense_matrix(galerkin, parallel, n_jobs)
+
+    def extract_diagonal(
+        self,
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> np.ndarray:
+        """
+        Overrides the base method to efficiently extract the main diagonal.
+        """
+        if galerkin == self.is_galerkin:
+            return self._matrix.diagonal(k=0)
+        else:
+            return super().extract_diagonal(
+                galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
+            )
+
+    def extract_diagonals(
+        self,
+        offsets: List[int],
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> Tuple[np.ndarray, List[int]]:
+        """
+        Overrides the base method for efficiency by extracting diagonals
+        directly from the stored sparse array.
+        """
+        if galerkin != self.is_galerkin:
+            return super().extract_diagonals(
+                offsets, galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
+            )
+
+        dim = self.domain.dim
+        diagonals_array = np.zeros((len(offsets), dim))
+
+        for i, k in enumerate(offsets):
+            # Use the sparse array's fast .diagonal() method
+            diag_k = self._matrix.diagonal(k=k)
+
+            # Place the raw diagonal into the padded output array
+            if k >= 0:
+                diagonals_array[i, k : k + len(diag_k)] = diag_k
+            else:
+                diagonals_array[i, : len(diag_k)] = diag_k
+
+        return diagonals_array, offsets
+
+
+class DiagonalSparseMatrixLinearOperator(SparseMatrixLinearOperator):
+    """
+    A highly specialized operator for matrices defined purely by a set of
+    non-zero diagonals.
+
+    This class internally stores the operator using a `scipy.sparse.dia_array`
+    for maximum efficiency in storage and matrix-vector products. It provides
+    extremely fast methods for extracting diagonals, as this is its native
+    storage format.
+
+    A key feature of this class is its support for **functional calculus**. It
+    dynamically proxies element-wise mathematical functions (e.g., `.sqrt()`,
+    `.log()`, `abs()`, `**`) to the underlying sparse array. For reasons of
+    mathematical correctness, these operations are restricted to operators that
+    are **strictly diagonal** (i.e., have only a non-zero main diagonal) and
+    will raise a `NotImplementedError` otherwise.
+
+    Aggregation methods that do not return a new operator (e.g., `.sum()`)
+    are not restricted and can be used on any multi-diagonal operator.
+
+    Class Methods
+    -------------
+    from_diagonal_values:
+        Constructs a strictly diagonal operator from a 1D array of values.
+    from_operator:
+        Creates a diagonal approximation of another LinearOperator.
+
+    Properties
+    ----------
+    offsets:
+        The array of stored diagonal offsets.
+    is_strictly_diagonal:
+        True if the operator only has a non-zero main diagonal.
+    inverse:
+        The inverse of a strictly diagonal operator.
+    sqrt:
+        The square root of a strictly diagonal operator.
+    """
+
+    def __init__(
+        self,
+        domain: HilbertSpace,
+        codomain: HilbertSpace,
+        diagonals: Tuple[np.ndarray, List[int]],
+        /,
+        *,
+        galerkin: bool = False,
+    ):
+        """
+        Args:
+            domain: The domain of the operator.
+            codomain: The codomain of the operator.
+            diagonals: A tuple `(data, offsets)` where `data` is a 2D array
+                       of diagonal values and `offsets` is a list of their
+                       positions. This is the native format for a dia_array.
+            galerkin: If True, the matrix is in Galerkin form. Defaults to False.
+        """
+        shape = (codomain.dim, domain.dim)
+        dia_array = sp.dia_array(diagonals, shape=shape)
+
+        MatrixLinearOperator.__init__(
+            self, domain, codomain, dia_array, galerkin=galerkin
+        )
+
+    @classmethod
+    def from_operator(
+        cls, operator: LinearOperator, offsets: List[int], /, *, galerkin: bool = True
+    ) -> DiagonalSparseMatrixLinearOperator:
+        """
+        Creates a diagonal approximation of another LinearOperator.
+
+        This factory method works by calling the source operator's
+        `.extract_diagonals()` method and using the result to construct a
+        new, highly efficient DiagonalSparseMatrixLinearOperator.
+
+        Args:
+            operator: The source operator to approximate.
+            offsets: The list of diagonal offsets to extract and keep.
+            galerkin: Specifies which matrix representation to use.
+
+        Returns:
+            A new DiagonalSparseMatrixLinearOperator.
+        """
+        diagonals_data, extracted_offsets = operator.extract_diagonals(
+            offsets, galerkin=galerkin
+        )
+        return cls(
             operator.domain,
             operator.codomain,
-            operator,
-            adjoint_mapping=operator.adjoint,
+            (diagonals_data, extracted_offsets),
+            galerkin=galerkin,
         )
 
-    @property
-    def diagonal_values(self) -> np.ndarray:
-        """The diagonal entries of the operator's matrix representation."""
-        return self._diagonal_values
-
-    def function(self, f: Callable[[float], float]) -> DiagonalLinearOperator:
-        """Applies a function to the operator via functional calculus.
+    @classmethod
+    def from_diagonal_values(
+        cls,
+        domain: HilbertSpace,
+        codomain: HilbertSpace,
+        diagonal_values: np.ndarray,
+        /,
+        *,
+        galerkin: bool = False,
+    ) -> "DiagonalSparseMatrixLinearOperator":
+        """
+        Constructs a purely diagonal operator from a 1D array of values.
 
-        This creates a new `DiagonalLinearOperator` where the function `f` has
-        been applied to each of the diagonal entries. For example,
-        `op.function(lambda x: 1/x)` computes the inverse.
+        This provides a convenient way to create an operator with non-zero
+        entries only on its main diagonal (offset k=0).
 
         Args:
-            f: A scalar function to apply to the diagonal values.
+            domain: The domain of the operator.
+            codomain: The codomain of the operator. Must have the same dimension.
+            diagonal_values: A 1D NumPy array of the values for the main diagonal.
+            galerkin: If True, the operator is in Galerkin form.
 
         Returns:
-            A new `DiagonalLinearOperator` with the transformed diagonal.
+            A new DiagonalSparseMatrixLinearOperator.
         """
-        diagonal_values = np.array([f(x) for x in self.diagonal_values])
-        return DiagonalLinearOperator(self.domain, self.codomain, diagonal_values)
+        if domain.dim != codomain.dim or domain.dim != len(diagonal_values):
+            raise ValueError(
+                "Domain, codomain, and diagonal_values must all have the same dimension."
+            )
+
+        # Reshape the 1D array of values into the 2D `data` array format
+        diagonals_data = diagonal_values.reshape(1, -1)
+        offsets = [0]
+
+        return cls(domain, codomain, (diagonals_data, offsets), galerkin=galerkin)
 
     @property
-    def inverse(self) -> DiagonalLinearOperator:
+    def offsets(self) -> np.ndarray:
+        """Returns the array of stored diagonal offsets."""
+        return self._matrix.offsets
+
+    @property
+    def is_strictly_diagonal(self) -> bool:
+        """
+        True if the operator only has a non-zero main diagonal (offset=0).
+        """
+        return len(self.offsets) == 1 and self.offsets[0] == 0
+
+    @property
+    def inverse(self) -> "DiagonalSparseMatrixLinearOperator":
         """
         The inverse of the operator, computed via functional calculus.
-        Requires all diagonal values to be non-zero.
+        Requires the operator to be strictly diagonal with no zero entries.
         """
-        assert all(val != 0 for val in self.diagonal_values)
-        return self.function(lambda x: 1 / x)
+        if not self.is_strictly_diagonal:
+            raise NotImplementedError(
+                "Inverse is only implemented for strictly diagonal operators."
+            )
+
+        if np.any(self._matrix.diagonal(k=0) == 0):
+            raise ValueError("Cannot invert an operator with zeros on the diagonal.")
+
+        return self**-1
 
     @property
-    def sqrt(self) -> DiagonalLinearOperator:
+    def sqrt(self) -> "DiagonalSparseMatrixLinearOperator":
         """
         The square root of the operator, computed via functional calculus.
-        Requires all diagonal values to be non-negative.
+        Requires the operator to be strictly diagonal with non-negative entries.
+        """
+
+        if np.any(self._matrix.data < 0):
+            raise ValueError(
+                "Cannot take the square root of an operator with negative entries."
+            )
+
+        return self.__getattr__("sqrt")()
+
+    def extract_diagonals(
+        self,
+        offsets: List[int],
+        /,
+        *,
+        galerkin: bool = True,
+        # parallel and n_jobs are ignored but kept for signature consistency
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> Tuple[np.ndarray, List[int]]:
+        """
+        Overrides the base method for extreme efficiency.
+
+        This operation is nearly free, as it involves selecting the requested
+        diagonals from the data already stored in the native format.
+        """
+        if galerkin != self.is_galerkin:
+            return super().extract_diagonals(offsets, galerkin=galerkin)
+
+        # Create a result array and fill it with the requested stored diagonals
+        result_diagonals = np.zeros((len(offsets), self.domain.dim))
+
+        # Create a mapping from stored offset to its data row for quick lookup
+        stored_diagonals = dict(zip(self.offsets, self._matrix.data))
+
+        for i, k in enumerate(offsets):
+            if k in stored_diagonals:
+                result_diagonals[i, :] = stored_diagonals[k]
+
+        return result_diagonals, offsets
+
+    def __getattr__(self, name: str):
+        """
+        Dynamically proxies method calls to the underlying dia_array.
+
+        For element-wise mathematical functions that return a new operator,
+        this method enforces that the operator must be strictly diagonal.
+        """
+        attr = getattr(self._matrix, name)
+
+        if callable(attr):
+
+            def wrapper(*args, **kwargs):
+                result = attr(*args, **kwargs)
+
+                if isinstance(result, sp.sparray):
+                    if not self.is_strictly_diagonal:
+                        raise NotImplementedError(
+                            f"Element-wise function '{name}' is only defined for "
+                            "strictly diagonal operators."
+                        )
+
+                    return DiagonalSparseMatrixLinearOperator(
+                        self.domain,
+                        self.codomain,
+                        (result.data, result.offsets),
+                        galerkin=self.is_galerkin,
+                    )
+                else:
+                    return result
+
+            return wrapper
+        else:
+            return attr
+
+    def __abs__(self):
+        """Explicitly handle the built-in abs() function."""
+        return self.__getattr__("__abs__")()
+
+    def __pow__(self, power):
+        """Explicitly handle the power operator (**)."""
+        return self.__getattr__("__pow__")(power)
+
+
+class NormalSumOperator(LinearOperator):
+    """
+    Represents a self-adjoint operator of the form N = A @ Q @ A.adjoint + B.
+
+    The operators Q and B are expected to be self-adjoint for the resulting
+    operator to be mathematically correct.
+
+    Q and B are optional. If Q is None, it defaults to the identity operator.
+    If B is None, it defaults to the zero operator.
+
+    This class uses operator algebra for a concise definition and provides an
+    optimized, parallelizable method for computing its dense Galerkin matrix.
+    """
+
+    def __init__(
+        self,
+        A: LinearOperator,
+        Q: Optional[LinearOperator] = None,
+        B: Optional[LinearOperator] = None,
+    ) -> None:
+
+        op_domain = A.codomain
+
+        if Q is None:
+            Q = A.domain.identity_operator()
+
+        if B is None:
+            B = op_domain.zero_operator()
+
+        if A.domain != Q.domain:
+            raise ValueError("The domain of A must match the domain of Q.")
+        if op_domain != B.domain:
+            raise ValueError("The domain of B must match the codomain of A.")
+
+        self._A = A
+        self._Q = Q
+        self._B = B
+
+        composite_op = self._A @ self._Q @ self._A.adjoint + self._B
+
+        super().__init__(
+            composite_op.domain,
+            composite_op.codomain,
+            composite_op,
+            adjoint_mapping=composite_op,
+        )
+
+    def _compute_dense_matrix(
+        self, galerkin: bool, parallel: bool, n_jobs: int
+    ) -> np.ndarray:
+        """
+        Overloaded method using the matrix-free approach for Q and a cleaner
+        implementation leveraging the base class's methods.
         """
-        assert all(val >= 0 for val in self._diagonal_values)
-        return self.function(np.sqrt)
+        if not galerkin:
+            return super()._compute_dense_matrix(galerkin, parallel, n_jobs)
+
+        domain_Y = self._A.codomain
+        dim = self.domain.dim
+        jobs = n_jobs if parallel else 1
+
+        a_star_mat = self._A.adjoint.matrix(
+            dense=True, galerkin=False, parallel=parallel, n_jobs=n_jobs
+        )
+
+        v_vectors = [domain_Y.from_components(a_star_mat[:, j]) for j in range(dim)]
+        w_vectors = Parallel(n_jobs=jobs)(delayed(self._Q)(v_j) for v_j in v_vectors)
+
+        def compute_row(i: int) -> np.ndarray:
+            """Computes the i-th row of the inner product matrix."""
+            v_i = v_vectors[i]
+            return np.array([domain_Y.inner_product(v_i, w_j) for w_j in w_vectors])
+
+        rows = Parallel(n_jobs=jobs)(delayed(compute_row)(i) for i in range(dim))
+        m_aqa_mat = np.vstack(rows)
+
+        b_mat = self._B.matrix(
+            dense=True, galerkin=True, parallel=parallel, n_jobs=n_jobs
+        )
+
+        return m_aqa_mat + b_mat
+
+    def extract_diagonal(
+        self,
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> np.ndarray:
+        """Overrides base method for efficiency."""
+        if not galerkin:
+            return super().extract_diagonal(
+                galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
+            )
+
+        diag_B = self._B.extract_diagonal(
+            galerkin=True, parallel=parallel, n_jobs=n_jobs
+        )
+
+        dim = self.domain.dim
+        jobs = n_jobs if parallel else 1
+
+        def compute_entry(i: int) -> float:
+            e_i = self.domain.basis_vector(i)
+            v_i = self._A.adjoint(e_i)
+            w_i = self._Q(v_i)
+            return self._A.domain.inner_product(v_i, w_i)
+
+        diag_AQA_T = Parallel(n_jobs=jobs)(
+            delayed(compute_entry)(i) for i in range(dim)
+        )
+
+        return np.array(diag_AQA_T) + diag_B
+
+    def extract_diagonals(
+        self,
+        offsets: List[int],
+        /,
+        *,
+        galerkin: bool = False,
+        parallel: bool = False,
+        n_jobs: int = -1,
+    ) -> Tuple[np.ndarray, List[int]]:
+        """Overrides base method for efficiency."""
+        if not galerkin:
+            return super().extract_diagonals(
+                offsets, galerkin=galerkin, parallel=parallel, n_jobs=n_jobs
+            )
+
+        diagonals_B, _ = self._B.extract_diagonals(
+            offsets, galerkin=True, parallel=parallel, n_jobs=n_jobs
+        )
+
+        dim = self.domain.dim
+        jobs = n_jobs if parallel else 1
+
+        # Pre-compute A*e_i for all i
+        v_vectors = Parallel(n_jobs=jobs)(
+            delayed(self._A.adjoint)(self.domain.basis_vector(i)) for i in range(dim)
+        )
+
+        def compute_column_entries(j: int) -> Dict[int, Dict[int, float]]:
+            col_results = defaultdict(dict)
+            v_j = v_vectors[j]
+            w_j = self._Q(v_j)
+
+            for k in offsets:
+                i = j - k
+                if 0 <= i < dim:
+                    v_i = v_vectors[i]
+                    val = self._A.domain.inner_product(v_i, w_j)
+                    col_results[k][i] = val
+            return col_results
+
+        column_data = Parallel(n_jobs=jobs)(
+            delayed(compute_column_entries)(j) for j in range(dim)
+        )
+
+        results: Dict[int, Dict[int, float]] = defaultdict(dict)
+        for col_dict in column_data:
+            for k, entries in col_dict.items():
+                results[k].update(entries)
+
+        diagonals_array = np.zeros((len(offsets), dim))
+        for idx, k in enumerate(offsets):
+            diag_entries = results[k]
+            for i, val in diag_entries.items():
+                j = i + k
+                if 0 <= j < dim:
+                    diagonals_array[idx, j] = val
+
+        return diagonals_array + diagonals_B, offsets