eztaox.kernels.quasisep#
Quasiseparable kernels.
Scalable kernels exploiting the quasiseparable structure in the relevant matrices to achieve a O(N) scaling.
This module extends the tinygp.kernels.quasisep module.
Classes#
An extension of the tinygp.kernels.quasisep.Quasisep kernel. |
|
A helper to represent the sum of two quasiseparable kernels. |
|
A helper to represent the product of two quasiseparable kernels. |
|
The product of a scalar and a quasiseparable kernel. |
|
Extends the tinygp.kernels.quasisep.Exp kernel, adding a power method. |
|
Extends the tinygp.kernels.quasisep.Cosine kernel, adding a power method. |
|
Extends the tinygp.kernels.quasisep.Celerite kernel, adding a power method. |
|
Extends the tinygp.kernels.quasisep.Matern32 kernel, adding a power method. |
|
Extends the tinygp.kernels.quasisep.Matern52 kernel, adding a power method. |
|
Extends the tinygp.kernels.quasisep.SHO kernel, adding a power method. |
|
The Lorentzian kernel. |
|
A continuous-time autoregressive moving-average kernel. |
|
A multiband kernel implementating a low-rank Kronecker covariance structure. |
Functions#
|
Compute the CARMA state stationary covariance from AR roots. |
|
Compute roots of a CARMA characteristic polynomial. |
|
Expand a product of CARMA quadratic factors into a full polynomial. |
|
Factorize a CARMA polynomial into quadratic and linear factors. |
|
Compute exponential-basis coefficients of the CARMA autocovariance. |
Module Contents#
- class Quasisep[source]#
Bases:
tinygp.kernels.quasisep.QuasisepAn extension of the tinygp.kernels.quasisep.Quasisep kernel.
tinygp.kernels.quasisep.Quasisep is the base class for all kernels that can be evaluated following an O(N) scaling. This extension adds a power method to return the power spectral density (PSD) of a quasiseparable kernel at an input frequency.
- class Sum[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.SumA helper to represent the sum of two quasiseparable kernels.
- class Product[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.ProductA helper to represent the product of two quasiseparable kernels.
- class Scale[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.ScaleThe product of a scalar and a quasiseparable kernel.
- class Exp[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.ExpExtends the tinygp.kernels.quasisep.Exp kernel, adding a power method.
- class Cosine[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.CosineExtends the tinygp.kernels.quasisep.Cosine kernel, adding a power method.
- class Celerite[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.CeleriteExtends the tinygp.kernels.quasisep.Celerite kernel, adding a power method.
- class Matern32[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.Matern32Extends the tinygp.kernels.quasisep.Matern32 kernel, adding a power method.
- class Matern52[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.Matern52Extends the tinygp.kernels.quasisep.Matern52 kernel, adding a power method.
- class SHO[source]#
Bases:
Quasisep,tinygp.kernels.quasisep.SHOExtends the tinygp.kernels.quasisep.SHO kernel, adding a power method.
- class Lorentzian[source]#
Bases:
QuasisepThe Lorentzian kernel.
The kernel takes the form:
\[k(\tau) = \sigma^2\,\exp(-b\,\tau)\,cos(\omega\,\tau)\]for \(\tau = |x_i - x_j|\) and \(b = \frac{\omega}{2\,Q}\).
- Parameters:
omega – The parameter \(\omega\).
quality – The parameter \(Q\).
sigma (optional) – The parameter \(\sigma\). Defaults to a value of 1. Specifying the explicit value here provides a slight performance boost compared to independently multiplying the kernel with a prefactor.
- get_scale() tuple[tinygp.helpers.JAXArray | float, tinygp.helpers.JAXArray | float][source]#
Scale of the Lorentzian.
- observation_model(X: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray[source]#
The observation model for the process
- class CARMA[source]#
Bases:
QuasisepA continuous-time autoregressive moving-average kernel.
This kernel represents a CARMA(\(p, q\)) process in companion-form state-space notation so it can be used with the quasiseparable solvers in tinygp.
The autoregressive polynomial is parameterized in ascending power order as
\[\alpha(D) = \alpha_0 + \alpha_1 D + \cdots + \alpha_{p-1} D^{p-1} + D^p,\]where the leading coefficient of \(D^p\) is fixed to 1 and supplied implicitly. The moving-average polynomial is
\[\beta(D) = \beta_0 + \beta_1 D + \cdots + \beta_q D^q.\]The corresponding power spectral density (PSD) is
\[P(\omega) = \sigma^2\,\frac{|\sum_{q} \beta_q\,(i\,\omega)^q|^2}{|\sum_{p} \alpha_p\,(i\,\omega)^p|^2}\]following Equation 1 in Kelly et al. (2014), where \(\alpha_p\) and \(\beta_0\) are set to 1. In this implementation, we absorb \(\sigma\) into the definition of the \(\beta\) parameters. That is, \(\beta_{\mathrm{new}} = \beta\,\sigma\).
- Parameters:
alpha – Autoregressive coefficients in ascending power order, excluding the leading coefficient fixed to 1.
beta – Moving-average coefficients \([\beta_0, \ldots, \beta_q]\) in ascending power order.
sigma_w – Standard deviation of the white-noise driving term used when constructing the stationary state covariance.
- classmethod from_quads(alpha_quads: tinygp.helpers.JAXArray | numpy.typing.NDArray, beta_quads: tinygp.helpers.JAXArray | numpy.typing.NDArray, beta_mult: tinygp.helpers.JAXArray | numpy.typing.NDArray) CARMA[source]#
Construct a CARMA kernel using the roots of its characteristic polynomials.
The roots can be parameterized as the 0th and 1st order coefficients of a set of quadratic equations (2nd order coefficient equals 1). The product of those quadratic equations gives the characteristic polynomials of CARMA. The input of this method are said coefficients of the quadratic equations. See Equation 30 in Kelly et al. (2014). for more detail.
- Parameters:
alpha_quads – Coefficients of the auto-regressive (AR) quadratic equations corresponding to the \(\alpha\) parameters. This should be an array of length p.
beta_quads – Coefficients of the moving-average (MA) quadratic equations corresponding to the \(\beta\) parameters. This should be an array of length q.
beta_mult – A multiplier of the MA coefficients, equivalent to \(\beta_q\)—the last entry of the \(\beta\) parameters input to the
init()method.
- property arroots: tinygp.helpers.JAXArray[source]#
Return the autoregressive roots sorted by real part.
- observation_model(X)[source]#
Return the observation vector that maps the state to the process value.
- stationary_covariance()[source]#
Return the stationary covariance of the latent companion-form state.
- carma_root_stationary_covariance(arroots: tinygp.helpers.JAXArray, sigma: tinygp.helpers.JAXArray | float = 1.0) tinygp.helpers.JAXArray[source]#
Compute the CARMA state stationary covariance from AR roots.
This implements the closed-form expression
\[V_{ij} = - \,\sigma^2 \, \sum_{k=1}^{p} \frac{r_k^i(-r_k)^j}{ 2\mathrm{Re}(r_k) \prod_{l=1, l\ne k}^{p}(r_l-r_k)(r_l^*+r_k) }\]where \(r_k\) are the autoregressive roots and \(i,j \in [0, p-1]\).
- Parameters:
arroots – The roots of the autoregressive characteristic polynomial.
sigma – The driving-noise amplitude \(\sigma\).
- Returns:
The \(p \times p\) matrix defined by the root-based covariance expression above.
- carma_roots(poly_coeffs: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray[source]#
Compute roots of a CARMA characteristic polynomial.
- Parameters:
poly_coeffs – Polynomial coefficients in ascending power order, so the first element is the constant term.
- Returns:
The polynomial roots sorted by their real part.
- carma_quads2poly(quads_coeffs: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray[source]#
Expand a product of CARMA quadratic factors into a full polynomial.
- Parameters:
quads_coeffs – Constant and linear coefficients of the quadratic factors used by the Kelly et al. parameterization. The last entry is the multiplier for the highest-order term of the reconstructed polynomial.
- Returns:
Polynomial coefficients in ascending power order.
- carma_poly2quads(poly_coeffs: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray[source]#
Factorize a CARMA polynomial into quadratic and linear factors.
- Parameters:
poly_coeffs – Coefficients of the input characteristic polynomial. The first entry corresponds to the constant term.
- Returns:
Constant and linear coefficients for the factorized quadratic blocks, followed by the multiplier for the highest-order term.
- carma_acvf(arroots: tinygp.helpers.JAXArray, arparam: tinygp.helpers.JAXArray, maparam: tinygp.helpers.JAXArray) tinygp.helpers.JAXArray[source]#
Compute exponential-basis coefficients of the CARMA autocovariance.
- Parameters:
arroots – The roots of the autoregressive characteristic polynomial.
arparam – Autoregressive coefficients in ascending power order.
maparam – Moving-average coefficients \([\beta_0, \ldots, \beta_q]\) in ascending power order.
- Returns:
The coefficients of the exponential expansion of the ACVF, with one coefficient per autoregressive root.
- class MultibandLowRank[source]#
Bases:
tinygp.kernels.quasisep.WrapperA multiband kernel implementating a low-rank Kronecker covariance structure.
The specific form of the cross-band Kronecker covariance matrix is given by Equation 13 of Gordon et al. (2020). The implementation is inspired by this tinygp tutorial.
- Parameters:
params – A dictionary of string and array pairs, which are used in the observational_model method to describe the cross-band covariance.