Rewriting the particle simulator in NumPy

In this section, we will optimize our particle simulator by rewriting some parts of it in NumPy. We found, from the profiling we did in Chapter 1, Benchmarking and Profiling, that 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_vel * v_x         d_y = timestep * p.ang_vel * v_y         p.x += d_x         p.y += d_y 

You may have noticed 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. ...

Get Python High Performance - Second Edition now with O’Reilly online learning.

O’Reilly members experience live online training, plus books, videos, and digital content from 200+ publishers.