Speeding up simulations with JAX #1072

corbinklett · 2024-12-03T16:42:38Z

corbinklett
Dec 3, 2024

Does it make sense to incorporate JAX in order to run simulations faster? It seems to speed up simulation significantly via just-in-time compilation. This could help with large models, long time horizons, and realtime testing, including software-in-the-loop. Just wondering if this type of effort is something has has been considered before.

murrayrm · 2024-12-03T18:52:03Z

murrayrm
Dec 3, 2024
Maintainer

Worth a shot if someone wants to code it up. A couple of thoughts:

This is probably going to matter most for nonlinear models and will likely require passing additional information to the simulators (eg, gradients) to take advantage of the structure, when available. In addition, we may have to modify the interconnect function and perhaps a few others to use JAX instead of NumPy for various operations.
If we implement this, we should probably use a x_method keyword (with 'x' as some appropriate prefix) so that users can choose to use 'jax' (if installed) versus 'numpy'. This would also let us maintain use of NumPy if JAX is not installed (so that JAX is an optional dependency, like pandas and cvxopt).
It would be good to put together some benchmarks so we can make sure things are faster. See the benchmarks/ directory for examples and the README file in that directory for more information.
Another possibility (and method) would be to allow the use of numba. This can possibly be done completely by the user, though again if you want to interconnect subsystems then some modifications to nlsys.py may be required.

1 reply

corbinklett Dec 3, 2024
Author

Thanks for the suggestions! Will post an update when I'm able to experiment more with this.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Speeding up simulations with JAX #1072

Uh oh!

{{title}}

Uh oh!

Replies: 1 comment 1 reply

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

Speeding up simulations with JAX #1072

Uh oh!

corbinklett Dec 3, 2024

Replies: 1 comment · 1 reply

Uh oh!

murrayrm Dec 3, 2024 Maintainer

Uh oh!

corbinklett Dec 3, 2024 Author

corbinklett
Dec 3, 2024

Replies: 1 comment 1 reply

murrayrm
Dec 3, 2024
Maintainer

corbinklett Dec 3, 2024
Author