base
====

.. py:module:: smolgp.solvers.base

.. autoapi-nested-parse::

   :class:`Solver` helps organize the common elements of all solvers, including
       1. the required methods (:meth:`Kalman`, :meth:`RTS`, :meth:`smoothing_gains`, :meth:`predict`),
       2. the shared state-order bookkeeping definitions,
       3. and a default implementation of the marginal likelihood in terms of the
          filter's innovations, which every Kalman filter produces. A subclass may
          override :meth:`log_probability` with a more efficient scan that only
          accumulates the likelihood contributions (rather than the full filter outputs).

   Subclasses of :class:`Solver` only inherit when the parent's method would be still
   be correct (though perhaps slower). As such, the inheritance tree looks like:

       Solver
       |-- StateSpaceSolver                        instantaneous, sequential Kalman/RTS
       |   `-- ParallelStateSpaceSolver            associative-scan Kalman/RTS
       `-- IntegratedStateSpaceSolver              exposure-aware, sequential Kalman/RTS (over K = 2N states)
           `-- ParallelIntegratedStateSpaceSolver  associative-scan exposure-aware Kalman/RTS (over K = 2N states)



Classes
-------

.. autoapisummary::

   smolgp.solvers.base.Solver


Functions
---------

.. autoapisummary::

   smolgp.solvers.base.log_prob_from_v_S


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

.. py:function:: log_prob_from_v_S(v: tinygp.helpers.JAXArray, S: tinygp.helpers.JAXArray) -> tinygp.helpers.JAXArray

   The Gaussian log probability from a Kalman filter's innovations.

   Every Kalman filter produces the innovation :math:`v_k = y_k - H_k m_k^-`
   and its covariance :math:`S_k`, from which the marginal log likelihood is

   .. math::

       \log p(y) = -\tfrac{1}{2} \sum_k \left(
           v_k^T S_k^{-1} v_k + \log\det S_k + D \log 2\pi \right),

   :param v: shape ``(N, D)``, the innovations.
   :param S: shape ``(N, D, D)``, the innovation covariances.


.. py:class:: Solver

   Bases: :py:obj:`equinox.Module`


   Base class for a smolgp solver.

   Subclasses must implement :meth:`Kalman`, :meth:`RTS`, :meth:`smoothing_gains`
   and :meth:`predict`. :meth:`condition` is provided here too, being only a
   sequencing of :meth:`Kalman` and :meth:`RTS`, but a solver whose smoother
   takes different arguments overrides it.

   The likelihood is implemented here in terms of the filter's innovations,
   which every Kalman filter produces, but a subclass may override
   :meth:`log_probability` with a more efficient scan that only accumulates
   the likelihood contributions (rather than the full filter outputs).

   .. attribute:: kernel

      The kernel defining the state space model.

      :type: StateSpaceModel

   .. attribute:: X

      The observed input coordinates.

      :type: JAXArray

   .. attribute:: noise

      Per-observation noise covariance, shape ``(N, D, D)``.

      :type: JAXArray

   .. attribute:: state_coords

      State-level bookkeeping. One state per
      observation for an instantaneous kernel, two (exposure start and
      end) for an integrated kernel.

      :type: StateCoords


   .. py:attribute:: X
      :type:  tinygp.helpers.JAXArray


   .. py:attribute:: kernel
      :type:  smolgp.kernels.base.StateSpaceModel


   .. py:attribute:: noise
      :type:  tinygp.helpers.JAXArray


   .. py:attribute:: state_coords
      :type:  smolgp.solvers.state_coords.StateCoords


   .. py:property:: t_states
      :type: tinygp.helpers.JAXArray


      The chronologically sorted time coordinate of every state.


   .. py:method:: _to_state_order(*arrays: tinygp.helpers.JAXArray) -> tuple[tinygp.helpers.JAXArray, Ellipsis]

      Gather per-observation arrays into the solver's (chronologically
      sorted) state order.

      ``self.X`` and everything derived from it (``y``, ``noise``) are kept
      in the caller's input order; the filter/smoother step chronologically,
      so they need the sorted order instead. ``state_coords.obsid`` is exactly
      that permutation. Results come back in state order and are mapped
      back by the usual sort-by-``obsid`` machinery.



   .. py:method:: Kalman(y, return_v_S: bool = False) -> Any
      :abstractmethod:


      Run this solver's Kalman filter.

      Returns ``(m_filtered, P_filtered, m_predicted, P_predicted)``,
          plus ``(v, S)`` when ``return_v_S`` is True.



   .. py:method:: RTS(kalman_results) -> Any
      :abstractmethod:


      Run this solver's RTS smoother over :meth:`Kalman`'s output.



   .. py:method:: smoothing_gains(P_filtered, P_predicted) -> tinygp.helpers.JAXArray
      :abstractmethod:


      The ``y``-independent RTS smoothing gains for this state timeline.



   .. py:method:: condition(y, return_v_S: bool = False) -> Any

      Filter then smooth, giving the posterior at the data.

      Implemented here rather than per solver: it is only a sequencing of
      :meth:`Kalman` and :meth:`RTS` plus packaging, so every solver whose
      ``RTS`` takes the filter's four outputs shares it verbatim. A solver
      whose smoother has a different signature overrides it -- see
      :class:`~smolgp.solvers.ParallelStateSpaceSolver`, whose parallel
      smoother consumes only the filtered pair.



   .. py:method:: predict(X_test, conditioned_results) -> Any
      :abstractmethod:


      The posterior at arbitrary test coordinates.



   .. py:method:: _log_probability_from_filter(y) -> tinygp.helpers.JAXArray

      The likelihood via the Kalman filter outputs.



   .. py:method:: log_probability(y) -> tinygp.helpers.JAXArray

      The marginal log likelihood of the data, ``y``.

      By default, runs the Kalman filter and reduces its innovations.
      However, the Kalman filter computes more than is necessary if one
      only wants the likelihood. Hence, a `Solver` can override this
      function with an optimized method, e.g.
      :meth:`~smolgp.solvers.solver.StateSpaceSolver.log_probability`