Ceyron/machine-learning-and-simulation

Lattice Boltzmann Simulation without Jax

mhechthz opened this issue · 1 comments

mhechthz commented

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

Ceyron commented

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 😊