kalman
======

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


Functions
---------

.. autoapisummary::

   smolgp.solvers.integrated.parallel.kalman.ParallelIntegratedKalmanFilter
   smolgp.solvers.integrated.parallel.kalman.make_associative_params
   smolgp.solvers.integrated.parallel.kalman._combine_per_pair
   smolgp.solvers.integrated.parallel.kalman.parallel_integrated_kalman_filter
   smolgp.solvers.integrated.parallel.kalman._calc_kf_predictions
   smolgp.solvers.integrated.parallel.kalman._calc_vS
   smolgp.solvers.integrated.parallel.kalman.postprocess


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

.. py:function:: ParallelIntegratedKalmanFilter(kernel, X, y, t_states, obsid, instid, stateid, R, return_v_S=False)

   Wrapper for parallel_integrated_kalman_filter function

   :param kernel: IntegratedStateSpaceModel kernel
   :param X: Array of size N, data coordinates (e.g. (time, texp, instid))
   :param y: Array of size (N, D), measurements at the data coordinates
   :param t_states: Array of size K, sorted time coordinate of all states (exposure starts and ends)
   :param obsid: Array of size N, which observation (0,...,N-1) is being made at each state k
   :param instid: Array of size N, which instrument (0,...,Ninst-1) recorded observation n
   :param stateid: Array of size K, 0 for exposure-start, 1 for exposure-end
   :param R: Observation noise covariance, shape (N, D, D)
   :param return_v_S: Whether to return innovation and its covariance (for likelihood computation)

   :returns: filtered means
             P_filtered : filtered covariances
             m_predicted: predicted means
             P_predicted: predicted covariances
   :rtype: m_filtered


.. py:function:: make_associative_params(Phi_aug, H_aug, Q_aug, RESET, R, X, y, t_states, obsid, instid, stateid, m0, P0)

   Generate the associative parameters needed for parallel Kalman.
   See Eqns. 10, 11, 12 from Sarkka & Garcia-Fernandez (2020)


.. py:function:: _combine_per_pair(left, right)

   See Eqn. 13 & 14 of Sarkka & Garcia-Fernandez (2020) for
   a the algorithm and notation.


.. py:function:: parallel_integrated_kalman_filter(asso_params)

   Jax implementation of the parallel Kalman filter algorithm
   for integrated measurements.

   See Section 4A of Sarkka & Garcia-Fernandez (2020) for
   a detailed description of the algorithm and notation,
   and section 3.2.4 of Rubenzahl & Hattori et al. (2025)
   for the integrated measurement case.

   Total runtime (span) complexity is O(N/T + logT) where N is the
   number of time steps and T is the number of parallel threads.


.. py:function:: _calc_kf_predictions(Phi_aug, Q_aug, H_aug, t_states, b, C, m0, P0)

.. py:function:: _calc_vS(H_aug, R, m_pred, P_pred, X, y, stateid, obsid)

.. py:function:: postprocess(Phi_aug, Q_aug, H_aug, R, X, y, t_states, obsid, stateid, b, C, m0, P0)