Parallel Lorenz Simulation in JAX

The JAX Deep Learning framework in Python is a powerful superset of familiar NumPy functions. In this video, we use its automatic vectorization via jax.vmap to simulate multiple Lorenz trajectories/butterfly shapes at the same time. Here is the code: https://github.com/Ceyron/machine-lea...

------

👉 This educational series is supported by the world-leaders in integrating machine learning and artificial intelligence with simulation and scientific computing, Pasteur Labs and Institute for Simulation Intelligence. Check out https://simulation.science/ for more on their pursuit of 'Nobel-Turing' technologies (https://arxiv.org/abs/2112.03235 ), and for partnership or career opportunities.

-------

📝 : Check out the GitHub Repository of the channel, where I upload all the handwritten notes and source-code files: https://github.com/Ceyron/machine-lea...

📢 : Follow me on LinkedIn or Twitter for updates on the channel and other cool Machine Learning & Simulation stuff:   / felix-koehler   and   / felix_m_koehler  

💸 : If you want to support my work on the channel, you can become a Patreon here:   / mlsim  

🪙: Or you can make a one-time donation via PayPal: https://www.paypal.com/paypalme/Felix...

-----

Timestamps:

00:00 Intro
01:08 Imports
01:30 Adapting Lorenz RHS and RK4 Simulator
02:44 Autoregressive Rollout (to get a trajectory)
08:06 Comparison of the trajectories (chaos due to single precision)
09:11 Lorenz Map
12:01 About the automatic vectorization in JAX
14:50 Multiple Initial Conditions
15:57 jax.vmap for parallel RK4 stepping
16:54 Parallel Rollout/Simulation for multiple trajectories
21:54 Visualize all 9 trajectories
25:33 Compute & Visualize all 9 trajectories
30:30 Outro