Added support for multiterm (imex solvers) by MatthieuDarcy · Pull Request #307 · BasisResearch/dynestyx

MatthieuDarcy · 2026-08-11T19:50:50Z

Requires diffrax==0.7.2 which is only possible when CD-Dynamax is updated, so marking this as a draft in the meantime.
#304

mattlevine22 · 2026-08-17T17:18:33Z

+    Checked structurally (not via `isinstance(..., ImExDrift)`) because the
+    concrete `ImExDrift` class lives in `state_evolution.py`, which imports
+    from `core.py`, which imports from this module -- importing it here
+    would cycle.


This doesn't seem like a good reason to create a dummy class. Can you take a look to see what sort of file refactor / creation would be needed to avoid this circular import issue?

mattlevine22 · 2026-08-17T17:59:48Z

        raise NotImplementedError()


+type _Drift = Callable[


How about models/types.py?

just trying to think about big picture organization...dynestyx/types.py are pretty high-level types, so seems weird to drop ImExDrift there.

actually, I guess models/core.py is already mostly a "types.py" file.

How about models/drifts.py and include Drift, AffineDrift, and ImExDrift there?

I'm going to be a bit verbose to make sure I understand things correctly, let me know if I'm off track.

Right now incheckers.py, we use _validate_continuous_state_evolution to check that we are not using both an ImExDrift AND a state_evolution.potential term (because we decided that combining these two is ambiguous), raising an error if so.

I think adding a drifts.py with Drift, AffineDrift, ImExDrift, Potential (?) and importing them into checkers.py could make sense.

One thing I'm not completely sure about is that diffusion.py is imported at the level of core.py, one level down so there is a difference between the two (and I would naively assume that they would mirror each other, I don't know if this is a concern). I could see if there is a way of checking the ImExDrift vs Potential at the level of core.py but it's definitely more complicated.

I want to add that I believe that we have 3 choices:

add the check in checkers.py, in _validate_continuous_state_evolution.

add the check in core.py, in ContinuousTimeStateEvolution (doing it during class instantiation(

add the check in odes.py, in _solve
I'm not a fan of 3, because the error will be raised much later.

I think it makes sense to include Potential too!

Not sure I understand the concern about diffusions.py. Both drifts.py and diffusions.py would be leaf modules imported by core.py; and checkers.py would additionally import ImExDrift

That should be OK? want to understand the complication here

I'm just wanted to make sure that it was fine that checkers.py depends on drifts.py but not on diffusions.py is fine. I might be overthinking this, if it's not an issue, then will implement.

yeah, checkers.py COULD depend on diffusions.py if we needed it to though right?

We just don't do any diffusion checks

dirfts.py now includes potential, drift, imexdrift. pure leaf file imported by checkers.py amd core.py

MatthieuDarcy · 2026-08-18T14:10:55Z

There now is a drifts.py, imported by core.py and checkers.py. It's a leaf, no internal Dynestyx imports. I think the solution works pretty well, let me know what you think @mattlevine22

mattlevine22 · 2026-08-18T15:19:52Z

I guess this file path should change to models/drifts/drift.md

mattlevine22 · 2026-08-18T15:20:02Z

I guess this file path should change to models/drifts/potential.md

mattlevine22 · 2026-08-18T15:20:16Z

change file path

mattlevine22 · 2026-08-18T15:20:56Z

There now is a drifts.py, imported by core.py and checkers.py. It's a leaf, no internal Dynestyx imports. I think the solution works pretty well, let me know what you think @mattlevine22

Yeah looks reasonable---I guess you need to change some of the docs .md file paths accordingly though

mattlevine22 · 2026-08-18T15:30:45Z

@MatthieuDarcy I did a quick iteration with Codex, which picked up 2 issues:

MultiTerm is not specific to IMEX. Diffrax also uses it for SDE and Langevin solvers, where the terms have different meanings. Checking only get_origin(solver.term_structure) is dfx.MultiTerm can therefore misclassify solvers like EulerHeun and pass them two ODE terms incorrectly.

Recommended fix: require both MultiTerm and dfx.AbstractImplicitSolver, or use an explicit supported-solver check:
```
is_imex_solver = (
    isinstance(solver, dfx.AbstractImplicitSolver)
    and get_origin(solver.term_structure) is dfx.MultiTerm
)
```
A potential drift can be silently discarded inside dsx.plate. The ImExDrift + potential validation is skipped in plate contexts, while the IMEX solve evaluates only explicit_term and implicit_term, bypassing total_drift and omitting the potential gradient.

Recommended fix: perform the incompatibility check unconditionally, outside plate-skipped shape validation. Alternatively, have ContinuousTimeStateEvolution construct the complete IMEX decomposition, including the potential contribution. A plated regression test would be helpful.

mattlevine22 · 2026-08-18T15:42:57Z

Also, the reaction diffusion notebook kinda uses IMEX at the end, but it doesn't show any plots or justify the choice

MatthieuDarcy · 2026-08-18T18:09:29Z

@MatthieuDarcy I did a quick iteration with Codex, which picked up 2 issues:

1. `MultiTerm` is not specific to IMEX. Diffrax also uses it for SDE and Langevin solvers, where the terms have different meanings. Checking only `get_origin(solver.term_structure) is dfx.MultiTerm` can therefore misclassify solvers like `EulerHeun` and pass them two ODE terms incorrectly.
   Recommended fix: require both `MultiTerm` and `dfx.AbstractImplicitSolver`, or use an explicit supported-solver check:
   ```python
   is_imex_solver = (
       isinstance(solver, dfx.AbstractImplicitSolver)
       and get_origin(solver.term_structure) is dfx.MultiTerm
   )
   ```

2. A potential drift can be silently discarded inside `dsx.plate`. The `ImExDrift + potential` validation is skipped in plate contexts, while the IMEX solve evaluates only `explicit_term` and `implicit_term`, bypassing `total_drift` and omitting the potential gradient.
   Recommended fix: perform the incompatibility check unconditionally, outside plate-skipped shape validation. Alternatively, have `ContinuousTimeStateEvolution` construct the complete IMEX decomposition, including the potential contribution. A plated regression test would be helpful.

I fixed both.

In addition to what you suggested, I added a check at the very beginning of solve_ode_state_path raising an error if the solver is of typedfx.AbstractStratonovichSolver ordfx.AbstractItoSolver. The reason for this is because it's technically possible to have implicit SDE solvers, they just don't exist in diffrax yet. Checking for dfx.AbstractImplicitSolver only could therefore lead to a bug in a potential future release? In principle this should not alter solve_ode_state_path as it should already error if given an SDE solver and the error message should now be clearer. Let me know if you approve of this (another option is to add this check only when there is an ImExDrift, more specific). I put the dfx.AbstractImplicitSolver check only.
I added a specific check for ImExDrift and Potential in checkers.py , called outside the _inside_plate check, only checks for continuous_time.
Will work on the docs later.

mattlevine22 · 2026-08-18T18:23:04Z

Yeah looks good, thanks!

MatthieuDarcy · 2026-08-18T18:27:01Z

Actually nevermind, ODE solvers can also be labelled as dfx.AbstractStratonovichSolver or dfx.AbstractItoSolver. I'm changing this to only check for AbstractImplicitSolver, might break in a future release but so be it.

DanWaxman · 2026-08-19T02:49:36Z

FYI I updated the PR to merge to main + made a merge commit, since the previous base branch landed on main.

mattlevine22

I'm happy with this PR, so approving it.

Would recommend updating the end of the notebook with a plot of the pde solution that used the imex solver though.

changed the order added a drifts section (similar to diffusion)

MatthieuDarcy · 2026-08-19T15:21:53Z

@mattlevine22 I modified some of the API documentation in /core

The order is changed (it seemed a bit random, I tried to have more coherent order).

There now is a dedicated Driftssection, similar to Diffusions, with Drift, Potential, AffineDrift, and ImExDrift.
Note that AffineDrifthas its own separate documentation in Specialized, I didn't want to touch it.

I updated the notebook with a plot and added a justification for using Imex.

Let me know what you think, didn't want to merge without your approval.

DanWaxman

Some small comments

DanWaxman · 2026-08-19T15:31:11Z

-class AffineDrift(eqx.Module):
-    """
-    Affine drift function for continuous-time models.
-
-    This implements an affine map of the form
-
-    $$f(x, u, t) = A x + B u + b,$$
-
-    where $A \\in \\mathbb{R}^{d_x \\times d_x}$, $B \\in \\mathbb{R}^{d_x \\times d_u}$
-    (optional), and $b \\in \\mathbb{R}^{d_x}$ (optional). The time argument $t$
-    is accepted for compatibility with the `Drift` protocol but is not used.
-
-    This is commonly used as the drift term $\\mu(x_t, u_t, t)$ inside
-    `ContinuousTimeStateEvolution`, and is a building block for LTI models such as
-    `LTI_continuous`.
-
-    Attributes:
-        A (jax.Array): Drift matrix with shape $(d_x, d_x)$.
-        B (jax.Array | None): Optional control matrix with shape $(d_x, d_u)$.
-        b (jax.Array | None): Optional additive bias with shape $(d_x,)$.
-    """


I think we have enough (even if the total number is small) users and this is a common enough class that we should be keeping this API entrypoint with a deprecation warning that it will be removed in v0.5.0.

state_evolution.py now imports AffineDrift from drifts.pyso this restores the API entrypoint (with a warning).

DanWaxman · 2026-08-19T15:35:19Z

Sorry I realized afterwards that Matt had approved it already lol. But these are tiny, I'm happy for it to be merged afterwards :)

Co-authored-by: Dan Waxman <dan.waxman1@gmail.com>

into md-imex_schemes

DanWaxman

Thanks!!

MatthieuDarcy added 3 commits August 11, 2026 15:49

Added support for multiterm (imex solvers)

8eee555

Merge branch 'main' into imex_schemes

0b65160

Merge branch 'main' into imex_schemes

f519001

MatthieuDarcy mentioned this pull request Aug 14, 2026

Add support for IMEX solvers #304

Open

Merge branch 'main' into imex_schemes

678da11

DanWaxman changed the base branch from main to staging-advanced-pdes August 14, 2026 18:13

Upgrade diffrax + cd-dynamax

8928665

MatthieuDarcy changed the base branch from staging-advanced-pdes to dw-upgrade-cd-dynamax August 14, 2026 19:43

MatthieuDarcy and others added 4 commits August 14, 2026 15:44

Merge branch 'dw-upgrade-cd-dynamax' into imex_schemes

ceeed15

Merge branch 'main' into dw-upgrade-cd-dynamax

7d4be9b

Added imex class

ebe7341

Merge branch 'dw-upgrade-cd-dynamax' into imex_schemes

5930523

mattlevine22 reviewed Aug 17, 2026

View reviewed changes

MatthieuDarcy added 5 commits August 17, 2026 13:30

added imex example in PDEs

c603e34

moved imexdrift to types

95aeb47

Update reaction_diffusion_inference.ipynb

54a05bd

updated notebook + cd-dynamx version

6d9a430

Update reaction_diffusion_inference.ipynb

8015014

mattlevine22 reviewed Aug 17, 2026

View reviewed changes

created drifts.py

af5cd6e

dirfts.py now includes potential, drift, imexdrift. pure leaf file imported by checkers.py amd core.py

MatthieuDarcy marked this pull request as ready for review August 18, 2026 14:08

fix lint stuff

ea9e393

mattlevine22 reviewed Aug 18, 2026

View reviewed changes

Comment thread docs/api_reference/public/models/specialized/affine_drift.md

mattlevine22 Aug 18, 2026

Copy link
Copy Markdown

Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

change file path

removed the ability to use an sde solver with imexdrift

17357ff

MatthieuDarcy added 3 commits August 18, 2026 13:35

added imex checker + compatible with plate

a2d0ea3

added extra check to exclude ito and stratonovich solvers

fa05a57

added sde solver rejection when solving an ode

8f9ea03

removed ito and stratonovich checks

4b9cb40

DanWaxman changed the base branch from dw-upgrade-cd-dynamax to main August 19, 2026 02:48

Merge branch 'main' into imex_schemes

397810d

mattlevine22 approved these changes Aug 19, 2026

View reviewed changes

mattlevine22 and others added 3 commits August 19, 2026 09:12

Merge branch 'main' into imex_schemes

a857db7

added plot of imex solution

eec7c10

modified API docs

133be75

changed the order added a drifts section (similar to diffusion)

MatthieuDarcy requested a review from mattlevine22 August 19, 2026 15:16

Update drifts.py

efba71a

added short justification for imex scheme and reaction diffusion

25f3048

DanWaxman requested changes Aug 19, 2026

View reviewed changes

MatthieuDarcy and others added 3 commits August 19, 2026 11:40

Update dynestyx/models/drifts.py

e54a9c1

Co-authored-by: Dan Waxman <dan.waxman1@gmail.com>

added the ability to import affine drfit from state_evolution

1abcf70

Merge branch 'imex_schemes' of https://github.com/BasisResearch/dynestyx

c9cfb89

into md-imex_schemes

DanWaxman mentioned this pull request Aug 19, 2026

Standardize Notation Across Docs #327

Open

DanWaxman approved these changes Aug 19, 2026

View reviewed changes

DanWaxman merged commit 8788ba6 into main Aug 19, 2026
3 checks passed

MatthieuDarcy deleted the imex_schemes branch August 19, 2026 17:09

Uh oh!

Conversation

MatthieuDarcy commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mattlevine22 Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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 Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattlevine22 Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MatthieuDarcy Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MatthieuDarcy Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

MatthieuDarcy commented Aug 18, 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!

mattlevine22 commented Aug 18, 2026

Uh oh!

mattlevine22 commented Aug 18, 2026

Uh oh!

mattlevine22 commented Aug 18, 2026

Uh oh!

MatthieuDarcy commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mattlevine22 commented Aug 18, 2026

Uh oh!

MatthieuDarcy commented Aug 18, 2026

Uh oh!

DanWaxman commented Aug 19, 2026

Uh oh!

mattlevine22 left a comment

Choose a reason for hiding this comment

Uh oh!

MatthieuDarcy commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DanWaxman left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DanWaxman commented Aug 19, 2026

Uh oh!

DanWaxman left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MatthieuDarcy commented Aug 11, 2026 •

edited

Loading

mattlevine22 Aug 17, 2026 •

edited

Loading

mattlevine22 Aug 17, 2026 •

edited

Loading

mattlevine22 Aug 17, 2026 •

edited

Loading

MatthieuDarcy Aug 17, 2026 •

edited

Loading

MatthieuDarcy Aug 17, 2026 •

edited

Loading

MatthieuDarcy commented Aug 18, 2026 •

edited

Loading

MatthieuDarcy commented Aug 19, 2026 •

edited

Loading