diff --git a/dynestyx/api.py b/dynestyx/api.py index 67d2b6c6..7b9b4574 100644 --- a/dynestyx/api.py +++ b/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) diff --git a/dynestyx/handlers.py b/dynestyx/handlers.py index 5b8bfe29..e4a9fb06 100644 --- a/dynestyx/handlers.py +++ b/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 diff --git a/dynestyx/models/checkers.py b/dynestyx/models/checkers.py index c778700f..a2ac7765 100644 --- a/dynestyx/models/checkers.py +++ b/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( diff --git a/dynestyx/models/core.py b/dynestyx/models/core.py index 4fa14852..643d5ddf 100644 --- a/dynestyx/models/core.py +++ b/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 diff --git a/dynestyx/simulation/discrete.py b/dynestyx/simulation/discrete.py index 45a60135..fa9ce445 100644 --- a/dynestyx/simulation/discrete.py +++ b/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, ) diff --git a/dynestyx/simulation/utils.py b/dynestyx/simulation/utils.py index 44bbb7f7..68eab8a4 100644 --- a/dynestyx/simulation/utils.py +++ b/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) diff --git a/dynestyx/types.py b/dynestyx/types.py index 283cd738..a189d994 100644 --- a/dynestyx/types.py +++ b/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 diff --git a/dynestyx/utils.py b/dynestyx/utils.py index ff122ded..4311e21a 100644 --- a/dynestyx/utils.py +++ b/dynestyx/utils.py @@ -334,6 +334,8 @@ def _validate_controls( ctrl_values: Real[Array, "*ctrl_value_plate ctrl_time control_dim"] | Real[Array, "*ctrl_value_plate ctrl_time"] | None, + *, + observation_control_alignment: str = "same_time", ) -> None: """ Validate control inputs against model time grids. @@ -344,11 +346,22 @@ def _validate_controls( - If both obs_times and predict_times are present, ctrl_times must match their union. - Otherwise ctrl_times must match whichever single grid is provided. - Matching is set-like (order-insensitive) and length-preserving. + - When observation_control_alignment is "previous_transition", ctrl_times must + instead match predict_times[:-1] (one control per transition); obs_times-based + conditioning is not supported yet under this convention (see issue #312). Raises: ValueError: If controls are partially provided or no time grid is provided. """ + if observation_control_alignment == "previous_transition" and obs_times is not None: + raise ValueError( + "observation_control_alignment='previous_transition' does not " + "support obs_times-based conditioning yet (Filter/Smoother/" + "LatentPathBuilder posterior rollout); only predict_times-only " + "generation is supported. See issue #312." + ) + if ctrl_times is None: if ctrl_values is not None: raise ValueError( @@ -365,7 +378,11 @@ def _validate_controls( if obs_times is None and predict_times is None: raise ValueError("At least one of obs_times or predict_times must be provided") - if obs_times is None: + if observation_control_alignment == "previous_transition": + # obs_times is None here -- already rejected above otherwise. + assert predict_times is not None + total_obs_pred_times = predict_times[:-1] + elif obs_times is None: total_obs_pred_times = predict_times elif predict_times is None: total_obs_pred_times = obs_times diff --git a/tests/test_discrete_control.py b/tests/test_discrete_control.py index 97ba1229..bbc473ff 100644 --- a/tests/test_discrete_control.py +++ b/tests/test_discrete_control.py @@ -943,13 +943,15 @@ def test_dsx_simulate_with_control_policy_rejects_simulator_config(): def test_dsx_simulate_without_control_policy_unchanged(): """No control_policy given -> falls back to today's type-based routing, - returning a plain SimulatedResult (no controls field at all), not a - ControlledSimulatedResult.""" + returning a plain SimulatedResult, not a ControlledSimulatedResult. (The + plain SimulatedResult does carry its own `controls` field now (#312 + follow-up), but it's None here since no ctrl_values were supplied.)""" dynamics = _lti_1d() predict_times = jnp.arange(0.0, 5.0) result = dsx.simulate(dynamics, rng_key=jr.PRNGKey(0), predict_times=predict_times) - assert not hasattr(result, "controls") + assert not isinstance(result, ControlledSimulatedResult) + assert result.controls is None def test_initial_policy_state_threads_through_dsx_simulate(): diff --git a/tests/test_models_core.py b/tests/test_models_core.py index a877e82e..a602e3f3 100644 --- a/tests/test_models_core.py +++ b/tests/test_models_core.py @@ -803,3 +803,60 @@ def bad_cov_fn(t_now, t_next): ), observation_model=dsx.LinearGaussianObservation(H=jnp.eye(2), R=jnp.eye(2)), ) + + +# --------------------------------------------------------------------------- +# observation_control_alignment field (#312) +# --------------------------------------------------------------------------- + + +def test_observation_control_alignment_defaults_to_same_time() -> None: + model = _simple_discrete_model() + assert model.observation_control_alignment == "same_time" + + +def test_observation_control_alignment_previous_transition_stored() -> None: + model = DynamicalModel( + initial_condition=dist.Normal(0.0, 1.0), + state_evolution=lambda x, u, t_now, t_next: dist.Normal(x, 0.1), + observation_model=lambda x, u, t: dist.Normal(x, 0.1), + control_dim=0, + observation_control_alignment="previous_transition", + ) + assert model.observation_control_alignment == "previous_transition" + + +def test_observation_control_alignment_rejects_invalid_literal() -> None: + # Under the pytest jaxtyping import hook (see pyproject.toml addopts), + # jaxtyping's own Literal[...] enforcement rejects an invalid value before + # DynamicalModel.__init__'s body -- and _validate_observation_control_alignment + # within it -- ever runs, raising jaxtyping.TypeCheckError (a TypeError + # subclass) rather than the ValueError _validate_observation_control_alignment + # raises outside that instrumented context. Accept either so this test is + # correct with or without the import hook active. + with pytest.raises((ValueError, TypeError)): + DynamicalModel( + initial_condition=dist.Normal(0.0, 1.0), + state_evolution=lambda x, u, t_now, t_next: dist.Normal(x, 0.1), + observation_model=lambda x, u, t: dist.Normal(x, 0.1), + control_dim=0, + observation_control_alignment="bogus", # ty: ignore[invalid-argument-type] + ) + + +def test_observation_control_alignment_previous_transition_rejects_continuous_time() -> ( + None +): + with pytest.raises( + ValueError, + match="observation_control_alignment='previous_transition' is only supported " + "for discrete-time models", + ): + DynamicalModel( + initial_condition=_initial_condition_2d(), + state_evolution=ContinuousTimeStateEvolution( + drift=lambda x, u, t: -0.3 * x + ), + observation_model=_observation_model_2d, + observation_control_alignment="previous_transition", + ) diff --git a/tests/test_simulate_standalone.py b/tests/test_simulate_standalone.py index 07db9cb8..da6ad897 100644 --- a/tests/test_simulate_standalone.py +++ b/tests/test_simulate_standalone.py @@ -406,3 +406,233 @@ def model(): "Simulator + Predictive, dsx.simulate, and pre-split Simulator.simulate " f"produced different values for {mismatched_fields}" ) + + +# --------------------------------------------------------------------------- +# observation_control_alignment="previous_transition" (#312) +# --------------------------------------------------------------------------- + + +def _make_previous_transition_dynamics() -> dsx.DynamicalModel: + """Deterministic 1-D discrete model whose observation reveals both the + state and (scaled) control it was conditioned on, so tests can check + exactly which control an observation used.""" + + def _state_evolution(x, u, t_now, t_next): + del t_now, t_next + u = jnp.zeros_like(x) if u is None else u + return dist.Delta(x + u).to_event(1) + + def _observation_model(x, u, t): + del t + u = jnp.zeros_like(x) if u is None else u + return dist.Delta(x + 100.0 * u).to_event(1) + + return dsx.DynamicalModel( + control_dim=1, + initial_condition=dist.Delta(jnp.array([0.0])).to_event(1), + state_evolution=_state_evolution, + observation_model=_observation_model, + observation_control_alignment="previous_transition", + ) + + +def test_discrete_simulator_previous_transition_aligns_ctrl_values_shorter_by_one(): + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_times = predict_times[:-1] + ctrl_values = jnp.array([[1.0], [2.0], [3.0]]) + + result = dsx.DiscreteTimeSimulator().simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=ctrl_times, + ctrl_values=ctrl_values, + ) + + states = jnp.asarray(result.states) + observations = jnp.asarray(result.observations) + controls = jnp.asarray(result.controls) + times = jnp.asarray(result.times) + + # x_0 is populated and states/times include it (length 4), like same_time. + assert result.x_0 is not None + assert jnp.array_equal(jnp.asarray(result.x_0), jnp.array([[0.0]])) + assert times.shape == (1, 4) + assert jnp.allclose(times[0], predict_times) + assert states.shape == (1, 4, 1) + # x_0=0, x_1=1, x_2=3, x_3=6 (cumulative sum of controls) + expected_states = jnp.array([[[0.0], [1.0], [3.0], [6.0]]]) + assert jnp.array_equal(states, expected_states) + # observations/controls are one shorter -- y_0 is never sampled. + assert observations.shape == (1, 3, 1) + assert controls.shape == (1, 3, 1) + assert jnp.array_equal(controls[0], ctrl_values) + # y_{k+1} = x_{k+1} + 100 * u_k + expected_observations = expected_states[:, 1:, :] + 100.0 * jnp.array( + [[[1.0], [2.0], [3.0]]] + ) + assert jnp.array_equal(observations, expected_observations) + + +def test_discrete_simulator_previous_transition_rejects_ctrl_times_matching_full_predict_times(): + """dsx.simulate's _validate_controls gate requires an exact-length match + against predict_times[:-1] for previous_transition; DiscreteTimeSimulator's + own _align_ctrl_values_to_times permits a superset ctrl_times, so this must + go through the public dsx.simulate entry point to observe the rejection.""" + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_values = jnp.array([[1.0], [2.0], [3.0], [4.0]]) + + with pytest.raises(Exception): + dsx.simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=predict_times, + ctrl_values=ctrl_values, + ) + + +def test_dsx_simulate_previous_transition_end_to_end(): + """Exercises dsx.simulate -> api.py -> utils.py::_validate_controls threading.""" + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_times = predict_times[:-1] + ctrl_values = jnp.array([[1.0], [2.0], [3.0]]) + + result = dsx.simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=ctrl_times, + ctrl_values=ctrl_values, + ) + + assert result.x_0 is not None + assert jnp.asarray(result.times).shape == (1, 4) + assert jnp.asarray(result.states).shape == (1, 4, 1) + assert jnp.asarray(result.observations).shape == (1, 3, 1) + assert jnp.asarray(result.controls).shape == (1, 3, 1) + + +def test_discrete_simulator_previous_transition_zero_length_predict_times_edge_case(): + """With a single prediction time, there are zero transitions/controls, so + observations/controls are empty -- but x_0/states/times are still the + (length-1) seed, same as same_time.""" + result = dsx.simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=jnp.arange(1.0), + ) + + assert result.x_0 is not None + assert jnp.asarray(result.times).shape == (1, 1) + assert jnp.asarray(result.states).shape == (1, 1, 1) + assert jnp.asarray(result.observations).shape == (1, 0, 1) + assert result.controls is None # no ctrl_values supplied + + +def test_discrete_simulator_previous_transition_rejects_obs_times(): + predict_times = jnp.array([0.0, 1.0, 2.0]) + + with pytest.raises( + ValueError, + match="observation_control_alignment='previous_transition' does not support " + "obs_times", + ): + dsx.condition( + "f", + _make_previous_transition_dynamics(), + obs_times=predict_times, + obs_values=jnp.zeros((3, 1)), + predict_times=predict_times, + ) + + +def test_previous_transition_observation_uses_previous_step_control_not_same_index(): + """y_{k+1} must use u_k (the control that produced x_{k+1}), not u_{k+1}.""" + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_times = predict_times[:-1] + ctrl_values = jnp.array([[1.0], [2.0], [3.0]]) + + result = dsx.DiscreteTimeSimulator().simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=ctrl_times, + ctrl_values=ctrl_values, + ) + + states = jnp.asarray(result.states) + observations = jnp.asarray(result.observations) + # observation_model reveals x + 100*u, so subtracting the realized state + # recovers exactly which (scaled) control each observation used. + # states includes x_0 (one longer than observations), so observations[k] + # pairs with states[k+1], not states[k]. + revealed_control = (observations - states[:, 1:, :]) / 100.0 + assert jnp.array_equal(revealed_control[0], ctrl_values) + + +def test_discrete_simulator_previous_transition_states_include_x0(): + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_times = predict_times[:-1] + ctrl_values = jnp.array([[1.0], [2.0], [3.0]]) + + result = dsx.DiscreteTimeSimulator().simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=ctrl_times, + ctrl_values=ctrl_values, + ) + + states = jnp.asarray(result.states) + observations = jnp.asarray(result.observations) + assert jnp.array_equal(states[:, 0, :], jnp.asarray(result.x_0)) + assert states.shape[1] == observations.shape[1] + 1 + + +def test_discrete_simulator_returns_controls_same_time(): + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_values = jnp.array([[1.0], [2.0], [3.0], [4.0]]) + + result = dsx.DiscreteTimeSimulator().simulate( + _make_controlled_deterministic_discrete_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=predict_times, + ctrl_values=ctrl_values, + ) + + assert result.controls is not None + assert jnp.asarray(result.controls).shape == (1, 4, 1) + assert jnp.array_equal(jnp.asarray(result.controls)[0], ctrl_values) + + +def test_discrete_simulator_returns_controls_previous_transition(): + predict_times = jnp.array([0.0, 1.0, 2.0, 3.0]) + ctrl_times = predict_times[:-1] + ctrl_values = jnp.array([[1.0], [2.0], [3.0]]) + + result = dsx.DiscreteTimeSimulator().simulate( + _make_previous_transition_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ctrl_times=ctrl_times, + ctrl_values=ctrl_values, + ) + + assert result.controls is not None + assert jnp.asarray(result.controls).shape == (1, 3, 1) + assert jnp.array_equal(jnp.asarray(result.controls)[0], ctrl_values) + + +def test_discrete_simulator_controls_none_when_uncontrolled(): + predict_times = jnp.arange(4.0) + + result = dsx.simulate( + _make_discrete_dynamics(), + rng_key=jr.PRNGKey(0), + predict_times=predict_times, + ) + + assert result.controls is None