dense
=====

.. py:module:: smolgp.kernels.dense

.. autoapi-nested-parse::

   This module uses tinygp to construct dense representations of
   kernels for testing/benchmarking performance and accuracy.

   .. danger::

      They are not meant for use on large datasets!

   Kernels defined here:

   - `SHOKernel`: A stochastic harmonic oscillator (SHO) kernel

   - `IntegratedSHOKernel`: A tinygp/JAX implementation of the integrated SHO kernel from `Luhn et al. 2026 <https://arxiv.org/abs/2601.02462>`__.



Classes
-------

.. autoapisummary::

   smolgp.kernels.dense.SHOKernel
   smolgp.kernels.dense.IntegratedSHOKernel


Functions
---------

.. autoapisummary::

   smolgp.kernels.dense.extract_leaf_kernels
   smolgp.kernels.dense.unpack_coordinates


Module Contents
---------------

.. py:function:: extract_leaf_kernels(kernel, all=False)

   Recursively extract all tinygp leaf kernels from a sum or product of tinygp kernels


.. py:function:: unpack_coordinates(X1, X2)

   Unpack the input coordinates X1 and X2 into time, instrument ID, and exposure time.

   X1 and X2 can either be
       t1 and t2 for single instrument, no exposure times
   or
       (t1, instid1) and (t2, instid2) for multiple instruments, no exposure times
   or
       (t1, instid1, delta1) and (t2, instid2, delta2) for exposure times, also requires instrument IDs

   Returns
       (t1, instid1, delta1), (t2, instid2, delta2)
       where instid and delta are arrays of zeros if not provided


.. py:class:: SHOKernel(S=None, w=None, Q=None, rho=None, tau=None, sig=None)

   Bases: :py:obj:`tinygp.kernels.Kernel`


   The base class for all kernel implementations

   This subclass provides default implementations to add and multiply kernels.
   Subclasses should accept parameters in their ``__init__`` and then override
   :func:`Kernel.evaluate` with custom behavior.


   .. py:attribute:: S
      :type:  jax.Array | float


   .. py:attribute:: w
      :type:  jax.Array | float


   .. py:attribute:: Q
      :type:  jax.Array | float


   .. py:attribute:: rho
      :type:  jax.Array | float


   .. py:attribute:: tau
      :type:  jax.Array | float


   .. py:attribute:: sig
      :type:  jax.Array | float


   .. py:method:: evaluate(X1, X2)

      Calculate the kernel for given pair of times X1 and X2.

      X1 and X2 can either be
          t1 and t2 for single instrument, no exposure times
      or
          (t1, instid1) and (t2, instid2) for multiple instruments, no exposure times
      or
          (t1, delta1, instid1) and (t2, delta2, instid2) for exposure times,
              also requires instrument IDs even if a single instrument



   .. py:method:: __repr__()


.. py:class:: IntegratedSHOKernel(S=None, w=None, Q=None, rho=None, tau=None, sig=None)

   Bases: :py:obj:`tinygp.kernels.Kernel`


   The base class for all kernel implementations

   This subclass provides default implementations to add and multiply kernels.
   Subclasses should accept parameters in their ``__init__`` and then override
   :func:`Kernel.evaluate` with custom behavior.


   .. py:attribute:: S
      :type:  jax.Array | float


   .. py:attribute:: w
      :type:  jax.Array | float


   .. py:attribute:: Q
      :type:  jax.Array | float


   .. py:attribute:: rho
      :type:  jax.Array | float


   .. py:attribute:: tau
      :type:  jax.Array | float


   .. py:attribute:: sig
      :type:  jax.Array | float


   .. py:attribute:: eta
      :type:  jax.Array | float


   .. py:attribute:: a
      :type:  jax.Array | float


   .. py:method:: __repr__()


   .. py:method:: Delta(X1, X2)

      Time difference (absolute value)
      between pairs of observations X1, X2.

      Delta = | t1 - t2 |



   .. py:method:: latent_process(Delta)

      Latent kernel for given time difference Delta



   .. py:method:: I0(y)

      Helper function for single integral (Eq. 11 in L26)



   .. py:method:: I1(lower, upper)

      Helper function for double integrals (Eq. 14 in L26)



   .. py:method:: I2(lower, upper)

      Helper function for double integrals (Eq. 15 in L26)



   .. py:method:: integrated_separate(Delta, delta1, delta2)

      The double integral for two non-overlapping observations.
      Depends on the time-lag (Delta) and the two exposure times (delta1, delta2)

      Eq. 17 in L26



   .. py:method:: integrated_overlap(delta)

      The double integral for two perfectly overlapping
      observations (i.e. zero time-lag). As such, this
      only depends on the exposure time (delta)

      Eq. 13 in L26



   .. py:method:: integrated_single(Delta, delta)

      The single integral for when one observation
      has zero exposure time (e.g. latent curve)

      Eq. 10 in L26



   .. py:method:: full_single_integral(Delta, delta)

      This helper function handles the logic for the single integral case.

      For pairs of observation coordinates separated by Delta, such that
      one is the 0-exposure test point (i.e., the latent curve) and the
      other has exposure time delta, this function handles the logic to
      break up the integral and call self.integrated_single() for the
      sub-exposures where appropriate



   .. py:method:: evaluate(X1, X2)

      Compute the integrated SHO kernel for a pair of observations X1 and X2.

      X1 and X2 can either be
          (t1, texp1) and (t2, texp2) for single instrument
      or
          (t1, texp1, instid1) and (t2, texp2, instid2) for multiple instruments