Performance refactor for Jax BDF Solver, fixes #4455

[BradyPlanden]

Description

This PR refactors the JAX BDF solver for performance updates. It also updates the JaxSolver's default method to be "BDF" as performance is much higher than "RK45". This PR also bug fixes the calculate_sensitivities error described in #4455. Lastly, this PR adds JAX vectorised example script to showcase the JaxSolvers performance in highly vectorised usage.

Fixes #4455

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

• New feature (non-breaking change which adds functionality)
• Optimization (back-end change that speeds up the code)
• Bug fix (non-breaking change which fixes an issue)

Key checklist:

• No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

• Code is commented, particularly in hard-to-understand areas
• Tests added that prove fix is effective or that feature works

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.46%. Comparing base (1390ea3) to head (3c34305).
Report is 1 commits behind head on develop.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #4456      +/-   ##
===========================================
- Coverage    99.46%   99.46%   -0.01%     
===========================================
  Files          293      293              
  Lines        22386    22332      -54     
===========================================
- Hits         22266    22212      -54     
  Misses         120      120              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[BradyPlanden]

Benchmarking for the BDF method is provided below. This was ran on a MacBook Pro M3Pro laptop, with the DFN discretised with var_pts = {"x_n": 10, "x_s": 10, "x_p": 10, "r_n": 10, "r_p": 10}. For the RK45 method, the SPMe had not solved within 5 minutes (conservatively estimated), so I did not proceed to the DFN.

Unless otherwise stated, the results below are in seconds. This is the average across ten simulations of the example script (with the number of simulations reduced from 1000 to 100), including the JIT compilation time.

Model	develop (v24.9)	#4456	RK45
SPM	0.97	0.85	1.11
SPMe	3.09	1.82	>5min
DFN	18.86	8.29	≫5min

[martinjrobins]

Great stuff, thanks @BradyPlanden. I'll have a closer look at the code in a tic, I'm really suprised how the timings have evolved, when I wrote this the jit compilation of the BDF solver was very slow compared with RK45, hence the RK45 default. But it looks like there have been many changes in JAX since then. Do you have timings including the solve by itself (without including JIT)?

[BradyPlanden]

I was also a bit surprised, as I didn't go into this refactor looking for performance, mostly just to clean up the code and understand the underlying JAX methods. I suspect there are areas for further improvement. Here are the post-JIT timings per 100 solves, now in milliseconds.

Model	develop (v24.9)	#4456	RK45
SPM	154	39.5	889
SPMe	1620	320	N/A
DFN	12372	2340	N/A

The JIT compilation for the BDF has improved quite a bit, there are a few areas I think it's most likely to come from. I think reducing the memory copies and using newer JAX methods (lax.cond, etc.) is doing a lot of the heavier lifting here.

Edit: it would be interesting to compare these numbers to the recent idaklu parallelisation

[martinjrobins]

This looks great, thanks @BradyPlanden, glad to see the BDF jit compiling so much faster now :) Interested to see the runtime differences to idaklu.

[martinjrobins]

thanks @BradyPlanden. I don't think we need to store additional data to implement this property, see below

[martinjrobins]

great, thanks @BradyPlanden :)

[BradyPlanden]

Ready for merge now :)