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.

Sign up for GitHub

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

Adds JAX IDAKLU solver integration #481

Merged
merged 40 commits into from
Dec 21, 2024
Merged

Adds JAX IDAKLU solver integration #481

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
40 commits
Select commit Hold shift + click to select a range
1e08d6e
working commit, adds [jax] optional dependency to pybamm, example ida…
BradyPlanden Aug 20, 2024
79234f9
adds experimental sub-directory, sync commit
BradyPlanden Aug 20, 2024
7d8aa9c
Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation
BradyPlanden Aug 20, 2024
ab861a9
Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation
BradyPlanden Aug 28, 2024
f2bf143
feat: adds Jax functionality via idaklu_jax solver, with example and …
BradyPlanden Aug 31, 2024
025c7d8
Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation
BradyPlanden Sep 2, 2024
d3a18d5
Post merge fixes, update base_model solver property
BradyPlanden Sep 2, 2024
eb08c00
examples: update benchmarking script
BradyPlanden Sep 2, 2024
92265b6
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Sep 6, 2024
f1d03a9
Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation
BradyPlanden Sep 24, 2024
3ecd5dd
updts docstrings, jax arg to models, adds tests.
BradyPlanden Sep 24, 2024
9794910
add coverage, refactor BaseModel solver setter
BradyPlanden Sep 25, 2024
df1ba0c
adds coverage
BradyPlanden Sep 25, 2024
d770315
convert BaseModel.calculate_sensitivites to property
BradyPlanden Sep 25, 2024
daae4ba
add changelog entry
BradyPlanden Sep 25, 2024
424699b
feat: removes Jax arg from BaseModel, BaseProblem.model is copied ins…
BradyPlanden Sep 30, 2024
2695fd7
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Oct 11, 2024
23b8e3d
fix: cost shape for CMAES, adds example comments for jax solver
BradyPlanden Oct 21, 2024
bba2e56
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Oct 21, 2024
f7672ce
fix: update model reparameterisation as model is now copied within Ba…
BradyPlanden Oct 21, 2024
0549d14
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Oct 23, 2024
44e3095
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Oct 23, 2024
e1b99d7
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Nov 2, 2024
8f11efe
refactor: Jax implementation, FittingProblem.evaluate, adds Fisher In…
BradyPlanden Nov 4, 2024
e362938
refactor: Jax costs and add Jax-based evaluator
BradyPlanden Nov 25, 2024
a6f987d
refactor: update init structure
BradyPlanden Dec 4, 2024
9dec62d
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Dec 4, 2024
b7c054a
tests: adds coverage
BradyPlanden Dec 6, 2024
0a54573
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Dec 6, 2024
8f74e33
Merge remote-tracking branch 'origin/jaxify-idaklu-implementation' in…
BradyPlanden Dec 6, 2024
8a49024
adds coverage, updates jaxified solver for current pybamm bug
BradyPlanden Dec 19, 2024
4e563db
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Dec 19, 2024
f729a95
Adds SciPyEvaluator, updates examples, tests.
BradyPlanden Dec 20, 2024
30be3a2
Adds coverage, updates evaluation, and improves integration between B…
BradyPlanden Dec 20, 2024
ca77081
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Dec 21, 2024
ccb7493
Upper pin BPX to <0.5
BradyPlanden Dec 21, 2024
762062c
Merge branch 'develop' into jaxify-idaklu-implementation
BradyPlanden Dec 21, 2024
5ae6543
refactor: changes post review
BradyPlanden Dec 21, 2024
99d0bb5
refactor: updates to docstrings, general improvements
BradyPlanden Dec 21, 2024
b93bd36
tests: increment sigma0
BradyPlanden Dec 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

## Features

- [#481](https://github.com/pybop-team/PyBOP/pull/481) - Adds experimental support for PyBaMM's jaxified IDAKLU solver. Includes Jax-specific cost functions `pybop.JaxSumSquareError` and `pybop.JaxLogNormalLikelihood`. Adds `Jax` optional dependency to PyBaMM dependency.
- [#597](https://github.com/pybop-team/PyBOP/pull/597) - Adds number of function evaluations `n_evaluations` to `OptimisationResult`.
- [#362](https://github.com/pybop-team/PyBOP/issues/362) - Adds the `classify_using_Hessian` functionality to classify the optimised result.
- [#584](https://github.com/pybop-team/PyBOP/pull/584) - Adds the `GroupedSPMe` model for parameter identification.
Expand Down Expand Up @@ -35,6 +36,7 @@

## Breaking Changes

- [#481](https://github.com/pybop-team/PyBOP/pull/481) - `problem.model` is now a copied instance of `model`
- [#598](https://github.com/pybop-team/PyBOP/pull/598) - Depreciated `Adam` optimiser has been removed, see `AdamW` for replacement.
- [#531](https://github.com/pybop-team/PyBOP/pull/531) - Plot methods moved to `pybop.plot` with mostly minimal renaming. For example, `pybop.plot_parameters` is now `pybop.plot.parameters`. Other breaking changes include: `pybop.plot2d` to `pybop.plot.contour`.
- [#526](https://github.com/pybop-team/PyBOP/pull/526) - Refactor `OptimisationResults` classes, with `optim.run()` now return the full object. Adds finite cost value check for optimised parameters.
Expand Down
Loading
Loading