Diffrax integration - #817
Open
caprilesport wants to merge 5 commits into
Open
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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...