Use NumPy & Avoid For Loops

I was given a Python script by a student for simulating the random motion of particles.  It was really good for a first attempt but he had ignored my pleas to avoid nested for loops and to use NumPy arrays and functions instead.  The following is a comparison of the original script and my two attempts to improve it.

This is the figure that the scripts output.  The student wanted to show the final position of all the particles and the trajectory of the first (red) and last (green) particle.  For this test I tracked 100 particles through 100,000 steps.

3D_100_Particles_100_Steps_v0

The following is a simplified version of the original script.  It takes 128 seconds to execute.  For each particle and time step, a random direction and distance is applied.

import math
import matplotlib.pyplot as pp
from mpl_toolkits.mplot3d import Axes3D as pp3d
from matplotlib import rcParams
import numpy as np
import random
from time import clock

## Input parameters
particles = 100
steps = 100000

calc_time = clock()

## Initialize all particles at zero
points = np.zeros((particles,3))

## For the first and last points, save the full path history
## Need one extra step value to store starting point
pointsR = np.zeros((steps+1,3))
pointsG = np.zeros((steps+1,3))

## Loop through each particle
for j in range(0,particles):

    ## Loop through each step
    for i in range(1,steps+1):
        ## For the angles, use the python uniform random distribution
        theta = random.uniform(0, 2 * math.pi)
        phi = random.uniform(0, 2 * math.pi)
        ## For the distance, use the python gaussian random distribution
        r = random.gauss(1.0, 0.01)

        ## Apply the motion to each particle for eacch step
        points[j,0] += r * np.sin(phi) * np.cos(theta)
        points[j,1] += r * np.sin(phi) * np.sin(theta)
        points[j,2] += r * np.cos(phi)

        ## Save the full trajectory for the first and last particle
        if j == 0:
            pointsR[i,0] = points[j,0]
            pointsR[i,1] = points[j,1]
            pointsR[i,2] = points[j,2]
        if j == particles - 1:
            pointsG[i,0] = points[j,0]
            pointsG[i,1] = points[j,1]
            pointsG[i,2] = points[j,2]

## Plot the final point locations
fig = pp.figure()
rcParams.update({'figure.autolayout': True}) # adjusts layout of graph
ax = fig.add_subplot(111, projection='3d')
ax.scatter(points[:,0], points[:,1], points[:,2], c='b')
ax.plot(pointsR[:,0], pointsR[:,1], pointsR[:,2], '-', c='r')
ax.plot(pointsG[:,0], pointsG[:,1], pointsG[:,2], '-', c='g')
print('Calc Time: {0}'.format(clock()-calc_time))
pp.show() 

 

As a first step, I thought that all the particles could be processed at once, instead of separately.  For each time step, I built an array of the random motions for all the particles and then applied these all at once.  This version takes 6.4 seconds to execute.

## Initialize all particles at zero
points = np.zeros((particles,3))

## For the first and last points, save the full path history
## Need one extra step value to store starting point
pointsR = np.zeros((steps+1,3))
pointsG = np.zeros((steps+1,3))

## Loop through each step
for i in range(steps):
    ## For the angles, use the NumPy uniform random distribution function
    theta = np.random.uniform(0, 2 * math.pi, particles)
    phi = np.random.uniform(0, 2 * math.pi, particles)

    ## For the distance, use the NumPy normal (Gaussian) random distribution
    r = np.random.normal(1.0, 0.01, particles)

    ## Apply the motion to each axis for all particles
    points[:,0] += r * np.sin(phi) * np.cos(theta)
    points[:,1] += r * np.sin(phi) * np.sin(theta)
    points[:,2] += r * np.cos(phi)

    ## Save the full trajectory of the first and last point
    ## Particles started at 0 so add 1 to index
    pointsR[i+1,:] = points[0,:]
    pointsG[i+1,:] = points[-1,:]

 

It also seemed like this should be possible without using any for loops.  To do that I built an array of the random motions for all particles and steps, and then used the NumPy cumulative sum function to apply the motions across the “time steps” axis of the array.  This approach was faster for some combinations of particles and steps.  With the 100 particles and 100,000 steps this takes 2.3 seconds to execute.  The main downside to this approach is that it can use a lot of memory since all the motions are stored in memory at once.

## Initialize all particles and steps at zero
points = np.zeros((3,steps,particles))

## Pre-calculate sin/cos for all particles and steps
## For the angles, use the NumPy uniform random distribution function
phi = np.random.uniform(0, 2 * math.pi, (steps,particles))
theta = np.random.uniform(0, 2 * math.pi, (steps,particles))

## For the distance, use the NumPy normal (Gaussian) random distribution
r = np.random.normal(1.0, 0.01, (steps,particles))

## Pre calculate the x,y,z motion for each particle and step
points[0,:,:] = r * np.sin(phi) * np.cos(theta)
points[1,:,:] = r * np.sin(phi) * np.sin(theta)
points[2,:,:] = r * np.cos(phi)

## For all particles, sum the motions for all steps at once
np.cumsum(points, axis=1, out=points)

## Save the full trajectory of the first and last point
## This doesn't include first point at the origin though
pointsR = points[:,:,0]
pointsG = points[:,:,particles-1]

 

 

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s