Running an ML-GPR fit

There are two ways to run ML-GPR. The configuration-driven interface provides the complete workflow in a reproducible TOML file and is the best choice for most users. The advanced interface exposes the kernels, data preparation, regressor, and sampler directly when more control is needed. Choose the one that best fits the analysis; both use the same underlying model and inference classes.

Configuration-driven interface

The high-level interface runs a complete analysis from a TOML configuration. A minimal file using the packaged defaults might be:

ml_gpr.toml

kern.uv_bins_du = 20
kern.fg = ["fg_int"]
kern.eor = ["eor"]
gp.use_uv_weight = true

sampler_method = "mcmc"
mcmc.n_steps = 500
mcmc.n_walkers = 50
mcmc.n_burn = 300

[fg_int]
type = "MRBF"
variance.prior = "Uniform(0.8, 1.2)"
lengthscale.prior = "Uniform(2, 60)"
ls_alpha.prior = "Fixed(0)"
var_alpha.prior = "Fixed(0)"

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

[kern.noise]
alpha.prior = "Uniform(0.9, 1.6)"
use_simulated_noise_cube = false
use_sefd_freqs_estimate = false
use_noise_std = false
estimate_baseline_noise_from_stokes = "V"
estimate_baseline_noise_remove_n_pca = 0
estimate_baseline_noise_from_channel_diff = false
scale_baseline_noise = false

The following executes that configuration. Kernel types, component composition, and priors are described in Building the covariance model; sampler choices and their settings are described in Choosing an inference 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)

The TOML configuration selects the covariance components and priors, UV binning, noise treatment, and inference method. The returned result keeps the configuration, posterior samples, input data, and sampler diagnostics together.

The data and noise cubes must cover compatible frequency and UV grids. process_noise_cube converts the supplied noise information into the form selected by the configuration. Depending on that configuration, it can apply differencing, estimate the noise level, scale it with baseline length, or simulate a noise realization for the fit.

After fitting, see Posterior components and power spectra for reconstructing components, residuals, and posterior power spectra from result.

Advanced interface

The covariance components can be constructed directly in Python, without a configuration file. The following kernels and priors reproduce the model in the TOML example above. Building the covariance model gives the available kernels, prior types, UV-dependent parameterizations, and composition rules.

from ps_eor.ml_gpr.kernels import (
    UVScaledKernel,
    WhiteHeteroscedasticKernel,
)
from ps_eor.ml_gpr.params import Components
from ps_eor.ml_gpr.priors import Log10UniformPrior, UniformPrior

foreground = UVScaledKernel(
    "rbf", variance=1, lengthscale=10, var_alpha=0, ls_alpha=0)
foreground.set_free("variance", UniformPrior(0.8, 1.2))
foreground.set_free("lengthscale", UniformPrior(2, 60))

signal = UVScaledKernel(
    "exponential", variance=1e-3, lengthscale=1, var_alpha=0, ls_alpha=0)
signal.set_free(
    "variance", Log10UniformPrior(-5, -1), log_scale=True)
signal.set_free("lengthscale", UniformPrior(0.2, 2))

noise = WhiteHeteroscedasticKernel(alpha=1)
noise.set_free("alpha", UniformPrior(0.9, 1.6))

components = Components({
    "fg_int": foreground,
    "eor": signal,
    "noise": noise,
})

When a configuration is already available, the equivalent component construction is simply:

components = config.get_kern()

The remaining steps prepare the binned data, build the regressor, and run the sampler. Here noise_cube is used directly because the example TOML disables additional noise preprocessing:

from ps_eor.ml_gpr.multidata import MultiData
from ps_eor.ml_gpr.regressor import MultiGPRegressor
from ps_eor.ml_gpr.samplers import MCMCSampler

multi_data = MultiData(
    data_cube,
    noise_cube=noise_cube,
    uv_bins_du=20,
    uv_bins_n_uni=0,
)
gp = MultiGPRegressor(
    multi_data,
    components,
    use_uv_weight=True,
)
sampler = MCMCSampler(
    gp,
    n_walkers=50,
    emcee_moves="kde",
)
samples = sampler.run(
    n_steps=500,
    n_burn=300,
)

The data and regression API explains MultiData and MultiGPRegressor in more detail. For choosing a sampler, checking convergence, and comparing inference methods, see Choosing an inference method.

The advanced interface is an alternative to configuration-driven fitting. It gives direct control over each step: preparing the data representation, constructing covariance components, creating the regressor, running inference, and producing posterior results. The numbered tutorials introduce these pieces before showing the configuration-driven interface; see Tutorial notebooks for the notebook sequence.