In this section, we will optimize our particle simulator by rewriting some parts of it in NumPy. From the profiling we did in Chapter 1, Benchmarking and Profiling, the slowest part of our program is the following loop contained in the ParticleSimulator.evolve method:
for i in range(nsteps):
for p in self.particles:
norm = (p.x**2 + p.y**2)**0.5
v_x = (-p.y)/norm
v_y = p.x/norm
d_x = timestep * p.ang_speed * v_x
d_y = timestep * p.ang_speed * v_y
p.x += d_x
p.y += d_yWe may notice that the body of the loop acts solely on the current particle. If we had an array containing the particle positions and angular speed, we could rewrite the loop using a broadcasted operation. In contrast, the loop over the time steps depends on the previous step and cannot be treated in a parallel fashion.
It's natural then, to store all the array coordinates in an array of shape (nparticles, 2) and the angular speed in an array of shape ...