sample

sample#

Functions#

sample_prior_trajectory(→ tinygp.helpers.JAXArray)

Draw one exact forward-simulation sample of the (possibly-augmented)

data_order_indices(→ tinygp.helpers.JAXArray)

The state-array indices (into state_coords' own K-length arrays)

project_trajectory_at_positions(→ tinygp.helpers.JAXArray)

Project a full state trajectory (shape (K_total, dim)) at

project_trajectory_at_data(→ tinygp.helpers.JAXArray)

Project a full state trajectory (shape (K, dim)) down to the N

merge_test_coords(...)

Merge new delta=0 (instantaneous) test times into an existing sorted

merge_exposure_test_coords(kernel, state_coords, ...)

Build an extended kernel and a merged, sorted timeline for jointly

project_exposure_test_points(→ tinygp.helpers.JAXArray)

Read out the prior sample at M (possibly exposure-integrated)

Module Contents#

smolgp.solvers.sample.sample_prior_trajectory(kernel, state_coords: smolgp.solvers.state_coords.StateCoords, key: jax.random.KeyArray) → tinygp.helpers.JAXArray[source]#

Draw one exact forward-simulation sample of the (possibly-augmented) latent state trajectory at every state in state_coords.

This is the “prior sampling” half of the residual/Matheron’s-rule conditional-sampling method (see smolgp.gp.GaussianProcess.sample()): a pure forward SDE simulation using the same transition_matrix/ process_noise/stationary_covariance/reset_matrix the Kalman filter itself uses, but with no Kalman gain/update logic.

Mirrors kalman_filter()’s own conventions exactly: the first state is drawn directly from the stationary distribution (matching that filter’s m0=zeros/P0=stationary_covariance() initialization) rather than via a transition from a nonexistent previous state, and reset_matrix is applied at stateid==0 states exactly as integrated_kalman_filter()’s update_start does. For non-integrated kernels, pass the same “unified” fallback used elsewhere (StateCoords.instantaneous(), i.e. instid=zeros, obsid=arange(K), stateid=ones(K)) – since StateSpaceModel.reset_matrix defaults to the identity, the stateid==0 branch is then a harmless no-op.

Parameters:

  • kernel – The (possibly-augmented) state space model.

  • state_coords – a StateCoords, as held by solver.state_coords / returned by solver.condition().

  • key – A jax random number key.

Returns:

The sampled trajectory, shape (K, kernel.dimension), in the same order as state_coords.

smolgp.solvers.sample.data_order_indices(state_coords: smolgp.solvers.state_coords.StateCoords, N: int) → tinygp.helpers.JAXArray[source]#

The state-array indices (into state_coords’ own K-length arrays) of the N real data points, sorted into data order (obsid 0..N-1).

Mirrors smolgp.gp.ConditionedStates.project_at_data()’s select-stateid==1/sort-by-obsid logic exactly.

smolgp.solvers.sample.project_trajectory_at_positions(X: tinygp.helpers.JAXArray, positions: tinygp.helpers.JAXArray, x_traj: tinygp.helpers.JAXArray, observation_model) → tinygp.helpers.JAXArray[source]#

Project a full state trajectory (shape (K_total, dim)) at explicit positions into it, applying observation_model at the corresponding X coordinates (X’s own leading dimension must match positions).

smolgp.solvers.sample.project_trajectory_at_data(X: tinygp.helpers.JAXArray, state_coords: smolgp.solvers.state_coords.StateCoords, x_traj: tinygp.helpers.JAXArray, observation_model, N: int) → tinygp.helpers.JAXArray[source]#

Project a full state trajectory (shape (K, dim)) down to the N real data points, in data order.

For the non-integrated fallback state_coords (stateid all 1, obsid=arange(K)), this reduces to the identity permutation.

smolgp.solvers.sample.merge_test_coords(state_coords: smolgp.solvers.state_coords.StateCoords, t_test: tinygp.helpers.JAXArray) → tuple[smolgp.solvers.state_coords.StateCoords, tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]#

Merge new delta=0 (instantaneous) test times into an existing sorted state_coords timeline, for a joint prior-trajectory sample covering both the original states and the new test points.

Test points are inserted with stateid=1 (no reset) and a dummy in-bounds obsid=0 (the reset branch is traced-but-discarded via jax.lax.cond even when unused, so it must not index out of bounds). Sorting uses the same (t, -stateid) tie-break convention IntegratedStateSpaceSolver already uses, so a test point exactly at an existing state’s time gets a zero-length transition to/from it and gets identically copied.

Parameters:

  • state_coords – the training timeline’s StateCoords (K states).

  • t_test – sortable test times to insert, length M.

Returns:

the merged, sorted StateCoords, with K+M

states. instid is passed through unchanged (length N): test points add new states, not new observations.

train_positions: length K, the new position of each original

state (in state_coords’ own order) within the merged arrays.

test_positions: length M, the new position of each t_test[i]

within the merged arrays.

Return type:

merged_state_coords

smolgp.solvers.sample.merge_exposure_test_coords(kernel, state_coords, t_test, delta_test, instid_test, num_test_insts)[source]#

Build an extended kernel and a merged, sorted timeline for jointly sampling the prior at (possibly) exposure-integrated (delta>0) test points. Generalizes merge_test_coords() to include exposures.

Each test window is simulated with a “virtual probe instrument”, the same idea predict_exposure() uses for filtering (simulation needs less machinery, since a prior draw has no data/update step at all). Concretely:

  • Probe dimensions. kernel_ext extends kernel by num_test_insts + 1 instruments. Test points sharing an instid_test share one probe dimension, exactly as repeated real exposures from one physical instrument already share one z slot, enabling overlapping test exposures to be simulated simultaneously with correct covariances.

  • Two states per test point. A reset at a_i = t_i - delta_i/2 (stateid=0), and a readout at b_i = t_i + delta_i/2 (stateid=1).

  • One joint simulation. Running sample_prior_trajectory() once over the merged timeline gives a single draw covering the real states and every probe, correctly correlated with each other and with any overlapping real observations (independent per-point simulations would not give).

Extending num_insts leaves the dynamics of the original n-dim block unchanged, so only the projection changes: use kernel_ext.observation_model (its extra probe columns are zero by construction) to match the extended trajectory’s dimension.

Parameters:

  • kernel – the original (non-extended) kernel.

  • state_coords – the training timeline’s StateCoords (K states, instid of length N).

  • t_test – exposure midpoints, length M.

  • delta_test – exposure widths (0 for instantaneous), length M.

  • instid_test – length M, which probe group (0..num_test_insts-1) each test point belongs to.

  • num_test_insts – the number of distinct probe groups, i.e. int(jnp.max(instid_test)) + 1. Must be a static Python int (concrete outside any enclosing jax.jit), since it determines kernel_ext’s dimension

Returns:

kernel with num_insts extended by num_test_insts + 1

(the +1 being the shared trash dimension).

merged_state_coords: the merged, sorted StateCoords, with

K + 2M states and an instid of length N + M (each test point contributes one new probe “observation”).

train_positions: length K, new position of each original state. b_positions: length M, new position of each test point’s b_i. probe_dims: length M, the state dimension to read each test

point’s probe value from at b_positions[i] (equal across test points sharing an instid_test).

Return type:

kernel_ext

smolgp.solvers.sample.project_exposure_test_points(X_test: tinygp.helpers.JAXArray, kernel_ext, x_traj_ext: tinygp.helpers.JAXArray, b_positions: tinygp.helpers.JAXArray, probe_dims: tinygp.helpers.JAXArray, delta_test: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Read out the prior sample at M (possibly exposure-integrated) test points from a trajectory produced by simulating over merge_exposure_test_coords()’s merged timeline.

For delta_test[i] == 0, uses the ordinary kernel_ext.observation_model-based readout at X_test[i] (matching project_trajectory_at_positions()). For delta_test[i] > 0, reads x_traj_ext[b_positions[i], probe_dims[i]] / delta_test[i] directly – the simulation analog of predict_exposure()’s z_mean = m_final[probe_idx] readout. Both are computed for every point and combined with jnp.where (rather than a data-dependent partition, which JAX’s static shapes don’t allow); the discarded branch for each point is finite but otherwise meaningless.