Diffrax integration - #817
Conversation
a656228 to
b43c245
Compare
| if first_step is None: | ||
| first_step = np.finfo(np.float64).eps |
There was a problem hiding this comment.
Let's either pass a given first step through or let the underlying algorithms decide it. I think it is simpler and better matches diffrax?
There was a problem hiding this comment.
Yeah I've made a bit of a mess of it I think. Now it's either set by the user, kept as None (in the case of diffrax, then it chooses one) or the old heuristic is applied. I think keeping the old behaviour at least for now is the call, but if you want to change this I'm happy to do it in this PR.
This is awesome and would close #771.
Let's put everything behind the I also noted that your change in |
b43c245 to
d600870
Compare
58fe5bb to
4201979
Compare
Initial idea for integrating diffrax, currently dispatches to a _get_y_diffrax function if the selected method is one of the provided by this integration
4201979 to
f76af87
Compare
Uses the constants defined in overreact.simulate to sync the possible choices with the simulate module. Improves the error message to also print the available scipy solvers that don't require jax. A supported_methods constant also is exposed by the module.
After avoiding differentiating 0**0 we basically solved the issue where we couldn't use jax in solve_ivp.
|
I think all the comments are adressed, If there's anything else you think is needed I'm happy to include in this PR.
Yes it does! I didn't quite know why we left it disabled, but just allowing it and running the test suite runs smoothly. So apparently after we merge we can also close #422 :) |
2 participants
Initial integration of the Diffrax solvers, mainly to add the stiff solvers it provides.
The diffrax solver requires Jax, so the first question is, we now require Jax as a core feature or hide diffrax behind the
fastfeature flag?In this initial implementation I've kept it behind the
fastfeature flag, but an option would be to replace the scipy solvers with the diffrax ones, as I think they are much more solid than the ones scipy exposes.The only problem is that if someone wants to reproduce old simulated system this would be a bit of a breaking update, so I think it would be better to still keep the old solvers in the current proposed change? Open to feedback regarding this.
I haven't yet throughtly tested this in a "real world scenario", but the plan is to try to get some stiff system this week to compare the scipy solvers and the newly added Kvaerno solvers to see if they can handle it better.
I've also noted that I accidently linked the diffrax issue (#771) when doing the uv migration (which should have been #759, oopsie) instead of linking the correct issue. So if you want to re-open that one and close the uv issue...