Predict Times by DanWaxman · Pull Request #135 · BasisResearch/dynestyx

DanWaxman · 2026-03-10T14:28:20Z

Addresses #50 #116

…mall typo.

…f_times, etc.

…crete-time); new demos in tutorials

…ick likelihood profile to quickstart

mattlevine22 · 2026-03-16T21:26:56Z

@DanWaxman I THINK this PR is more-or-less ready to be looked at.

DanWaxman · 2026-03-17T04:44:30Z

This looks great to a first read! I'll read in more depth tomorrow. In the mean time, we should write up & update the PR description. This should be squash & merged, so the description should be fairly complete as to sufficiently describe the progress & design choices on its own.

DanWaxman · 2026-03-17T14:00:28Z

dynestyx/inference/integrations/cd_dynamax/discrete.py

+    if ctrl_values is None:
        inputs = jnp.zeros((T1, control_dim))
+    elif ctrl_values.shape[0] > T1:
+        inds = jnp.searchsorted(ctrl_times, obs_times, side="left")
+        inputs = ctrl_values[inds]
+    else:
+        inputs = ctrl_values


Suggested change

if ctrl_values is None:

inputs = jnp.zeros((T1, control_dim))

elif ctrl_values.shape[0] > T1:

inds = jnp.searchsorted(ctrl_times, obs_times, side="left")

inputs = ctrl_values[inds]

else:

inputs = ctrl_values

if ctrl_values is None:

inputs = jnp.zeros((T1, control_dim))

elif ctrl_values.shape[0] > T1:

# Find controls aligned to obs_times

inds = jnp.searchsorted(ctrl_times, obs_times, side="left")

inputs = ctrl_values[inds]

else:

# Controls should align exactly with obs_times

inputs = ctrl_values

DanWaxman · 2026-03-17T14:02:18Z

dynestyx/inference/integrations/cd_dynamax/discrete.py

+    if ctrl_values is None:
+        inputs = jnp.zeros((T1, control_dim))
+    elif ctrl_values.shape[0] > T1:
+        # ctrl spans union of obs_times and predict_times; filter needs ctrl at obs_times only
+        inds = jnp.searchsorted(ctrl_times, obs_times, side="left")
+        inputs = ctrl_values[inds]
+    else:
+        inputs = ctrl_values


Suggested change

if ctrl_values is None:

inputs = jnp.zeros((T1, control_dim))

elif ctrl_values.shape[0] > T1:

# ctrl spans union of obs_times and predict_times; filter needs ctrl at obs_times only

inds = jnp.searchsorted(ctrl_times, obs_times, side="left")

inputs = ctrl_values[inds]

else:

inputs = ctrl_values

if ctrl_values is None:

inputs = jnp.zeros((T1, control_dim))

elif ctrl_values.shape[0] > T1:

# Find controls aligned to obs_times

inds = jnp.searchsorted(ctrl_times, obs_times, side="left")

inputs = ctrl_values[inds]

else:

# Controls should align exactly with obs_times

inputs = ctrl_values

DanWaxman · 2026-03-17T14:27:40Z

dynestyx/inference/integrations/cd_dynamax/discrete.py

-    Args:
-        name: Name of the factor.
-        dynamics: Dynamical model to filter.
-        filter_config: KFConfig, EKFConfig, or UKFConfig.
-        obs_times: Observation times.
-        obs_values: Observed values.
-        ctrl_times: Control times (optional).
-        ctrl_values: Control values (optional).


Why removed?

DanWaxman · 2026-03-17T14:28:23Z

dynestyx/inference/integrations/cuthbert/discrete.py

    record_kwargs = _config_to_record_kwargs(filter_config)

    ys = obs_values
    t1 = int(ys.shape[0])  # this is T+1 in cuthbert's convention


Suggested change

t1 = int(ys.shape[0]) # this is T+1 in cuthbert's convention

T1 = int(ys.shape[0]) # this is T+1 in cuthbert's convention

Let's use T1 to mirror the cddynamax conventions

DanWaxman · 2026-03-17T14:28:33Z

dynestyx/inference/integrations/cuthbert/discrete.py

    if ctrl_values is None:
        control_dim = dynamics.control_dim
        ctrl_values = jnp.zeros((t1, control_dim), dtype=ys.dtype)
+    elif ctrl_values.shape[0] > t1:


Suggested change

elif ctrl_values.shape[0] > t1:

elif ctrl_values.shape[0] > T1:

DanWaxman · 2026-03-17T15:38:22Z

dynestyx/simulators.py

+
+        return jax.vmap(_obs_step)(jnp.arange(T))
+    else:
+


Suggested change

DanWaxman · 2026-03-17T15:41:21Z

dynestyx/simulators.py

+            # TODO: Handle this case.
+            raise NotImplementedError(
+                "this is to-be-implemented. Should pass forward whatever is from previous operator in **kwargs."
+            )


DanWaxman · 2026-03-17T15:42:25Z

dynestyx/simulators.py

-                "Construct dynamics via DynamicalModel before simulation."
+        t0 = dynamics.t0 if dynamics.t0 is not None else times[0]
+
+        def _run_one_from_x0(key: Array, x0: Array) -> tuple[Array, Array]:


Run one what?

DanWaxman · 2026-03-17T15:43:58Z

dynestyx/simulators.py

-                f"y_{t_idx + 1}",
-                dynamics.observation_model(x=x_t, u=u_next, t=t_next),
-                obs=_get_val_or_None(obs_values, t_idx + 1),
+        # n_simulations > 1: vmap over scan with dist.sample (no numpyro.sample in body)


Should exaplain this more -- is it because numpyro doesn't place nice with the vmap?

DanWaxman · 2026-03-17T15:45:04Z

dynestyx/utils.py

    """
    Validate that ctrl_times/ctrl_values align with obs_times if provided.

    Raises:
        ValueError: If control times length doesn't match observation times length.
    """


Explain all the rules that are being tested

DanWaxman · 2026-03-17T15:45:46Z

@mattlevine22 left feedback

DanWaxman and others added 8 commits March 5, 2026 12:05

wip

0369682

wip formatted

2bb8f63

removing unneeded validations since they are done at handler level. s…

205665f

…mall typo.

add wip filtered predictions for sde

e606fd1

fmt

2570629

Merge main

b01ff1f

wip (use dynamicalmodel.t0 to start simulations)

8aca61e

dealing with length-1 simulations

ca22296

DanWaxman mentioned this pull request Mar 10, 2026

Adding predict_times option #119

Closed

mattlevine22 added 21 commits March 10, 2026 15:31

updating tests to use predict times and map onto new sites f_states, …

710f2a1

…f_times, etc.

more test restrictions and fix the shape of deltas

2eaa3aa

small bugfix/simplification

61d2507

giant/bold cursor-based commit. a lot of new predict_times functionality

47d78c6

update simulator predict-mode to get conntrol segments (needed in dis…

559e930

…crete-time); new demos in tutorials

final linter fixes and clean up notebooks more

b46a1d6

fixing tests with missing predict_times

a640eed

remove unneeded script

3874cb6

add interpolation between obs_times via predict_times

7c59c72

make sure all table cases work

b8a77ad

support num_samples>1 and n_simulations>1 and update notebooks

c89a033

re-run MLL nb

971c80d

re-ran notebook

45ee1bd

Merge branch 'main' into dw-ml-predict-times

50940db

quickstart uses more data for better MCMC (slower now); also added qu…

93463cd

…ick likelihood profile to quickstart

simplify code and add trailing dim for HMMs

41d506a

remove unneeded fallback

d77331d

please the lint

6563842

simplifying simulators.py

62aab00

modularizing _step in DiscreteTimeSimulator

b1d4821

remove unused function

bb9def6

mattlevine22 added 2 commits March 16, 2026 16:51

only check n_sim once;

99f0577

simplify BaseSimulator _sample_ds

d164210

mattlevine22 marked this pull request as ready for review March 16, 2026 21:25

mattlevine22 assigned DanWaxman Mar 16, 2026

update api reference documentation

0ebccc1

This was linked to issues Mar 17, 2026

Forecast evaluations will need a posterior "rollout" mode #50

Open

Support non-filtered predict_times for trajectory-based evaluation #116

Open

DanWaxman changed the title ~~Predict Times [WIP]~~ Predict Times Mar 17, 2026

DanWaxman requested a review from mattlevine22 March 17, 2026 13:33

DanWaxman commented Mar 17, 2026

View reviewed changes

Merge branch 'main' into dw-ml-predict-times

837abcc

	t1 = int(ys.shape[0]) # this is T+1 in cuthbert's convention
	T1 = int(ys.shape[0]) # this is T+1 in cuthbert's convention

	elif ctrl_values.shape[0] > t1:
	elif ctrl_values.shape[0] > T1:

Conversation

DanWaxman commented Mar 10, 2026

Uh oh!

mattlevine22 commented Mar 16, 2026

Uh oh!

DanWaxman commented Mar 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DanWaxman commented Mar 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants