state_coords

Classes

StateCoords

The state-level bookkeeping shared by every solver, filter, smoother,

Module Contents

class smolgp.solvers.state_coords.StateCoords

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

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

instid

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

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.

stateid

shape (K,). 0 for an exposure-start state (where 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.

t_states: tinygp.helpers.JAXArray

instid: tinygp.helpers.JAXArray

obsid: tinygp.helpers.JAXArray

stateid: tinygp.helpers.JAXArray

classmethod instantaneous(t_states: tinygp.helpers.JAXArray, *, sort: bool = True) → StateCoords

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 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 (instid_per_state(), project_at_data(), 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.

Parameters:

• t_states – the sortable coordinate of each observation, in input order (not necessarily sorted).

• sort – set False only if t_states is known to be sorted and the identity obsid is wanted verbatim.

property num_states: int

K, the number of states in the timeline.

property num_obs: int

N, the number of observations.

instid_per_state() → tinygp.helpers.JAXArray

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