Langevin PR #453

andyElking · 2024-07-01T10:56:08Z

Hi Patrick,

Thanks for the heads up, here's the reuploaded PR (with another quick fix).

This PR contains all the new Langevin solvers. All of these inherit from AbstractLangevinSRK in langevin_srk.py.
Another important addition is LangevinTerm in _term.py. I explained why it is needed in a comment bellow.

I haven't added the new solvers to the docs and autocite yet, because 1) the relevant paper is not on arxiv yet, but might be in a month or two and 2) I expect you might suggest several changes, so might as well write the docs once the rest is stationary. Still, I think the docstrings and comments are quite comprehensive.

I'm making this PR now so you have ample time to have a look at it, but I will be away for the next few weeks, so there is aboslutely no hurry.

Best,
Andraž

diffrax/_term.py

andyElking · 2024-07-20T13:02:35Z

Hi @patrick-kidger, just bumping this in case you didn't notice that the tests passed. No hurry though :)

patrick-kidger

Haha, sorry, it's taken me a while to get to this!

Here's a first review. Many of my comments on align apply to the others too, I think.

diffrax/_integrate.py

diffrax/_solver/align.py

diffrax/_solver/langevin_srk.py

diffrax/_term.py

andyElking · 2024-07-21T21:01:58Z

Hi Patrick,

Thanks so much for your review! I addressed almost all of your comments. Two of them I will address in a later commit (it's getting late today haha).

I think you didn't quite understand the reason why the LangevinTerm and the changes in _integrate.py are needed, so I tried to give an explanation, but if that is unclear we can have a call at some point. Let me know :)

Also you might want to read my reply here, but it's a bit hidden way up the conversation: #453 (comment)

andyElking · 2024-07-25T21:00:44Z

Quick heads up: I now made all the edits you suggested and the tests all passed :)

andyElking · 2024-07-29T11:42:21Z

@patrick-kidger a note about interpolation for ALIGN:
to maintain a 2nd order of convergence at interpolated points, the interpolated value cannot depend just on t0, t, t1, y0, y1, but must also depend on W0, H0, W1, H1. Do you think it would be feasible to modify the interpolation code in order to allow for that? Maybe I could store W and H as part of the solution? Do you have any ideas how to do this in a way that fits into Diffrax naturally?

patrick-kidger · 2024-08-02T12:00:03Z

to maintain a 2nd order of convergence at interpolated points, the interpolated value cannot depend just on t0, t, t1, y0, y1, but must also depend on W0, H0, W1, H1. Do you think it would be feasible to modify the interpolation code in order to allow for that? Maybe I could store W and H as part of the solution? Do you have any ideas how to do this in a way that fits into Diffrax naturally?

I think should be totally fine -- it can go into the dense_info output of a solver step. Take a look at how Runge--Kutta methods output the intermediate stages, for example.

andyElking · 2024-08-02T14:47:42Z

I think should be totally fine -- it can go into the dense_info output of a solver step. Take a look at how Runge--Kutta methods output the intermediate stages, for example.

That's good to know, thanks! I'll still make it a separate PR - I have a few other tasks to complete beforehand. Otherwise I think the only comment that remains unresolved is #453 (comment). Let me know if there is anything else you'd like me to improve.

andyElking · 2024-08-04T15:02:14Z

@patrick-kidger I now made a temporary fix, but there are two things that remain to be solved:

For some reason eqx.filter_eval_shape(term.vf, 0.0, y0, args) doesn't work when term is a LangevinTerm (the stacktrace is in a comment above). I genuinely do not undertand what is going wrong, so please help.
I think we are still not on the same page as to why I introduced LangevinTerm in the first place. Please let's have a call at some point to discuss it. And then hopefully it will make sense why using MultiTerm[LangevinDriftTerm, LangevinDiffusionTerm] could lead to incorrect results.

andyElking · 2024-08-09T23:53:20Z

Great news @patrick-kidger @lockwo: what we discussed worked! Thanks for your advice!

Patrick, if there is nothing else about the code itself that you'd like me to change, then I'll add all of the new things to the docs. Do you think I should add a short example of how to use the langevin solvers as well? Maybe a simple Langevin Monte Carlo example?

patrick-kidger

Nice! I really like this approach -- I think this is now looking really elegant.

I've gone through and left comments but they're mostly nits and tidy ups. I think we've now cracked the structure of this problem! :)

diffrax/_term.py

test/helpers.py

diffrax/_solver/langevin_srk.py

patrick-kidger · 2024-08-17T07:36:49Z

Do you think I should add a short example of how to use the langevin solvers as well? Maybe a simple Langevin Monte Carlo example?

I like the sound of that! An example would be great.

I'll emphasise 'short' -- I really try to keep the examples pedagogical.

andyElking · 2024-08-17T22:24:24Z

Thanks so much for the review, Patrick! And sorry for all my dummy mistakes 😅. I made all the smaller edits already and tomorrow I'll write a short example, a test for the backward solve and put everything into the docs.

diffrax/_term.py

patrick-kidger

Okay, this basically LGTM!

I've gone through and left what I think is a final round of nits, but I think they should all be super minor.

diffrax/_solver/foster_langevin_srk.py

diffrax/_term.py

docs/api/solvers/sde_solvers.md

docs/devdocs/srk_example.ipynb

examples/underdamped_langevin_example.ipynb

andyElking · 2024-09-01T12:13:25Z

Quick question @patrick-kidger:

Should I make AbstractFosterLangevinSRK public and add it under Abstract Solvers in the docs? I haven't done it so far, because, unlike RK and SRK, where the user just has to specify a tableau, making a custom child of AbstractFosterLangevinSRK is much more involved, so I doubt users who are just using a packaged version of Diffrax would do that. WDYT?

patrick-kidger · 2024-09-01T12:43:20Z

I think it probably should be public + in the abstract solvers. I agree that writing your own here is incredibly niche, but I think it's useful for inquisitive users to be able to poke at such things.

diffrax/_term.py

andyElking · 2024-09-01T14:46:09Z

So I added the check that drift and diffusion have the same arguments in AbstractFosterLangevinSRK.init and also added a short test for this in test_underdamped_langevin.py.

I will now do the scan trick.

andyElking · 2024-09-01T15:28:51Z

I think I now addressed everything you suggested, including the scan trick in both QUICSORT and ShOULD. I'll do another quick check and then I pass it back to you.

andyElking · 2024-09-01T16:20:13Z

I went through the code and the docs again and now I think I addressed everything you mentioned. Also please take a look at this comment here and let me know if I should revert it to how it was before. Other than that there are no major changes.

Also please take a look at the conversations I left unresolved, namely this and this. Thanks!

andyElking · 2024-09-01T16:51:07Z

Sorry, I left a tiny problem in the test I added, I fixed it now.

diffrax/_solver/foster_langevin_srk.py

docs/api/terms.md

patrick-kidger · 2024-09-01T19:31:39Z

Aaaaand... merged! 🎉
Great job getting this one done, I'm really happy to have it in! :)

andyElking · 2024-09-01T19:41:15Z

Thanks for bearing with me and taking your time Patrick!! I really appreciate it!

* Langevin PR * Minor fixes * removed the SORT solver (superseded by QUICSORT) * made LangevinTerm.term a static field * temporary fix for _term_compatible and LangevinTerm * Fixed LangevinTerm YAAAYYYYY * Nits * Added Langevin docs, a Langevin example and backwards in time test * Fixed Patrick's comments * langevin -> underdamped_langevin * round of small fixes * check langevin drift term and diffusion term have same args * added scan_trick in QUICSORT and ShOULD * using RuntimeError for when ULD args have wrong structure * small fixes

* Langevin PR (#453) * Langevin PR * Minor fixes * removed the SORT solver (superseded by QUICSORT) * made LangevinTerm.term a static field * temporary fix for _term_compatible and LangevinTerm * Fixed LangevinTerm YAAAYYYYY * Nits * Added Langevin docs, a Langevin example and backwards in time test * Fixed Patrick's comments * langevin -> underdamped_langevin * round of small fixes * check langevin drift term and diffusion term have same args * added scan_trick in QUICSORT and ShOULD * using RuntimeError for when ULD args have wrong structure * small fixes * tidy-ups * Split SDE tests in half, to try and avoid GitHub runner issues? * Added effects_barrier to fix test issue with JAX 0.4.33+ * small fix of docs in all three and a return type in quicsort * bump doc building pipeline * Compatibility with JAX 0.4.36, which removes ConcreteArray * using a fori_loop to save states in edge case t0==t1 * added case for saving t0 data, which was also not getting updated. Added a test * using while_loop, ran into issues with reverse-mode diff using the fori_loop * bug fix for cases when t0=True * simplified logic for saving, no loop necessary * added vmap test * using a fori_loop to save states in edge case t0==t1 * added case for saving t0 data, which was also not getting updated. Added a test * using while_loop, ran into issues with reverse-mode diff using the fori_loop * bug fix for cases when t0=True * simplified logic for saving, no loop necessary * added vmap test * fix t1 out of bounds issue * fix for steps: don't want to update those values if t0==t1 since we didn't take any steps. Added test --------- Co-authored-by: Andraž Jelinčič <[email protected]> Co-authored-by: Patrick Kidger <[email protected]> Co-authored-by: andyElking <[email protected]>

* Langevin PR * Minor fixes * removed the SORT solver (superseded by QUICSORT) * made LangevinTerm.term a static field * temporary fix for _term_compatible and LangevinTerm * Fixed LangevinTerm YAAAYYYYY * Nits * Added Langevin docs, a Langevin example and backwards in time test * Fixed Patrick's comments * langevin -> underdamped_langevin * round of small fixes * check langevin drift term and diffusion term have same args * added scan_trick in QUICSORT and ShOULD * using RuntimeError for when ULD args have wrong structure * small fixes

* Langevin PR (#453) * Langevin PR * Minor fixes * removed the SORT solver (superseded by QUICSORT) * made LangevinTerm.term a static field * temporary fix for _term_compatible and LangevinTerm * Fixed LangevinTerm YAAAYYYYY * Nits * Added Langevin docs, a Langevin example and backwards in time test * Fixed Patrick's comments * langevin -> underdamped_langevin * round of small fixes * check langevin drift term and diffusion term have same args * added scan_trick in QUICSORT and ShOULD * using RuntimeError for when ULD args have wrong structure * small fixes * tidy-ups * Split SDE tests in half, to try and avoid GitHub runner issues? * Added effects_barrier to fix test issue with JAX 0.4.33+ * small fix of docs in all three and a return type in quicsort * bump doc building pipeline * Compatibility with JAX 0.4.36, which removes ConcreteArray * using a fori_loop to save states in edge case t0==t1 * added case for saving t0 data, which was also not getting updated. Added a test * using while_loop, ran into issues with reverse-mode diff using the fori_loop * bug fix for cases when t0=True * simplified logic for saving, no loop necessary * added vmap test * using a fori_loop to save states in edge case t0==t1 * added case for saving t0 data, which was also not getting updated. Added a test * using while_loop, ran into issues with reverse-mode diff using the fori_loop * bug fix for cases when t0=True * simplified logic for saving, no loop necessary * added vmap test * fix t1 out of bounds issue * fix for steps: don't want to update those values if t0==t1 since we didn't take any steps. Added test --------- Co-authored-by: Andraž Jelinčič <[email protected]> Co-authored-by: Patrick Kidger <[email protected]> Co-authored-by: andyElking <[email protected]>

andyElking commented Jul 1, 2024

View reviewed changes

diffrax/_term.py Outdated Show resolved Hide resolved

andyElking force-pushed the langevin_pr branch 2 times, most recently from 2d5dcad to 21a5730 Compare July 11, 2024 15:07

Langevin PR

881cd84

andyElking force-pushed the langevin_pr branch from 21a5730 to 881cd84 Compare July 11, 2024 17:08

patrick-kidger reviewed Jul 21, 2024

View reviewed changes

lockwo mentioned this pull request Jul 22, 2024

KL Divergence for Latent SDEs #463

Open

andyElking force-pushed the langevin_pr branch from d84d187 to c4fb51b Compare July 24, 2024 16:21

Minor fixes

cd0ec5e

andyElking force-pushed the langevin_pr branch from c4fb51b to cd0ec5e Compare July 24, 2024 19:33

removed the SORT solver (superseded by QUICSORT)

15c21a6

andyElking added 2 commits August 2, 2024 15:53

made LangevinTerm.term a static field

34c51a9

temporary fix for _term_compatible and LangevinTerm

34aa6cb

andyElking force-pushed the langevin_pr branch 2 times, most recently from 2719cc9 to 67011f0 Compare August 9, 2024 21:03

Fixed LangevinTerm YAAAYYYYY

3e3656b

andyElking force-pushed the langevin_pr branch from 67011f0 to 3e3656b Compare August 9, 2024 21:03

patrick-kidger reviewed Aug 17, 2024

View reviewed changes

Nits

2b80e3b

Added Langevin docs, a Langevin example and backwards in time test

e38a932

andyElking commented Aug 31, 2024

View reviewed changes

diffrax/_term.py Outdated Show resolved Hide resolved

patrick-kidger reviewed Sep 1, 2024

View reviewed changes

andyElking commented Sep 1, 2024

View reviewed changes

diffrax/_term.py Outdated Show resolved Hide resolved

andyElking force-pushed the langevin_pr branch from 2eafdbd to aa5cbd0 Compare September 1, 2024 15:31

round of small fixes

0a5ffeb

andyElking force-pushed the langevin_pr branch 4 times, most recently from 914ef13 to 5af3e06 Compare September 1, 2024 16:00

andyElking added 2 commits September 1, 2024 17:48

check langevin drift term and diffusion term have same args

3aab39c

added scan_trick in QUICSORT and ShOULD

b760c65

andyElking force-pushed the langevin_pr branch from 5af3e06 to b760c65 Compare September 1, 2024 16:50

patrick-kidger reviewed Sep 1, 2024

View reviewed changes

diffrax/_solver/foster_langevin_srk.py Outdated Show resolved Hide resolved

docs/api/terms.md Outdated Show resolved Hide resolved

andyElking added 2 commits September 1, 2024 18:08

using RuntimeError for when ULD args have wrong structure

dff726b

small fixes

be892ae

patrick-kidger changed the base branch from main to dev September 1, 2024 19:28

patrick-kidger merged commit fa7417f into patrick-kidger:dev Sep 1, 2024
2 checks passed

andyElking deleted the langevin_pr branch September 26, 2024 10:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Langevin PR #453

Langevin PR #453

andyElking commented Jul 1, 2024 •

edited

Loading

andyElking commented Jul 20, 2024

patrick-kidger left a comment

andyElking commented Jul 21, 2024 •

edited

Loading

andyElking commented Jul 25, 2024

andyElking commented Jul 29, 2024

patrick-kidger commented Aug 2, 2024

andyElking commented Aug 2, 2024

andyElking commented Aug 4, 2024

andyElking commented Aug 9, 2024

patrick-kidger left a comment

patrick-kidger commented Aug 17, 2024

andyElking commented Aug 17, 2024

patrick-kidger left a comment

andyElking commented Sep 1, 2024

patrick-kidger commented Sep 1, 2024

andyElking commented Sep 1, 2024

andyElking commented Sep 1, 2024

andyElking commented Sep 1, 2024 •

edited

Loading

andyElking commented Sep 1, 2024

patrick-kidger commented Sep 1, 2024 •

edited

Loading

andyElking commented Sep 1, 2024

Langevin PR #453

Langevin PR #453

Conversation

andyElking commented Jul 1, 2024 • edited Loading

andyElking commented Jul 20, 2024

patrick-kidger left a comment

Choose a reason for hiding this comment

andyElking commented Jul 21, 2024 • edited Loading

andyElking commented Jul 25, 2024

andyElking commented Jul 29, 2024

patrick-kidger commented Aug 2, 2024

andyElking commented Aug 2, 2024

andyElking commented Aug 4, 2024

andyElking commented Aug 9, 2024

patrick-kidger left a comment

Choose a reason for hiding this comment

patrick-kidger commented Aug 17, 2024

andyElking commented Aug 17, 2024

patrick-kidger left a comment

Choose a reason for hiding this comment

andyElking commented Sep 1, 2024

patrick-kidger commented Sep 1, 2024

andyElking commented Sep 1, 2024

andyElking commented Sep 1, 2024

andyElking commented Sep 1, 2024 • edited Loading

andyElking commented Sep 1, 2024

patrick-kidger commented Sep 1, 2024 • edited Loading

andyElking commented Sep 1, 2024

andyElking commented Jul 1, 2024 •

edited

Loading

andyElking commented Jul 21, 2024 •

edited

Loading

andyElking commented Sep 1, 2024 •

edited

Loading

patrick-kidger commented Sep 1, 2024 •

edited

Loading