integrated
==========

.. py:module:: smolgp.solvers.integrated

.. autoapi-nested-parse::

   This subpackage implements the Kalman filter and RTS smoother for an augmented
   state space model which includes an integral state as a
   :class:`IntegratedStateSpaceSolver`. This is intended to be used with
   :class:`smolgp.kernels.integrated` state space models to properly account
   for integrated (e.g. exposure-averaged) measurements.

   See :ref:`integrated` for a tutorial on using the integrated solvers.



Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/smolgp/solvers/integrated/kalman/index
   /autoapi/smolgp/solvers/integrated/parallel/index
   /autoapi/smolgp/solvers/integrated/rts/index
   /autoapi/smolgp/solvers/integrated/solver/index


Classes
-------

.. autoapisummary::

   smolgp.solvers.integrated.IntegratedStateSpaceSolver
   smolgp.solvers.integrated.ParallelIntegratedStateSpaceSolver


Package Contents
----------------

.. py:class:: IntegratedStateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)

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


   A solver that uses ``jax.lax.scan`` to implement Kalman filtering
   and RTS smoothing for integrated measurements


   .. 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:  tinygp.helpers.JAXArray


   .. py:method:: normalization() -> tinygp.helpers.JAXArray


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

      Wrapper for Kalman filter used with this solver



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

      Wrapper for RTS smoother used with this solver



   .. py:method:: condition(y, return_v_S=False) -> tinygp.helpers.JAXArray

      Compute the Kalman predicted, filtered, and RTS smoothed
      means and covariances at each of the input coordinates



   .. py:method:: predict(X_test, conditioned_results) -> tinygp.helpers.JAXArray

      Algorithm for making predictions at arbitrary coordinates X_test

      :param X_test: The test coordinates.
      :param conditioned_results: The output of self.condition()

      :returns: Predicted means of the states at X_test
                pred_var  : Predicted variances of the states at X_test
      :rtype: pred_mean

      There are three cases:
          1. Retrodiction  : smoothing from the first data point
                             using the prior as the prediction
          2. Interpolation : filtering from most recent data point
                             and smoothing from next future point
          3. Extrapolation : predicting from final filtered point



.. py:class:: ParallelIntegratedStateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)

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


   A solver that uses ``jax.lax.associative_scan`` to implement
   parallel Kalman filtering and RTS smoothing for integrated measurements


   .. 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:  tinygp.helpers.JAXArray


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


   .. py:method:: normalization() -> tinygp.helpers.JAXArray


   .. py:method:: Kalman(y, return_v_S=True) -> Any

      Wrapper for Kalman filter used with this solver



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

      Wrapper for RTS smoother used with this solver



   .. py:method:: condition(y, return_v_S=True) -> tinygp.helpers.JAXArray

      Compute the Kalman predicted, filtered, and RTS smoothed
      means and covariances at each of the input coordinates



   .. py:method:: predict(X_test, conditioned_results) -> tinygp.helpers.JAXArray

      Algorithm for making predictions at arbitrary coordinates X_test

      :param X_test: The test coordinates.
      :param conditioned_results: The output of self.condition()
      :param observation_model: (optional) H for the test points
                                should be a function just like
                                self.kernel.observation_model

      There are three cases:
          1. Retrodiction  : smoothing from the first data point
                             using the prior as the prediction
          2. Interpolation : filtering from most recent data point
                             and smoothing from next future point
          3. Extrapolation : predicting from final filtered point