[naillo]

You should check out taichi: https://github.com/taichi-dev/taichi They have a ton of great demos for doing physics but check out this example in particular for something related to your project: https://github.com/taichi-dev/quantaichi#game-of-life-gol. (Taichi also makes it super easy to write things for the GPU and the kernels are differentiable :).)