state_coords
============

.. py:module:: smolgp.solvers.state_coords


Classes
-------

.. autoapisummary::

   smolgp.solvers.state_coords.StateCoords


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

.. py:class:: StateCoords

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


   The state-level bookkeeping shared by every solver, filter, smoother,
   and sampler in ``smolgp``.

   A state space model is solved by stepping through a sorted timeline of
   *states*. For an instantaneous kernel there is exactly one state per
   observation, but for an integrated (exposure-averaged) kernel each
   observation contributes *two* states: the start and end of its exposure
   window. Therefore, the number of states ``K`` and the number of observations
   ``N`` are different, and the mapping between them has to be carried
   around explicitly. Additionally, the integrated solver needs to know for
   a given state what observation it belongs to and whether that state was
   an exposure start or end, as these are handled differently.
   That mapping is what this object holds.

   .. attribute:: t_states

      shape ``(K,)``. The sortable coordinate (e.g. time) of each
      state, in ascending order. Ties are broken so that exposure *ends*
      (``stateid==1``) precede exposure *starts* (``stateid==0``) at the
      same instant, via a ``jnp.lexsort((-stateid, t_states))``.

   .. attribute:: instid

      shape ``(N,)`` **per observation, not per state**. Which
      instrument recorded each observation. Index it through ``obsid``
      (or use :meth:`instid_per_state`) to get a state's instrument.

   .. attribute:: obsid

      shape ``(K,)``. Which observation ``0..N-1`` each state belongs
      to. For an integrated kernel, the start and end states of one
      exposure share an ``obsid``.

   .. attribute:: stateid

      shape ``(K,)``. ``0`` for an exposure-start state (where
      :meth:`~smolgp.kernels.base.StateSpaceModel.reset_matrix` is
      applied), ``1`` for an exposure-end state or a plain instantaneous
      observation (i.e. the states that carry data).

   For an instantaneous kernel the "integrated" fields degenerate to a trivial
   convention: ``K == N``, ``obsid == arange(N)``, ``stateid == 1``
   everywhere (every state carries data). A single instrument dataset likewise
   trivially has ``instid == 0`` everywhere.


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


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


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


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


   .. py:method:: instantaneous(t_states: tinygp.helpers.JAXArray, *, sort: bool = True) -> StateCoords
      :classmethod:


      The degenerate ``StateCoords`` for an instantaneous kernel: one
      state per observation, all carrying data, all on one instrument.

      The state timeline is sorted (every solver steps forward in time, and
      :meth:`predict`'s ``searchsorted`` bracketing assumes it), with
      ``obsid`` recording which observation each state came from -- the same
      role it plays for an integrated kernel, so the usual
      ``sort-by-obsid`` machinery (:meth:`instid_per_state`,
      :meth:`~smolgp.gp.ConditionedStates.project_at_data`,
      :func:`~smolgp.solvers.sample.data_order_indices`) maps results back
      to input order unchanged. Already-sorted input yields the identity
      permutation, so this is a no-op in the common case.

      The sort is *stable*, so observations sharing a timestamp keep their
      input order relative to each other.

      :param t_states: the sortable coordinate of each observation, in input
                       order (not necessarily sorted).
      :param sort: set ``False`` only if ``t_states`` is known to be sorted and
                   the identity ``obsid`` is wanted verbatim.



   .. py:property:: num_states
      :type: int


      ``K``, the number of states in the timeline.


   .. py:property:: num_obs
      :type: int


      ``N``, the number of observations.


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

      ``instid`` gathered to shape ``(K,)``: the instrument id of the
      observation each state belongs to.