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.

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]