Add fast SDE solver backends and filter-simulator coverage. by mattlevine22 · Pull Request #178 · BasisResearch/dynestyx

mattlevine22 · 2026-03-26T23:32:21Z

Introduce solver modules for ODE and SDE paths, add an Euler-Maruyama scan source for SDESimulator, and update filter+simulator tests and continuous-time tutorial outputs to cover the new backend behavior.

Made-with: Cursor

Introduce solver modules for ODE and SDE paths, add an Euler-Maruyama scan source for SDESimulator, and update filter+simulator tests and continuous-time tutorial outputs to cover the new backend behavior. Made-with: Cursor

mattlevine22 · 2026-03-27T00:14:07Z

This PR seems ready, and I think its tests are actually passing because main is failing right now.

mattlevine22 · 2026-03-27T17:52:55Z

Let's finish the hierarchical PR first #173

DanWaxman · 2026-04-13T19:00:46Z

With some Codex help, merged main into here. Changes seem reasonable + seems to work still.

n.b. unlike the NumPyro calls, we can actually vmap over these scan calls instead of doing everything in a big for loop. Worth considering.

mattlevine22 · 2026-04-14T14:57:29Z

@DanWaxman looks great thanks!

I added a speed tip to the SDESimulator docs
Regarding vmap vs for loop over the scans...I believe we already vmap; SDESimulator line 838 has states, emissions = jax.vmap(_sim_one_trajectory)(...)

DanWaxman · 2026-04-14T15:26:44Z

Regarding vmap vs for loop over the scans...I believe we already vmap; SDESimulator line 838 has states, emissions = jax.vmap(_sim_one_trajectory)(...)

Ah, you're right, thanks! I thought we didn't do this for ODE for some reason, I guess we're just slow for discrete-time systems. :)

Copilot

Pull request overview

Adds dedicated ODE/SDE solver backends and expands simulator+filter test coverage to exercise both Diffrax and a faster Euler–Maruyama scan backend for SDE simulation.

Changes:

Introduces dynestyx.solvers.solve_ode and dynestyx.solvers.solve_sde (with "diffrax" and "em_scan" sources for SDEs).
Updates SDESimulator to dispatch SDE integration via the new solver backend and exposes a source selector.
Expands test parametrization to run key simulator/filter shape and smoke tests against both SDE backends.

Reviewed changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`dynestyx/simulators.py`	Routes ODE/SDE integration through new solver backends; adds `SDESimulator(source=...)`.
`dynestyx/solvers/sde.py`	New SDE backend module with Diffrax and Euler–Maruyama (`lax.scan`) implementations.
`dynestyx/solvers/odes.py`	New ODE backend module wrapping Diffrax ODE solves.
`dynestyx/solvers/__init__.py`	Exposes `solve_ode` / `solve_sde` public API.
`tests/test_predictive_filter_simulator_shapes.py`	Parametrizes SDE shape tests across `"diffrax"` and `"em_scan"`.
`tests/test_hierarchical_simulator_discretizer_smokes.py`	Parametrizes plated SDE smoke test across both SDE sources.
`tests/test_filter_simulator.py`	Parametrizes SDE simulator+filter integration tests across both SDE sources.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-14T15:32:50Z

+    if key is None:
+        raise ValueError("PRNG key is required for em_scan SDE solves.")
+    if dt0 <= 0:
+        raise ValueError(f"EM scan requires dt0 > 0, got dt0={dt0}.")


_solve_sde_scan raises when key is None even if no integration would be performed (e.g., all saveat_times <= t0, where the output is just x0 repeated). For parity with the diffrax backend (which early-returns without requiring a key), consider moving the key is None check to only the code path that actually advances time, or adding an early-return branch before validating the key.

Done. Resulted in a larger refactor to not duplicate early_return. This will need further refactor to support #199, but I suggest we complete this PR as is and shift the subsequent load onto #199.

Copilot · 2026-04-14T15:32:51Z

        assert self.tol_vbt < dt0, (
            "tol_vbt must be smaller than dt0 for statistically correct simulation."
        )
+        if self.source not in {"diffrax", "em_scan"}:
+            raise ValueError(


SDESimulator enforces tol_vbt < dt0 unconditionally, but tol_vbt is only relevant for the source="diffrax" backend. With source="em_scan", this assertion can fail even though the value is unused. Consider guarding the assertion (and any tol_vbt-specific validation) behind if self.source == "diffrax", and optionally ignoring/forbidding non-None tol_vbt when source="em_scan" to avoid confusing configuration errors.

Copilot · 2026-04-14T15:32:51Z

+        if not isinstance(dt0_setting, (int, float)):
+            raise ValueError(
+                "solve_sde(source='em_scan') requires a numeric dt0 in "
+                "diffeqsolve_settings."
+            )
+        return _solve_sde_scan(
+            dynamics,
+            t0,
+            saveat_times,
+            x0,
+            control_path_eval,
+            float(dt0_setting),


The em_scan branch rejects dt0 unless it is a Python int/float. This will raise for common scalar types like numpy.float64 (and potentially 0-d JAX arrays), even though float(dt0_setting) would work. Consider accepting any real scalar by attempting float(dt0_setting) in a try/except (or checking against numbers.Real / np.floating) so solve_sde(source="em_scan") is robust to scalar-like inputs.

Suggested change

if not isinstance(dt0_setting, (int, float)):

raise ValueError(

"solve_sde(source='em_scan') requires a numeric dt0 in "

"diffeqsolve_settings."

)

return _solve_sde_scan(

dynamics,

t0,

saveat_times,

x0,

control_path_eval,

float(dt0_setting),

try:

dt0 = float(dt0_setting)

except (TypeError, ValueError) as exc:

raise ValueError(

"solve_sde(source='em_scan') requires a numeric dt0 in "

"diffeqsolve_settings."

) from exc

return _solve_sde_scan(

dynamics,

t0,

saveat_times,

x0,

control_path_eval,

dt0,

Done---introduced a _coerce_dt utility.

DanWaxman

Minor comments only. Other remaining thing, we should probably have a test that compares statistics of our hand-written Euler-Maruyama, on a non-trivial SDE with controls, to the results of diffrax using the Euler solver and a sufficiently small tol_vbt parameter.

DanWaxman · 2026-04-21T18:06:04Z

+def _coerce_dt(dt: object, *, name: str = "dt") -> float:
+    """Coerce a numeric scalar timestep-like value to Python float."""
+    if isinstance(dt, bool):
+        raise ValueError(
+            f"{name} must be a numeric scalar (Python/NumPy real or scalar JAX array), "
+            f"got {type(dt)}."
+        )
+    if isinstance(dt, Real):
+        return float(dt)
+    if (
+        isinstance(dt, Array)
+        and dt.ndim == 0
+        and not jnp.issubdtype(dt.dtype, jnp.bool_)
+        and (
+            jnp.issubdtype(dt.dtype, jnp.integer)
+            or jnp.issubdtype(dt.dtype, jnp.floating)
+        )
+    ):
+        return float(dt)
+    raise ValueError(
+        f"{name} must be a numeric scalar (Python/NumPy real or scalar JAX array), "
+        f"got {type(dt)}."
+    )


I kinda don't understand why we need this. If we need a float, we probably need a float everywhere and should require that from the beginning.

I think the main use case is when it is given as a jax.Array (perhaps due to some batching needs).

The bool/Real checks are definitely a bit much, want me to remove them?

As discussed, I guess my main objection is that we are concretizing a jax array in the first place. I think this will break things if traced. If we can at all (i) store as a float from the beginning, and only convert to an array as-needed, or (ii) always store as a jax array and deal with the consequences of that, I think the code will be more robust.

I looked more closely, and I think the main reason we have floats anywhere is for friendly user API (no one wants to write dt0=jnp.array(0.01) ).

What do you think about doing float -> Array at a high-level early on and then do our internals all on Arrays?

@DanWaxman

Apply suggestion from @DanWaxman Co-authored-by: Dan Waxman <dan.waxman1@gmail.com>

mattlevine22

Fixed the 1 minor correction and responded to use case of _coerce_dt

mattlevine22 · 2026-04-21T18:48:11Z

+def _coerce_dt(dt: object, *, name: str = "dt") -> float:
+    """Coerce a numeric scalar timestep-like value to Python float."""
+    if isinstance(dt, bool):
+        raise ValueError(
+            f"{name} must be a numeric scalar (Python/NumPy real or scalar JAX array), "
+            f"got {type(dt)}."
+        )
+    if isinstance(dt, Real):
+        return float(dt)
+    if (
+        isinstance(dt, Array)
+        and dt.ndim == 0
+        and not jnp.issubdtype(dt.dtype, jnp.bool_)
+        and (
+            jnp.issubdtype(dt.dtype, jnp.integer)
+            or jnp.issubdtype(dt.dtype, jnp.floating)
+        )
+    ):
+        return float(dt)
+    raise ValueError(
+        f"{name} must be a numeric scalar (Python/NumPy real or scalar JAX array), "
+        f"got {type(dt)}."
+    )


I think the main use case is when it is given as a jax.Array (perhaps due to some batching needs).

The bool/Real checks are definitely a bit much, want me to remove them?

DanWaxman · 2026-04-21T19:50:28Z

I'd be happy with conversion at init :).

…

On Tue, Apr 21, 2026 at 3:46 PM Matt Levine ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In dynestyx/utils.py <#178?email_source=notifications&email_token=AC2DPON4AGQNCRSQSWCXY2D4W7FYRA5CNFSNUABKM5UWIORPF5TWS5BNNB2WEL2QOVWGYUTFOF2WK43UKJSXM2LFO4XTIMJVGAZTQMBWGQ22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJL3QOJPXEZLWNFSXOX3DNRUWG2Y#discussion_r3119935099> : > +def _coerce_dt(dt: object, *, name: str = "dt") -> float: + """Coerce a numeric scalar timestep-like value to Python float.""" + if isinstance(dt, bool): + raise ValueError( + f"{name} must be a numeric scalar (Python/NumPy real or scalar JAX array), " + f"got {type(dt)}." + ) + if isinstance(dt, Real): + return float(dt) + if ( + isinstance(dt, Array) + and dt.ndim == 0 + and not jnp.issubdtype(dt.dtype, jnp.bool_) + and ( + jnp.issubdtype(dt.dtype, jnp.integer) + or jnp.issubdtype(dt.dtype, jnp.floating) + ) + ): + return float(dt) + raise ValueError( + f"{name} must be a numeric scalar (Python/NumPy real or scalar JAX array), " + f"got {type(dt)}." + ) I looked more closely, and I think the main reason we have floats *anywhere* is for friendly user API (no one wants to write dt0=jnp.array(0.01) ). What do you think about doing float -> Array at a high-level early on and then do our internals all on Arrays? — Reply to this email directly, view it on GitHub <#178?email_source=notifications&email_token=AC2DPOKMCTZPSIISZU5BTY34W7FYRA5CNFSNUABKM5UWIORPF5TWS5BNNB2WEL2QOVWGYUTFOF2WK43UKJSXM2LFO4XTIMJVGAZTQMBWGQ22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJPLQOJPXEZLWNFSXOX3ON52GSZTJMNQXI2LPNZZV6Y3MNFRWW#discussion_r3119935099>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AC2DPOJ2FOSM4VWCJTNND534W7FYRAVCNFSM6AAAAACXBE7HN2VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHM2DCNJQGM4DANRUGU> . You are receiving this because you were mentioned.Message ID: ***@***.***>

Add fast SDE solver backends and filter-simulator coverage.

2dff055

Introduce solver modules for ODE and SDE paths, add an Euler-Maruyama scan source for SDESimulator, and update filter+simulator tests and continuous-time tutorial outputs to cover the new backend behavior. Made-with: Cursor

mattlevine22 mentioned this pull request Mar 27, 2026

Add a locally hosted lax.scan-based SDE solver #177

Closed

mattlevine22 marked this pull request as ready for review March 27, 2026 15:44

Merge branch 'main' into ml-fast-sde-solver2

3c90be4

Merge remote-tracking branch 'origin/main' into ml-fast-sde-solver2

a28eb99

add a speed-tip for SDESimulator documentation

62b1e01

DanWaxman requested a review from Copilot April 14, 2026 15:26

Copilot started reviewing on behalf of DanWaxman April 14, 2026 15:27 View session

Copilot AI reviewed Apr 14, 2026

View reviewed changes

respond to copilot suggestions

69496e9

DanWaxman requested changes Apr 15, 2026

View reviewed changes

Comment thread dynestyx/solvers/sde.py Outdated

Comment thread dynestyx/solvers/sde.py Outdated

Comment thread dynestyx/solvers/sde.py Outdated

Comment thread dynestyx/solvers/sde.py

mattlevine22 added 2 commits April 15, 2026 17:08

simplify key / tol_vbt types and checking

5485d87

oops, forgot tol_vbt in the new diffrax test

058d68b

mattlevine22 mentioned this pull request Apr 16, 2026

Faster SDE solver for simple (non-inference-related) unrolling #170

Open

DanWaxman requested changes Apr 21, 2026

View reviewed changes

mattlevine22 and others added 2 commits April 21, 2026 14:44

Remove duplicate line

80adbcb

Apply suggestion from @DanWaxman Co-authored-by: Dan Waxman <dan.waxman1@gmail.com>

Merge branch 'main' into ml-fast-sde-solver2

5fbec58

mattlevine22 commented Apr 21, 2026

View reviewed changes

Normalize time scalars to JAX arrays across model and solver internals

724596d

DanWaxman self-requested a review April 22, 2026 14:14

DanWaxman approved these changes Apr 22, 2026

View reviewed changes

DanWaxman merged commit 1ef0782 into main Apr 22, 2026
3 checks passed

DanWaxman mentioned this pull request Apr 22, 2026

Make Discretizer use GaussianStateEvolution for increased filter compatibility #199

Open

Conversation

mattlevine22 commented Mar 26, 2026

Uh oh!

mattlevine22 commented Mar 27, 2026

Uh oh!

mattlevine22 commented Mar 27, 2026

Uh oh!

DanWaxman commented Apr 13, 2026

Uh oh!

mattlevine22 commented Apr 14, 2026

Uh oh!

DanWaxman commented Apr 14, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

DanWaxman left a comment

Choose a reason for hiding this comment

Uh oh!

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!

mattlevine22 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DanWaxman commented Apr 21, 2026 via email

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants