pygeoinf 1.3.3__py3-none-any.whl → 1.3.5__py3-none-any.whl

pygeoinf/linear_optimisation.py

@@ -16,6 +16,8 @@ Key Classes
 - `LinearMinimumNormInversion`: Finds the model with the smallest norm that
   fits the data to a statistically acceptable degree using the discrepancy
   principle.
+- `ConstrainedLinearLeastSquaresInversion`: Solves a linear inverse problem
+  subject to an affine subspace constraint.
 """
 
 from __future__ import annotations
@@ -30,6 +32,7 @@ from .forward_problem import LinearForwardProblem
 from .linear_operators import LinearOperator
 from .linear_solvers import LinearSolver, IterativeLinearSolver
 from .hilbert_space import Vector
+from .subspaces import AffineSubspace
 
 
 class LinearLeastSquaresInversion(LinearInversion):
@@ -160,6 +163,113 @@ class LinearLeastSquaresInversion(LinearInversion):
             return inverse_normal_operator @ forward_operator.adjoint
 
 
+class ConstrainedLinearLeastSquaresInversion(LinearInversion):
+    """
+    Solves a linear inverse problem subject to an affine subspace constraint.
+
+    Problem:
+        Minimize J(u) = || A(u) - d ||_D^2 + alpha * || u ||_M^2
+        Subject to u in A (Affine Subspace)
+
+    Method:
+        The problem is reduced to an unconstrained minimization in the subspace.
+        We decompose the model as u = u_base + w, where u_base is the element
+        of the affine subspace closest to the origin (orthogonal to the tangent space),
+        and w is a perturbation in the tangent space.
+
+        The cost function separates (due to orthogonality) into:
+        J(w) = || A(w) - (d - A(u_base)) ||^2 + alpha * || w ||^2 + (alpha * ||u_base||^2)
+
+        This is solved using the standard LinearLeastSquaresInversion on a
+        reduced forward problem.
+    """
+
+    def __init__(
+        self, forward_problem: LinearForwardProblem, constraint: AffineSubspace
+    ) -> None:
+        """
+        Args:
+            forward_problem: The original unconstrained forward problem.
+            constraint: The affine subspace A where the solution must lie.
+        """
+        super().__init__(forward_problem)
+        self._constraint = constraint
+
+        # 1. Compute the Orthogonal Base Vector (u_base)
+        # u_base = (I - P) * translation
+        # This is the unique vector in the affine space that is orthogonal to the tangent space.
+        # It ensures ||u||^2 = ||u_base||^2 + ||w||^2, decoupling the regularization.
+        self._u_base = constraint.domain.subtract(
+            constraint.translation, constraint.projector(constraint.translation)
+        )
+
+        # 2. Construct Reduced Forward Problem
+        # Operator: A_tilde = A @ P
+        reduced_operator = forward_problem.forward_operator @ constraint.projector
+
+        # The error measure on the data remains valid for the reduced problem
+        # because the noise model is additive and independent of the model parameters.
+        self._reduced_forward_problem = LinearForwardProblem(
+            reduced_operator,
+            data_error_measure=(
+                forward_problem.data_error_measure
+                if forward_problem.data_error_measure_set
+                else None
+            ),
+        )
+
+        # 3. Initialize the internal unconstrained solver
+        self._unconstrained_inversion = LinearLeastSquaresInversion(
+            self._reduced_forward_problem
+        )
+
+    def least_squares_operator(
+        self,
+        damping: float,
+        solver: LinearSolver,
+        /,
+        **kwargs,
+    ) -> NonLinearOperator:
+        """
+        Returns an operator that maps data to the constrained least-squares solution.
+
+        Args:
+            damping: The Tikhonov damping parameter.
+            solver: The linear solver for the reduced normal equations.
+            **kwargs: Additional arguments passed to the solver (e.g., preconditioner).
+
+        Returns:
+            A NonLinearOperator mapping d -> u_constrained.
+        """
+
+        # Get the operator L_tilde such that w = L_tilde(d_tilde)
+        reduced_op = self._unconstrained_inversion.least_squares_operator(
+            damping, solver, **kwargs
+        )
+
+        # Precompute A(u_base) to shift the data efficiently
+        # This represents the data predicted by the "base" model.
+        data_offset = self.forward_problem.forward_operator(self._u_base)
+
+        domain = self.data_space
+        codomain = self.model_space
+
+        def mapping(d: Vector) -> Vector:
+            # 1. Shift Data: d_tilde = d - A(u_base)
+            d_tilde = domain.subtract(d, data_offset)
+
+            # 2. Solve for perturbation w in the tangent space
+            # w = (P A* A P + alpha I)^-1 P A* d_tilde
+            w = reduced_op(d_tilde)
+
+            # 3. Reconstruct full model: u = u_base + w
+            # Note: w is guaranteed to be in the tangent space (Range of P)
+            # because of the structure of the reduced normal equations.
+            return codomain.add(self._u_base, w)
+
+        return NonLinearOperator(domain, codomain, mapping)
+
+
 class LinearMinimumNormInversion(LinearInversion):
     """
     Finds a regularized solution using the discrepancy principle.
@@ -309,3 +419,111 @@ class LinearMinimumNormInversion(LinearInversion):
             normal_operator = forward_operator @ forward_operator.adjoint
             inverse_normal_operator = solver(normal_operator)
             return forward_operator.adjoint @ inverse_normal_operator
+
+
+class ConstrainedLinearMinimumNormInversion(LinearInversion):
+    """
+    Finds the minimum-norm solution subject to an affine subspace constraint
+    using the discrepancy principle.
+
+    Problem:
+        Minimize ||u||
+        Subject to u in A (Affine Subspace)
+        And chi_squared(u, d) <= critical_value
+
+    Method:
+        We decompose the model as u = u_base + w, where u_base is the element
+        of the affine subspace with the smallest norm (orthogonal to the tangent
+        space), and w is a perturbation in the tangent space.
+
+        Because u_base and w are orthogonal, ||u||^2 = ||u_base||^2 + ||w||^2.
+        Minimizing ||u|| is therefore equivalent to minimizing ||w||.
+
+        The problem reduces to finding the minimum norm w such that:
+        || A(w) - (d - A(u_base)) ||_D^2 <= critical_value
+
+        This is solved using the standard LinearMinimumNormInversion on a
+        reduced forward problem.
+    """
+
+    def __init__(
+        self,
+        forward_problem: LinearForwardProblem,
+        constraint: AffineSubspace,
+    ) -> None:
+        """
+        Args:
+            forward_problem: The original unconstrained forward problem.
+            constraint: The affine subspace A where the solution must lie.
+        """
+        super().__init__(forward_problem)
+        if self.forward_problem.data_error_measure_set:
+            self.assert_inverse_data_covariance()
+
+        self._constraint = constraint
+
+        # 1. Compute the Orthogonal Base Vector (u_base)
+        # u_base = (I - P) * translation
+        # This is the vector in the affine space closest to the origin.
+        self._u_base = constraint.domain.subtract(
+            constraint.translation, constraint.projector(constraint.translation)
+        )
+
+        # 2. Construct Reduced Forward Problem
+        # Operator: A_tilde = A @ P
+        reduced_operator = forward_problem.forward_operator @ constraint.projector
+
+        self._reduced_forward_problem = LinearForwardProblem(
+            reduced_operator,
+            data_error_measure=(
+                forward_problem.data_error_measure
+                if forward_problem.data_error_measure_set
+                else None
+            ),
+        )
+
+        # 3. Initialize the internal unconstrained solver
+        self._unconstrained_inversion = LinearMinimumNormInversion(
+            self._reduced_forward_problem
+        )
+
+    def minimum_norm_operator(
+        self,
+        solver: LinearSolver,
+        /,
+        **kwargs,
+    ) -> NonLinearOperator:
+        """
+        Returns an operator that maps data to the constrained minimum-norm solution.
+
+        Args:
+            solver: The linear solver for the reduced normal equations.
+            **kwargs: Arguments passed to LinearMinimumNormInversion (e.g.,
+                      significance_level, rtol, maxiter).
+
+        Returns:
+            A NonLinearOperator mapping d -> u_constrained.
+        """
+
+        # Get the operator L_tilde such that w = L_tilde(d_tilde)
+        reduced_op = self._unconstrained_inversion.minimum_norm_operator(
+            solver, **kwargs
+        )
+
+        # Precompute A(u_base) to shift the data
+        data_offset = self.forward_problem.forward_operator(self._u_base)
+
+        domain = self.data_space
+        codomain = self.model_space
+
+        def mapping(d: Vector) -> Vector:
+            # 1. Shift Data: d_tilde = d - A(u_base)
+            d_tilde = domain.subtract(d, data_offset)
+
+            # 2. Solve for perturbation w in the tangent space
+            w = reduced_op(d_tilde)
+
+            # 3. Reconstruct full model: u = u_base + w
+            return codomain.add(self._u_base, w)
+
+        return NonLinearOperator(domain, codomain, mapping)

pygeoinf/plot.py

@@ -247,7 +247,7 @@ def plot_corner_distributions(
                 sigma = np.sqrt(cov_posterior[i, i])
                 
                 # Create x-axis range
-                x = np.linspace(mu - 4 * sigma, mu + 4 * sigma, 200)
+                x = np.linspace(mu - 3.75 * sigma, mu + 3.75 * sigma, 200)
                 pdf = stats.norm.pdf(x, mu, sigma)
                 
                 # Plot the PDF
@@ -276,10 +276,10 @@ def plot_corner_distributions(
                 sigma_j = np.sqrt(cov_posterior[j, j])
                 sigma_i = np.sqrt(cov_posterior[i, i])
                 
-                x_range = np.linspace(mean_2d[0] - 3.5 * sigma_j, 
-                                    mean_2d[0] + 3.5 * sigma_j, 100)
-                y_range = np.linspace(mean_2d[1] - 3.5 * sigma_i, 
-                                    mean_2d[1] + 3.5 * sigma_i, 100)
+                x_range = np.linspace(mean_2d[0] - 3.75 * sigma_j, 
+                                    mean_2d[0] + 3.75 * sigma_j, 100)
+                y_range = np.linspace(mean_2d[1] - 3.75 * sigma_i, 
+                                    mean_2d[1] + 3.75 * sigma_i, 100)
                 
                 X, Y = np.meshgrid(x_range, y_range)
                 pos = np.dstack((X, Y))

pygeoinf/subspaces.py

@@ -0,0 +1,311 @@
+"""
+Defines classes for representing affine and linear subspaces.
+
+The primary abstraction is the `AffineSubspace`, which represents a subset of
+a Hilbert space defined by a translation and a closed linear tangent space.
+`LinearSubspace` is a specialization where the translation is zero.
+"""
+
+from __future__ import annotations
+from typing import List, Optional, Any, Callable, TYPE_CHECKING
+import numpy as np
+
+from .linear_operators import LinearOperator
+from .hilbert_space import HilbertSpace, Vector, EuclideanSpace
+from .linear_solvers import LinearSolver, CholeskySolver, IterativeLinearSolver
+
+if TYPE_CHECKING:
+    # Avoid circular imports for type checking
+    pass
+
+
+class OrthogonalProjector(LinearOperator):
+    """
+    Internal engine for subspace projections.
+
+    Represents an orthogonal projection operator P = P* = P^2.
+    While this class can be used directly, it is generally recommended to use
+    `AffineSubspace` or `LinearSubspace` for high-level problem definitions.
+    """
+
+    def __init__(
+        self,
+        domain: HilbertSpace,
+        mapping: Callable[[Any], Any],
+        complement_projector: Optional[LinearOperator] = None,
+    ) -> None:
+        super().__init__(domain, domain, mapping, adjoint_mapping=mapping)
+        self._complement_projector = complement_projector
+
+    @property
+    def complement(self) -> LinearOperator:
+        """Returns the projector onto the orthogonal complement (I - P)."""
+        if self._complement_projector is None:
+            identity = self.domain.identity_operator()
+            self._complement_projector = identity - self
+        return self._complement_projector
+
+    @classmethod
+    def from_basis(
+        cls,
+        domain: HilbertSpace,
+        basis_vectors: List[Vector],
+        orthonormalize: bool = True,
+    ) -> OrthogonalProjector:
+        """Constructs P from a basis spanning the range."""
+        if not basis_vectors:
+            # Return zero operator if basis is empty
+            return domain.zero_operator(domain)
+
+        if orthonormalize:
+            e_vectors = domain.gram_schmidt(basis_vectors)
+        else:
+            e_vectors = basis_vectors
+
+        # P = sum (v_i x v_i)
+        tensor_op = LinearOperator.self_adjoint_from_tensor_product(domain, e_vectors)
+        return cls(domain, tensor_op)
+
+
+class AffineSubspace:
+    """
+    Represents an affine subspace A = x0 + V.
+    """
+
+    def __init__(
+        self,
+        projector: OrthogonalProjector,
+        translation: Optional[Vector] = None,
+        constraint_operator: Optional[LinearOperator] = None,
+        constraint_value: Optional[Vector] = None,
+    ) -> None:
+        self._projector = projector
+
+        if translation is None:
+            self._translation = projector.domain.zero
+        else:
+            if not projector.domain.is_element(translation):
+                raise ValueError("Translation vector not in domain.")
+            self._translation = translation
+
+        self._constraint_operator = constraint_operator
+        self._constraint_value = constraint_value
+
+    @property
+    def domain(self) -> HilbertSpace:
+        return self._projector.domain
+
+    @property
+    def translation(self) -> Vector:
+        return self._translation
+
+    @property
+    def projector(self) -> OrthogonalProjector:
+        return self._projector
+
+    @property
+    def tangent_space(self) -> LinearSubspace:
+        return LinearSubspace(self._projector)
+
+    @property
+    def has_constraint_equation(self) -> bool:
+        return self._constraint_operator is not None
+
+    @property
+    def constraint_operator(self) -> LinearOperator:
+        if self._constraint_operator is None:
+            raise AttributeError("This subspace is not defined by a linear equation.")
+        return self._constraint_operator
+
+    @property
+    def constraint_value(self) -> Vector:
+        if self._constraint_value is None:
+            raise AttributeError("This subspace is not defined by a linear equation.")
+        return self._constraint_value
+
+    def project(self, x: Vector) -> Vector:
+        diff = self.domain.subtract(x, self.translation)
+        proj_diff = self.projector(diff)
+        return self.domain.add(self.translation, proj_diff)
+
+    def is_element(self, x: Vector, rtol: float = 1e-6) -> bool:
+        proj = self.project(x)
+        diff = self.domain.subtract(x, proj)
+        norm_diff = self.domain.norm(diff)
+        norm_x = self.domain.norm(x)
+        scale = norm_x if norm_x > 1e-12 else 1.0
+        return norm_diff <= rtol * scale
+
+    @classmethod
+    def from_linear_equation(
+        cls,
+        operator: LinearOperator,
+        value: Vector,
+        solver: Optional[LinearSolver] = None,
+        preconditioner: Optional[LinearOperator] = None,
+    ) -> AffineSubspace:
+        """Constructs the subspace {u | B(u) = w}."""
+        domain = operator.domain
+        G = operator @ operator.adjoint
+
+        if solver is None:
+            solver = CholeskySolver(galerkin=True)
+
+        if isinstance(solver, IterativeLinearSolver):
+            G_inv = solver(G, preconditioner=preconditioner)
+        else:
+            G_inv = solver(G)
+
+        intermediate = G_inv(value)
+        translation = operator.adjoint(intermediate)
+        P_perp_op = operator.adjoint @ G_inv @ operator
+
+        def mapping(x: Any) -> Any:
+            return domain.subtract(x, P_perp_op(x))
+
+        projector = OrthogonalProjector(domain, mapping, complement_projector=P_perp_op)
+
+        return cls(
+            projector, translation, constraint_operator=operator, constraint_value=value
+        )
+
+    @classmethod
+    def from_tangent_basis(
+        cls,
+        domain: HilbertSpace,
+        basis_vectors: List[Vector],
+        translation: Optional[Vector] = None,
+        orthonormalize: bool = True,
+    ) -> AffineSubspace:
+        """
+        Constructs the subspace passing through 'translation' with the given
+        tangent basis.
+
+        Note: This does not define a constraint equation B(u)=w, so it cannot
+        be used directly with ConstrainedLinearBayesianInversion.
+        """
+        projector = OrthogonalProjector.from_basis(
+            domain, basis_vectors, orthonormalize=orthonormalize
+        )
+        return cls(projector, translation)
+
+    @classmethod
+    def from_complement_basis(
+        cls,
+        domain: HilbertSpace,
+        basis_vectors: List[Vector],
+        translation: Optional[Vector] = None,
+        orthonormalize: bool = True,
+    ) -> AffineSubspace:
+        """
+        Constructs the subspace orthogonal to the given basis, passing through
+        'translation'.
+
+        This automatically constructs the constraint operator B such that
+        the subspace is {u | B(u) = B(translation)}.
+        """
+        # 1. Orthonormalize basis for stability
+        if orthonormalize:
+            e_vectors = domain.gram_schmidt(basis_vectors)
+        else:
+            e_vectors = basis_vectors
+
+        # 2. Construct Projector P_perp
+        complement_projector = OrthogonalProjector.from_basis(
+            domain, e_vectors, orthonormalize=False  # Already done
+        )
+
+        # 3. Construct Projector P = I - P_perp
+        def mapping(x: Any) -> Any:
+            return domain.subtract(x, complement_projector(x))
+
+        projector = OrthogonalProjector(
+            domain, mapping, complement_projector=complement_projector
+        )
+
+        # 4. Construct Constraint Operator B implicitly defined by the basis
+        # B: E -> R^k, u -> [<e_1, u>, ..., <e_k, u>]
+        # Since e_i are orthonormal, BB* = I, which is perfect for solvers.
+        codomain = EuclideanSpace(len(e_vectors))
+
+        def constraint_mapping(u: Vector) -> np.ndarray:
+            return np.array([domain.inner_product(e, u) for e in e_vectors])
+
+        def constraint_adjoint(c: np.ndarray) -> Vector:
+            # sum c_i e_i
+            res = domain.zero
+            for i, e in enumerate(e_vectors):
+                domain.axpy(c[i], e, res)
+            return res
+
+        B = LinearOperator(
+            domain, codomain, constraint_mapping, adjoint_mapping=constraint_adjoint
+        )
+
+        # 5. Determine Constraint Value w = B(translation)
+        # If translation is None (zero), w is zero.
+        if translation is None:
+            _translation = domain.zero
+            w = codomain.zero
+        else:
+            _translation = translation
+            w = B(_translation)
+
+        return cls(projector, _translation, constraint_operator=B, constraint_value=w)
+
+
+class LinearSubspace(AffineSubspace):
+    """
+    Represents a linear subspace (an affine subspace passing through zero).
+    """
+
+    def __init__(self, projector: OrthogonalProjector) -> None:
+        super().__init__(projector, translation=None)
+
+    @property
+    def complement(self) -> LinearSubspace:
+        op_perp = self.projector.complement
+        if isinstance(op_perp, OrthogonalProjector):
+            return LinearSubspace(op_perp)
+        p_perp = OrthogonalProjector(self.domain, op_perp._mapping)
+        return LinearSubspace(p_perp)
+
+    @classmethod
+    def from_kernel(
+        cls, operator: LinearOperator, solver: Optional[LinearSolver] = None
+    ) -> LinearSubspace:
+        affine = AffineSubspace.from_linear_equation(
+            operator, operator.codomain.zero, solver
+        )
+        instance = cls(affine.projector)
+        instance._constraint_operator = operator
+        instance._constraint_value = operator.codomain.zero
+        return instance
+
+    @classmethod
+    def from_basis(
+        cls,
+        domain: HilbertSpace,
+        basis_vectors: List[Vector],
+        orthonormalize: bool = True,
+    ) -> LinearSubspace:
+        projector = OrthogonalProjector.from_basis(
+            domain, basis_vectors, orthonormalize=orthonormalize
+        )
+        return cls(projector)
+
+    @classmethod
+    def from_complement_basis(
+        cls,
+        domain: HilbertSpace,
+        basis_vectors: List[Vector],
+        orthonormalize: bool = True,
+    ) -> LinearSubspace:
+        affine = AffineSubspace.from_complement_basis(
+            domain, basis_vectors, translation=None, orthonormalize=orthonormalize
+        )
+        # Copy constraint info from the affine instance
+        instance = cls(affine.projector)
+        instance._constraint_operator = affine.constraint_operator
+        instance._constraint_value = affine.constraint_value
+        return instance

pygeoinf/symmetric_space/sh_tools.py

@@ -0,0 +1,95 @@
+"""
+Module collecting some small helper functions or classes.
+"""
+
+from typing import Optional
+
+import numpy as np
+
+
+class SHVectorConverter:
+    """
+    Handles conversion between pyshtools 3D coefficient arrays and 1D vectors.
+
+    This class provides a bridge between the pyshtools data structure for
+    spherical harmonic coefficients, a 3D array of shape (2, lmax+1, lmax+1),
+    and the 1D vector format often used in linear algebra and inverse problems.
+
+    The vector is ordered by degree l, and within each degree, by order m,
+    from -l to +l.
+
+    Args:
+        lmax (int): The maximum spherical harmonic degree to include.
+        lmin (int): The minimum spherical harmonic degree to include. Defaults to 2.
+    """
+
+    def __init__(self, lmax: int, lmin: int = 0):
+        if not isinstance(lmax, int) or not isinstance(lmin, int):
+            raise TypeError("lmax and lmin must be integers.")
+        if lmin > lmax:
+            raise ValueError("lmin cannot be greater than lmax.")
+
+        self.lmax = lmax
+        self.lmin = lmin
+        self.vector_size = (self.lmax + 1) ** 2 - self.lmin**2
+
+    def to_vector(self, coeffs: np.ndarray) -> np.ndarray:
+        """Converts a pyshtools 3D coefficient array to a 1D vector.
+
+        If the input coefficients have a smaller lmax than the converter,
+        the missing high-degree coefficients in the output vector will be zero.
+
+        Args:
+            coeffs (np.ndarray): A pyshtools-compatible coefficient array
+                of shape (2, l_in+1, l_in+1).
+
+        Returns:
+            np.ndarray: A 1D vector of the coefficients from lmin to lmax.
+        """
+        lmax_in = coeffs.shape[1] - 1
+        vec = np.zeros(self.vector_size)
+        loop_lmax = min(self.lmax, lmax_in)
+
+        for l in range(self.lmin, loop_lmax + 1):
+            start_idx = l**2 - self.lmin**2
+            sin_part = coeffs[1, l, 1 : l + 1][::-1]
+            cos_part = coeffs[0, l, 0 : l + 1]
+            vec[start_idx : start_idx + l] = sin_part
+            vec[start_idx + l : start_idx + 2 * l + 1] = cos_part
+
+        return vec
+
+    def from_vector(
+        self, vec: np.ndarray, output_lmax: Optional[int] = None
+    ) -> np.ndarray:
+        """Converts a 1D vector back to a pyshtools 3D coefficient array.
+
+        This method can create an array that is larger (zero-padding) or
+        smaller (truncating) than the lmax of the converter.
+
+        Args:
+            vec (np.ndarray): A 1D vector of coefficients.
+            output_lmax (Optional[int]): The desired lmax for the output array.
+                If None, defaults to the converter's lmax.
+
+        Returns:
+            np.ndarray: A pyshtools-compatible coefficient array.
+        """
+        if vec.size != self.vector_size:
+            raise ValueError("Input vector has incorrect size.")
+
+        # If output_lmax is not specified, default to the converter's lmax
+        lmax_out = output_lmax if output_lmax is not None else self.lmax
+
+        # Create the output array of the desired size, initialized to zeros
+        coeffs = np.zeros((2, lmax_out + 1, lmax_out + 1))
+
+        # Determine the loop range: iterate up to the minimum of the two lmax values
+        loop_lmax = min(self.lmax, lmax_out)
+
+        for l in range(self.lmin, loop_lmax + 1):
+            start_idx = l**2 - self.lmin**2
+            coeffs[1, l, 1 : l + 1] = vec[start_idx : start_idx + l][::-1]
+            coeffs[0, l, 0 : l + 1] = vec[start_idx + l : start_idx + 2 * l + 1]
+
+        return coeffs