As you note, the particles are independent of each other in the current example so it’s hard to give advice. Can you post a simple version of the model you’re actually interested in, even if it’s in numpy code? Scan would be the most general solution, but you might also be able to get away with something more clever depending on the dependency structure between the particles.
1 Like