3 ms·
Jax is super useful for scientific computing. Although nbody sims might not be the best application. A naive nbody sim is very easy to implement and accelerate
by tehsauce 3y ago
Jax is super useful for scientific computing. Although nbody sims might not be the best application. A naive nbody sim is very easy to implement and accelerate in jax (here’s my version: https://github.com/PWhiddy/jax-experiments/blob/main/nbody.ipynb https://github.com/PWhiddy/jax-experiments/blob/main/nbody.i...), but it can be tricky to scale it. This is because efficient nbody sims usually either rely on trees or spatial hashing/sorting which are tricky to efficiently implement with jax.
- mattjjatgoogle 3y agoHave you seen JAX MD? https://github.com/jax-md/jax-md https://github.com/jax-md/jax-md
- tehsauce 3y agoI've seen it although haven't dived deep into it. It looks like they have some interesting support for particle cell data structures, but is fairly complicated and carries limitations: https://jax-md.readthedocs.io/en/main/_modules/jax_md/partition.html#cell_list https://jax-md.readthedocs.io/en/main/_modules/jax_md/partit...
- dekhn 3y agoLast time I looked at JAX MD it didn't support most of the force field terms necessary for simulating proteins and DNA. For example, it could do n-body simulations with some potential, but not the bonds/torsions between atoms. It's unclear if they added support, but that's a huge gap in functionality compared to other systems.