solver
======

.. py:module:: smolgp.solvers.integrated.parallel.solver


Classes
-------

.. autoapisummary::

   smolgp.solvers.integrated.parallel.solver.ParallelIntegratedStateSpaceSolver


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

.. 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