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 Functions --------- .. autoapisummary:: eztaox.kernels.quasisep.carma_root_stationary_covariance 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 Bases: :py:obj:`Quasisep` A continuous-time autoregressive moving-average kernel. This kernel represents a CARMA(:math:`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 .. math:: \alpha(D) = \alpha_0 + \alpha_1 D + \cdots + \alpha_{p-1} D^{p-1} + D^p, where the leading coefficient of :math:`D^p` is fixed to 1 and supplied implicitly. The moving-average polynomial is .. math:: \beta(D) = \beta_0 + \beta_1 D + \cdots + \beta_q D^q. The corresponding power spectral density (PSD) is .. math:: 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 :math:`\alpha_p` and :math:`\beta_0` are set to 1. In this implementation, we absorb :math:`\sigma` into the definition of the :math:`\beta` parameters. That is, :math:`\beta_{\mathrm{new}} = \beta\,\sigma`. :param alpha: Autoregressive coefficients in ascending power order, excluding the leading coefficient fixed to 1. :param beta: Moving-average coefficients :math:`[\beta_0, \ldots, \beta_q]` in ascending power order. :param sigma_w: Standard deviation of the white-noise driving term used when constructing the stationary state covariance. .. 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:property:: arroots :type: tinygp.helpers.JAXArray Return the autoregressive roots sorted by real part. .. py:method:: design_matrix() Return the companion-form drift matrix for the latent CAR state. .. py:method:: observation_model(X) Return the observation vector that maps the state to the process value. .. py:method:: stationary_covariance() Return the stationary covariance of the latent companion-form state. .. py:method:: transition_matrix(X1, X2) Return the state transition matrix between two one-dimensional inputs. .. 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_root_stationary_covariance(arroots: tinygp.helpers.JAXArray, sigma: tinygp.helpers.JAXArray | float = 1.0) -> tinygp.helpers.JAXArray Compute the CARMA state stationary covariance from AR roots. This implements the closed-form expression .. math:: 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 :math:`r_k` are the autoregressive roots and :math:`i,j \in [0, p-1]`. :param arroots: The roots of the autoregressive characteristic polynomial. :param sigma: The driving-noise amplitude :math:`\sigma`. :returns: The :math:`p \times p` matrix defined by the root-based covariance expression above. .. py:function:: carma_roots(poly_coeffs: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Compute roots of a CARMA characteristic polynomial. :param 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. .. py:function:: carma_quads2poly(quads_coeffs: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Expand a product of CARMA quadratic factors into a full polynomial. :param 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. .. py:function:: carma_poly2quads(poly_coeffs: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Factorize a CARMA polynomial into quadratic and linear factors. :param 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. .. py:function:: carma_acvf(arroots: tinygp.helpers.JAXArray, arparam: tinygp.helpers.JAXArray, maparam: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray Compute exponential-basis coefficients of the CARMA autocovariance. :param arroots: The roots of the autoregressive characteristic polynomial. :param arparam: Autoregressive coefficients in ascending power order. :param maparam: Moving-average coefficients :math:`[\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. .. 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