patrick-kidger / patrick-kidger/diffrax

How does diffrax handle state variables becoming inf?

Open
#223 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

question
Dominant language
Python
Stars
2.1k
Forks
189
Avg merge
3d 18h
Merged PRs (30d)
1

Description

I'm trying to port some code from pure python to JAX. The code is a bit complicated with time-varying interpolated inputs but essentially involves solving a system of 8 coupled ODEs. The state variables and independent time variable span a large dynamic range so the system is solved in logarithmic form (i.e., read in 10^logt and 10^logy, do calculations in physical units, return dlogy/dlogt). The system is stiff but my pure python implementation based on scipy.integrate.solve_ivp leads to reasonable solutions for multiple solvers (the implicit methods BDF and Radau being the best, but RK23 and DOP853 also finish and give similar solutions though require lower tolerances and more steps; all of these agree with my manual Euler integrations with extremely small fixed timesteps, increasing my confidence in the solutions).

However, all of the adaptive diffrax solvers I've tried fail and JAX's odeint also fails. Both packages only get 4 steps in before they return NaN's (and diffrax hits up against max_steps=10K whereas scipy only needs 300 for BDF and 3000 for RK23). I used print/jax.debug.print() statements inside my integrator function to print the values of the state variables and their associated derivatives for both the pure python and JAX versions of my code. I find that early on, the derivatives for one or two of my state variables blows up (dlogy/dlogt=10^6) so in the next step, the value of that state variable becomes inf, and then that leads to NaN's. solve_ivp is able to deal with this by repeatedly going back to an earlier time and taking ever smaller timesteps, eventually getting over this initial hump. But both jax-odeint and diffrax remain stuck with NaN's. I was hoping diffrax's implicit solvers and adaptive step-size controller would help compared to the default jax-odeint one, but it has the same problem.

Should diffrax/jax-odeint be able to adapt to state variables becoming inf? Why is solve_ivp successful?

I can probably share a self-contained github repo with two example Jupyter notebooks comparing solve_ivp to diffrax/jax-odeint if that would help.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the reported comparisons between scipy.integrate.solve_ivp, diffrax's adaptive solvers, and JAX's odeint. If the author provides the proposed self-contained Jupyter notebooks, reproduce the first few steps and inspect how each solver handles the state variables becoming inf. Done means explaining whether recovery from inf is supported and why the solvers behave differently.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.