Code

Adds JAX IDAKLU solver integration#481

Code

Merged

BradyPlanden merged 40 commits intodeveloppybop-team/PyBOP:developfrom jaxify-idaklu-implementationpybop-team/PyBOP:jaxify-idaklu-implementationCopy head branch name to clipboardDec 21, 2024

+909-1,546

Commits

Commits on Aug 20, 2024

• working commit, adds [jax] optional dependency to pybamm, example idaklu-jax solver

  BradyPlanden

  committed
  1e08d6e

  Copy full SHA for 1e08d6e

• adds experimental sub-directory, sync commit

  BradyPlanden

  committed
  79234f9

  Copy full SHA for 79234f9

• Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation

  BradyPlanden

  committed
  7d8aa9c

  Copy full SHA for 7d8aa9c

Commits on Aug 28, 2024

• Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation

  BradyPlanden

  committed
  ab861a9

  Copy full SHA for ab861a9

Commits on Aug 31, 2024

• feat: adds Jax functionality via idaklu_jax solver, with example and benchmarks. Adds JaxSumSquaredError and JaxLogNormalLikelihood.

  BradyPlanden

  committed
  f2bf143

  Copy full SHA for f2bf143

Commits on Sep 2, 2024

• Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation

  Show description for 025c7d8
  BradyPlanden

  committed
  025c7d8

  Copy full SHA for 025c7d8

• Post merge fixes, update base_model solver property

  BradyPlanden

  committed
  d3a18d5

  Copy full SHA for d3a18d5

• examples: update benchmarking script

  BradyPlanden

  committed
  eb08c00

  Copy full SHA for eb08c00

Commits on Sep 6, 2024

• Merge branch 'develop' into jaxify-idaklu-implementation

  BradyPlanden

  authored
  92265b6

  Copy full SHA for 92265b6

Commits on Sep 24, 2024

• Merge branch 'refs/heads/develop' into jaxify-idaklu-implementation

  Show description for f1d03a9
  BradyPlanden

  committed
  f1d03a9

  Copy full SHA for f1d03a9

• updts docstrings, jax arg to models, adds tests.

  BradyPlanden

  committed
  3ecd5dd

  Copy full SHA for 3ecd5dd

Commits on Sep 25, 2024

• add coverage, refactor BaseModel solver setter

  BradyPlanden

  committed
  9794910

  Copy full SHA for 9794910

• adds coverage

  BradyPlanden

  committed
  df1ba0c

  Copy full SHA for df1ba0c

• convert BaseModel.calculate_sensitivites to property

  BradyPlanden

  committed
  d770315

  Copy full SHA for d770315

• add changelog entry

  BradyPlanden

  committed
  daae4ba

  Copy full SHA for daae4ba

Commits on Sep 30, 2024

• feat: removes Jax arg from BaseModel, BaseProblem.model is copied instance

  BradyPlanden

  committed
  424699b

  Copy full SHA for 424699b

Commits on Oct 11, 2024

• Merge branch 'develop' into jaxify-idaklu-implementation

  BradyPlanden

  authored
  2695fd7

  Copy full SHA for 2695fd7

Commits on Oct 21, 2024

• fix: cost shape for CMAES, adds example comments for jax solver

  BradyPlanden

  committed
  23b8e3d

  Copy full SHA for 23b8e3d

• Merge branch 'develop' into jaxify-idaklu-implementation

  Show description for bba2e56
  BradyPlanden

  committed
  bba2e56

  Copy full SHA for bba2e56

• fix: update model reparameterisation as model is now copied within BaseProblem()

  BradyPlanden

  committed
  f7672ce

  Copy full SHA for f7672ce

Commits on Oct 23, 2024

• Merge branch 'develop' into jaxify-idaklu-implementation

  Show description for 0549d14
  BradyPlanden

  committed
  0549d14

  Copy full SHA for 0549d14

• Merge branch 'develop' into jaxify-idaklu-implementation

  Show description for 44e3095
  BradyPlanden

  committed
  44e3095

  Copy full SHA for 44e3095

Commits on Nov 2, 2024

• Merge branch 'develop' into jaxify-idaklu-implementation

  BradyPlanden

  committed
  e1b99d7

  Copy full SHA for e1b99d7

Commits on Nov 4, 2024

• refactor: Jax implementation, FittingProblem.evaluate, adds Fisher Information to OptimisationResult

  BradyPlanden

  committed
  8f11efe

  Copy full SHA for 8f11efe

Commits on Nov 25, 2024

• refactor: Jax costs and add Jax-based evaluator

  BradyPlanden

  committed
  e362938

  Copy full SHA for e362938

Commits on Dec 4, 2024

• refactor: update init structure

  BradyPlanden

  committed
  a6f987d

  Copy full SHA for a6f987d

• Merge branch 'develop' into jaxify-idaklu-implementation

  Show description for 9dec62d
  BradyPlanden

  committed
  9dec62d

  Copy full SHA for 9dec62d

Commits on Dec 6, 2024

• tests: adds coverage

  BradyPlanden

  committed
  b7c054a

  Copy full SHA for b7c054a

• Merge branch 'develop' into jaxify-idaklu-implementation

  BradyPlanden

  authored
  0a54573

  Copy full SHA for 0a54573

• Merge remote-tracking branch 'origin/jaxify-idaklu-implementation' into jaxify-idaklu-implementation

  BradyPlanden

  committed
  8f74e33

  Copy full SHA for 8f74e33

Commits on Dec 19, 2024

• adds coverage, updates jaxified solver for current pybamm bug

  BradyPlanden

  committed
  8a49024

  Copy full SHA for 8a49024

• Merge branch 'develop' into jaxify-idaklu-implementation

  BradyPlanden

  committed
  4e563db

  Copy full SHA for 4e563db

Commits on Dec 20, 2024

• Adds SciPyEvaluator, updates examples, tests.

  BradyPlanden

  committed
  f729a95

  Copy full SHA for f729a95

Commits on Dec 21, 2024

• Adds coverage, updates evaluation, and improves integration between BaseCost & BaseJaxCost

  BradyPlanden

  committed
  30be3a2

  Copy full SHA for 30be3a2

• Merge branch 'develop' into jaxify-idaklu-implementation

  Show description for ca77081
  BradyPlanden

  committed
  ca77081

  Copy full SHA for ca77081

• Upper pin BPX to <0.5

  BradyPlanden

  committed
  ccb7493

  Copy full SHA for ccb7493

• Merge branch 'develop' into jaxify-idaklu-implementation

  Show description for 762062c
  BradyPlanden

  committed
  762062c

  Copy full SHA for 762062c

• refactor: changes post review

  BradyPlanden

  committed
  5ae6543

  Copy full SHA for 5ae6543

• refactor: updates to docstrings, general improvements

  BradyPlanden

  committed
  99d0bb5

  Copy full SHA for 99d0bb5

• tests: increment sigma0

  BradyPlanden

  committed
  b93bd36

  Copy full SHA for b93bd36