pygeoinf 1.3.8__py3-none-any.whl → 1.4.0__py3-none-any.whl

pygeoinf/__init__.py

@@ -104,8 +104,27 @@ from .nonlinear_optimisation import (
 
 from .subspaces import OrthogonalProjector, AffineSubspace, LinearSubspace
 
+from .subsets import (
+    Subset,
+    EmptySet,
+    UniversalSet,
+    Complement,
+    Intersection,
+    Union,
+    SublevelSet,
+    LevelSet,
+    ConvexSubset,
+    Ellipsoid,
+    NormalisedEllipsoid,
+    EllipsoidSurface,
+    Ball,
+    Sphere,
+)
+
 from .plot import plot_1d_distributions, plot_corner_distributions
 
+from .utils import configure_threading
+
 __all__ = [
     # random_matrix
     "fixed_rank_random_range",
@@ -184,7 +203,24 @@ __all__ = [
     "OrthogonalProjector",
     "AffineSubspace",
     "LinearSubspace",
+    # Subsets
+    "Subset",
+    "EmptySet",
+    "UniversalSet",
+    "Complement",
+    "Intersection",
+    "Union",
+    "SublevelSet",
+    "LevelSet",
+    "ConvexSubset",
+    "Ellipsoid",
+    "NormalisedEllipsoid",
+    "EllipsoidSurface",
+    "Ball",
+    "Sphere",
     # plot
     "plot_1d_distributions",
     "plot_corner_distributions",
+    # utils
+    "configure_threading",
 ]

pygeoinf/gaussian_measure.py

@@ -27,7 +27,7 @@ import numpy as np
 from scipy.linalg import eigh
 from scipy.sparse import diags
 from scipy.stats import multivariate_normal
-
+from joblib import Parallel, delayed
 
 from .hilbert_space import EuclideanSpace, HilbertModule, Vector
 
@@ -44,7 +44,6 @@ from .direct_sum import (
 # This block is only processed by type checkers, not at runtime.
 if TYPE_CHECKING:
     from .hilbert_space import HilbertSpace
-    from .typing import Vector
 
 
 class GaussianMeasure:
@@ -402,24 +401,52 @@ class GaussianMeasure:
             raise NotImplementedError("A sample method is not set for this measure.")
         return self._sample()
 
-    def samples(self, n: int) -> List[Vector]:
-        """Returns a list of n random samples from the measure."""
+    def samples(
+        self, n: int, /, *, parallel: bool = False, n_jobs: int = -1
+    ) -> List[Vector]:
+        """
+        Returns a list of n random samples from the measure.
+
+        Args:
+            n: Number of samples to draw.
+            parallel: If True, draws samples in parallel.
+            n_jobs: Number of CPU cores to use. -1 means all available.
+        """
         if n < 1:
             raise ValueError("Number of samples must be a positive integer.")
-        return [self.sample() for _ in range(n)]
 
-    def sample_expectation(self, n: int) -> Vector:
-        """Estimates the expectation by drawing n samples."""
+        if not parallel:
+            return [self.sample() for _ in range(n)]
+
+        return Parallel(n_jobs=n_jobs)(delayed(self.sample)() for _ in range(n))
+
+    def sample_expectation(
+        self, n: int, /, *, parallel: bool = False, n_jobs: int = -1
+    ) -> Vector:
+        """
+        Estimates the expectation by drawing n samples.
+
+        Args:
+            n: Number of samples to draw.
+            parallel: If True, draws samples in parallel.
+            n_jobs: Number of CPU cores to use. -1 means all available.
+        """
         if n < 1:
             raise ValueError("Number of samples must be a positive integer.")
-        return self.domain.sample_expectation(self.samples(n))
+        return self.domain.sample_expectation(
+            self.samples(n, parallel=parallel, n_jobs=n_jobs)
+        )
 
-    def sample_pointwise_variance(self, n: int) -> Vector:
+    def sample_pointwise_variance(
+        self, n: int, /, *, parallel: bool = False, n_jobs: int = -1
+    ) -> Vector:
         """
         Estimates the pointwise variance by drawing n samples.
 
-        This method is only available if the domain supports vector
-        multiplication.
+        Args:
+            n: Number of samples to draw.
+            parallel: If True, draws samples in parallel.
+            n_jobs: Number of CPU cores to use. -1 means all available.
         """
         if not isinstance(self.domain, HilbertModule):
             raise NotImplementedError(
@@ -428,7 +455,10 @@ class GaussianMeasure:
         if n < 1:
             raise ValueError("Number of samples must be a positive integer.")
 
-        samples = self.samples(n)
+        # Draw samples
+        samples = self.samples(n, parallel=parallel, n_jobs=n_jobs)
+
+        # Compute variance using vector arithmetic
         expectation = self.expectation
         variance = self.domain.zero
 
@@ -439,6 +469,42 @@ class GaussianMeasure:
 
         return variance
 
+    def sample_pointwise_std(
+        self, n: int, /, *, parallel: bool = False, n_jobs: int = -1
+    ) -> Vector:
+        """
+        Estimates the pointwise standard deviation by drawing n samples.
+
+        Args:
+            n: Number of samples to draw.
+            parallel: If True, draws samples in parallel.
+            n_jobs: Number of CPU cores to use. -1 means all available.
+        """
+        variance = self.sample_pointwise_variance(n, parallel=parallel, n_jobs=n_jobs)
+        return self.domain.vector_sqrt(variance)
+
+    def with_dense_covariance(self, parallel: bool = False, n_jobs: int = -1):
+        """
+        Forms a new Gaussian measure equivalent to the existing one, but
+        with its covariance matrix stored in dense form. The dense matrix
+        calculation can optionally be parallelised.
+
+        Args:
+            parallel: If True, computes the covariance in parallel.
+            n_jobs: Number of CPU cores to use. -1 means all available.
+
+        Returns:
+            The new Gaussian measure.
+        """
+
+        covariance_matrix = self.covariance.matrix(
+            dense=True, galerkin=True, parallel=parallel, n_jobs=n_jobs
+        )
+
+        return GaussianMeasure.from_covariance_matrix(
+            self.domain, covariance_matrix, expectation=self.expectation
+        )
+
     def affine_mapping(
         self, /, *, operator: LinearOperator = None, translation: Vector = None
     ) -> GaussianMeasure:
@@ -516,7 +582,7 @@ class GaussianMeasure:
 
         # Pass the parallelization arguments directly to the matrix creation method
         cov_matrix = self.covariance.matrix(
-            dense=True, parallel=parallel, n_jobs=n_jobs
+            dense=True, galerkin=True, parallel=parallel, n_jobs=n_jobs
         )
 
         try:

pygeoinf/hilbert_space.py

@@ -547,6 +547,12 @@ class HilbertModule(HilbertSpace, ABC):
             The product of the two vectors.
         """
 
+    @abstractmethod
+    def vector_sqrt(self, x: Vector) -> Vector:
+        """
+        Returns the square root of a vector.
+        """
+
 
 class EuclideanSpace(HilbertSpace):
     """
@@ -829,3 +835,12 @@ class MassWeightedHilbertModule(MassWeightedHilbertSpace, HilbertModule):
         is itself an instance of `HilbertModule`.
         """
         return self.underlying_space.vector_multiply(x1, x2)
+
+    def vector_sqrt(self, x: Vector) -> Vector:
+        """
+        Computes vector multiplication by delegating to the underlying space.
+
+        Note: This assumes the underlying space provided during initialization
+        is itself an instance of `HilbertModule`.
+        """
+        return self.underlying_space.vector_sqrt(x)

pygeoinf/linear_operators.py

@@ -102,10 +102,13 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         self.__adjoint_mapping: Callable[[Any], Any]
         self.__dual_mapping: Callable[[Any], Any]
 
+        self.__using_default_dual_and_adjoint = False
+
         if dual_mapping is None:
             if adjoint_mapping is None:
                 self.__dual_mapping = self._dual_mapping_default
                 self.__adjoint_mapping = self._adjoint_mapping_from_dual
+                self.__using_default_dual_and_adjoint = True
             else:
                 self.__adjoint_mapping = adjoint_mapping
                 self.__dual_mapping = self._dual_mapping_from_adjoint
@@ -919,6 +922,35 @@ class LinearOperator(NonLinearOperator, LinearOperatorAxiomChecks):
         self, galerkin: bool, parallel: bool, n_jobs: int
     ) -> np.ndarray:
 
+        # Optimization: If the codomain is smaller than the domain, it is cheaper
+        # to compute the matrix of the adjoint/dual (which has fewer columns)
+        # and transpose the result.
+
+        # Note: This recursion naturally terminates because the adjoint/dual
+        # swaps the domain and codomain. In the recursive call,
+        # (codomain.dim < domain.dim) will be False, forcing the standard path.
+
+        # If the operator has its dual and adjoint actions done using the
+        # default implementation, this optimisation is skipped.
+        if (
+            self.codomain.dim < self.domain.dim
+            and not self.__using_default_dual_and_adjoint
+        ):
+            if galerkin:
+                # For Galerkin representations: Matrix(L) = Matrix(L*).T
+                return self.adjoint.matrix(
+                    dense=True, galerkin=True, parallel=parallel, n_jobs=n_jobs
+                ).T
+            else:
+                # For Standard representations: Matrix(L) = Matrix(L').T
+                return self.dual.matrix(
+                    dense=True, galerkin=False, parallel=parallel, n_jobs=n_jobs
+                ).T
+
+        # --- Standard Column-wise Construction ---
+        # This block executes if optimization is not applicable (or in the
+        # recursive base case).
+
         scipy_op_wrapper = self.matrix(galerkin=galerkin)
 
         if not parallel:

pygeoinf/plot.py

@@ -220,6 +220,8 @@ def plot_corner_distributions(
     show_plot: bool = True,
     include_sigma_contours: bool = True,
     colormap: str = "Blues",
+    parallel: bool = False,
+    n_jobs: int = -1,
 ):
     """
     Create a corner plot for multi-dimensional posterior distributions.
@@ -233,6 +235,8 @@ def plot_corner_distributions(
         show_plot: Whether to display the plot
         include_sigma_contours: Whether to include 1-sigma contour lines
         colormap: Colormap for 2D plots
+        parallel: Compute dense covariance matrix in parallel, default False.
+        n_jobs: Number of cores to use in parallel calculations, default -1.
 
     Returns:
         fig, axes: Figure and axes array
@@ -243,7 +247,9 @@ def plot_corner_distributions(
         posterior_measure, "covariance"
     ):
         mean_posterior = posterior_measure.expectation
-        cov_posterior = posterior_measure.covariance.matrix(dense=True, parallel=True)
+        cov_posterior = posterior_measure.covariance.matrix(
+            dense=True, parallel=parallel, n_jobs=n_jobs
+        )
     else:
         raise ValueError(
             "posterior_measure must have 'expectation' and 'covariance' attributes"

pygeoinf/preconditioners.py

@@ -36,7 +36,7 @@ class JacobiPreconditioningMethod(LinearSolver):
         method: str = "variable",
         rtol: float = 1e-2,
         block_size: int = 10,
-        parallel: bool = True,
+        parallel: bool = False,
         n_jobs: int = -1,
     ) -> None:
         # Damping is removed: the operator passed to __call__ is already damped