pygeoinf 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

pygeoinf/{operators.py → linear_operators.py}

@@ -1,22 +1,20 @@
 """
-Defines a class hierarchy for operators between Hilbert spaces.
+Provides classes for linear operators between Hilbert spaces.
 
-This module provides the tools for defining and manipulating mappings between
-`HilbertSpace` objects. It distinguishes between general non-linear operators
-and the more structured linear operators, which are the primary focus and
-support a rich algebra.
+This module is the primary tool for defining and manipulating linear mappings
+between `HilbertSpace` objects. It provides a powerful `LinearOperator` class
+that supports a rich algebra and includes numerous factory methods for
+convenient construction from matrices, forms, or tensor products.
 
 Key Classes
 -----------
-- `Operator`: A general, potentially non-linear operator defined by a simple
-  mapping function.
 - `LinearOperator`: The main workhorse for linear algebra. It represents a
-  linear map and provides rich functionality, including composition (`@`),
-  adjoints (`.adjoint`), duals (`.dual`), and matrix representations (`.matrix`).
-  It includes numerous factory methods for convenient construction.
+  linear map `L(x) = Ax` and provides rich functionality, including composition
+  (`@`), adjoints (`.adjoint`), duals (`.dual`), and matrix representations
+  (`.matrix`).
 - `DiagonalLinearOperator`: A specialized, efficient implementation for linear
-  operators that are diagonal in their component representation, which supports
-  functional calculus.
+  operators that are diagonal in their component representation, notable for
+  supporting functional calculus (e.g., `.inverse`, `.sqrt`).
 """
 
 from __future__ import annotations
@@ -26,10 +24,11 @@ import numpy as np
 from scipy.sparse.linalg import LinearOperator as ScipyLinOp
 from scipy.sparse import diags
 
+# from .operators import Operator
+from .nonlinear_operators import NonLinearOperator
 
 from .random_matrix import (
-    fixed_rank_random_range,
-    variable_rank_random_range,
+    random_range,
     random_svd as rm_svd,
     random_cholesky as rm_chol,
     random_eig as rm_eig,
@@ -43,94 +42,18 @@ if TYPE_CHECKING:
     from .linear_forms import LinearForm
 
 
-def _parallel_scipy_op_col(scipy_op: ScipyLinOp, j: int, domain_dim: int) -> np.ndarray:
-    """
-    A top-level helper that applies a scipy.LinearOperator to a basis vector.
-
-    This function is simple and serializable ("picklable").
-
-    Args:
-        scipy_op: The SciPy LinearOperator wrapper for the matrix action.
-        j: The index of the basis vector (column) to compute.
-        domain_dim: The dimension of the domain space.
-
-    Returns:
-        The j-th column of the dense matrix as a NumPy array.
-    """
-    # Create the j-th component basis vector
-    cx = np.zeros(domain_dim)
-    cx[j] = 1.0
-
-    # Apply the SciPy wrapper, which handles all necessary conversions
-    return scipy_op @ cx
+class LinearOperator(NonLinearOperator):
+    """A linear operator between two Hilbert spaces.
 
+    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.
 
-class Operator:
-    """
-    A general, potentially non-linear operator between two Hilbert spaces.
-    """
-
-    def __init__(
-        self,
-        domain: HilbertSpace,
-        codomain: HilbertSpace,
-        mapping: Callable[[Any], Any],
-    ) -> None:
-        """
-        Initializes the Operator.
-
-        Args:
-            domain (HilbertSpace): Domain of the operator.
-            codomain (HilbertSpace): Codomain of the operator.
-            mapping (callable): The function defining the mapping from the
-                domain to the codomain.
-        """
-        self._domain: HilbertSpace = domain
-        self._codomain: HilbertSpace = codomain
-        self.__mapping: Callable[[Any], Any] = mapping
-
-    @property
-    def domain(self) -> HilbertSpace:
-        """The domain of the operator."""
-        return self._domain
-
-    @property
-    def codomain(self) -> HilbertSpace:
-        """The codomain of the operator."""
-        return self._codomain
-
-    @property
-    def is_automorphism(self) -> bool:
-        """True if the operator maps a space into itself."""
-        return self.domain == self.codomain
-
-    @property
-    def is_square(self) -> bool:
-        """True if the operator's domain and codomain have the same dimension."""
-        return self.domain.dim == self.codomain.dim
-
-    @property
-    def linear(self) -> bool:
-        """False for a general operator. Overridden by LinearOperator."""
-        return False
-
-    def __call__(self, x: Any) -> Any:
-        """Applies the operator's mapping to a vector."""
-        return self.__mapping(x)
-
-
-class LinearOperator(Operator):
-    """
-    A linear operator between two Hilbert spaces.
-
-    This class is the primary workhorse for linear algebraic operations. An
-    operator can be defined "on the fly" from a callable mapping. The class
-    automatically derives the associated `adjoint` and `dual` operators,
-    which are fundamental for solving linear systems and for optimization.
-
-    It supports a rich algebra, including composition (`@`), addition (`+`),
-    and scalar multiplication (`*`). Operators can also be represented as
-    dense or matrix-free (`scipy`) matrices for use with numerical solvers.
+    Key features include operator algebra (`@`, `+`, `*`), automatic
+    derivation of adjoint (`.adjoint`) and dual (`.dual`) operators, and
+    multiple matrix representations (`.matrix()`) for use with numerical
+    solvers.
     """
 
     def __init__(
@@ -159,7 +82,10 @@ class LinearOperator(Operator):
             dual_base (LinearOperator, optional): Internal use for duals.
             adjoint_base (LinearOperator, optional): Internal use for adjoints.
         """
-        super().__init__(domain, codomain, mapping)
+        super().__init__(
+            domain, codomain, self._mapping_impl, derivative=self._derivative_impl
+        )
+        self._mapping = mapping
         self._dual_base: Optional[LinearOperator] = dual_base
         self._adjoint_base: Optional[LinearOperator] = adjoint_base
         self._thread_safe: bool = thread_safe
@@ -329,47 +255,25 @@ class LinearOperator(Operator):
         """
         Creates a LinearOperator from its matrix representation.
 
-        This factory method allows you to define a `LinearOperator` using a
-        concrete matrix (like a `numpy.ndarray`) that acts on the component
-        vectors of the abstract Hilbert space vectors. The `galerkin` flag
-        determines how this matrix action is interpreted.
+        This factory defines a `LinearOperator` using a concrete matrix that
+        acts on the component vectors of the abstract Hilbert space vectors.
 
         Args:
-            domain (HilbertSpace): The operator's domain.
-            codomain (HilbertSpace): The operator's codomain.
-            matrix (MatrixLike): The matrix representation, which can be a dense
-                NumPy array or a SciPy LinearOperator. Its shape must be
-                (codomain.dim, domain.dim).
-            galerkin (bool): Specifies the interpretation of the matrix.
-
-                - **`galerkin=False` (Default): Standard Component Mapping**
-                  This is the most direct interpretation. The matrix `M` maps the
-                  component vector `c_x` of an input vector `x` directly to the
-                  component vector `c_y` of the output vector `y`.
-
-                - **`galerkin=True`: Galerkin (or "Weak Form") Representation**
-                  This interpretation is standard in the finite element method (FEM)
-                  and other variational techniques. The matrix `M` maps the component
-                  vector `c_x` of an input `x` to the component vector `c_yp` of the
-                  *dual* of the output vector `y`.
-
-                  - **Matrix Entries**: The matrix elements are defined by inner
-                    products with basis vectors: `M_ij = inner_product(A(b_j), b_i)`,
-                    where `b_j` are domain basis vectors and `b_i` are codomain
-                    basis vectors.
-                  - **Use Case**: This is critically important for preserving the
-                    mathematical properties of an operator. For example, if an operator
-                    `A` is self-adjoint, its Galerkin matrix `M` will be **symmetric**
-                    (`M.T == M`). This allows the use of highly efficient numerical
-                    methods like the Conjugate Gradient solver or Cholesky
-                    factorization, which rely on symmetry. The standard component
-                    matrix of a self-adjoint operator is generally not symmetric
-                    unless the basis is orthonormal.
+            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.
 
         Returns:
-            LinearOperator: A new `LinearOperator` instance whose action is
-                defined by the provided matrix and interpretation.
+            A new `LinearOperator` defined by the matrix action.
         """
+
         assert matrix.shape == (codomain.dim, domain.dim)
 
         if galerkin:
@@ -521,53 +425,28 @@ class LinearOperator(Operator):
         parallel: bool = False,
         n_jobs: int = -1,
     ) -> Union[ScipyLinOp, np.ndarray]:
-        """
-        Returns a matrix representation of the operator.
+        """Returns a matrix representation of the operator.
 
-        This method provides a concrete matrix that represents the abstract
-        linear operator's action on the underlying component vectors.
+        This provides a concrete matrix that represents the operator's action
+        on the underlying component vectors.
 
         Args:
-            dense (bool): Determines the format of the returned matrix.
-                - If `True`, this method computes and returns a dense `numpy.ndarray`.
-                  Be aware that this can be very memory-intensive for
-                  high-dimensional spaces.
-                - If `False` (default), it returns a matrix-free
-                  `scipy.sparse.linalg.LinearOperator`. This object encapsulates
-                  the operator's action (`matvec`) and its transpose action
-                  (`rmatvec`) without ever explicitly forming the full matrix in memory,
-                  making it ideal for large-scale problems.
-
-            galerkin (bool): Specifies the interpretation of the matrix representation. This
-                flag is crucial for correctly using the matrix with numerical solvers.
-
-                - **`galerkin=False` (Default): Standard Component Mapping**
-                  The returned matrix `M` performs a standard component-to-component
-                  mapping.
-                  - **`matvec` action**: Takes the component vector `c_x` of an input `x`
-                    and returns the component vector `c_y` of the output `y`.
-                  - **`rmatvec` action**: Corresponds to the matrix of the **dual operator**, `A'`.
-
-                - **`galerkin=True`: Galerkin (or "Weak Form") Representation**
-                  The returned matrix `M` represents the operator in a weak form, mapping
-                  components of a vector to components of a dual vector.
-                  - **`matvec` action**: Takes the component vector `c_x` of an input `x`
-                    and returns the component vector `c_yp` of the *dual* of the output `y`.
-                  - **`rmatvec` action**: Corresponds to the matrix of the **adjoint operator**, `A*`.
-                  - **Key Property**: This representation is designed to preserve fundamental
-                    mathematical properties. For instance, if the `LinearOperator` is
-                    self-adjoint, its Galerkin matrix will be **symmetric**, which is a
-                    prerequisite for algorithms like the Conjugate Gradient method.
-
-            parallel (bool): If True, use parallel computing. Defaults to False.
-                This is only relevant for dense matrices.
-            n_jobs (int): Number of parallel jobs. Defaults to -1.
-                This is only relevant for dense matrices.
+            dense: If `True`, returns a dense `numpy.ndarray`. If `False`
+                (default), returns a memory-efficient, matrix-free
+                `scipy.sparse.linalg.LinearOperator`.
+            galerkin: If `True`, the returned matrix is the Galerkin
+                representation, whose `rmatvec` corresponds to the
+                **adjoint** operator. If `False` (default), the `rmatvec`
+                corresponds to the **dual** operator. The Galerkin form is
+                essential for algorithms that rely on symmetry/self-adjointness.
+            parallel: If `True` and `dense=True`, computes the matrix columns
+                in parallel.
+            n_jobs: Number of parallel jobs to use. `-1` uses all available cores.
 
         Returns:
-            Union[ScipyLinOp, np.ndarray]: The matrix representation of the
-                operator, either as a dense array or a matrix-free object.
+            The matrix representation, either dense or matrix-free.
         """
+
         if dense:
             return self._compute_dense_matrix(galerkin, parallel, n_jobs)
         else:
@@ -625,31 +504,38 @@ class LinearOperator(Operator):
 
     def random_svd(
         self,
-        rank: int,
+        size_estimate: int,
         /,
         *,
-        power: int = 0,
         galerkin: bool = False,
-        rtol: float = 1e-3,
-        method: str = "fixed",
+        method: str = "variable",
+        max_rank: int = None,
+        power: int = 2,
+        rtol: float = 1e-4,
+        block_size: int = 10,
         parallel: bool = False,
         n_jobs: int = -1,
-    ) -> Tuple[LinearOperator, "DiagonalLinearOperator", LinearOperator]:
+    ) -> Tuple[LinearOperator, DiagonalLinearOperator, LinearOperator]:
         """
         Computes an approximate SVD using a randomized algorithm.
 
         Args:
-            rank (int): The desired rank of the SVD.
-            power (int): The power of the random matrix.
+            size_estimate: For 'fixed' method, the exact target rank. For 'variable'
+                       method, this is the initial rank to sample.
             galerkin (bool): If True, use the Galerkin representation.
-            rtol (float): The relative tolerance for the SVD.
-            method (str): The method to use for the SVD.
-                - "fixed": Use a fixed rank SVD.
-                - "variable": Use a variable rank SVD.
-            parallel (bool): If True, use parallel computing. Defaults to False.
-                Only used with fixed rank method.
-            n_jobs (int): Number of parallel jobs. Defaults to -1.
-                Only used with fixed rank method.
+            method ({'variable', 'fixed'}): The algorithm to use.
+            - 'variable': (Default) Progressively samples to find the rank needed
+                          to meet tolerance `rtol`, stopping at `max_rank`.
+            - 'fixed': Returns a basis with exactly `size_estimate` columns.
+            max_rank: For 'variable' method, a hard limit on the rank. Ignored if
+                    method='fixed'. Defaults to min(m, n).
+            power: Number of power iterations to improve accuracy.
+            rtol: Relative tolerance for the 'variable' method. Ignored if
+                method='fixed'.
+            block_size: Number of new vectors to sample per iteration in 'variable'
+                        method. Ignored if method='fixed'.
+            parallel: Whether to use parallel matrix multiplication.
+            n_jobs: Number of jobs for parallelism.
 
         Returns:
             left (LinearOperator): The left singular vector matrix.
@@ -659,25 +545,25 @@ class LinearOperator(Operator):
         Notes:
             The right factor is in transposed form. This means the original
             operator can be approximated as:
-                A = left @ singular_values @ right
-
+            A = left @ singular_values @ right
         """
         from .hilbert_space import EuclideanSpace
 
         matrix = self.matrix(galerkin=galerkin)
         m, n = matrix.shape
         k = min(m, n)
-        rank = rank if rank <= k else k
 
-        qr_factor: np.ndarray
-        if method == "fixed":
-            qr_factor = fixed_rank_random_range(
-                matrix, rank, power=power, parallel=parallel, n_jobs=n_jobs
-            )
-        elif method == "variable":
-            qr_factor = variable_rank_random_range(matrix, rank, power=power, rtol=rtol)
-        else:
-            raise ValueError("Invalid method selected")
+        qr_factor = random_range(
+            matrix,
+            size_estimate if size_estimate < k else k,
+            method=method,
+            max_rank=max_rank,
+            power=power,
+            rtol=rtol,
+            block_size=block_size,
+            parallel=parallel,
+            n_jobs=n_jobs,
+        )
 
         left_factor_mat, singular_values, right_factor_transposed = rm_svd(
             matrix, qr_factor
@@ -705,39 +591,40 @@ class LinearOperator(Operator):
 
     def random_eig(
         self,
-        rank: int,
+        size_estimate: int,
         /,
         *,
-        power: int = 0,
-        rtol: float = 1e-3,
-        method: str = "fixed",
+        method: str = "variable",
+        max_rank: int = None,
+        power: int = 2,
+        rtol: float = 1e-4,
+        block_size: int = 10,
         parallel: bool = False,
         n_jobs: int = -1,
-    ) -> Tuple[LinearOperator, "DiagonalLinearOperator"]:
+    ) -> Tuple[LinearOperator, DiagonalLinearOperator]:
         """
-        Computes an approximate eigendecomposition for a self-adjoint
-        operator using a randomized algorithm.
+        Computes an approximate eigen-decomposition using a randomized algorithm.
 
         Args:
-            rank (int): The desired rank of the eigendecomposition.
-            power (int): The power of the random matrix.
-            rtol (float): The relative tolerance for the eigendecomposition.
-            method (str): The method to use for the eigendecomposition.
-                - "fixed": Use a fixed rank eigendecomposition.
-                - "variable": Use a variable rank eigendecomposition.
-            parallel (bool): If True, use parallel computing. Defaults to False.
-                Only used with fixed rank method.
-            n_jobs (int): Number of parallel jobs. Defaults to -1.
-                Only used with fixed rank method.
+            size_estimate: For 'fixed' method, the exact target rank. For 'variable'
+                       method, this is the initial rank to sample.
+            method ({'variable', 'fixed'}): The algorithm to use.
+            - 'variable': (Default) Progressively samples to find the rank needed
+                          to meet tolerance `rtol`, stopping at `max_rank`.
+            - 'fixed': Returns a basis with exactly `size_estimate` columns.
+            max_rank: For 'variable' method, a hard limit on the rank. Ignored if
+                    method='fixed'. Defaults to min(m, n).
+            power: Number of power iterations to improve accuracy.
+            rtol: Relative tolerance for the 'variable' method. Ignored if
+                method='fixed'.
+            block_size: Number of new vectors to sample per iteration in 'variable'
+                        method. Ignored if method='fixed'.
+            parallel: Whether to use parallel matrix multiplication.
+            n_jobs: Number of jobs for parallelism.
 
         Returns:
-            expansion (LinearOperator): A linear operator that maps coefficients
-                in the eigen-basis to the resulting vector.
-            eigenvalues (DiagonalLinearOperator): The eigenvalues.
-
-        Notes:
-            The original operator can be approximated as:
-                A = expansion @ eigenvalues @ expansion.adjoint
+            expansion (LinearOperator): Mapping from coefficients in eigen-basis to vectors.
+            eigenvaluevalues (DiagonalLinearOperator): The eigenvalues values.
 
         """
         from .hilbert_space import EuclideanSpace
@@ -746,17 +633,18 @@ class LinearOperator(Operator):
         matrix = self.matrix(galerkin=True)
         m, n = matrix.shape
         k = min(m, n)
-        rank = rank if rank <= k else k
 
-        qr_factor: np.ndarray
-        if method == "fixed":
-            qr_factor = fixed_rank_random_range(
-                matrix, rank, power=power, parallel=parallel, n_jobs=n_jobs
-            )
-        elif method == "variable":
-            qr_factor = variable_rank_random_range(matrix, rank, power=power, rtol=rtol)
-        else:
-            raise ValueError("Invalid method selected")
+        qr_factor = random_range(
+            matrix,
+            size_estimate if size_estimate < k else k,
+            method=method,
+            max_rank=max_rank,
+            power=power,
+            rtol=rtol,
+            block_size=block_size,
+            parallel=parallel,
+            n_jobs=n_jobs,
+        )
 
         eigenvectors, eigenvalues = rm_eig(matrix, qr_factor)
         euclidean = EuclideanSpace(qr_factor.shape[1])
@@ -770,12 +658,14 @@ class LinearOperator(Operator):
 
     def random_cholesky(
         self,
-        rank: int,
+        size_estimate: int,
         /,
         *,
-        power: int = 0,
-        rtol: float = 1e-3,
-        method: str = "fixed",
+        method: str = "variable",
+        max_rank: int = None,
+        power: int = 2,
+        rtol: float = 1e-4,
+        block_size: int = 10,
         parallel: bool = False,
         n_jobs: int = -1,
     ) -> LinearOperator:
@@ -784,16 +674,21 @@ class LinearOperator(Operator):
         self-adjoint operator using a randomized algorithm.
 
         Args:
-            rank (int): The desired rank of the Cholesky decomposition.
-            power (int): The power of the random matrix.
-            rtol (float): The relative tolerance for the Cholesky decomposition.
-            method (str): The method to use for the Cholesky decomposition.
-                - "fixed": Use a fixed rank Cholesky decomposition.
-                - "variable": Use a variable rank Cholesky decomposition.
-        parallel (bool): If True, use parallel computing. Defaults to False.
-            Only used with fixed rank method.
-        n_jobs (int): Number of parallel jobs. Defaults to -1.
-                Only used with fixed rank method.
+            size_estimate: For 'fixed' method, the exact target rank. For 'variable'
+                       method, this is the initial rank to sample.
+            method ({'variable', 'fixed'}): The algorithm to use.
+            - 'variable': (Default) Progressively samples to find the rank needed
+                          to meet tolerance `rtol`, stopping at `max_rank`.
+            - 'fixed': Returns a basis with exactly `size_estimate` columns.
+            max_rank: For 'variable' method, a hard limit on the rank. Ignored if
+                    method='fixed'. Defaults to min(m, n).
+            power: Number of power iterations to improve accuracy.
+            rtol: Relative tolerance for the 'variable' method. Ignored if
+                method='fixed'.
+            block_size: Number of new vectors to sample per iteration in 'variable'
+                        method. Ignored if method='fixed'.
+            parallel: Whether to use parallel matrix multiplication.
+            n_jobs: Number of jobs for parallelism.
 
         Returns:
             factor (LinearOperator): A linear operator from a Euclidean space
@@ -810,17 +705,18 @@ class LinearOperator(Operator):
         matrix = self.matrix(galerkin=True)
         m, n = matrix.shape
         k = min(m, n)
-        rank = rank if rank <= k else k
 
-        qr_factor: np.ndarray
-        if method == "fixed":
-            qr_factor = fixed_rank_random_range(
-                matrix, rank, power=power, parallel=parallel, n_jobs=n_jobs
-            )
-        elif method == "variable":
-            qr_factor = variable_rank_random_range(matrix, rank, power=power, rtol=rtol)
-        else:
-            raise ValueError("Invalid method selected")
+        qr_factor = random_range(
+            matrix,
+            size_estimate if size_estimate < k else k,
+            method=method,
+            max_rank=max_rank,
+            power=power,
+            rtol=rtol,
+            block_size=block_size,
+            parallel=parallel,
+            n_jobs=n_jobs,
+        )
 
         cholesky_factor = rm_chol(matrix, qr_factor)
 
@@ -831,6 +727,12 @@ class LinearOperator(Operator):
             galerkin=True,
         )
 
+    def _mapping_impl(self, x: Any) -> Any:
+        return self._mapping(x)
+
+    def _derivative_impl(self, _: Any) -> LinearOperator:
+        return self
+
     def _dual_mapping_default(self, yp: Any) -> LinearForm:
         from .linear_forms import LinearForm
 
@@ -899,55 +801,126 @@ class LinearOperator(Operator):
     def __truediv__(self, a: float) -> LinearOperator:
         return self * (1.0 / a)
 
-    def __add__(self, other: LinearOperator) -> LinearOperator:
-        domain = self.domain
-        codomain = self.codomain
+    def __add__(
+        self, other: NonLinearOperator | LinearOperator
+    ) -> NonLinearOperator | LinearOperator:
+        """Returns the sum of this operator and another.
 
-        def mapping(x: Any) -> Any:
-            return codomain.add(self(x), other(x))
+        If `other` is also a `LinearOperator`, this performs an optimized
+        addition that preserves linearity and correctly defines the new
+        operator's `adjoint`. Otherwise, it delegates to the general
+        implementation in the `NonLinearOperator` base class.
 
-        def adjoint_mapping(y: Any) -> Any:
-            return domain.add(self.adjoint(y), other.adjoint(y))
+        Args:
+            other: The operator to add to this one.
 
-        return LinearOperator(
-            domain, codomain, mapping, adjoint_mapping=adjoint_mapping
-        )
+        Returns:
+            A new `LinearOperator` if adding two linear operators, otherwise
+            a `NonLinearOperator`.
+        """
 
-    def __sub__(self, other: LinearOperator) -> LinearOperator:
-        domain = self.domain
-        codomain = self.codomain
+        if isinstance(other, LinearOperator):
+            domain = self.domain
+            codomain = self.codomain
 
-        def mapping(x: Any) -> Any:
-            return codomain.subtract(self(x), other(x))
+            def mapping(x: Any) -> Any:
+                return codomain.add(self(x), other(x))
 
-        def adjoint_mapping(y: Any) -> Any:
-            return domain.subtract(self.adjoint(y), other.adjoint(y))
+            def adjoint_mapping(y: Any) -> Any:
+                return domain.add(self.adjoint(y), other.adjoint(y))
 
-        return LinearOperator(
-            domain, codomain, mapping, adjoint_mapping=adjoint_mapping
-        )
+            return LinearOperator(
+                domain, codomain, mapping, adjoint_mapping=adjoint_mapping
+            )
+        else:
+            return super().__add__(other)
 
-    def __matmul__(self, other: LinearOperator) -> LinearOperator:
-        domain = other.domain
-        codomain = self.codomain
+    def __sub__(
+        self, other: NonLinearOperator | LinearOperator
+    ) -> NonLinearOperator | LinearOperator:
+        """Returns the difference between this operator and another.
 
-        def mapping(x: Any) -> Any:
-            return self(other(x))
+        If `other` is also a `LinearOperator`, this performs an optimized
+        subtraction that preserves linearity and correctly defines the new
+        operator's `adjoint`. Otherwise, it delegates to the general
+        implementation in the `NonLinearOperator` base class.
 
-        def adjoint_mapping(y: Any) -> Any:
-            return other.adjoint(self.adjoint(y))
+        Args:
+            other: The operator to subtract from this one.
 
-        return LinearOperator(
-            domain, codomain, mapping, adjoint_mapping=adjoint_mapping
-        )
+        Returns:
+            A new `LinearOperator` if subtracting two linear operators,
+            otherwise a `NonLinearOperator`.
+        """
+
+        if isinstance(other, LinearOperator):
+
+            domain = self.domain
+            codomain = self.codomain
+
+            def mapping(x: Any) -> Any:
+                return codomain.subtract(self(x), other(x))
+
+            def adjoint_mapping(y: Any) -> Any:
+                return domain.subtract(self.adjoint(y), other.adjoint(y))
+
+            return LinearOperator(
+                domain, codomain, mapping, adjoint_mapping=adjoint_mapping
+            )
+        else:
+            return super().__sub__(other)
+
+    def __matmul__(
+        self, other: NonLinearOperator | LinearOperator
+    ) -> NonLinearOperator | LinearOperator:
+        """Composes this operator with another using the @ symbol.
+
+        The composition `(self @ other)` results in a new operator that
+        first applies `other` and then applies `self`, i.e.,
+        `(self @ other)(x) = self(other(x))`.
+
+        If `other` is also a `LinearOperator`, this creates a new `LinearOperator`
+        whose adjoint is correctly defined using the composition rule:
+        `(L1 @ L2)* = L2* @ L1*`. Otherwise, it delegates to the general
+        `NonLinearOperator` implementation.
+
+        Args:
+            other: The operator to compose with (the right-hand operator).
+
+        Returns:
+            A new `LinearOperator` if composing two linear operators,
+            otherwise a `NonLinearOperator`.
+        """
+
+        if isinstance(other, LinearOperator):
+            domain = other.domain
+            codomain = self.codomain
+
+            def mapping(x: Any) -> Any:
+                return self(other(x))
+
+            def adjoint_mapping(y: Any) -> Any:
+                return other.adjoint(self.adjoint(y))
+
+            return LinearOperator(
+                domain, codomain, mapping, adjoint_mapping=adjoint_mapping
+            )
+
+        else:
+            return super().__matmul__(other)
 
     def __str__(self) -> str:
         return self.matrix(dense=True).__str__()
 
 
 class DiagonalLinearOperator(LinearOperator):
-    """
-    A LinearOperator that is diagonal in its component representation.
+    """A LinearOperator that is diagonal in its component representation.
+
+    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
+
+    square root (`.sqrt`) by applying the function to the diagonal entries.
     """
 
     def __init__(
@@ -989,21 +962,24 @@ class DiagonalLinearOperator(LinearOperator):
         """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.
+    def function(self, f: Callable[[float], float]) -> DiagonalLinearOperator:
+        """Applies a function to the operator via functional calculus.
 
-        This creates a new DiagonalLinearOperator where each diagonal entry `d_i`
-        is replaced by `f(d_i)`.
+        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.
 
         Args:
-            f: A function that maps a float to a float.
+            f: A scalar function to apply to the diagonal values.
+
+        Returns:
+            A new `DiagonalLinearOperator` with the transformed diagonal.
         """
         diagonal_values = np.array([f(x) for x in self.diagonal_values])
         return DiagonalLinearOperator(self.domain, self.codomain, diagonal_values)
 
     @property
-    def inverse(self) -> "DiagonalLinearOperator":
+    def inverse(self) -> DiagonalLinearOperator:
         """
         The inverse of the operator, computed via functional calculus.
         Requires all diagonal values to be non-zero.
@@ -1012,7 +988,7 @@ class DiagonalLinearOperator(LinearOperator):
         return self.function(lambda x: 1 / x)
 
     @property
-    def sqrt(self) -> "DiagonalLinearOperator":
+    def sqrt(self) -> DiagonalLinearOperator:
         """
         The square root of the operator, computed via functional calculus.
         Requires all diagonal values to be non-negative.