Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
numpy performance improvement for a particle simulator program (from Python High Performance)
"""
This program is my implementation of G.Lanaro particle simulator.
The objective is to demonstrate efficiency of numpy operations in a simulation problem.
https://www.packtpub.com/application-development/python-high-performance-second-edition
"""
import matplotlib.pyplot as plt
from matplotlib import animation
from typing import Sequence
from random import uniform
import timeit
import cProfile
import numpy as np
class Particle:
def __init__(self, x, y, ang_v):
self.x = x
self.y = y
self.ang_v = ang_v
class ParticleSimualtor:
def __init__(self, particles: Sequence[Particle]):
self.particles = particles
def evolve_numpy(self, dt):
timestep = 0.00001
nsteps = int(dt / timestep)
r_i = np.array([[p.x, p.y] for p in self.particles])
ang_vel_i = np.array([p.ang_v for p in self.particles])
for i in range(nsteps):
norm_i = np.sqrt((r_i ** 2).sum(axis=1))
v_i = r_i[:, [1, 0]]
v_i[:, 0] *= -1
v_i /= norm_i[:, np.newaxis]
d_i = timestep * ang_vel_i[:, np.newaxis] * v_i
r_i += d_i
for i, p in enumerate(self.particles):
p.x, p.y = r_i[i]
def evolve(self, dt):
timestep = .00001
nstep = int(dt / timestep)
for i in range(nstep):
for p in self.particles:
t_x_ang = timestep * p.ang_v
norm = (p.x ** 2 + p.y ** 2) ** .5
p.x, p.y = (p.x - t_x_ang * p.y / norm, p.y + t_x_ang * p.x / norm)
# repeat for all particles
# repeat for all time-steps
def visualize(simulator: ParticleSimualtor):
X = [p.x for p in simulator.particles]
Y = [p.y for p in simulator.particles]
fig = plt.figure()
ax = plt.subplot(111, aspect='equal')
line, = ax.plot(X, Y, 'ro')
plt.xlim(-1, 1)
plt.ylim(-1, 1)
plt.grid(True)
def init():
line.set_data([], [])
return line,
def animate(i):
simulator.evolve(.01)
X = [p.x for p in simulator.particles]
Y = [p.y for p in simulator.particles]
line.set_data(X, Y)
return line,
anim = animation.FuncAnimation(fig, animate, init_func=init, blit=True, interval=10)
plt.show()
def test_visualize():
particles = [Particle(.3, .5, 1), Particle(0, -.5, -1),
Particle(-.1, -.4, 3)]
sim = ParticleSimualtor(particles)
visualize(sim)
def benchmark2(nparts, m):
particles = [Particle(uniform(-1, 1), uniform(-1, 1), uniform(-1, 1))
for i in range(nparts)]
sim = ParticleSimualtor(particles)
if m == 1:
sim.evolve(.1)
elif m == 2:
sim.evolve_numpy(.1)
else:
raise Exception('not a valid method')
if __name__ == '__main__':
# test_visualize()
nloop = 10 # number of executions
result = timeit.timeit('benchmark2(100, 2)',
setup='from __main__ import benchmark2',
number=nloop)
print('{:.2f} seconds to run {} loops'.format(result, nloop))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.