Bayesian Inference (geobrain.bayes)

Bayesian Inference (geobrain.bayes)#

The bayes module provides gradient-informed sampling methods for uncertainty quantification in geophysical inverse problems.

Overview#

Rather than finding a single “best” model, Bayesian inference characterizes the full posterior distribution of model parameters, providing uncertainty estimates that are critical for decision-making.

GeoBrain provides four samplers, all supporting automatic differentiation through the forward model:

Sampler

Method

Best For

SVGD

Stein Variational Gradient Descent

Multi-modal posteriors, large ensembles

HMC

Hamiltonian Monte Carlo

High-dimensional with tuned step size

NUTS

No-U-Turn Sampler

Auto-tuned HMC

LDS

Langevin Dynamics (ULA/MALA)

Simple gradient-based sampling

SVGD (Stein Variational Gradient Descent)#

SVGD maintains an ensemble of particles that are iteratively transported to approximate the posterior distribution:

from geobrain import InverseProblem, SVGD, Posterior, Gaussian
import torch

# Define problem
problem = InverseProblem(
    forward_fn=forward,
    observed=data,
    noise_std=0.01,
)

# Define prior (Gaussian takes mean tensor and covariance matrix)
prior = Gaussian(mean=m_prior, cov=500.0**2 * torch.eye(n_params))

# Create posterior
posterior = problem.as_posterior(log_prior=prior.log_prob)

# Run SVGD
svgd = SVGD(target=posterior, lr=0.005)
result = svgd.run(
    n_samples=50,    # Number of particles
    n_steps=1000,    # Number of iterations
)

# Access results
mean_model = result.samples.mean(dim=0)   # Ensemble mean
std_model = result.samples.std(dim=0)     # Ensemble standard deviation
samples = result.samples                   # All particles [n_samples, dim]
print(result.summary())
../_images/13_bayesian_avo_comparison.png

Fig. 25 Bayesian AVO inversion: posterior mean and uncertainty from four samplers.#

Distributions#

Gaussian#

Multivariate Gaussian parameterized by mean vector and covariance (or precision) matrix:

from geobrain import Gaussian
import torch

# Isotropic Gaussian
mean = torch.zeros(10)
cov = torch.eye(10)
prior = Gaussian(mean=mean, cov=cov)

# Or with precision matrix
prior = Gaussian(mean=mean, precision=torch.eye(10))

log_p = prior.log_prob(samples)   # [batch, 1]
score = prior.score(samples)      # [batch, dim]
drawn = prior.sample(100)         # [100, dim]

Gaussian Mixture#

from geobrain import GaussianMixture
import torch

prior = GaussianMixture(
    means=[torch.tensor([2000.0]), torch.tensor([3000.0])],
    covs=[torch.tensor([[90000.0]]), torch.tensor([[160000.0]])],
    weights=torch.tensor([0.6, 0.4]),
)

Custom Posterior#

from geobrain import Posterior

posterior = Posterior(
    log_likelihood=log_likelihood_fn,
    log_prior=log_prior_fn,
    data=d_obs,
    dim=n_params,
)

Kernels#

Kernels measure particle similarity for SVGD repulsion:

from geobrain import RBFKernel, IMQKernel

kernel = RBFKernel()            # Radial Basis Function (default)
kernel = IMQKernel()            # Inverse Multi-Quadratic

Hamiltonian Monte Carlo (HMC)#

from geobrain.bayes import HMC

hmc = HMC(
    target=posterior,
    step_size=1e-5,
    n_leapfrog=20,
    adapt_step_size=True,
    target_accept_rate=0.65,
)
result = hmc.run(
    n_samples=10,
    n_steps=300,
    initial_samples=init_samples,
    n_burnin=100,
)

No-U-Turn Sampler (NUTS)#

from geobrain.bayes import NUTS

nuts = NUTS(
    target=posterior,
    target_accept_rate=0.65,
    max_tree_depth=10,
)
result = nuts.run(
    n_samples=10,
    n_steps=300,
    initial_samples=init_samples,
    n_burnin=100,
)

Langevin Dynamics Sampler (LDS)#

Supports both Unadjusted Langevin Algorithm (ULA) and Metropolis-Adjusted Langevin Algorithm (MALA):

from geobrain.bayes import LDS

lds = LDS(
    target=posterior,
    step_size=1e-5,
    use_mh_correction=True,      # MALA (False for ULA)
    adapt_step_size=True,
    target_accept_rate=0.574,
)
result = lds.run(
    n_samples=50,
    n_steps=500,
    initial_samples=init_samples,
    n_burnin=200,
)

SamplingResult#

All samplers return a SamplingResult dataclass:

result.samples              # Final samples [n_samples, dim]
result.log_prob_history     # Mean log prob per step
result.algorithm            # Algorithm name string
result.n_steps              # Steps completed
result.total_time           # Wall time (seconds)
result.summary()            # Formatted text summary
result.to_dict()            # Serializable dictionary
result.to_numpy()           # Samples as numpy array
../_images/14_posterior_stats.png

Fig. 26 IBDP case study: posterior mean and standard deviation of porosity.#

Utility Functions#

from geobrain.bayes import pairwise_distances, median_heuristic, mmd, energy_distance

# Compute MMD between two sample sets
distance = mmd(samples_p, samples_q)
Contents