Define observation/control alignment for discrete-time models

dynestyx/api.py

@@ -55,7 +55,9 @@ def simulate(
         dynamics: Dynamical model to simulate.
         rng_key: JAX pseudorandom number generator key.
         ctrl_times: Times associated with `ctrl_values`. If controls are
-            provided, these times must match `predict_times`.
+            provided, these times must match `predict_times` for models with
+            `dynamics.observation_control_alignment="same_time"` (default), or
+            `predict_times[:-1]` for `"previous_transition"`.
         ctrl_values: Control values, or `None` for an uncontrolled model.
         predict_times: Times at which to simulate states and observations.
         n_simulations: Number of independent trajectories to simulate.
@@ -97,7 +99,13 @@ def simulate(
 
     _validate_site_sorting(ctrl_times, name="ctrl_times")
     _validate_site_sorting(predict_times, name="predict_times")
-    _validate_controls(None, predict_times, ctrl_times, ctrl_values)
+    _validate_controls(
+        None,
+        predict_times,
+        ctrl_times,
+        ctrl_values,
+        observation_control_alignment=dynamics.observation_control_alignment,
+    )
     _validate_control_dim(dynamics, ctrl_values)
 
     dynamics_with_t0 = _get_dynamics_with_t0(dynamics, None, predict_times)

dynestyx/handlers.py

@@ -100,7 +100,13 @@ def _validate_and_prepare(
     _validate_site_sorting(ctrl_times, name="ctrl_times")
     _validate_site_sorting(predict_times, name="predict_times")
 
-    _validate_controls(obs_times, predict_times, ctrl_times, ctrl_values)
+    _validate_controls(
+        obs_times,
+        predict_times,
+        ctrl_times,
+        ctrl_values,
+        observation_control_alignment=dynamics.observation_control_alignment,
+    )
     _validate_control_dim(dynamics, ctrl_values)
 
     # Initial dynamics may not have t0, which is then inferred from obs_times

dynestyx/models/checkers.py

@@ -235,6 +235,32 @@ def _validate_categorical_state(
         )
 
 
+_VALID_OBSERVATION_CONTROL_ALIGNMENTS = ("same_time", "previous_transition")
+
+
+def _validate_observation_control_alignment(
+    observation_control_alignment: str, continuous_time: bool
+) -> None:
+    """Validate the observation/control alignment convention.
+
+    Unlike continuous_time/categorical_state, this field has no inferred
+    counterpart -- it is purely user-supplied with a default, so this only
+    checks (1) the value is one of the two supported literals and (2)
+    'previous_transition' is only used with discrete-time state evolution.
+    """
+    if observation_control_alignment not in _VALID_OBSERVATION_CONTROL_ALIGNMENTS:
+        raise ValueError(
+            "observation_control_alignment must be one of "
+            f"{_VALID_OBSERVATION_CONTROL_ALIGNMENTS}, got "
+            f"{observation_control_alignment!r}."
+        )
+    if observation_control_alignment == "previous_transition" and continuous_time:
+        raise ValueError(
+            "observation_control_alignment='previous_transition' is only "
+            "supported for discrete-time models (continuous_time=False)."
+        )
+
+
 def _inside_numpyro_plate_context() -> bool:
     """Return True when currently executing inside any active numpyro.plate frame."""
     return any(

dynestyx/models/core.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Protocol, cast, runtime_checkable
+from typing import Any, Literal, Protocol, cast, runtime_checkable
 
 import equinox as eqx
 import jax
@@ -22,6 +22,7 @@
     _validate_continuous_time_flag,
     _validate_discrete_state_evolution_output_shape,
     _validate_imex_potential_conflict,
+    _validate_observation_control_alignment,
     _validate_observation_dim,
     _validate_state_dim,
 )
@@ -103,7 +104,16 @@ class DynamicalModel(eqx.Module):
             exactly; a mismatch raises a ``ValueError`` at simulation time.
         continuous_time (bool): Whether the model uses continuous-time state evolution (SDE) or discrete-time.
             Gets set automatically from the concrete type of `state_evolution`.
-
+        observation_control_alignment ("same_time" | "previous_transition"): Convention for how
+            observations pair with controls in discrete time. `"same_time"` (default) pairs
+            $y_k$ with $u_k$, matching today's `dsx.simulate` behavior. `"previous_transition"`
+            pairs $y_{k+1}$ with $u_k$ (the control that produced $x_{k+1}$), matching
+            `DiscreteControlLoopSimulator`'s closed-loop convention; under this convention
+            $y_0$ is never sampled. Only `"same_time"` is honored outside the plain
+            `Simulator`/`DiscreteTimeSimulator`/`dsx.simulate` generation path (not yet by
+            Filter/Smoother/`LatentPathBuilder` posterior rollout, `DiscreteControlLoopSimulator`,
+            or `mppi.py` -- see [issue #312](https://github.com/BasisResearch/dynestyx/issues/312)).
+
     Note:
         - `continuous_time`, `state_dim`, `observation_dim`, and `categorical_state` are inferred automatically; do not pass them to the constructor.
         - Logic for control_model is not implemented yet.
@@ -127,6 +137,7 @@ class DynamicalModel(eqx.Module):
     observation_dim: int
     categorical_state: bool
     continuous_time: bool
+    observation_control_alignment: Literal["same_time", "previous_transition"]
 
     def __init__(
         self,
@@ -141,12 +152,19 @@ def __init__(
         observation_dim: int | None = None,
         categorical_state: bool | None = None,
         continuous_time: bool | None = None,
+        observation_control_alignment: Literal[
+            "same_time", "previous_transition"
+        ] = "same_time",
     ):
         inferred_continuous_time = isinstance(
             state_evolution, ContinuousTimeStateEvolution
         )
         _validate_continuous_time_flag(continuous_time, inferred_continuous_time)
         self.continuous_time = inferred_continuous_time
+        _validate_observation_control_alignment(
+            observation_control_alignment, self.continuous_time
+        )
+        self.observation_control_alignment = observation_control_alignment
         self.initial_condition = initial_condition
         self.state_evolution = state_evolution
         self.observation_model = observation_model

dynestyx/simulation/discrete.py

@@ -49,7 +49,9 @@ def _sample_discrete_state_path_from_initial_state(
     initial_state: Real[Array, " state_dim"] | Real[Array, ""],
     rng_key: PRNGKeyArray,
     times: Real[Array, " time"],
-    ctrl_values: Real[Array, "time control_dim"] | Real[Array, " time"] | None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"]
+    | Real[Array, " ctrl_time"]
+    | None,
 ) -> Real[Array, "time state_dim"] | Real[Array, " time"]:
     """Sample one canonical discrete state path from a fixed initial state."""
     if len(times) == 1:
@@ -108,20 +110,14 @@ class DiscreteTimeSimulator(BaseSimulator):
     r"""Generate trajectories from a discrete-time dynamical model.
 
     For prediction times \(t_0,\ldots,t_{T-1}\), this simulator draws
-    `n_simulations` independent paths according to
-
-    \[
-    x_0^{(m)} \sim p_0(x_0), \qquad
-    x_{k+1}^{(m)}
-      \sim p\!\left(x_{k+1}\mid x_k^{(m)},u_k,t_k,t_{k+1}\right),
-    \qquad
-    y_k^{(m)} \sim p(y_k\mid x_k^{(m)},u_k,t_k).
-    \]
-
-    The first state in the returned path is the initial-condition draw at
-    `predict_times[0]`; the simulator then makes one transition draw for each
-    adjacent pair of prediction times and samples one observation conditional
-    on every realized state. See
+    `n_simulations` independent paths. The observation/control pairing
+    depends on `dynamics.observation_control_alignment`:
+
+    For `"same_time"` (default): y_{k} is paired with u_{k} and x_{k} (including k=0). States, times, observations, and controls are all of length \(T\).
+    For `"previous_transition"`: y_{k+1} is paired with u_{k} and x_{k+1} (y_0 is never sampled). States and times are of length \(T\), but observations and controls are of length \(T-1\).
+
+
+    See
     [DiscreteTimeStateEvolution][dynestyx.models.core.DiscreteTimeStateEvolution]
     for how a discrete transition model is represented in a `DynamicalModel`.
 
@@ -165,10 +161,14 @@ class DiscreteTimeSimulator(BaseSimulator):
     not be uniformly spaced, provided the model's transition accepts those
     intervals.
 
-    If controls are supplied, `ctrl_times` must contain every prediction time
-    exactly. `ctrl_values[k]` is used for the transition beginning at \(t_k\)
-    and for the observation at \(t_k\). The paired control arrays are validated
-    before simulation.
+    If controls are supplied, `ctrl_times` must exactly match the grid
+    required by `dynamics.observation_control_alignment`: the full
+    `predict_times` for `"same_time"` (default) -- `ctrl_values[k]` is used
+    for the transition beginning at \(t_k\) and for the observation at
+    \(t_k\) -- or `predict_times[:-1]` for `"previous_transition"` --
+    `ctrl_values[k]` drives the transition into \(x_{k+1}\) and the
+    observation \(y_{k+1}\). The paired control arrays are validated before
+    simulation.
 
     This handler is generation-only and does not condition on `obs_times` or
     `obs_values`. Use
@@ -198,7 +198,13 @@ class DiscreteTimeSimulator(BaseSimulator):
     - `"f_states"`: latent states, shape
       `(*plate_shape, n_simulations, T, state_dim)`;
     - `"f_observations"`: sampled observations, shape
-      `(*plate_shape, n_simulations, T, observation_dim)`.
+      `(*plate_shape, n_simulations, T, observation_dim)` for `"same_time"`,
+      or `(*plate_shape, n_simulations, T-1, observation_dim)` for
+      `"previous_transition"`;
+    - `"f_controls"`: the (aligned) controls used, when the model is
+      controlled, shape `(*plate_shape, n_simulations, T, control_dim)` for
+      `"same_time"` or `(*plate_shape, n_simulations, T-1, control_dim)` for
+      `"previous_transition"`; absent when the model is uncontrolled.
 
     Here `"f"` is replaced by the `name` passed to `dsx.sample`. Under
     `Predictive(..., num_samples=N)`, NumPyro prepends an `N` axis to each
@@ -250,15 +256,22 @@ def _simulate_forward_from_initial_state(
         | Real[Array, " n_simulations"],
         rng_key: PRNGKeyArray,
         times: Real[Array, " time"],
-        ctrl_values: Real[Array, "time control_dim"] | Real[Array, " time"] | None,
+        ctrl_values: Real[Array, "ctrl_time control_dim"]
+        | Real[Array, " ctrl_time"]
+        | None,
     ) -> SimulatedResult:
-        """Run pure forward simulation for a discrete-time model."""
+        """Run pure forward simulation for a discrete-time model.
+
+        ctrl_values has its own length ("ctrl_time"), decoupled from `times`
+        (always the full predict_times grid): len(times) for same_time or
+        len(times) - 1 for previous_transition. States always include x_0
+        (length matches `times`) for both conventions; only observations (and
+        the returned controls) are one shorter for previous_transition.
+        """
         n_sim = initial_state.shape[0]
         sim_keys = jr.split(rng_key, n_sim)
-        ctrl_eval = (
-            (lambda t: ctrl_values[jnp.searchsorted(times, t, side="left")])
-            if ctrl_values is not None
-            else None
+        include_initial_condition = (
+            dynamics.observation_control_alignment != "previous_transition"
         )
 
         def _sim_one_trajectory(
@@ -273,23 +286,38 @@ def _sim_one_trajectory(
                 times=times,
                 ctrl_values=ctrl_values,
             )
+            # For previous_transition, drop x_0/t_0 before sampling
+            # observations -- y_0 is never sampled under that convention.
+            # Otherwise (same_time) this is a no-op.
+            obs_states, obs_times = (
+                (states, times)
+                if include_initial_condition
+                else (states[1:], times[1:])
+            )
+            ctrl_eval = (
+                (lambda t: ctrl_values[jnp.searchsorted(obs_times, t, side="left")])
+                if ctrl_values is not None
+                else None
+            )
             observations = self._emit_observations(
-                "",
-                dynamics,
-                states,
-                times,
-                None,
-                ctrl_eval,
-                key=key_obs,
+                "", dynamics, obs_states, obs_times, None, ctrl_eval, key=key_obs
             )
             return states, observations
 
         states, observations = jax.vmap(_sim_one_trajectory)(sim_keys, initial_state)
+
+        controls = None
+        if ctrl_values is not None:
+            controls = _ensure_trailing_dim(
+                jnp.broadcast_to(ctrl_values[None], (n_sim, *ctrl_values.shape))
+            )
+
         return SimulatedResult(
             times=_tile_times(times, n_sim),
             x_0=initial_state,
             states=_ensure_trailing_dim(states),
             observations=_ensure_trailing_dim(observations),
+            controls=controls,
         )
 
     def simulate(
@@ -315,8 +343,13 @@ def simulate(
         if predict_times is None:
             raise ValueError("predict_times must be provided")
 
+        align_times = (
+            predict_times[:-1]
+            if dynamics.observation_control_alignment == "previous_transition"
+            else predict_times
+        )
         aligned_ctrl_values = _align_ctrl_values_to_times(
-            times=predict_times,
+            times=align_times,
             ctrl_times=ctrl_times,
             ctrl_values=ctrl_values,
         )

dynestyx/simulation/utils.py

@@ -9,7 +9,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import numpyro
-from jaxtyping import Array, Bool, Int, PRNGKeyArray, Real
+from jaxtyping import Array, Bool, PRNGKeyArray, Real
 
 from dynestyx.models import DynamicalModel
 from dynestyx.types import SimulatedResult, chain_numpyro_site_registrations
@@ -115,10 +115,14 @@ def _sample_observation_path(
     ctrl = control_path_eval if control_path_eval is not None else (lambda t: None)
     obs_keys = jr.split(rng_key, len(times))
 
-    def _sample_at_time(t_idx: Int[Array, ""]):
-        x_t = states[t_idx]
-        t = times[t_idx]
+    # Map directly over states/times/obs_keys rather than indexing by a
+    # scanned integer inside the mapped body: jax.vmap traces its body once
+    # regardless of batch size, so indexing into a genuinely zero-length
+    # array (e.g. a previous_transition observation path sliced down from a
+    # single-timepoint prediction grid) would raise immediately. Mapping over
+    # the arrays directly lets vmap's own batching handle the zero-size case.
+    def _sample_at(x_t, t, key):
         obs_dist = dynamics.observation_model(x=x_t, u=ctrl(t), t=t)
-        return obs_dist.sample(obs_keys[t_idx])
+        return obs_dist.sample(key)
 
-    return jax.vmap(_sample_at_time)(jnp.arange(len(times)))
+    return jax.vmap(_sample_at)(states, times, obs_keys)

dynestyx/types.py

@@ -143,8 +143,30 @@ class SimulatedResult(eqx.Module):
     posterior rollout, the same result object instead carries
     ``predicted_times``, ``predicted_states``, and
     ``predicted_observations``.
+
+    ``controls`` carries the (aligned) control values used to produce this
+    result, when the model was controlled -- ``None`` otherwise. It is
+    populated by ``DiscreteTimeSimulator`` for both
+    ``observation_control_alignment`` conventions; ODE/SDE simulators leave it
+    ``None`` for now.
+
+    For a discrete-time model with
+    ``dynamics.observation_control_alignment="previous_transition"``, ``x_0``
+    is populated and ``states`` includes it (length :math:`T`, matching
+    ``times``), exactly like ``"same_time"``. ``observations`` and
+    ``controls``, however, are one shorter (length :math:`T-1`:
+    :math:`y_1,\\dots,y_{T-1}` and :math:`u_0,\\dots,u_{T-2}`), since
+    :math:`y_0` is never sampled under this convention -- there is no control
+    that produced it. So ``states`` is intentionally one longer than
+    ``observations``/``controls``: ``states[k+1]`` pairs with
+    ``observations[k]``/``controls[k]``, not ``states[k]``. See
+    [DiscreteTimeSimulator][dynestyx.simulation.discrete.DiscreteTimeSimulator].
     """
 
+    # observations/controls use their own axis names ("obs_time"/"ctrl_time")
+    # rather than sharing "time" with times/states: under
+    # observation_control_alignment="previous_transition" they are one
+    # shorter than times/states, so jaxtyping must not enforce equal length.
     times: Real[Array, "*plate n_simulations time"] | None = None
     x_0: (
         Real[Array, "*plate n_simulations state_dim"]
@@ -157,8 +179,13 @@ class SimulatedResult(eqx.Module):
         | None
     ) = None
     observations: (
-        Real[Array, "*plate n_simulations time observation_dim"]
-        | Real[Array, "*plate n_simulations time"]
+        Real[Array, "*plate n_simulations obs_time observation_dim"]
+        | Real[Array, "*plate n_simulations obs_time"]
+        | None
+    ) = None
+    controls: (
+        Real[Array, "*plate n_simulations ctrl_time control_dim"]
+        | Real[Array, "*plate n_simulations ctrl_time"]
         | None
     ) = None
     predicted_times: Real[Array, "*plate n_simulations predict_time"] | None = None