pygeoinf 1.3.7__py3-none-any.whl → 1.3.9__py3-none-any.whl

pygeoinf/__init__.py

@@ -104,6 +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",
@@ -182,4 +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)
+        # Step 1: Draw samples (Parallelized)
+        samples = self.samples(n, parallel=parallel, n_jobs=n_jobs)
+
+        # Step 2: Compute variance using vector arithmetic
         expectation = self.expectation
         variance = self.domain.zero
 

pygeoinf/plot.py

@@ -4,18 +4,21 @@ import numpy as np
 import scipy.stats as stats
 from typing import Union, List, Optional
 
+
 def plot_1d_distributions(
     posterior_measures: Union[object, List[object]],
+    /,
+    *,
     prior_measures: Optional[Union[object, List[object]]] = None,
     true_value: Optional[float] = None,
     xlabel: str = "Property Value",
     title: str = "Prior and Posterior Probability Distributions",
     figsize: tuple = (12, 7),
-    show_plot: bool = True
+    show_plot: bool = True,
 ):
     """
     Plot 1D probability distributions for prior and posterior measures using dual y-axes.
-    
+
     Args:
         posterior_measures: Single measure or list of measures for posterior distributions
         prior_measures: Single measure or list of measures for prior distributions (optional)
@@ -24,26 +27,44 @@ def plot_1d_distributions(
         title: Title for the plot
         figsize: Figure size tuple
         show_plot: Whether to display the plot
-        
+
     Returns:
         fig, (ax1, ax2): Figure and axes objects
     """
-    
+
     # Convert single measures to lists for uniform handling
     if not isinstance(posterior_measures, list):
         posterior_measures = [posterior_measures]
-    
+
     if prior_measures is not None and not isinstance(prior_measures, list):
         prior_measures = [prior_measures]
-    
+
     # Define color sequences
-    prior_colors = ['green', 'orange', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan']
-    posterior_colors = ['blue', 'red', 'darkgreen', 'orange', 'purple', 'brown', 'pink', 'gray']
-    
+    prior_colors = [
+        "green",
+        "orange",
+        "purple",
+        "brown",
+        "pink",
+        "gray",
+        "olive",
+        "cyan",
+    ]
+    posterior_colors = [
+        "blue",
+        "red",
+        "darkgreen",
+        "orange",
+        "purple",
+        "brown",
+        "pink",
+        "gray",
+    ]
+
     # Calculate statistics for all distributions
     posterior_stats = []
     for measure in posterior_measures:
-        if hasattr(measure, 'expectation') and hasattr(measure, 'covariance'):
+        if hasattr(measure, "expectation") and hasattr(measure, "covariance"):
             # For pygeoinf measures
             mean = measure.expectation[0]
             var = measure.covariance.matrix(dense=True)[0, 0]
@@ -53,11 +74,11 @@ def plot_1d_distributions(
             mean = measure.mean[0]
             std = np.sqrt(measure.cov[0, 0])
         posterior_stats.append((mean, std))
-    
+
     prior_stats = []
     if prior_measures is not None:
         for measure in prior_measures:
-            if hasattr(measure, 'expectation') and hasattr(measure, 'covariance'):
+            if hasattr(measure, "expectation") and hasattr(measure, "covariance"):
                 # For pygeoinf measures
                 mean = measure.expectation[0]
                 var = measure.covariance.matrix(dense=True)[0, 0]
@@ -67,99 +88,106 @@ def plot_1d_distributions(
                 mean = measure.mean[0]
                 std = np.sqrt(measure.cov[0, 0])
             prior_stats.append((mean, std))
-    
+
     # Determine plot range to include all distributions
     all_means = [stat[0] for stat in posterior_stats]
     all_stds = [stat[1] for stat in posterior_stats]
-    
+
     if prior_measures is not None:
         all_means.extend([stat[0] for stat in prior_stats])
         all_stds.extend([stat[1] for stat in prior_stats])
-    
+
     if true_value is not None:
         all_means.append(true_value)
         all_stds.append(0)  # No std for true value
-    
+
     # Calculate x-axis range (6 sigma coverage)
     x_min = min([mean - 6 * std for mean, std in zip(all_means, all_stds) if std > 0])
     x_max = max([mean + 6 * std for mean, std in zip(all_means, all_stds) if std > 0])
-    
+
     # Add some padding around true value if needed
     if true_value is not None:
         range_size = x_max - x_min
         x_min = min(x_min, true_value - 0.1 * range_size)
         x_max = max(x_max, true_value + 0.1 * range_size)
-    
+
     x_axis = np.linspace(x_min, x_max, 1000)
-    
+
     # Create the plot with two y-axes
     fig, ax1 = plt.subplots(figsize=figsize)
-    
+
     # Plot priors on the first axis (left y-axis) if provided
     if prior_measures is not None:
-        color1 = prior_colors[0] if len(prior_measures) > 0 else 'green'
+        color1 = prior_colors[0] if len(prior_measures) > 0 else "green"
         ax1.set_xlabel(xlabel)
-        ax1.set_ylabel('Prior Probability Density', color=color1)
-        
+        ax1.set_ylabel("Prior Probability Density", color=color1)
+
         for i, (measure, (mean, std)) in enumerate(zip(prior_measures, prior_stats)):
             color = prior_colors[i % len(prior_colors)]
-            
+
             # Calculate PDF values using scipy.stats
             pdf_values = stats.norm.pdf(x_axis, loc=mean, scale=std)
-            
+
             # Determine label
             if len(prior_measures) == 1:
-                label = f'Prior PDF (Mean: {mean:.5f})'
+                label = f"Prior PDF (Mean: {mean:.5f})"
             else:
-                label = f'Prior {i+1} (Mean: {mean:.5f})'
-            
-            ax1.plot(x_axis, pdf_values, color=color, lw=2, linestyle=':', label=label)
+                label = f"Prior {i+1} (Mean: {mean:.5f})"
+
+            ax1.plot(x_axis, pdf_values, color=color, lw=2, linestyle=":", label=label)
             ax1.fill_between(x_axis, pdf_values, color=color, alpha=0.15)
-        
-        ax1.tick_params(axis='y', labelcolor=color1)
-        ax1.grid(True, linestyle='--')
+
+        ax1.tick_params(axis="y", labelcolor=color1)
+        ax1.grid(True, linestyle="--")
     else:
         # If no priors, use the left axis for posteriors
         ax1.set_xlabel(xlabel)
-        ax1.set_ylabel('Probability Density')
-        ax1.grid(True, linestyle='--')
-    
+        ax1.set_ylabel("Probability Density")
+        ax1.grid(True, linestyle="--")
+
     # Create second y-axis for posteriors (or use first if no priors)
     if prior_measures is not None:
         ax2 = ax1.twinx()
-        color2 = posterior_colors[0] if len(posterior_measures) > 0 else 'blue'
-        ax2.set_ylabel('Posterior Probability Density', color=color2)
-        ax2.tick_params(axis='y', labelcolor=color2)
+        color2 = posterior_colors[0] if len(posterior_measures) > 0 else "blue"
+        ax2.set_ylabel("Posterior Probability Density", color=color2)
+        ax2.tick_params(axis="y", labelcolor=color2)
         ax2.grid(False)
         plot_ax = ax2
     else:
         plot_ax = ax1
-        color2 = posterior_colors[0] if len(posterior_measures) > 0 else 'blue'
-    
+        color2 = posterior_colors[0] if len(posterior_measures) > 0 else "blue"
+
     # Plot posteriors
-    for i, (measure, (mean, std)) in enumerate(zip(posterior_measures, posterior_stats)):
+    for i, (measure, (mean, std)) in enumerate(
+        zip(posterior_measures, posterior_stats)
+    ):
         color = posterior_colors[i % len(posterior_colors)]
-        
+
         # Calculate PDF values using scipy.stats
         pdf_values = stats.norm.pdf(x_axis, loc=mean, scale=std)
-        
+
         # Determine label
         if len(posterior_measures) == 1:
-            label = f'Posterior PDF (Mean: {mean:.5f})'
+            label = f"Posterior PDF (Mean: {mean:.5f})"
         else:
-            label = f'Posterior {i+1} (Mean: {mean:.5f})'
-        
+            label = f"Posterior {i+1} (Mean: {mean:.5f})"
+
         plot_ax.plot(x_axis, pdf_values, color=color, lw=2, label=label)
         plot_ax.fill_between(x_axis, pdf_values, color=color, alpha=0.2)
-    
+
     # Plot true value if provided
     if true_value is not None:
-        ax1.axvline(true_value, color='black', linestyle='-', lw=2, 
-                   label=f'True Value: {true_value:.5f}')
-    
+        ax1.axvline(
+            true_value,
+            color="black",
+            linestyle="-",
+            lw=2,
+            label=f"True Value: {true_value:.5f}",
+        )
+
     # Create combined legend
     handles1, labels1 = ax1.get_legend_handles_labels()
-    
+
     if prior_measures is not None:
         handles2, labels2 = ax2.get_legend_handles_labels()
         all_handles = handles1 + handles2
@@ -167,14 +195,14 @@ def plot_1d_distributions(
     else:
         all_handles = handles1
         all_labels = [h.get_label() for h in all_handles]
-    
-    fig.legend(all_handles, all_labels, loc='upper right', bbox_to_anchor=(0.9, 0.9))
+
+    fig.legend(all_handles, all_labels, loc="upper right", bbox_to_anchor=(0.9, 0.9))
     fig.suptitle(title, fontsize=16)
     fig.tight_layout(rect=[0, 0, 1, 0.96])
-    
+
     if show_plot:
         plt.show()
-    
+
     if prior_measures is not None:
         return fig, (ax1, ax2)
     else:
@@ -183,17 +211,21 @@ def plot_1d_distributions(
 
 def plot_corner_distributions(
     posterior_measure: object,
+    /,
+    *,
     true_values: Optional[Union[List[float], np.ndarray]] = None,
     labels: Optional[List[str]] = None,
     title: str = "Joint Posterior Distribution",
     figsize: Optional[tuple] = None,
     show_plot: bool = True,
     include_sigma_contours: bool = True,
-    colormap: str = "Blues"
+    colormap: str = "Blues",
+    parallel: bool = False,
+    n_jobs: int = -1,
 ):
     """
     Create a corner plot for multi-dimensional posterior distributions.
-    
+
     Args:
         posterior_measure: Multi-dimensional posterior measure (pygeoinf object)
         true_values: True values for each dimension (optional)
@@ -203,148 +235,184 @@ 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
     """
-    
+
     # Extract statistics from the measure
-    if hasattr(posterior_measure, 'expectation') and hasattr(posterior_measure, 'covariance'):
+    if hasattr(posterior_measure, "expectation") and hasattr(
+        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")
-    
+        raise ValueError(
+            "posterior_measure must have 'expectation' and 'covariance' attributes"
+        )
+
     n_dims = len(mean_posterior)
-    
+
     # Set default labels if not provided
     if labels is None:
         labels = [f"Dimension {i+1}" for i in range(n_dims)]
-    
+
     # Set figure size based on dimensions if not provided
     if figsize is None:
         figsize = (3 * n_dims, 3 * n_dims)
-    
+
     # Create subplots
     fig, axes = plt.subplots(n_dims, n_dims, figsize=figsize)
     fig.suptitle(title, fontsize=16)
-    
+
     # Ensure axes is always 2D array
     if n_dims == 1:
         axes = np.array([[axes]])
     elif n_dims == 2:
         axes = axes.reshape(2, 2)
-    
+
     # Initialize pcm variable for colorbar
     pcm = None
-    
+
     for i in range(n_dims):
         for j in range(n_dims):
             ax = axes[i, j]
-            
+
             if i == j:  # Diagonal plots (1D marginal distributions)
                 mu = mean_posterior[i]
                 sigma = np.sqrt(cov_posterior[i, i])
-                
+
                 # Create x-axis range
                 x = np.linspace(mu - 3.75 * sigma, mu + 3.75 * sigma, 200)
                 pdf = stats.norm.pdf(x, mu, sigma)
-                
+
                 # Plot the PDF
                 ax.plot(x, pdf, "darkblue", label="Posterior PDF")
                 ax.fill_between(x, pdf, color="lightblue", alpha=0.6)
-                
+
                 # Add true value if provided
                 if true_values is not None:
                     true_val = true_values[i]
-                    ax.axvline(true_val, color="black", linestyle="-", 
-                              label=f"True: {true_val:.2f}")
-                
+                    ax.axvline(
+                        true_val,
+                        color="black",
+                        linestyle="-",
+                        label=f"True: {true_val:.2f}",
+                    )
+
                 ax.set_xlabel(labels[i])
                 ax.set_ylabel("Density" if i == 0 else "")
                 ax.set_yticklabels([])
-            
+
             elif i > j:  # Lower triangle: 2D joint distributions
                 # Extract 2D mean and covariance
                 mean_2d = np.array([mean_posterior[j], mean_posterior[i]])
-                cov_2d = np.array([
-                    [cov_posterior[j, j], cov_posterior[j, i]],
-                    [cov_posterior[i, j], cov_posterior[i, i]]
-                ])
-                
+                cov_2d = np.array(
+                    [
+                        [cov_posterior[j, j], cov_posterior[j, i]],
+                        [cov_posterior[i, j], cov_posterior[i, i]],
+                    ]
+                )
+
                 # Create 2D grid
                 sigma_j = np.sqrt(cov_posterior[j, j])
                 sigma_i = np.sqrt(cov_posterior[i, i])
-                
-                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_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))
-                
+
                 # Calculate PDF values
                 rv = stats.multivariate_normal(mean_2d, cov_2d)
                 Z = rv.pdf(pos)
-                
+
                 # Create filled contour plot using pcolormesh like the original
                 pcm = ax.pcolormesh(
-                    X, Y, Z, shading="auto", cmap=colormap,
-                    norm=colors.LogNorm(vmin=Z.min(), vmax=Z.max())
+                    X,
+                    Y,
+                    Z,
+                    shading="auto",
+                    cmap=colormap,
+                    norm=colors.LogNorm(vmin=Z.min(), vmax=Z.max()),
                 )
-                
+
                 # Add contour lines
                 ax.contour(X, Y, Z, colors="black", linewidths=0.5, alpha=0.6)
-                
+
                 # Add 1-sigma contour if requested
                 if include_sigma_contours:
                     # Calculate 1-sigma level (approximately 39% of peak for 2D Gaussian)
                     sigma_level = rv.pdf(mean_2d) * np.exp(-0.5)
-                    ax.contour(X, Y, Z, levels=[sigma_level], colors="red", 
-                              linewidths=1, linestyles="--", alpha=0.8)
-                
+                    ax.contour(
+                        X,
+                        Y,
+                        Z,
+                        levels=[sigma_level],
+                        colors="red",
+                        linewidths=1,
+                        linestyles="--",
+                        alpha=0.8,
+                    )
+
                 # Plot mean point
-                ax.plot(mean_posterior[j], mean_posterior[i], "r+", 
-                       markersize=10, mew=2, label="Posterior Mean")
-                
+                ax.plot(
+                    mean_posterior[j],
+                    mean_posterior[i],
+                    "r+",
+                    markersize=10,
+                    mew=2,
+                    label="Posterior Mean",
+                )
+
                 # Plot true value if provided
                 if true_values is not None:
-                    ax.plot(true_values[j], true_values[i], "kx", 
-                           markersize=10, mew=2, label="True Value")
-                
+                    ax.plot(
+                        true_values[j],
+                        true_values[i],
+                        "kx",
+                        markersize=10,
+                        mew=2,
+                        label="True Value",
+                    )
+
                 ax.set_xlabel(labels[j])
                 ax.set_ylabel(labels[i])
-            
+
             else:  # Upper triangle: hide these plots
                 ax.axis("off")
-    
+
     # Create legend similar to the original
     handles, labels_leg = axes[0, 0].get_legend_handles_labels()
     if n_dims > 1:
         handles2, labels2 = axes[1, 0].get_legend_handles_labels()
         handles.extend(handles2)
         labels_leg.extend(labels2)
-    
+
     # Clean up labels by removing values after colons
     cleaned_labels = [label.split(":")[0] for label in labels_leg]
-    
-    fig.legend(
-        handles, cleaned_labels,
-        loc="upper right", 
-        bbox_to_anchor=(0.9, 0.95)
-    )
-    
+
+    fig.legend(handles, cleaned_labels, loc="upper right", bbox_to_anchor=(0.9, 0.95))
+
     # Adjust main plot layout to make room on the right for the colorbar
     plt.tight_layout(rect=[0, 0, 0.88, 0.96])
-    
+
     # Add a colorbar if we have 2D plots
     if n_dims > 1 and pcm is not None:
         cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.7])
         cbar = fig.colorbar(pcm, cax=cbar_ax)
         cbar.set_label("Probability Density", size=12)
-    
+
     if show_plot:
         plt.show()
-    
-    return fig, axes
+
+    return fig, axes

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