sample ====== .. py:module:: smolgp.solvers.sample Functions --------- .. autoapisummary:: smolgp.solvers.sample.sample_prior_trajectory smolgp.solvers.sample.data_order_indices smolgp.solvers.sample.project_trajectory_at_positions smolgp.solvers.sample.project_trajectory_at_data smolgp.solvers.sample.merge_test_coords smolgp.solvers.sample.merge_exposure_test_coords smolgp.solvers.sample.project_exposure_test_points Module Contents --------------- .. py:function:: sample_prior_trajectory(kernel, state_coords: smolgp.solvers.state_coords.StateCoords, key: jax.random.KeyArray) -> tinygp.helpers.JAXArray 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 :meth:`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 :func:`~smolgp.solvers.kalman.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 :func:`~smolgp.solvers.integrated.kalman.integrated_kalman_filter`'s ``update_start`` does. For non-integrated kernels, pass the same "unified" fallback used elsewhere (:meth:`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. :param kernel: The (possibly-augmented) state space model. :param state_coords: a :class:`~smolgp.solvers.state_coords.StateCoords`, as held by ``solver.state_coords`` / returned by ``solver.condition()``. :param key: A ``jax`` random number key. :returns: The sampled trajectory, shape ``(K, kernel.dimension)``, in the same order as ``state_coords``. .. py:function:: data_order_indices(state_coords: smolgp.solvers.state_coords.StateCoords, N: int) -> tinygp.helpers.JAXArray 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 :meth:`smolgp.gp.ConditionedStates.project_at_data`'s select-``stateid==1``/sort-by-``obsid`` logic exactly. .. py:function:: project_trajectory_at_positions(X: tinygp.helpers.JAXArray, positions: tinygp.helpers.JAXArray, x_traj: tinygp.helpers.JAXArray, observation_model) -> tinygp.helpers.JAXArray 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``). .. py:function:: 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 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. .. py:function:: 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] 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 :class:`~smolgp.solvers.integrated.solver.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. :param state_coords: the training timeline's :class:`~smolgp.solvers.state_coords.StateCoords` (``K`` states). :param 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. :rtype: merged_state_coords .. py:function:: merge_exposure_test_coords(kernel, state_coords, t_test, delta_test, instid_test, num_test_insts) Build an extended kernel and a merged, sorted timeline for jointly sampling the prior at (possibly) exposure-integrated (``delta>0``) test points. Generalizes :func:`merge_test_coords` to include exposures. Each test window is simulated with a "virtual probe instrument", the same idea :func:`~smolgp.solvers.integrated.predict_exposure.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 :func:`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. :param kernel: the original (non-extended) kernel. :param state_coords: the training timeline's :class:`~smolgp.solvers.state_coords.StateCoords` (``K`` states, ``instid`` of length ``N``). :param t_test: exposure midpoints, length ``M``. :param delta_test: exposure widths (``0`` for instantaneous), length ``M``. :param instid_test: length ``M``, which probe group (``0..num_test_insts-1``) each test point belongs to. :param 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``). :rtype: kernel_ext .. py:function:: 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 Read out the prior sample at ``M`` (possibly exposure-integrated) test points from a trajectory produced by simulating over :func:`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 :func:`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 :func:`~smolgp.solvers.integrated.predict_exposure.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.