pybamm.
JaxSolver
(method='RK45', root_method=None, rtol=1e-06, atol=1e-06, extra_options=None)[source]¶Solve a discretised model using a JAX compiled solver.
Raises: |
|
---|---|
Parameters: |
|
create_solve
(model, t_eval)[source]¶Return a compiled JAX function that solves an ode model with input arguments.
Parameters: |
|
---|---|
Returns: | A function with signature f(inputs), where inputs are a dict containing any input parameters to pass to the model when solving |
Return type: | function |
get_solve
(model, t_eval)[source]¶Return a compiled JAX function that solves an ode model with input arguments.
Parameters: |
|
---|---|
Returns: | A function with signature f(inputs), where inputs are a dict containing any input parameters to pass to the model when solving |
Return type: | function |
pybamm.
jax_bdf_integrate
(func, y0, t_eval, *args, rtol=1e-06, atol=1e-06, mass=None)[source]¶Backward Difference formula (BDF) implicit multistep integrator. The basic algorithm is derived in [2]. This particular implementation follows that implemented in the Matlab routine ode15s described in [1] and the SciPy implementation [3], which features the NDF formulas for improved stability, with associated differences in the error constants, and calculates the jacobian at J(t_{n+1}, y^0_{n+1}). This implementation was based on that implemented in the scipy library [3], which also mainly follows [1] but uses the more standard jacobian update.
Parameters: |
|
---|---|
Returns: | y – calculated state vector at each of the m time points |
Return type: | ndarray with shape (n, m) |
References
[1] | (1, 2) L. F. Shampine, M. W. Reichelt, “THE MATLAB ODE SUITE”, SIAM J. SCI. COMPUTE., Vol. 18, No. 1, pp. 1-22, January 1997. |
[2] | G. D. Byrne, A. C. Hindmarsh, “A Polyalgorithm for the Numerical Solution of Ordinary Differential Equations”, ACM Transactions on Mathematical Software, Vol. 1, No. 1, pp. 71-96, March 1975. |
[3] | (1, 2) Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., … & van der Walt, S. J. (2020). SciPy 1.0: fundamental algorithms for scientific computing in Python. Nature methods, 17(3), 261-272. |