Choosing an inference method#

Select the inference method with sampler_method in the TOML configuration:

sampler_method = "mcmc"

All methods return the same SamplerResult interface. Their configuration parameters and diagnostics differ.

Ensemble MCMC#

mcmc uses the emcee ensemble sampler. Its trace and autocorrelation diagnostics help assess whether the retained samples represent the posterior.

Parameter	Default	Meaning
`mcmc.n_steps`	`500`	Number of steps run by each walker.
`mcmc.n_walkers`	`50`	Number of ensemble walkers. This should exceed the number of free covariance parameters.
`mcmc.n_burn`	`300`	Initial steps discarded from every walker.
`mcmc.move`	`"kde"`	Proposal move: `"kde"` or `"stretch"`.

MAP and Laplace approximation#

map finds a maximum-a-posteriori solution with L-BFGS. It can also draw from a local Laplace approximation. These draws describe a Gaussian approximation around one optimum, not a non-Gaussian or multimodal posterior.

Parameter	Default	Meaning
`map.n_restarts`	`4`	Number of optimizations from different initial parameter values; the solution with the highest posterior probability is retained.
`map.n_steps`	`100`	Maximum L-BFGS steps per restart.
`map.lr`	`1.0`	L-BFGS learning rate.
`map.n_laplace_samples`	`2000`	Number of draws from the local Gaussian approximation.
`map.validate_feasible`	`false`	Test every Laplace draw by constructing and factorizing its covariance. This adds a potentially substantial validation cost.

Dynesty nested sampling#

nested uses Dynesty and returns both posterior samples and the Bayesian evidence.

Parameter	Default	Meaning
`nested.nlive`	`500`	Number of live points.
`nested.bound`	`"multi"`	Dynesty’s live-point bounding strategy.
`nested.sample`	`"auto"`	Dynesty’s proposal used to sample within the bounds.

NUTS#

nuts uses Pyro’s gradient-based No-U-Turn Sampler. Its result includes divergence, effective-sample-size, and \(\hat R\) diagnostics.

Parameter	Default	Meaning
`nuts.num_samples`	`1000`	Retained posterior draws per chain.
`nuts.warmup_steps`	`500`	Adaptation steps discarded before collecting posterior draws.
`nuts.num_chains`	`1`	Number of independent chains.
`nuts.max_tree_depth`	`8`	Maximum NUTS trajectory-tree depth.
`nuts.target_accept_prob`	`0.8`	Target acceptance probability used during step-size adaptation.
`nuts.full_mass`	`false`	Adapt a dense rather than diagonal mass matrix.
`nuts.warm_start_map`	`false`	Initialize sampling from a MAP solution instead of a prior draw.

UltraNest#

ultranest uses UltraNest’s reactive nested sampler and requires the ultranest installation extra.

Parameter	Default	Meaning
`ultranest.nlive`	Not set	Minimum number of live points. This field must currently be supplied when selecting `"ultranest"`.

These are exactly the sampler parameters read from TOML by run(). Its live_update and verbose controls are method arguments rather than configuration fields. The sampler classes in the API reference expose additional run controls; the Dynesty and UltraNest classes also accept backend-specific keyword arguments.

Running inference#

Once the covariance model and inference method have been selected in the configuration, the same high-level call runs every method:

from ps_eor.ml_gpr import MLGPRConfigFile, MLGPRForegroundFitter

config = MLGPRConfigFile.load_with_defaults("ml_gpr.toml")
fitter = MLGPRForegroundFitter(config)

noise_for_fit = fitter.process_noise_cube(noise_cube)
result = fitter.run(
    data_cube,
    noise_for_fit,
    live_update=True,
    verbose=True,
)

live_update displays notebook progress for MCMC and the nested samplers. verbose enables the textual output supported by the selected method. Both are optional and do not change the samples.

Running sampler classes directly#

The advanced interface separates sampler definition from execution. Iterative methods use run(n_steps=...) consistently:

from ps_eor.ml_gpr.samplers import MCMCSampler, MAPOptimizer, NUTSSampler

mcmc_result = MCMCSampler(
    gp, n_walkers=50, emcee_moves="kde"
).run(
    n_steps=500, n_burn=300, verbose=False
)

map_result = MAPOptimizer(
    gp, n_restarts=4, n_laplace_samples=2000
).run(
    n_steps=100, lr=1.0, verbose=False
)

nuts_result = NUTSSampler(
    gp, warmup_steps=500, num_chains=4
).run(
    n_steps=1000, verbose=False
)

For MCMC, n_steps counts steps per walker and n_burn selects the initial steps omitted from the returned result. The same sampled chain can be reprocessed later with sampler.get_result(n_burn=...). For NUTS, n_steps is the number of retained draws per chain; adaptation uses the additional warmup_steps configured on the sampler.

Dynesty and UltraNest remain convergence-driven: their run() methods take backend stopping criteria such as dlogz rather than n_steps.

The returned MLGPRResult holds the common high-level interface. The sampler-specific result is available as result.sampler_result:

sampler_result = result.sampler_result
print(sampler_result)
print(sampler_result.diagnostics)

result.save("results", "fit")

Saving a result also saves its diagnostics. See Running an ML-GPR fit for data preparation and the inference API for running sampler classes directly.

Assessing inference quality#

The following plots are available for every SamplerResult:

sampler_result.plot_samples()
sampler_result.plot_samples_likelihood()
sampler_result.plot_corner()

plot_samples() is most useful for inspecting MCMC and NUTS sampling history. plot_samples_likelihood() reveals low-probability branches and parameters that are poorly identified. plot_corner() shows marginal constraints and parameter degeneracies; by default it also shades intervals from the priors.

The following trace and corner plot come from a small simulated data set with a smooth foreground, a 21-cm-like component, and white noise. The model has four free hyperparameters and is sampled with ensemble MCMC.

MCMC traces for foreground and 21-cm-like covariance parameters. — The walker traces expose the initial transient, mixing between walkers, and the chain length relative to the estimated autocorrelation time.#

Corner plot of the ML-GPR MCMC posterior. — Marginal posterior distributions and correlations between covariance hyperparameters. Green shading shows intervals implied by the priors.#

Diagnostics are stored either in sampler_result.diagnostics or, for MCMC and nested sampling, on sampler_result.samples:

Method	Available diagnostics	Interpretation
`mcmc`	`samples.autocorr_time`, raw chains, log probabilities	The trace shows whether walkers have reached the same stationary region. The retained chain should extend over several autocorrelation times after burn-in. The autocorrelation estimate alone does not mark a short chain as converged.
`map`	`laplace_valid`, Hessian eigenvalues and regularization, Laplace-draw counts, Cholesky jitter	These describe whether the local Gaussian approximation is numerically meaningful. They do not turn a local MAP solution into a full posterior exploration.
`nested`	`samples.logz`, `samples.logzerr`, `samples.niter`, `samples.ncall`	`logz` and `logzerr` are the evidence estimate and its uncertainty. Evidence comparisons require the same data and explicitly specified priors. Iteration and likelihood-call counts describe the work done.
`ultranest`	`samples.logz`, `samples.logzerr`, `samples.niter`, `samples.ncall`	UltraNest exposes the same stored evidence summary. Its convergence and stopping tests are handled by the backend during the run.
`nuts`	`n_chains`, `divergences`, `n_eff`, `r_hat`, `converged`	These test chain agreement, effective sample count, and problematic Hamiltonian trajectories. The combined `converged` flag requires no divergences and finite \(\hat R \leq 1.01\) for every parameter.

MCMC#

The autocorrelation time is stored once per free parameter:

tau = sampler_result.samples.autocorr_time
for name, value in zip(sampler_result.get_parameter_names(), tau):
    print(name, value)

The raw trace includes burn-in. plot_samples() marks one and five autocorrelation times, while posterior access through samples.get() applies the configured burn-in and filters stuck walkers and extreme outliers.

MAP and Laplace approximation#

MAP diagnostics are available directly as a dictionary:

diagnostics = sampler_result.diagnostics
print(diagnostics["laplace_valid"])
print(diagnostics["n_material_negative_eigenvalues"])
print(
    diagnostics["n_laplace_produced"],
    diagnostics["n_laplace_requested"],
)

laplace_valid=False means that the negative-log-posterior Hessian has a materially negative direction. In that case the reported Laplace uncertainty does not describe a well-defined local maximum. hessian_regularized records whether small or negative eigenvalues were floored before drawing, and cholesky_jitter records numerical stabilization used by the covariance factorization.

Nested sampling#

Dynesty and UltraNest store their evidence results on the sample container:

import numpy as np

samples = sampler_result.samples
log_evidence = np.asarray(samples.logz).reshape(-1)[-1]
log_evidence_error = np.asarray(samples.logzerr).reshape(-1)[-1]

The complete Dynesty-compatible result remains available through samples. The evidence is primarily useful for comparing covariance models; posterior plots are still needed to identify unconstrained parameters or separated modes.

NUTS#

NUTS records both an overall status and per-parameter diagnostics:

diagnostics = sampler_result.diagnostics
print("Converged:", diagnostics["converged"])
print("Divergences:", diagnostics["divergences"])

for name in sampler_result.get_parameter_names():
    print(
        name,
        diagnostics["n_eff"][name],
        diagnostics["r_hat"][name],
    )

r_hat measures agreement between chains and n_eff estimates the effective number of independent draws. A single-chain split \(\hat R\) is a weaker test than agreement between multiple independent chains. Divergences indicate that NUTS could not reliably follow part of the posterior geometry.

Assessing the covariance model#

Sampler diagnostics assess the numerical inference. Component reconstructions, residuals, and posterior predictive power spectra assess whether the covariance model represents the data and whether competing components are identifiable. Numerical convergence alone cannot establish that physical separation.

The complete script used to simulate the data, run MCMC, and generate these figures and the posterior figures is included below.

Show the simulated ML-GPR script

"""Generate the ML-GPR MCMC figures used in the user guide."""

from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import torch

from ps_eor import pspec, psutil
from ps_eor.ml_gpr import MLGPRConfigFile, MLGPRForegroundFitter
from ps_eor.ml_gpr.kernels import UVScaledKernel
from ps_eor.simu import SimuMultiGPCube

CONFIG = """
sampler_method = "mcmc"

[mcmc]
n_walkers = 50
n_steps = 200
n_burn = 100
move = "kde"

[gp]
use_uv_weight = false

[kern]
uv_bins_du = 40
fg = ["foreground"]
eor = ["signal"]

[foreground]
type = "MRBF"
variance.prior = "Log10Uniform(-0.2, 0.2)"
variance.log_scale = true
lengthscale.prior = "Uniform(2, 16)"
var_alpha.prior = "Fixed(0)"
ls_alpha.prior = "Fixed(0)"

[signal]
type = "MExponential"
variance.prior = "Log10Uniform(-5, -2)"
variance.log_scale = true
lengthscale.prior = "Uniform(0.1, 2)"
var_alpha.prior = "Fixed(0)"
ls_alpha.prior = "Fixed(0)"

[kern.noise]
alpha.prior = "Fixed(1)"
"""


def simulate():
    """Return simulated data and its foreground, signal, and noise parts."""

    freqs = np.arange(122, 132, 0.2) * 1e6
    image_fov_deg = 4
    image_umax = 1000
    image_res = 1 / (2 * image_umax)
    n_pixels = round(np.deg2rad(image_fov_deg) / image_res)

    simulation = SimuMultiGPCube.new(
        res=image_res,
        n_pix=n_pixels,
        freqs=freqs,
        umin=20,
        umax=250,
        uv_bins_du=40,
    )
    foreground = simulation.get_from_kern(
        UVScaledKernel(
            family="rbf",
            variance=1,
            lengthscale=10,
            use_uv_ps=False,
        )
    )
    signal = simulation.get_from_kern(
        UVScaledKernel(
            family="exponential",
            variance=1e-4,
            lengthscale=0.7,
            use_uv_ps=False,
        )
    )
    noise = simulation.get_noise(1e-3)
    return foreground + signal + noise, foreground, signal, noise


def save_sampler_figures(sampler_result, output_dir):
    """Save the MCMC trace and posterior corner plots."""

    fig = sampler_result.plot_samples()
    fig.savefig(output_dir / "ml_gpr_mcmc_trace.png", dpi=160)
    plt.close(fig)

    fig = sampler_result.plot_corner()
    fig.savefig(
        output_dir / "ml_gpr_mcmc_corner.png",
        dpi=160,
        bbox_inches="tight",
    )
    plt.close(fig)


def save_component_figure(
    sampler_result,
    data,
    foreground,
    signal,
    output_dir,
):
    """Compare true and posterior component spectra for one UV cell."""

    foreground_draws = np.stack(
        [
            cube.data.real
            for cube in sampler_result.generate_data_cubes(
                30,
                kern_name="fg*",
            )
        ]
    )
    signal_draws = np.stack(
        [
            cube.data.real
            for cube in sampler_result.generate_data_cubes(
                30,
                kern_name="eor*",
            )
        ]
    )

    mode = np.argsort(data.ru)[len(data.ru) // 2]
    frequency = data.freqs / 1e6

    fig, axes = plt.subplots(
        1,
        2,
        figsize=(10, 3.8),
        sharex=True,
        layout="compressed",
    )
    for ax, truth, draws, title in (
        (axes[0], foreground, foreground_draws, "Foreground component"),
        (axes[1], signal, signal_draws, "21-cm-like component"),
    ):
        q16, median, q84 = np.quantile(
            draws[:, :, mode],
            [0.16, 0.5, 0.84],
            axis=0,
        )
        ax.plot(frequency, truth.data.real[:, mode], label="Input", color="C1")
        ax.plot(frequency, median, label="Posterior median", color="C0")
        ax.fill_between(
            frequency,
            q16,
            q84,
            color="C0",
            alpha=0.25,
            label="68% interval",
        )
        ax.set_title(title)
        ax.set_xlabel("Frequency [MHz]")
        ax.grid(alpha=0.25)

    axes[0].set_ylabel("Real visibility [K]")
    axes[0].legend()
    fig.savefig(output_dir / "ml_gpr_posterior_components.png", dpi=160)
    plt.close(fig)


def save_power_spectrum_figure(
    sampler_result,
    signal,
    noise,
    output_dir,
):
    """Compare the input and posterior 21-cm-like power spectra."""

    ps_gen = pspec.PowerSpectraBuilder().get(
        signal,
        fmhz_range=(122, 132),
        umin=20,
        umax=250,
        du=10,
        rmean_freqs=False,
        window_fct="blackmanharris",
        ps2d_pos_only=False,
        ft_method="nudft",
    )
    # Start above the sparsely populated lowest modes so every spherical bin
    # contains data in this deliberately small simulation.
    kbins = np.logspace(np.log10(0.05), np.log10(1), 7)
    recovered = sampler_result.get_ps_stack(
        ps_gen,
        kbins,
        n_pick=30,
        kern_name="eor*",
    )

    fig, axes = plt.subplots(
        1,
        2,
        figsize=(8.5, 3.7),
        layout="compressed",
    )
    for cube, label, color in (
        (noise, "Noise", "C0"),
        (signal, "21-cm-like input", "C1"),
    ):
        ps_gen.get_variance(cube).plot(
            ax=axes[0],
            label=label,
            c=color,
        )
        ps_gen.get_ps3d(kbins, cube).plot(
            ax=axes[1],
            label=label,
            c=color,
            nsigma=0,
        )

    recovered.get_variance().plot(
        ax=axes[0],
        label="21-cm-like posterior",
        c=psutil.black,
    )
    recovered.get_ps3d().plot(
        ax=axes[1],
        label="21-cm-like posterior",
        c=psutil.black,
    )
    fig.legend(
        *axes[0].get_legend_handles_labels(),
        loc="upper center",
        ncol=3,
        bbox_to_anchor=(0.5, 1.08),
    )
    fig.savefig(
        output_dir / "ml_gpr_posterior_power_spectra.png",
        dpi=160,
        bbox_inches="tight",
    )
    plt.close(fig)


def make_figures(output_dir):
    """Simulate a small data set, fit it with MCMC, and save all figures."""

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(42)

    data, foreground, signal, noise = simulate()
    config = MLGPRConfigFile.load_from_string_with_defaults(CONFIG)
    result = MLGPRForegroundFitter(config).run(data, noise)
    sampler_result = result.sampler_result

    save_sampler_figures(sampler_result, output_dir)
    save_component_figure(
        sampler_result,
        data,
        foreground,
        signal,
        output_dir,
    )
    save_power_spectrum_figure(
        sampler_result,
        signal,
        noise,
        output_dir,
    )


if __name__ == "__main__":
    make_figures(Path(__file__).resolve().parents[1] / "_static")