Idaklu-Jax Interface

[jsbrittain]

Description

Provides an IDAKLU-JAX interface by exposing a JAX expression that performs an IDAKLU solve. There is a new IDAKLUJax API page in the docs, and usage examples have been provided in a new jupyter notebook.

Fixes #2282

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

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (84eb098) 99.59% compared to head (d201899) 99.60%.
Report is 1 commits behind head on develop.

Additional details and impacted files

@@            Coverage Diff            @@
##           develop    #3658    +/-   ##
=========================================
  Coverage    99.59%   99.60%            
=========================================
  Files          257      258     +1     
  Lines        20806    21226   +420     
=========================================
+ Hits         20722    21142   +420     
  Misses          84       84            

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

[jsbrittain]

Hi @martinjrobins - I'm ready for the IDAKLU-Jax interface to be reviewed. A good starting point would be the example jupyter notebook that I have included as this walks through a usage example.

PS: Benchmarks have run twice and failed on different tests each time (which do not appear related to this PR); I've re-started them again.

[jsbrittain]

Ready to merge. A note on the failing tests:

• Lychee is failing because of a documentation link in a new jupyter notebook that references a docs page that does not exist yet in the stable deployment (not sure how to avoid this).
• doctests run fine on my machine on py3.11, but fail on py3.12. I am seeing the same issue in other recent PRs too (the pydata_sphinx_theme extension is not safe for parallel writing), which appears to be a broader problem with that test.

On a side note, some of the unit tests required multiple runs to pass. This should not happen as they were made deterministic in #2844 - it appears as though some test classes have been introduced that inherit from unittest.TestCase instead of PyBaMM's custom TestCase class (which fixes the random seed). I will open a separate issue to address this.

[agriyakhetarpal]

• doctests run fine on my machine on py3.11, but fail on py3.12. I am seeing the same issue in other recent PRs too (the pydata_sphinx_theme extension is not safe for parallel writing), which appears to be a broader problem with that test.

Thanks, I noticed this too. I don't think there are any larger problems at play since it's just a warning message – I will take a look in a separate PR

[agriyakhetarpal]

Yet another note: since this PR has worked on a lot of compiled code, I would suggest running the wheel builds workflow on this branch once, just so that we don't face any problems later.

P.S. my suggestion is to make such checks automatic in some cases through our infrastructure in the long term, such as those being discussed in #3662.

[jsbrittain]

since this PR has worked on a lot of compiled code, I would suggest running the wheel builds workflow on this branch once, just so that we don't face any problems later.

@agriyakhetarpal Good idea - I just want to check how to do this without breaking anything: I can run the Build and publish package to PyPI Action on this branch (and with publish to testpypi) to check this without breaking anything, can't I? I think I've done this before, but wanted to double-check.

[agriyakhetarpal]

@agriyakhetarpal Good idea - I just want to check how to do this without breaking anything: I can run the Build and publish package to PyPI Action on this branch (and with publish to testpypi) to check this without breaking anything, can't I? I think I've done this before, but wanted to double-check.

@jsbrittain yes – the configuration by default is designed to not break anything. You can just set the workflow input as blank to avoid publishing on TestPyPI either. The PyPI publishing is guarded by the condition

PyBaMM/.github/workflows/publish_pypi.yml

Line 197 in b84f7ad

if: github.event.inputs.target == 'pypi' || github.event_name == 'release'

which helps automatically skip the job regardless of the input on non-release runs.

[jsbrittain]

since this PR has worked on a lot of compiled code, I would suggest running the wheel builds workflow on this branch once, just so that we don't face any problems later.

@agriyakhetarpal wheels build, see https://github.com/jsbrittain/PyBaMM/actions/runs/7632700433.

[arjxn-py]

see https://github.com/jsbrittain/PyBaMM/actions/runs/7632700433.

For some reason, I can't access this build. Although it should be alright but would it be nicer to test the wheels too?
Since this is such a big PR and involves crucial infrastructural development, wanted to appreciate and thanks a lot @jsbrittain

[agriyakhetarpal]

I can't access the link either but I trust the universe with it haha. The wheels are currently being tested by cibuildwheel for all platforms through the instantiation of pybamm.IDAKLUSolver (#3569). After we migrate to a better test infrastructure by the end of the year (or by fall if sooner), the idea is to run the integration tests there

Edit: was able to access the link earlier when it was posted, I forgot I had seen it lol

[jsbrittain]

For some reason, I can't access this build. Although it should be alright but would it be nicer to test the wheels too? Since this is such a big PR and involves crucial infrastructural development, wanted to appreciate and thanks a lot @jsbrittain

@arjxn-py @agriyakhetarpal I disabled Actions on my fork after running the workflow - I didn't realise that would prevent you from seeing the report / artifacts. I've re-enabled them and the link is working for me again: https://github.com/jsbrittain/PyBaMM/actions/runs/7632700433

[jsbrittain]

@martinjrobins fixed the merge conflicts; all tests passing (apart from lychee).

[BradyPlanden]

Any updates on merging this? It would be a great addition from the parameter identification workflow side. @jsbrittain @martinjrobins @pybamm-team/maintainers

[jsbrittain]

@BradyPlanden @martinjrobins @pybamm-team/maintainers I've updated the branch with the latest changes from develop, but do not have permissions to merge the other way, so will require someone else to authorise that.