I'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

I'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...