This notebook contains a brief introduction to JAX and a demonstration of how it can be used to define and solve the equations of motion of a simple mass-spring system. I wrote this up in the process of learning to create wave propagation simulations from scratch.
Relevant features of JAX, as described in the JAX repository:
- "automatically differentiate native Python and NumPy code" with the
gradfunction (I used this to automatically generate the equations of motion without having to compute them by hand) - "automatic vectorization" with
vmapfunction (I used this to simplify the function definitions and overall code structure)
I also looked into using the jit ("just-in-time") decorator to speed up function calls, but jit is not compatible with control flow operations like if-else statements, which I wanted to use for the boundary conditions.
View the Jupyter notebook on nbviewer