Code

Issue 1031 jax#1038

Code

Merged

martinjrobins merged 56 commits intodeveloppybamm-team/PyBaMM:developfrom issue-1031-jaxpybamm-team/PyBaMM:issue-1031-jaxCopy head branch name to clipboardJul 7, 2020

+1,748-105

Commits

Commits on May 22, 2020

• started experimenting with jax evaluator

  martinjrobins

  committed
  7b5551d

  Copy full SHA for 7b5551d

Commits on May 30, 2020

• #1031 added an evaluator class that compiles using JAX

  martinjrobins

  committed
  23adeaf

  Copy full SHA for 23adeaf

• #1031 add jacobian evaluation using jax

  martinjrobins

  committed
  34005bf

  Copy full SHA for 34005bf

• #1031 add jax convert to target

  martinjrobins

  committed
  d66f2d4

  Copy full SHA for d66f2d4

• #1031 add test to use jax conversion in solver

  martinjrobins

  committed
  8a7c850

  Copy full SHA for 8a7c850

Commits on May 31, 2020

• #1031 improved statevector and constant number handling in evaluators

  martinjrobins

  committed
  218ffa3

  Copy full SHA for 218ffa3

• #1031 fix introduced bugs in evaluators

  martinjrobins

  committed
  fb61c14

  Copy full SHA for fb61c14

• #1031 tidy up selfs in generated code

  martinjrobins

  committed
  9843c87

  Copy full SHA for 9843c87

Commits on Jun 3, 2020

• #1031 add jax solver

  martinjrobins

  committed
  adeba73

  Copy full SHA for adeba73

• #1031 docstrings for jax solver

  martinjrobins

  committed
  540a462

  Copy full SHA for 540a462

• #1031 add docs for jax solver

  martinjrobins

  committed
  d91da37

  Copy full SHA for d91da37

• #1031 fail if any sparse matrices are in model

  martinjrobins

  committed
  6eb0765

  Copy full SHA for 6eb0765

• #1031 add additional runtimeerrors for jax solver

  martinjrobins

  committed
  2fa6359

  Copy full SHA for 2fa6359

• #1031 update changelog

  martinjrobins

  committed
  aa15b62

  Copy full SHA for aa15b62

• #1031 fix flake8

  martinjrobins

  committed
  06d4538

  Copy full SHA for 06d4538

• #1031 handle numpy functions in jax evaluator

  martinjrobins

  committed
  708e00a

  Copy full SHA for 708e00a

• #1031 tidy imports

  martinjrobins

  committed
  5544caa

  Copy full SHA for 5544caa

• #1031 events should return scalars

  martinjrobins

  committed
  8488106

  Copy full SHA for 8488106

• #1031 forgot to add doc file

  martinjrobins

  committed
  6720ac3

  Copy full SHA for 6720ac3

• #1031 fix docstring format

  martinjrobins

  committed
  9abc7b6

  Copy full SHA for 9abc7b6

Commits on Jun 16, 2020

• #1031 default to 64bit doubles, allow conversion of sparse to dense for jax compilation

  martinjrobins

  committed
  eb09109

  Copy full SHA for eb09109

Commits on Jun 17, 2020

• #1031 start on bdf solver for jax

  martinjrobins

  committed
  9c390eb

  Copy full SHA for 9c390eb

Commits on Jun 18, 2020

• #1031 use bdf stepper in integrate function

  martinjrobins

  committed
  7835aa0

  Copy full SHA for 7835aa0

• #1031 create a simple test

  martinjrobins

  committed
  ab0565e

  Copy full SHA for ab0565e

• #1031 update index assignments to jax

  martinjrobins

  committed
  fea7901

  Copy full SHA for fea7901

• #1031 start with a sensible time step

  martinjrobins

  committed
  8a09091

  Copy full SHA for 8a09091

• #1031 solver works on simple test

  martinjrobins

  committed
  81845ae

  Copy full SHA for 81845ae

Commits on Jun 19, 2020

• #1031 convert OO stuff to pure functions for jax

  martinjrobins

  committed
  a2bba78

  Copy full SHA for a2bba78

• #1031 working on flow control for jax

  martinjrobins

  committed
  be8a40a

  Copy full SHA for be8a40a

Commits on Jun 20, 2020

• #1031 restructuring to avoid shape changes and dynamic indexing

  martinjrobins

  committed
  f6c8623

  Copy full SHA for f6c8623

Commits on Jun 21, 2020

• #1031 jax bdf solver working on simple test problem

  martinjrobins

  committed
  1f8f10d

  Copy full SHA for 1f8f10d

Commits on Jun 22, 2020

• #1031 add fun parameterised by inputs

  martinjrobins

  committed
  7656e51

  Copy full SHA for 7656e51

Commits on Jun 23, 2020

• #1031 add bdf to jax solver

  martinjrobins

  committed
  dc7a31b

  Copy full SHA for dc7a31b

• #1031 merge in develop

  martinjrobins

  committed
  0efc099

  Copy full SHA for 0efc099

Commits on Jun 24, 2020

• #1031 test bdf for jax solver

  martinjrobins

  committed
  7fe74e4

  Copy full SHA for 7fe74e4

• #1031 fix a failing test, make sure jax solvers convert back to numpy array

  martinjrobins

  committed
  e7d1797

  Copy full SHA for e7d1797

• #1031 jax jacobian return correct shape

  martinjrobins

  committed
  db3ad98

  Copy full SHA for db3ad98

Commits on Jun 26, 2020

• #1031 update some docstrings + comments, fix tests

  martinjrobins

  committed
  28ae09f

  Copy full SHA for 28ae09f

• #1031 fix flake8

  martinjrobins

  committed
  8e0dbfb

  Copy full SHA for 8e0dbfb

• #1031 improve docstrings

  martinjrobins

  committed
  d4be017

  Copy full SHA for d4be017

• #1031 move main jax_bdf_integrate function to top of file to make it more obvious

  martinjrobins

  committed
  7b9a757

  Copy full SHA for 7b9a757

• #1031 fix test

  martinjrobins

  committed
  25bb757

  Copy full SHA for 25bb757

• #1031 switch off jax for windows

  martinjrobins

  committed
  93c2386

  Copy full SHA for 93c2386

Commits on Jul 2, 2020

• #1031 merge develop

  martinjrobins

  committed
  bf76212

  Copy full SHA for bf76212

• #1031 fix syntax error

  martinjrobins

  committed
  fc48be7

  Copy full SHA for fc48be7

• #1031 fix docstrings

  martinjrobins

  committed
  ff529ca

  Copy full SHA for ff529ca

• #1031 fix docstrings

  martinjrobins

  committed
  7c201e6

  Copy full SHA for 7c201e6

• #1031 fix import error on windows

  martinjrobins

  committed
  bdbac39

  Copy full SHA for bdbac39

• #1031 fix windows import errors

  martinjrobins

  committed
  779543e

  Copy full SHA for 779543e

Commits on Jul 3, 2020

• #1031 improve test coverage

  martinjrobins

  committed
  fcb6269

  Copy full SHA for fcb6269

• #1031 test evaluator jax jacobian and debug

  martinjrobins

  committed
  f294039

  Copy full SHA for f294039

• #1031 improve coverage

  martinjrobins

  committed
  284521b

  Copy full SHA for 284521b

Commits on Jul 7, 2020

• #1031 fix typo in docstring

  martinjrobins

  committed
  74d4ba7

  Copy full SHA for 74d4ba7

• #1031 fix doctring for bdf init function

  martinjrobins

  committed
  b8e9030

  Copy full SHA for b8e9030

• #1031 remove commented out code

  martinjrobins

  committed
  44df73b

  Copy full SHA for 44df73b

• #1031 give more details when jax solver fails due to terminate event

  martinjrobins

  committed
  a30ddf0

  Copy full SHA for a30ddf0