Kernels and time covariance#

Covariance kernels used by ML-GPR.

The module provides baseline-dependent stationary kernels, measured-noise kernels, VAE-based 21-cm kernels, and time covariance for multi-epoch fits. Kernels may be combined with + or * after configuring their parameters.

ps_eor.ml_gpr.kernels.make_kernel(name, batch_shape=())[source]#

Create a stock GPyTorch kernel, optionally batched over UV bins.

class ps_eor.ml_gpr.kernels.UVScaledKernel(family, variance=1.0, lengthscale=1.0, var_alpha=0.0, ls_alpha=0.0, theta_rad=0.1, delay_buffer_us=0.0, power=2.0, wedge_parametrization=False, ls_is_period=False, latitude_deg=90.0, uv_min=None, uv_max=None, use_uv_ps=False, l_max=100000000.0, uv_break=0.1, **kwargs)[source]#

Bases: Parameterized, Kernel

A stationary frequency kernel whose scale varies with baseline.

Per uv bin, variance and lengthscale follow power laws in baseline length:

variance(u)    = variance * (u / u_min) ** var_alpha
lengthscale(u) = lengthscale / (1 + ls_alpha * 1e-3 * lengthscale * (u - u_min))

Setting both exponents to zero gives the same kernel in every UV bin.

Parameters:

  • family – one of FAMILY_TO_BASE_KERNEL (‘rbf’, ‘matern32’, ‘matern52’, ‘exponential’, ‘cosine’, ‘ratquad’).

  • variance – Covariance amplitude at the shortest baseline, in the normalized-data units used internally by the regressor.

  • lengthscale – Frequency correlation scale at the shortest baseline, in MHz.

  • var_alpha – how each scales with baseline.

  • ls_alpha – how each scales with baseline.

  • power – rational-quadratic shape parameter; ignored otherwise.

  • wedge_parametrization – Derive the lengthscale from theta_rad and delay_buffer_us instead of fitting it directly.

  • theta_rad – Wedge angle in radians.

  • delay_buffer_us – Additional wedge delay in microseconds.

  • use_uv_ps – take the variance law from the data’s own measured angular power spectrum instead of var_alpha.

  • uv_min – Zero the kernel outside this baseline range, in wavelengths.

  • uv_max – Zero the kernel outside this baseline range, in wavelengths.

  • l_max – Cap on the per-bin lengthscale, in MHz.

A MultiGPRegressor binds the kernel to its UV bins before evaluation.

bind_data(uv_bins, mean_fmhz, uv_ps_fct=None)[source]#

Bind the kernel to the data’s UV geometry.

uv_ps_fct is required when use_uv_ps=True. Called by MultiGPRegressor.

forward(x1, x2, diag=False, **params)[source]#

Computes the covariance between \(\mathbf x_1\) and \(\mathbf x_2\). This method should be implemented by all Kernel subclasses.

Parameters:

  • x1 – First set of data (… x N x D).

  • x2 – Second set of data (… x M x D).

  • diag – Should the Kernel compute the whole kernel, or just the diag? If True, it must be the case that x1 == x2. (Default: False.)

  • last_dim_is_batch – If True, treat the last dimension of x1 and x2 as another batch dimension. (Useful for additive structure over the dimensions). (Default: False.)

Returns:

The kernel matrix or vector. The shape depends on the kernel’s evaluation mode:

  • full_covar: … x N x M

  • full_covar with last_dim_is_batch=True: … x K x N x M

  • diag: … x N

  • diag with last_dim_is_batch=True: … x K x N

class ps_eor.ml_gpr.kernels.WhiteHeteroscedasticKernel(alpha=1.0, **kwargs)[source]#

Bases: Parameterized, Kernel

Diagonal noise covariance measured from a noise cube.

bind_data() supplies the variance per frequency and UV bin. alpha scales that measured level. Noise remains independent between frequencies and, for multi-epoch data, between epochs.

bind_data(freq_grid, per_bin_variance)[source]#

Bind measured noise variance to the frequency grid.

per_bin_variance has shape (n_bins, n_freq) or (n_bins, n_epochs, n_freq).

forward(x1, x2, diag=False, **params)[source]#

Computes the covariance between \(\mathbf x_1\) and \(\mathbf x_2\). This method should be implemented by all Kernel subclasses.

Parameters:

  • x1 – First set of data (… x N x D).

  • x2 – Second set of data (… x M x D).

  • diag – Should the Kernel compute the whole kernel, or just the diag? If True, it must be the case that x1 == x2. (Default: False.)

  • last_dim_is_batch – If True, treat the last dimension of x1 and x2 as another batch dimension. (Useful for additive structure over the dimensions). (Default: False.)

Returns:

The kernel matrix or vector. The shape depends on the kernel’s evaluation mode:

  • full_covar: … x N x M

  • full_covar with last_dim_is_batch=True: … x K x N x M

  • diag: … x N

  • diag with last_dim_is_batch=True: … x K x N

class ps_eor.ml_gpr.kernels.VAEKernTorch(vae_fitter, variance=1.0, **kwargs)[source]#

Bases: Parameterized, Kernel

A 21-cm signal kernel driven by a trained VAE decoder.

Latent parameters x1, x2, ... are decoded into a power spectrum and converted to frequency covariance. variance scales the result.

Parameters:

vae_fitter – trained VAEFitter or VAEFitterPreProc. The kernel obtains the decoder, training k grid, latent dimension, and preprocessing directly from the fitter.

Decoder weights may be float32 or float64 independently of the fixed float64 precision used by ML-GPR.

ps_eor.ml_gpr.config can load the fitter from a saved file.

bind_data(uv_bins)[source]#

Bind the kernel to the data’s UV bins.

forward(x1, x2, diag=False, **params)[source]#

Computes the covariance between \(\mathbf x_1\) and \(\mathbf x_2\). This method should be implemented by all Kernel subclasses.

Parameters:

  • x1 – First set of data (… x N x D).

  • x2 – Second set of data (… x M x D).

  • diag – Should the Kernel compute the whole kernel, or just the diag? If True, it must be the case that x1 == x2. (Default: False.)

  • last_dim_is_batch – If True, treat the last dimension of x1 and x2 as another batch dimension. (Useful for additive structure over the dimensions). (Default: False.)

Returns:

The kernel matrix or vector. The shape depends on the kernel’s evaluation mode:

  • full_covar: … x N x M

  • full_covar with last_dim_is_batch=True: … x K x N x M

  • diag: … x N

  • diag with last_dim_is_batch=True: … x K x N

class ps_eor.ml_gpr.kernels.TimeCovariance[source]#

Bases: Parameterized, Module

Base class for correlation between observing epochs.

Subclasses return a unit-diagonal time correlation matrix; component amplitude remains in the frequency kernel.

matrix(times1, times2=None)[source]#

Correlation matrix between the given epoch times, unit-diagonal.

class ps_eor.ml_gpr.kernels.IndependentTimeCovariance[source]#

Bases: TimeCovariance

Treat each epoch as an independent realization.

matrix(times1, times2=None)[source]#

Correlation matrix between the given epoch times, unit-diagonal.

class ps_eor.ml_gpr.kernels.CoherentTimeCovariance[source]#

Bases: TimeCovariance

Treat a component as identical in every epoch.

matrix(times1, times2=None)[source]#

Correlation matrix between the given epoch times, unit-diagonal.

class ps_eor.ml_gpr.kernels.SharedFractionTimeCovariance(rho=0.5)[source]#

Bases: TimeCovariance

Share a fraction rho of the variance between every epoch.

The remaining fraction is an independent realization in each epoch. The total covariance is rho * J + (1 - rho) * I.

part_matrix(times1, times2=None, part='total')[source]#

Time covariance of the total, coherent, or independent process.

matrix(times1, times2=None)[source]#

Correlation matrix between the given epoch times, unit-diagonal.

class ps_eor.ml_gpr.kernels.ExponentialTimeCovariance(tau=1.0)[source]#

Bases: TimeCovariance

The component decorrelates over a coherence time: B[i,j] = exp(-|t_i - t_j| / tau).

tau uses the same units as the epoch times and may be fitted.

matrix(times1, times2=None)[source]#

Correlation matrix between the given epoch times, unit-diagonal.

class ps_eor.ml_gpr.kernels.TimeKron(freq_kern, time_cov, **kwargs)[source]#

Bases: Parameterized, Kernel

A component with covariance in time as well as frequency.

Combines a frequency kernel with a TimeCovariance as K = B_time ⊗ K_freq. For example:

Components({
    'fg':      TimeKron(k_fg,  CoherentTimeCovariance()),
    'sys':     TimeKron(k_sys, ExponentialTimeCovariance(tau=1.0)),
    'eor':     TimeKron(k_eor, CoherentTimeCovariance()),
    'noise':   WhiteHeteroscedasticKernel(),
})

set_free() and set_fixed() route parameters to the frequency kernel or time covariance, so the component can be configured before or after wrapping. SharedFractionTimeCovariance components expose their latent coherent and independent processes through part().

set_free(name, prior, log_scale=False)[source]#

Sample a frequency or time parameter under prior.

set_fixed(name, value)[source]#

Fix a frequency or time parameter to value.

free_param_names()[source]#

Names of the parameters currently being sampled.

flat_param_entries(prefix)[source]#

ParamEntry per free parameter, named f'{prefix}.{name}'.

Raises if a parameter was left neither free nor fixed.

describe()[source]#

Describe every candidate parameter as free or fixed.

part(name)[source]#

Prediction view of the coherent or independent latent process.

forward(x1, x2, diag=False, **params)[source]#

Computes the covariance between \(\mathbf x_1\) and \(\mathbf x_2\). This method should be implemented by all Kernel subclasses.

Parameters:

  • x1 – First set of data (… x N x D).

  • x2 – Second set of data (… x M x D).

  • diag – Should the Kernel compute the whole kernel, or just the diag? If True, it must be the case that x1 == x2. (Default: False.)

  • last_dim_is_batch – If True, treat the last dimension of x1 and x2 as another batch dimension. (Useful for additive structure over the dimensions). (Default: False.)

Returns:

The kernel matrix or vector. The shape depends on the kernel’s evaluation mode:

  • full_covar: … x N x M

  • full_covar with last_dim_is_batch=True: … x K x N x M

  • diag: … x N

  • diag with last_dim_is_batch=True: … x K x N

class ps_eor.ml_gpr.kernels.CombinedAdditiveKern(*kernels)[source]#

Bases: _CombinedKern, AdditiveKernel

A sum of kernels that retains parameter metadata.

class ps_eor.ml_gpr.kernels.CombinedProductKern(*kernels)[source]#

Bases: _CombinedKern, ProductKernel

A product of kernels that retains parameter metadata.

ps_eor.ml_gpr.kernels.combine_kernels(kernels, operation='add')[source]#

Combine kernels while preserving sampleable-parameter metadata.

Parameters:

  • kernels – ordered iterable of GPyTorch kernels. A single kernel is returned unchanged.

  • operation – 'add' or 'multiply' ('+', 'x', '*' and '.' are accepted aliases).

Nested combinations using the same operation are flattened.

Prior distributions shared by all ML-GPR inference methods.

Priors operate on natural kernel values. site_* methods provide the parameterization required by gradient samplers; stored-chain transformations are handled by ps_eor.ml_gpr.params.

class ps_eor.ml_gpr.priors.UniformPrior(lower, upper)[source]#

Bases: object

Uniform prior over [lower, upper].

log_prob_torch(x)[source]#

Return the differentiable log density inside the support.

Bounds are enforced separately by the MAP optimizer.

site_dist(dtype=None)[source]#

Distribution a gradient sampler draws the site from, as a torch.distributions object.

class ps_eor.ml_gpr.priors.Log10UniformPrior(lower, upper)[source]#

Bases: object

Prior uniform in log10(x) between lower and upper.

For example, Log10Uniform(-8, 0) covers natural values from 1e-8 to 1.

log_prob_torch(x)[source]#

Return the differentiable natural-space log density.

site_dist(dtype=None)[source]#

Return a uniform distribution over log10(natural).

class ps_eor.ml_gpr.priors.GaussianPrior(mu, sigma)[source]#

Bases: object

Normal prior with mean mu and standard deviation sigma.

ps_eor.ml_gpr.priors.make_prior(name, *args)[source]#

Build a prior from its TOML name and arguments.

make_prior('Log10Uniform', -2, 2) -> Log10UniformPrior(-2, 2). Fixed is handled by Parameterized.