Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Idaklu-Jax Interface #3658

Merged
merged 58 commits into from
Feb 21, 2024
Merged

Idaklu-Jax Interface #3658

merged 58 commits into from
Feb 21, 2024

Conversation

jsbrittain
Copy link
Contributor

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

Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

Copy link

codecov bot commented Dec 23, 2023

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 jsbrittain marked this pull request as ready for review January 10, 2024 16:50
@jsbrittain
Copy link
Contributor Author

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
Copy link
Contributor Author

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
Copy link
Member

  • 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
Copy link
Member

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
Copy link
Contributor Author

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
Copy link
Member

@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

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

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

agriyakhetarpal added a commit to agriyakhetarpal/PyBaMM that referenced this pull request Jan 23, 2024
@jsbrittain
Copy link
Contributor Author

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
Copy link
Member

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
Copy link
Member

agriyakhetarpal commented Jan 24, 2024

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
Copy link
Contributor Author

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
Copy link
Contributor Author

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

@BradyPlanden
Copy link
Member

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

@jsbrittain
Copy link
Contributor Author

@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.

@martinjrobins martinjrobins merged commit cf686e7 into pybamm-team:develop Feb 21, 2024
46 of 47 checks passed
js1tr3 pushed a commit to js1tr3/PyBaMM that referenced this pull request Aug 12, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

new JAX solver based on IDAKLU for parameter inference
5 participants