Lattice Boltzmann Simulation without Jax

Ceyron / machine-learning-and-simulation

All the handwritten notes 📝 and source code files 🖥️ used in my YouTube Videos on Machine Learning & Simulation (https://www.youtube.com/channel/UCh0P7KwJhuQ4vrzc3IRuw4Q)

MIT License

801 stars 176 forks source link

Lattice Boltzmann Simulation without Jax #10

Closed mhechthz closed 12 months ago

mhechthz commented 1 year ago

Could you add the lattice boltzmann simulation without the use of jax?

Ceyron commented 1 year ago

Hi, I think it should be relatively straightforward to translate it since JAX essentially implements the numpy API. The three points that need changes are:

Change the import from jax.numpy as jnp to numpy as np and subsequently change all jnp to np
Remove the JAX jit compilation
Change to in-place array mutations

Would be a nice thing if you could give it a try and open a PR 😊