Code

Idaklu-Jax Interface#3658

Code

Merged

martinjrobins merged 58 commits intopybamm-team:developpybamm-team/PyBaMM:developfrom jsbrittain:idaklu_jaxjsbrittain/PyBaMM:idaklu_jaxCopy head branch name to clipboardFeb 21, 2024

+2,982-43

Commits

Commits on Dec 11, 2023

• Add IDAKLUJax class, tests, and run black

  jsbrittain

  committed
  dca5e3f

Commits on Dec 14, 2023

• Jax primitive can lower with correct dimensions (arbitrary return vector)

  jsbrittain

  committed
  9b3f776

Commits on Dec 18, 2023

• Return np_array to cpp

  jsbrittain

  committed
  924c673

Commits on Dec 19, 2023

• JIT implemented f, jvp, vjp; 2 input case only

  jsbrittain

  committed
  11a0b62

Commits on Dec 20, 2023

• JIT implemented for arbitrary number of inputs; tests pass, but jaxexpr-solver is state-dependent on class

  jsbrittain

  committed
  87133a0

• Refactor c_solver to include idaklu.jax cpp/hpp

  jsbrittain

  committed
  5268349

Commits on Dec 21, 2023

• Support multiple idaklu-jax instantiations; provide jupyter notebook

  jsbrittain

  committed
  4ab5cb2

Commits on Dec 23, 2023

• Increased annotations and fix jupyter notebook

  jsbrittain

  committed
  67452a6

• Modify examples to be compatible with output_variables declared in Solver object

  jsbrittain

  committed
  a8afe47

• Merge upstream and fix conflicts

  jsbrittain

  committed
  91de54b

• Tidy-up docs

  jsbrittain

  committed
  8e7dc25

• Fix tests; add changelog

  jsbrittain

  committed
  27e1f24

Commits on Dec 24, 2023

• Improve test coverage

  jsbrittain

  committed
  a0ca1fd

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  4f3bcbe

Commits on Jan 8, 2024

• Fix tests

  jsbrittain

  committed
  2c97a9f

• Fix merge conflict

  jsbrittain

  committed
  7434ee6

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  094482f

• Lychee fix

  jsbrittain

  committed
  69acbdd

• Fix tests

  jsbrittain

  committed
  ef718c2

• Fix tests

  jsbrittain

  committed
  fe37d8a

Commits on Jan 10, 2024

• Python3.12 compatibilty

  jsbrittain

  committed
  f297ee9

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  50a7068

• Improve coverage

  jsbrittain

  committed
  4d6f157

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  fd22424

• Improve coverage

  jsbrittain

  committed
  56dd5f8

Commits on Jan 12, 2024

• Sync idaklu source lists in CMakeLists.txt and setup.py

  jsbrittain

  committed
  2c4dd49

• Move paramaterized to [dev]; add [dev] to nox test environments

  jsbrittain

  committed
  e4edf72

Commits on Jan 13, 2024

• [fix] turn off debug statements in idaklu-jax

  jsbrittain

  committed
  ce2fec6

• docstring change

  jsbrittain

  committed
  ae55921

• Remove print statements from idaklu-jax interface tests

  jsbrittain

  committed
  aea7541

• Remove function input from idaklu-jax helper functions (get_var, get_vars)

  jsbrittain

  committed
  54773d4

Commits on Jan 15, 2024

• Jaxify on initialisation of idaklu-jax interface

  jsbrittain

  committed
  3a9c222

• Ensure calculate_sensitivities is functional

  jsbrittain

  committed
  c029d7c

• Add call signatures to idaklu-jax helper functions

  jsbrittain

  committed
  4f169e3

• Ensure idaklu-jax logger is named

  jsbrittain

  committed
  af98682

• Replace DFN model with simpler expontential decay model for unit testing the idaklu-jax interface

  jsbrittain

  committed
  f445115

• Pre-commit

  jsbrittain

  committed
  cf91be6

• Merge branch 'develop' into idaklu_jax

  jsbrittain

  authored
  7775642

• Remove from idaklu-jax interface set-up call

  jsbrittain

  committed
  97f56eb

• Update idaklu-jax jupyter notebook

  jsbrittain

  committed
  8e0eaa7

• Fix docs build

  jsbrittain

  committed
  dd1c9e7

• Ensure inputs are updated and in-order in idaklu-jax

  jsbrittain

  committed
  bdf445d

Commits on Jan 16, 2024

• Allow override of function in get_var() and get_vars() with either a JAX expression or an array

  jsbrittain

  committed
  0252064

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  80eff0d

• Fix docs and coverage

  jsbrittain

  committed
  391d121

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  1c76fa8

Commits on Jan 17, 2024

• Update CHANGELOG.md

  Show description for cad5f9a
  

  jsbrittain

  and

  martinjrobins

  authored
  cad5f9a

Commits on Jan 18, 2024

• Add parameterized to the optional dependencies table in the installation guide

  jsbrittain

  committed
  484ff96

• Add test to ensure that IDAKLUJax fails if Idaklu and/or Jax are not available

  jsbrittain

  committed
  2df1e9e

• style: pre-commit fixes

  pre-commit-ci[bot]

  committed
  ebd2b95

• Add idaklu-jax requirements to jupyter notebook

  jsbrittain

  committed
  ae82b8f

Commits on Jan 19, 2024

• Rephrase idaklu-jax requirements in jupyter notebook

  jsbrittain

  committed
  f2bc81b

• Merge branch 'develop' into idaklu_jax

  jsbrittain

  committed
  4896c11

Commits on Jan 22, 2024

• Merge branch 'develop' into idaklu_jax

  jsbrittain

  authored
  97af1ad

Commits on Feb 9, 2024

• Merge develop

  jsbrittain

  committed
  1fd7366

• Merge develop [fix]

  jsbrittain

  committed
  71aca14

• Add dev to nox test environments

  jsbrittain

  committed
  65e71e8

Commits on Feb 21, 2024

• Merge branch 'develop' into idaklu_jax

  jsbrittain

  committed
  d201899