Predict Times

docs/api_reference/developer/simulators.md

@@ -1,15 +1,28 @@
 # Overview
 
 Simulators (also called *unrollers*) turn a `DynamicalModel` into explicit NumPyro
-`sample` sites for latent states and observations on a provided time grid.
+`sample` sites for latent states and observations on a chosen time grid.
 
-!!! note "Context"
+!!! note "When to use each time argument"
+    - **`obs_times` and `obs_values` must be provided together**:
+      - `obs_times` defines where observation sample sites (`y_t`) live.
+      - `obs_values` provides conditioning values for those sites via `obs=...`.
+      - Typical use: observed-data simulation/inference on a known observation grid.
+    - **`predict_times`**: use this when you want rollout trajectories at specific
+      times for simulation and/or post-filter rollout.
+      - In filter-rollout mode, predictions are generated at `predict_times` from
+        filtered posteriors.
+      - Typical use: forward simulation, forecasting, or dense trajectories for
+        visualization.
+    - **If both are provided**:
+      - `obs_times` controls filtering/conditioning points.
+      - `predict_times` controls where predicted trajectories are reported.
+    - **If both are omitted**: simulator does not run and adds no deterministic sites.
+
+!!! note "Context and caveats"
     - **NumPyro context required**: simulators call `numpyro.sample(...)` and draw
       randomness via NumPyro PRNG keys, so they must run inside a NumPyro model
       (or a `numpyro.handlers.seed(...)` context).
-    - **`obs_times` is required**: simulators only run when observation times are
-      provided (e.g. `dsx.sample(..., DynamicalModel(...), obs_times=...)`), because
-      those times define the trajectory grid.
     - **Conditioning is optional**: if `obs_values` is provided (e.g.
       `dsx.sample(..., DynamicalModel(...), obs_times=..., obs_values=...)`),
       simulators pass these values as `obs=...` to the observation `numpyro.sample`
@@ -20,13 +33,25 @@ Simulators (also called *unrollers*) turn a `DynamicalModel` into explicit NumPy
       `SDESimulator` is usually a poor inference strategy.
 
 !!! note "Deterministic sites"
-    When a simulator runs (i.e., when `obs_times` is provided), it records:
-    - `"times"`: the observation-time grid used for unrolling,
-    - `"states"`: the latent trajectory on that grid,
-    - `"observations"`: sampled (or conditioned) emissions on that grid.
+    When simulator trajectories are produced, sites are recorded as `"{name}_{key}"`
+    where `name` is the first
+    argument to `dsx.sample(name, dynamics, ...)` (conventionally `"f"`):
+
+    - `"f_times"`: trajectory time grid, shape `(n_sim, T)`,
+    - `"f_states"`: latent trajectory, shape `(n_sim, T, state_dim)`,
+    - `"f_observations"`: sampled or conditioned emissions, shape `(n_sim, T, obs_dim)`.
+
+    In filter-rollout mode (`predict_times` with filtered posteriors), additional
+    keys `"f_predicted_states"`, `"f_predicted_times"`, and
+    `"f_predicted_observations"` are recorded.
+
+    Under `numpyro.infer.Predictive(model, num_samples=N)`, NumPyro prepends a leading
+    `num_samples` axis, giving final shapes `(num_samples, n_sim, T, dim)`.
+    Use `dynestyx.flatten_draws` to collapse the `(num_samples, n_sim)` prefix into one
+    axis for plotting or downstream analysis.
 
-    If `obs_times` is omitted, no simulation is performed and these deterministic
-    sites are not added.
+    If both `obs_times` and `predict_times` are omitted, no simulation is performed
+    and these sites are not added.
 
 ## Simulators
 

docs/api_reference/public/simulators/discrete_time_simulator.md

@@ -43,10 +43,10 @@
     with DiscreteTimeSimulator():
         prior_pred = Predictive(model, num_samples=5)(
             jr.PRNGKey(0),
-            obs_times=obs_times,
+            predict_times=obs_times,
         )
-    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f', 'observations', 'phi', 'states', 'times', ...]
-    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # e.g. first axis is num_samples=5
+    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f_observations', 'f_states', 'f_times', 'phi', ...]
+    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # trajectory arrays: (num_samples, n_sim, T, dim); here num_samples=5, n_sim=1
     ```
 
 ??? example "NUTS with DiscreteTimeSimulator"
@@ -66,11 +66,11 @@
     print("Posterior sample keys:", sorted(posterior.keys()))  # stochastic sites (often includes latent x_* and parameters like 'phi')
     print("Posterior sample shapes:", {k: v.shape for k, v in posterior.items()})
 
-    # Deterministic trajectory keys like 'states'/'observations' are in posterior predictive output.
+    # Deterministic trajectory keys like 'f_states'/'f_observations' are in posterior predictive output.
     with DiscreteTimeSimulator():
         post_pred = Predictive(model, posterior_samples=posterior)(
-            jr.PRNGKey(2), obs_times=obs_times
+            jr.PRNGKey(2), predict_times=obs_times
         )
-    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'states', 'observations', 'times'
+    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'f_states', 'f_observations', 'f_times'
     print("Posterior predictive shapes:", {k: v.shape for k, v in post_pred.items()})
     ```

docs/api_reference/public/simulators/ode_simulator.md

@@ -41,9 +41,9 @@
 
     obs_times = jnp.linspace(0.0, 5.0, 51)
     with ODESimulator():
-        prior_pred = Predictive(model, num_samples=5)(jr.PRNGKey(0), obs_times=obs_times)
-    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f', 'observations', 'sigma_y', 'states', 'theta', 'times', ...]
-    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # e.g. first axis is num_samples=5
+        prior_pred = Predictive(model, num_samples=5)(jr.PRNGKey(0), predict_times=obs_times)
+    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f_observations', 'f_states', 'f_times', 'sigma_y', 'theta', ...]
+    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # trajectory arrays: (num_samples, n_sim, T, dim); here num_samples=5, n_sim=1
     ```
 
 ??? example "NUTS with ODESimulator"
@@ -63,11 +63,11 @@
     print("Posterior sample keys:", sorted(posterior.keys()))  # stochastic sites (typically parameters and x_0)
     print("Posterior sample shapes:", {k: v.shape for k, v in posterior.items()})
 
-    # Deterministic trajectories are exposed as 'states'/'observations' in posterior predictive output.
+    # Deterministic trajectories are exposed as 'f_states'/'f_observations' in posterior predictive output.
     with ODESimulator():
         post_pred = Predictive(model, posterior_samples=posterior)(
-            jr.PRNGKey(2), obs_times=obs_times
+            jr.PRNGKey(2), predict_times=obs_times
         )
-    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'states', 'observations', 'times'
+    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'f_states', 'f_observations', 'f_times'
     print("Posterior predictive shapes:", {k: v.shape for k, v in post_pred.items()})
     ```

docs/api_reference/public/simulators/overview.md

@@ -3,13 +3,26 @@
 Simulators (also called *unrollers*) turn a `DynamicalModel` into explicit NumPyro
 `sample` sites for latent states and observations on a provided time grid.
 
-!!! note "Context"
+!!! note "When to use each time argument"
+    - **`obs_times` and `obs_values` must be provided together**:
+      - `obs_times` defines where observation sample sites (`y_t`) live.
+      - `obs_values` provides conditioning values for those sites via `obs=...`.
+      - Typical use: observed-data simulation/inference on a known observation grid.
+    - **`predict_times`**: use this when you want rollout trajectories at specific
+      times for simulation and/or post-filter rollout.
+      - In filter-rollout mode, predictions are generated at `predict_times` from
+        filtered posteriors.
+      - Typical use: forward simulation, forecasting, or dense trajectories for
+        visualization.
+    - **If both are provided**:
+      - `obs_times` controls filtering/conditioning points.
+      - `predict_times` controls where predicted trajectories are reported.
+    - **If both are omitted**: simulator does not run and adds no deterministic sites.
+
+!!! note "Context and caveats"
     - **NumPyro context required**: simulators call `numpyro.sample(...)` and draw
       randomness via NumPyro PRNG keys, so they must run inside a NumPyro model
       (or a `numpyro.handlers.seed(...)` context).
-    - **`obs_times` is required**: simulators only run when observation times are
-      provided (e.g. `dsx.sample(..., DynamicalModel(...), obs_times=...)`), because
-      those times define the trajectory grid.
     - **Conditioning is optional**: if `obs_values` is provided (e.g.
       `dsx.sample(..., DynamicalModel(...), obs_times=..., obs_values=...)`),
       simulators pass these values as `obs=...` to the observation `numpyro.sample`
@@ -20,13 +33,25 @@ Simulators (also called *unrollers*) turn a `DynamicalModel` into explicit NumPy
       `SDESimulator` is usually a poor inference strategy.
 
 !!! note "Deterministic sites"
-    When a simulator runs (i.e., when `obs_times` is provided), it records:
-    - `"times"`: the observation-time grid used for unrolling,
-    - `"states"`: the latent trajectory on that grid,
-    - `"observations"`: sampled (or conditioned) emissions on that grid.
+    When simulator trajectories are produced, sites are recorded as `"{name}_{key}"`
+    where `name` is the first
+    argument to `dsx.sample(name, dynamics, ...)` (conventionally `"f"`):
+
+    - `"f_times"`: trajectory time grid, shape `(n_sim, T)`,
+    - `"f_states"`: latent trajectory, shape `(n_sim, T, state_dim)`,
+    - `"f_observations"`: sampled or conditioned emissions, shape `(n_sim, T, obs_dim)`.
+
+    In filter-rollout mode (`predict_times` with filtered posteriors), additional
+    keys `"f_predicted_states"`, `"f_predicted_times"`, and
+    `"f_predicted_observations"` are recorded.
+
+    Under `numpyro.infer.Predictive(model, num_samples=N)`, NumPyro prepends a leading
+    `num_samples` axis, giving final shapes `(num_samples, n_sim, T, dim)`.
+    Use `dynestyx.flatten_draws` to collapse the `(num_samples, n_sim)` prefix into one
+    axis for plotting or downstream analysis.
 
-    If `obs_times` is omitted, no simulation is performed and these deterministic
-    sites are not added.
+    If both `obs_times` and `predict_times` are omitted, no simulation is performed
+    and these sites are not added.
 
 ## BaseSimulator
 

docs/api_reference/public/simulators/sde_simulator.md

@@ -44,9 +44,9 @@
 
     obs_times = jnp.linspace(0.0, 5.0, 51)
     with SDESimulator():
-        prior_pred = Predictive(model, num_samples=5)(jr.PRNGKey(0), obs_times=obs_times)
-    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f', 'observations', 'sigma_x', 'sigma_y', 'states', 'theta', 'times', ...]
-    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # e.g. first axis is num_samples=5
+        prior_pred = Predictive(model, num_samples=5)(jr.PRNGKey(0), predict_times=obs_times)
+    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f_observations', 'f_states', 'f_times', 'sigma_x', 'sigma_y', 'theta', ...]
+    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # trajectory arrays: (num_samples, n_sim, T, dim); here num_samples=5, n_sim=1
     ```
 
 ??? example "NUTS with SDESimulator (small demonstration)"
@@ -67,11 +67,11 @@
     print("Posterior sample keys:", sorted(posterior.keys()))  # stochastic sites (typically parameters and x_0)
     print("Posterior sample shapes:", {k: v.shape for k, v in posterior.items()})
 
-    # Deterministic trajectories are exposed as 'states'/'observations' in posterior predictive output.
+    # Deterministic trajectories are exposed as 'f_states'/'f_observations' in posterior predictive output.
     with SDESimulator():
         post_pred = Predictive(model, posterior_samples=posterior)(
-            jr.PRNGKey(2), obs_times=obs_times
+            jr.PRNGKey(2), predict_times=obs_times
         )
-    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'states', 'observations', 'times'
+    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'f_states', 'f_observations', 'f_times'
     print("Posterior predictive shapes:", {k: v.shape for k, v in post_pred.items()})
     ```

docs/api_reference/public/simulators/simulator_wrapper.md

@@ -43,10 +43,10 @@
     with Simulator():
         prior_pred = Predictive(model, num_samples=5)(
             jr.PRNGKey(0),
-            obs_times=obs_times,
+            predict_times=obs_times,
         )
-    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f', 'observations', 'phi', 'states', 'times', ...]
-    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # e.g. first axis is num_samples=5
+    print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f_observations', 'f_states', 'f_times', 'phi', ...]
+    print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # trajectory arrays: (num_samples, n_sim, T, dim); here num_samples=5, n_sim=1
     ```
 
 ??? example "NUTS inference with auto-routing"
@@ -67,11 +67,11 @@
     print("Posterior sample keys:", sorted(posterior.keys()))  # stochastic sites (e.g. parameters, and possibly latent x_* sites)
     print("Posterior sample shapes:", {k: v.shape for k, v in posterior.items()})  # each shape starts with num_samples (here 100)
 
-    # Deterministic trajectory keys like 'states'/'observations' are in posterior predictive output.
+    # Deterministic trajectory keys like 'f_states'/'f_observations' are in posterior predictive output.
     with Simulator():
         post_pred = Predictive(model, posterior_samples=posterior)(
-            jr.PRNGKey(2), obs_times=obs_times
+            jr.PRNGKey(2), predict_times=obs_times
         )
-    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'states', 'observations', 'times'
+    print("Posterior predictive keys:", sorted(post_pred.keys()))  # includes 'f_states', 'f_observations', 'f_times'
     print("Posterior predictive shapes:", {k: v.shape for k, v in post_pred.items()})
     ```

docs/deep_dives/fhn_sparse_id.ipynb

@@ -253,10 +253,10 @@
         "\n",
         "predictive = Predictive(model_with_true_drift, num_samples=1, exclude_deterministic=False)\n",
         "with SDESimulator():\n",
-        "    synthetic = predictive(k_data, obs_times=obs_times)\n",
+        "    synthetic = predictive(k_data, predict_times=obs_times)\n",
         "\n",
-        "obs_values = synthetic[\"observations\"][0]  # y_k  shape (n_obs, observation_dim)\n",
-        "states = synthetic[\"states\"][0]            # x_{t_k}  shape (n_obs, state_dim)\n",
+        "obs_values = synthetic[\"f_observations\"][0]  # y_k  shape (n_obs, observation_dim)\n",
+        "states = synthetic[\"f_states\"][0]            # x_{t_k}  shape (n_obs, state_dim)\n",
         "times_1d = jnp.asarray(obs_times).squeeze()"
       ]
     },
@@ -761,7 +761,7 @@
   ],
   "metadata": {
     "kernelspec": {
-      "display_name": "dynestyx (3.12.11)",
+      "display_name": ".venv (3.12.11)",
       "language": "python",
       "name": "python3"
     },

docs/deep_dives/gp_drift.ipynb

@@ -214,7 +214,7 @@
    "outputs": [],
    "source": [
     "# True system: pipe in drift, **kwargs the rest\n",
-    "def model_with_true_drift(obs_times=None, obs_values=None):\n",
+    "def model_with_true_drift(obs_times=None, obs_values=None, predict_times=None):\n",
     "    return dsx.sample(\n",
     "        \"f\",\n",
     "        DynamicalModel(\n",
@@ -225,18 +225,20 @@
     "        ),\n",
     "        obs_times=obs_times,\n",
     "        obs_values=obs_values,\n",
+    "        predict_times=predict_times,\n",
     "    )\n",
     "\n",
     "\n",
     "predictive = Predictive(\n",
     "    model_with_true_drift, num_samples=1, exclude_deterministic=False\n",
     ")\n",
     "with SDESimulator():\n",
-    "    synthetic = predictive(k_data, obs_times=obs_times)\n",
+    "    synthetic = predictive(k_data, predict_times=obs_times)\n",
     "\n",
-    "obs_values = synthetic[\"observations\"][0]\n",
-    "states = synthetic[\"states\"][0]\n",
-    "times_1d = jnp.asarray(obs_times).squeeze()"
+    "# f_observations / f_states: (num_samples, n_sim, T, ...)\n",
+    "obs_values = synthetic[\"f_observations\"][0, 0]\n",
+    "states = synthetic[\"f_states\"][0, 0]\n",
+    "times_1d = jnp.asarray(obs_times).squeeze()\n"
    ]
   },
   {
@@ -401,7 +403,7 @@
     "from dynestyx.inference.filter_configs import ContinuousTimeEnKFConfig\n",
     "\n",
     "\n",
-    "def model_with_gp_drift(obs_times=None, obs_values=None):\n",
+    "def model_with_gp_drift(obs_times=None, obs_values=None, predict_times=None):\n",
     "    beta = numpyro.sample(\n",
     "        \"beta\", dist.Normal(0.0, 1.0).expand((MSTAR, state_dim)).to_event(2)\n",
     "    )\n",
@@ -411,6 +413,7 @@
     "        DynamicalModel(state_evolution=make_state_evolution(drift), **dynamics_kwargs),\n",
     "        obs_times=obs_times,\n",
     "        obs_values=obs_values,\n",
+    "        predict_times=predict_times,\n",
     "    )\n",
     "\n",
     "\n",
@@ -429,7 +432,7 @@
     "        k_svi, num_steps=num_steps, obs_times=obs_times, obs_values=obs_values\n",
     "    )\n",
     "\n",
-    "beta_map = guide.median(svi_result.params)[\"beta\"]"
+    "beta_map = guide.median(svi_result.params)[\"beta\"]\n"
    ]
   },
   {
@@ -1331,7 +1334,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": ".venv (3.12.11)",
    "language": "python",
    "name": "python3"
   },
@@ -1345,7 +1348,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.10"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,