eztaox.kernels.quasisep ======================= .. py:module:: eztaox.kernels.quasisep .. autoapi-nested-parse:: 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 ------- .. autoapisummary:: eztaox.kernels.quasisep.Quasisep eztaox.kernels.quasisep.Sum eztaox.kernels.quasisep.Product eztaox.kernels.quasisep.Scale eztaox.kernels.quasisep.Exp eztaox.kernels.quasisep.Cosine eztaox.kernels.quasisep.Celerite eztaox.kernels.quasisep.Matern32 eztaox.kernels.quasisep.Matern52 eztaox.kernels.quasisep.SHO eztaox.kernels.quasisep.Lorentzian eztaox.kernels.quasisep.CARMA eztaox.kernels.quasisep.MultibandLowRank eztaox.kernels.quasisep.LaguerreSeries Functions --------- .. autoapisummary:: eztaox.kernels.quasisep.carma_roots eztaox.kernels.quasisep.carma_quads2poly eztaox.kernels.quasisep.carma_poly2quads eztaox.kernels.quasisep.carma_acvf Module Contents --------------- .. py:class:: Quasisep Bases: :py:obj:`tinygp.kernels.quasisep.Quasisep` An 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. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Sum Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Sum` A helper to represent the sum of two quasiseparable kernels. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Product Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Product` A helper to represent the product of two quasiseparable kernels. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Scale Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Scale` The product of a scalar and a quasiseparable kernel. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Exp Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Exp` Extends the `tinygp.kernels.quasisep.Exp` kernel, adding a power method. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Cosine Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Cosine` Extends the `tinygp.kernels.quasisep.Cosine` kernel, adding a power method. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Celerite Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Celerite` Extends the `tinygp.kernels.quasisep.Celerite` kernel, adding a power method. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Matern32 Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Matern32` Extends the `tinygp.kernels.quasisep.Matern32` kernel, adding a power method. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Matern52 Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.Matern52` Extends the `tinygp.kernels.quasisep.Matern52` kernel, adding a power method. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: SHO Bases: :py:obj:`Quasisep`, :py:obj:`tinygp.kernels.quasisep.SHO` Extends the `tinygp.kernels.quasisep.SHO` kernel, adding a power method. .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: Lorentzian Bases: :py:obj:`Quasisep` The Lorentzian kernel. The kernel takes the form: .. math:: k(\tau) = \sigma^2\,\exp(-b\,\tau)\,cos(\omega\,\tau) for :math:`\tau = |x_i - x_j|` and :math:`b = \frac{\omega}{2\,Q}`. :param omega: The parameter :math:`\omega`. :param quality: The parameter :math:`Q`. :param sigma: The parameter :math:`\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. :type sigma: optional .. py:method:: get_scale() -> tuple[tinygp.helpers.JAXArray | float, tinygp.helpers.JAXArray | float] Scale of the Lorentzian. .. py:method:: design_matrix() -> tinygp.helpers.JAXArray The design matrix for the process .. py:method:: stationary_covariance() -> tinygp.helpers.JAXArray The variance of the kernel at :math:`t=0`. .. py:method:: observation_model(X: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray The observation model for the process .. py:method:: transition_matrix(X1: tinygp.helpers.JAXArray, X2: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray The transition matrix between two coordinates .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:class:: CARMA(alpha: tinygp.helpers.JAXArray | numpy.typing.NDArray, beta: tinygp.helpers.JAXArray | numpy.typing.NDArray) Bases: :py:obj:`Quasisep` A continuous-time autoregressive moving average (CARMA) process kernel. This process has the power spectrum density (PSD) .. math:: P(\omega) = \sigma^2\,\frac{|\sum_{q} \beta_q\,(i\,\omega)^q|^2}{|\sum_{p} \alpha_p\,(i\,\omega)^p|^2} defined following Equation 1 in `Kelly et al. (2014) `_, where :math:`\alpha_p` and :math:`\beta_0` are set to 1. In this implementation, we absorb :math:`\sigma` into the definition of :math:`\beta` parameters. That is :math:`\beta_{new}` = :math:`\beta * \sigma`. .. note:: To construct a stationary CARMA kernel/process, the roots of the characteristic polynomials for Equation 1 in `Kelly et al. (2014)` must have negative real parts. This condition can be met automatically by requiring positive input parameters when instantiating the kernel using the :func:`init` method for CARMA(1,0), CARMA(2,0), and CARMA(2,1) models or by requiring positive input parameters when instantiating the kernel using the :func:`from_quads` method. .. note:: Implementation details The logic behind this implementation is simple---finding the correct combination of real/complex exponential kernels that resembles the autocovariance function of the CARMA model. Note that the order also matters. This task is achieved using the `acvf` method. Then the rest is copied from the `Exp` and `Celerite` kernel. Given the requirement of negative roots for stationarity, the `from_quads` method is implemented to facilitate consturcting stationary higher-order CARMA models beyond CARMA(2,1). The inputs for `from_quads` are the coefficients of the quadratic equations factorized out of the full characteristic polynomial. `poly2quads` is used to factorize a polynomial into a product of said quadractic equations, and `quads2poly` is used for the reverse process. One last trick is the use of `_real_mask`, `_complex_mask`, and `complex_select`, which are arrays of 0s and 1s. They are implemented to avoid control flows. More specifically, some intermediate quantities are computed regardless, but are only used if there is a matching real or complex exponential kernel for the specific CARMA kernel. :param alpha: The parameter :math:`\alpha` in the definition above, exlcuding :math:`\alpha_p`. This should be an array of length `p`. :param beta: The product of parameters :math:`\beta` and parameter :math:`\sigma` in the definition above. This should be an array of length `q+1`, where `q+1 <= p`. .. py:method:: init(alpha: tinygp.helpers.JAXArray, beta: tinygp.helpers.JAXArray) -> CARMA :classmethod: .. py:method:: 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 :classmethod: 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. :param alpha_quads: Coefficients of the auto-regressive (AR) quadratic equations corresponding to the :math:`\alpha` parameters. This should be an array of length `p`. :param beta_quads: Coefficients of the moving-average (MA) quadratic equations corresponding to the :math:`\beta` parameters. This should be an array of length `q`. :param beta_mult: A multiplier of the MA coefficients, equivalent to :math:`\beta_q`---the last entry of the :math:`\beta` parameters input to the :func:`init` method. .. py:method:: design_matrix() -> tinygp.helpers.JAXArray The design matrix for the process .. py:method:: stationary_covariance() -> tinygp.helpers.JAXArray The stationary covariance of the process .. py:method:: observation_model(X: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray The observation model for the process .. py:method:: transition_matrix(X1: tinygp.helpers.JAXArray, X2: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray The transition matrix between two coordinates .. py:method:: power(f: float | tinygp.helpers.JAXArray, df: float | tinygp.helpers.JAXArray | None = None) -> tinygp.helpers.JAXArray Compute the power spectral density (PSD) at frequency `f`. .. py:function:: carma_roots(poly_coeffs: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Compute the CARMA polynomial coefficient roots. :param poly_coeffs: coefficients of the polynomial :returns: roots of the coefficients .. py:function:: carma_quads2poly(quads_coeffs: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Expand a product of quadractic equations into a polynomial. :param quads_coeffs: The 0th and 1st order coefficients of the quadractic equations. The last entry is a multiplier, which corresponds to the coefficient of the highest order term in the output full polynomial. :returns: Coefficients of the full polynomial. The first entry corresponds to the lowest order term. .. py:function:: carma_poly2quads(poly_coeffs: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Factorize a polynomial into a product of quadratic equations. :param poly_coeffs: Coefficients of the input characteristic polynomial. The first entry corresponds to the lowest order term. :returns: The 0th and 1st order coefficients of the quadractic equations. The last entry is a multiplier, which corresponds to the coefficient of the highest order term in the full polynomial. .. py:function:: carma_acvf(arroots: tinygp.helpers.JAXArray, arparam: tinygp.helpers.JAXArray, maparam: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Compute the coefficients of the autocovariance function (ACVF). :param arroots: The roots of the autoregressive characteristic polynomial. :param arparam: :math:`\alpha` parameters :param maparam: :math:`\beta` parameters :returns: ACVF coefficients, each entry corresponds to one root. .. py:class:: MultibandLowRank Bases: :py:obj:`tinygp.kernels.quasisep.Wrapper` A 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 `_. :param params: A dictionary of string and array pairs, which are used in the `observational_model` method to describe the cross-band covariance. .. py:method:: coord_to_sortable(X) -> tinygp.helpers.JAXArray Extract the time-sortable component of the coordinates. .. py:method:: observation_model(X) -> tinygp.helpers.JAXArray The observation model for the process .. py:class:: LaguerreSeries Bases: :py:obj:`Quasisep` Laguerre series approximation of a stationary kernel. Wraps a kernel and approximates its autocovariance using a truncated Laguerre series expansion, enabling O(N) GP computations. The coefficients are derived from the wrapped kernel's parameters at construction time. :param kernel: The kernel to approximate (must have a `scale` attribute). :param order: Maximum Laguerre polynomial order. :param n_quad: Number of quadrature points for coefficient computation. .. py:method:: design_matrix() -> jax.Array Block diagonal of individual design matrices. .. py:method:: stationary_covariance() -> jax.Array Block diagonal of scaled stationary covariances. .. py:method:: observation_model(X: jax.Array) -> jax.Array Concatenation of observation models. .. py:method:: transition_matrix(X1: jax.Array, X2: jax.Array) -> jax.Array Block diagonal of transition matrices.