nifty.re.multi_grid.matern module

class FourierIntegralGrid(mode_lengths: jax.Array, mode_binbounds: jax.Array, min_dist: float, max_dist: float, num: int, ndim: int, weights: jax.Array)

Bases: object

__init__(mode_lengths: Array, mode_binbounds: Array, min_dist: float, max_dist: float, num: int, ndim: int, weights: Array) → None

max_dist: float

min_dist: float

mode_binbounds: Array

mode_lengths: Array

ndim: int

num: int

weights: Array

class IsotropicPowerSpectrumTransform(Ndim: int, Nint: int | str = 1024, h: float | str = 'auto')

Bases: object

Transform between isotropic power spectrum P(k) and covariance kernel \(\mathrm{Cov}(r)\) in \(N_\mathrm{dim}\) dimensions.

The general formula is:

\[\mathrm{Cov}(r) = \frac{1}{(2 \pi)^{N_\mathrm{dim}/2}} \int_0^\infty P(k) \, k^{(N_\mathrm{dim}-1)} \frac{J_\nu(k r)}{(k r)^\nu} dk\]

\[= \frac{1}{(2 \pi)^{N_\mathrm{dim}/2}} \int_0^\infty P\left(\frac{x}{r}\right) r^{-N_\mathrm{dim}} x^{N_\mathrm{dim}/2} J_\nu(x) dx\]

with \(x=kr\) and \(\nu = (N_\mathrm{dim}-2)/2\).

We use a modified version of Ogata quadrature (eqn. 5.2 in https://ems.press/journals/prims/articles/2319) to compute the integral.

The initialisation precomputes the Ogata quadrature nodes and weights for given \(N_\mathrm{dim}\). This needs to be done once at initialisation and depends on scipy. The application uses pure JAX for differentiability.

Some advice on modelling power spectra: Without a high-k cutoff, it can quickly happen that the covariance function diverges. We therefore recommend to use power spectra that either decay sufficiently fast at high k, or have compact support.

Parameters:

• Ndim (int) – Number of spatial dimensions

• Nint (int) – Number of integration nodes for Ogata quadrature. Chosen automatically if set to "auto" and h is provided.

• h (float) – Step-size. Controls accuracy. Chosen automatically if set to "auto" and Nint is provided.

__init__(Ndim: int, Nint: int | str = 1024, h: float | str = 'auto')

k_range_at_r(r)

Return (k_min, k_max) as functions of r.

pk_to_cov(P_func: Callable[[Array], Array], r: float | Array) → Array

Compute covariance kernel from power spectrum.

Parameters:

• P_func (callable) – Power spectrum P(k), must accept JAX arrays

• r (array) – Separation values at which to evaluate the covariance kernel

Returns:

cov – Covariance kernel values

Return type:

array

class MaternCovarianceKernel(Ndim: int, r_min: float, r_max: float, Ninterp: int | str = 'auto', jitter: float = 1e-05, Nint: int = 1024, h: float | str = 'auto', enforce_nonnegativity: bool = True, enforce_monotonicity: bool = True)

Bases: IsotropicPowerSpectrumTransform

Computes the covariance kernel for a Matern-like power spectrum with a high-k cutoff. The power spectrum is given by:

\[P(k) = \frac{1}{\left(1 + \left(\alpha k\right)^2\right)^{\beta / 2}} \cdot e^{-(k / k_\mathrm{cutoff})^2}\]

where \(\beta\) (negloglogslope) controls the logarithmic slope at high k, \(\alpha\) (lengthscale) controls the correlation length scale, and kcutoff sets the high-k cutoff.

Parameters:

• Ndim (int) – Number of spatial dimensions

• r_min (float) – Minimum r value for covariance kernel interpolation. This should correspond to the smallest length scale you want to resolve.

• r_max (float) – Maximum r value for covariance kernel interpolation. This should correspond to the largest length scale you want to resolve.

• Ninterp (int or "auto") – Number of interpolation/evaluation points for the covariance kernel. If "auto", it is chosen as 100 points per decade between r_min and r_max.

• jitter (float) – Small value added to the kernel at r=0 for numerical stability. Increase this if you encounter numerical issues in Gaussian process inference.

• Nint (int) – Number of integration nodes for Ogata quadrature. Chosen automatically if set to "auto" and h is provided.

• h (float or "auto") – Step-size. Controls accuracy. Chosen automatically if set to "auto" and Nint is provided.

• enforce_nonnegativity (bool) – If True, the covariance values are clipped to be non-negative. This is not mathematically necessary but can help with numerical stability, especially when used together with enforce_monotonicity.

• enforce_monotonicity (bool) – If True, the covariance values are post-processed to ensure they are non-increasing with r. This is not mathematically necessary but can help with numerical stability, especially when using ICR or graphgp.

static Pk(k, *, lengthscale, negloglogslope, kcutoff)

Matern-like power spectrum with high-k cutoff:

\[P(k) = \frac{1}{\left(1 + \left(\alpha k\right)^2\right)^{\beta / 2}} \cdot e^{-(k / k_\mathrm{cutoff})^2}\]

Parameters:

• k (array) – Wavenumber values

• lengthscale (float) – Length scale parameter of the Matern-like power spectrum

• negloglogslope (float) – Log-log slope parameter of the Matern-like power spectrum

• kcutoff (float) – High-k cutoff parameter of the Matern-like power spectrum

Returns:

P(k) – Power spectrum values

Return type:

array

__init__(Ndim: int, r_min: float, r_max: float, Ninterp: int | str = 'auto', jitter: float = 1e-05, Nint: int = 1024, h: float | str = 'auto', enforce_nonnegativity: bool = True, enforce_monotonicity: bool = True)

get_covariance_kernel_interpolator(variance, lengthscale, negloglogslope, kcutoff)

Returns a function cov_fn(r) that interpolates the covariance kernel for given Matern-like power spectrum parameters.

Note: The covariance is normalised by cov(0) = variance at r = 0. Afterwards, jitter is added at r = 0 for numerical stability.

Parameters:

• variance (float) – Scaling factor defining cov(0) = variance * (1 + jitter)

• lengthscale (float) – Length scale parameter of the Matern-like power spectrum

• negloglogslope (float) – Log-log slope parameter of the Matern-like power spectrum

• kcutoff (float) – High-k cutoff parameter of the Matern-like power spectrum

Returns:

cov_fn – Function that takes r values and returns linearly interpolated covariance values

Return type:

callable

get_covariance_kernel_pair(variance, lengthscale, negloglogslope, kcutoff)

Returns (cov_rs, cov_vals) where cov_rs are the r values (including r=0) and cov_vals are the corresponding covariance values.

Note: The covariance is normalised by cov(0) = variance. Afterwards, jitter is added at r = 0 for numerical stability.

Parameters:

• variance (float) – Scaling factor defining cov(0) = variance * (1 + jitter)

• lengthscale (float) – Length scale parameter of the Matern-like power spectrum

• negloglogslope (float) – Log-log slope parameter of the Matern-like power spectrum

• kcutoff (float) – High-k cutoff parameter of the Matern-like power spectrum

Returns:

• cov_rs (array) – r values including r=0

• cov_vals (array) – Corresponding covariance values

class MaternCovarianceModel(Ndim: int, r_min: float, r_max: float, variance: float | tuple | Model, lengthscale: float | tuple | Model, negloglogslope: float | tuple | Model, kcutoff: float | tuple | Model | str = 'auto', Ninterp: int | str = 'auto', jitter: float = 1e-05, Nint: int | str = 1024, h: float | str = 'auto', prefix: str = 'MaternCovarianceModel_', mode: str = 'ICR', enforce_nonnegativity: bool = True, enforce_monotonicity: bool = True)

Bases: MaternCovarianceKernel, Model

NIFTy-wrapper for the Matern covariance kernel. Computes the covariance kernel for a Matern-like power spectrum with a high-k cutoff. The power spectrum is given by:

\[P(k) = \frac{1}{\left(1 + \left(\alpha k\right)^2\right)^{\beta / 2}} \cdot e^{-(k / k_\mathrm{cutoff})^2}\]

where \(\beta\) (negloglogslope) controls the logarithmic slope at high k, \(\alpha\) (lengthscale) controls the correlation length scale, and \(k_\mathrm{cutoff}\) sets the high-k cutoff.

Parameters:

• Ndim (int) – Number of spatial dimensions

• r_min (float) – Minimum r value for covariance kernel interpolation. This should correspond to the smallest length scale you want to resolve.

• r_max (float) – Maximum r value for covariance kernel interpolation. This should correspond to the largest length scale you want to resolve.

• variance (float or tuple or Model) – This sets the variance of the Gaussian process, i.e. \(\mathrm{Cov}(0) = variance * (1 + jitter)\). If a float is provided, it is treated as a constant. If a tuple (mean, stddev) is provided, it is treated as a log-normal prior. If a Model is provided, it is used directly and must have a scalar target.

• lengthscale (float or tuple or Model) – This sets roughly the largest coherent structure size in the Gaussian process. If a float is provided, it is treated as a constant. If a tuple (mean, stddev) is provided, it is treated as a log-normal prior. If a Model is provided, it is used directly and must have a scalar target.

• negloglogslope (float or tuple or Model) – This controls the logarithmic slope of the power spectrum at intermediate to high k. If a float is provided, it is treated as a constant. If a tuple (mean, stddev) is provided, it is treated as a log-normal prior. If a Model is provided, it is used directly and must have a scalar target.

• kcutoff (float or tuple or Model or "auto") – This sets the high-k cutoff of the power spectrum. If a float is provided, it is treated as a constant. If a tuple (mean, stddev) is provided, it is treated as a log-normal prior. If a Model is provided, it is used directly and must have a scalar target. If "auto", it is set to \(\pi/r_\mathrm{min}\).

• Ninterp (int or "auto") – Number of interpolation/evaluation points for the covariance kernel. If "auto", it is chosen as 100 points per decade between r_min and r_max.

• jitter (float) – Small value added to the kernel at r=0 for numerical stability. Increase this if you encounter numerical issues in Gaussian process inference.

• Nint (int or "auto") – Number of integration nodes for Ogata quadrature. Chosen automatically if set to "auto" and h is provided.

• h (float or "auto") – Step-size. Controls accuracy. Chosen automatically if set to "auto" and Nint is provided.

• prefix (str) – Prefix for the names of the priors if tuple inputs are provided. This is useful to distinguish multiple instances of MaternCovarianceModel. The final prior names will be prefix + attribute name, e.g. "MaternCovarianceModel_lengthscale".

• mode (str) – Mode of operation. If "ICR", the model is intended to be used with ICR and the __call__ method will return a callable covariance kernel function. If "graphgp", the model is intended to be used with GraphGP and the __call__ method will return a tuple (cov_rs, cov_vals) with the r values and corresponding covariance values for interpolation.

• enforce_nonnegativity (bool) – If True, the covariance values are clipped to be non-negative. This is not mathematically necessary but can help with numerical stability, especially when used together with enforce_monotonicity.

• enforce_monotonicity (bool) – If True, the covariance values are post-processed to ensure they are non-increasing with r. This is not mathematically necessary but can help with numerical stability, especially when using ICR or graphgp.

__init__(Ndim: int, r_min: float, r_max: float, variance: float | tuple | Model, lengthscale: float | tuple | Model, negloglogslope: float | tuple | Model, kcutoff: float | tuple | Model | str = 'auto', Ninterp: int | str = 'auto', jitter: float = 1e-05, Nint: int | str = 1024, h: float | str = 'auto', prefix: str = 'MaternCovarianceModel_', mode: str = 'ICR', enforce_nonnegativity: bool = True, enforce_monotonicity: bool = True)

cov_rs_no_zero: Array

kcutoff: Callable

lengthscale: Callable

negloglogslope: Callable

nodes: Array

variance: Callable

weights: Array

class MaternHarmonicCovariance(scale: tuple | Callable | float, cutoff: tuple | nifty.re.model.Model | float, loglogslope: tuple | Callable | float, *, ndim: int, n_integrate=2048, n_interpolate=512, integration_dists_min_max=None, interpolation_dists_min_max=None, kind: str = 'amplitude', prefix: str = '')

Bases: Model

__init__(scale: tuple | Callable | float, cutoff: tuple | Model | float, loglogslope: tuple | Callable | float, *, ndim: int, n_integrate=2048, n_interpolate=512, integration_dists_min_max=None, interpolation_dists_min_max=None, kind: str = 'amplitude', prefix: str = '')

Compute the Matérn-kernel using its harmonic representation.

See also

Causal, Bayesian, & non-parametric modeling of the SARS-CoV-2 viral load vs. patient’s age, Guardiani, Matteo and Frank, Philipp and Kostić, Andrija and Edenhofer, Gordian and Roth, Jakob and Uhlmann, Berit and Enßlin, Torsten, https://arxiv.org/abs/2105.13483 https://doi.org/10.1371/journal.pone.0275011

cutoff: Model | float

loglogslope: Model | float

normalized_spectrum(x)

scale: Model | float

default_Nint(h, nu, t_max=2.0)

default_Ninterp(r_min, r_max, num_per_decade=200)

default_h(Nint, nu, t_max=2.0)

get_bessel_zeros(nu, N_zeros, n_iter=3)

Compute the first \(N_\mathrm{zeros}\) zeros of \(\mathcal{J}_\nu(x)\) (Bessel function of the first kind).

Parameters:

• nu (float) – Order of the Bessel function

• N_zeros (int) – Number of zeros to compute

• n_iter (int) – Number of Newton-Raphson refinement iterations

Returns:

zeros – First \(N_\mathrm{zeros}\) zeros of \(\mathcal{J}_\nu(x)\)

Return type:

array

j1(x)

Bessel function of order one - using the implementation from CEPHES, translated to Jax.

Pure JAX implementation for Bessel function from benjaminpope/sibylla .

make_integral_grid(min_dist: float, max_dist: float, num: int, ndim: int, normalize: bool = True) → FourierIntegralGrid

spectrum2covariance(fig: FourierIntegralGrid, spec: NDArray, *, ref_distance=1.0, normalize=True)