[Numpy-discussion] numpy/numexpr performance (particle simulation)

Kurt Smith kwmsmith at gmail.com
Mon Jun 29 11:13:09 EDT 2009

On Mon, Jun 29, 2009 at 8:20 AM, Prashant Saxena<animator333 at yahoo.com> wrote:
> Hi,
>
> I am doing a little test using numpy and numexpr to do a particle
> simulation. I never used either of them much and this is the first time I
> have to go deeper. Here is the code:
>
> import numpy as np
> import numexpr as nexpr
>
> class Particle( object ):
>
>     def __init__( self, id ):
>         self.position = [0.0, 0.0, 0.0]
>         self.color = [0.5, 0.5, 0.5]
>         self.direction = [0.0, 0.0, 0.0]
>         self.id = id
>         self.lifeSpan = 0
>         self.type = 0
>
> class Emitter( object ):
>
>     def __init__( self ):
>         self.particles = np.empty([], dtype=Particle())
>         self.currentParticleId = 0
>         self.numParticles = 1000
>         self.emissionRate = 10
>         self.position = [0.0, 0.0, 0.0]
>         self.rotation = [0.0, 0.0, 0.0]
>
>     # Add a single particle
>         """
>         Add a single particle in the emitter.
>         """
>         if self.currentParticleId < self.numParticles:
>             self.particles = np.append( self.particles, Particle(
> self..currentParticleId ) )
>             self.currentParticleId+=1
>
>
> #######################################################
> Problem 1:
> self.particles = np.empty([], dtype=Particle())
> In "Emitter" class, how do I initialize a numpy array of "Particle" type?
>
> Problem 2:
> If problem 1 can be solved, is it possible to use numexpr to alter the
> position of each particle's position
> in emitter.particles array. For example, multiply x,y and z of each particle
> by using three random values. In other words
>
> self.particles.position*= random(1.0)
> self.particles.position*= random(2.0)
> self.particles.position*= random(3.0)
>
> If problem 1 cannot be solved then may be I could use something like this
> using simple list:
>
> a = [Particle(0), Particle(1)]
>
> How do I modify the position of each particle in array"a" using numexpr?
>
> I would be using 1-5 million particles where each particle's attribute
> postion, rotation and color will be modified using some complex
> equations. The above code is just the begining hence any tips for
> performance would be greatly apperciated.

1-5 million particles means that 1) the 'particle' object shouldn't be
a pure python object, with lists for the position field, etc.  You
won't get very good cache locality (although this might be hard to do
anyway, depending on the sort of operations you're doing), all the
ints, floats, etc are wrapped up inside python objects, making simple
arithmetic very expensive, and you'll want to start thinking about
memory usage and minimizing pointer lookups.  If you want to go the
pure numpy route (that is, not resorting to putting it into an
extension module via Pyrex/Cython, f2py, etc.), you could try
something like this:

import numpy as np

DTYPE = np.float32
position_dtype = np.dtype({'names':('x','y','z'), 'formats':(DTYPE,
DTYPE, DTYPE)})

# this is a nested dtype, for illustration.  It's a nicer data representation,
# but you might not want the nesting.
particle_dtype = np.dtype(
{
'names':('position', 'color', 'direction',
'id', 'lifeSpan', 'type'),
'formats':(position_dtype, position_dtype, position_dtype,
np.int32, np.int32, np.int32)
}
)

# This is the flat version of the above.  Should be obvious what's going on.
flat_particle_dtype = np.dtype(
{
'names':('posx', 'posy', 'posz', 'cx', 'cy', 'cz',
'dirx', 'diry', 'dirz', 'id', 'lifespan', 'type'),
'formats':[DTYPE]*9+[np.int32]*3
}
)

class Emitter(object):

def __init__(self, N):
self.numParticles = N
self.particles = np.zeros((self.numParticles,),
dtype=flat_particle_dtype)
# or you can use the nested version:
# self.particles = np.zeros((self.numParticles,), dtype=particle_dtype)
self.emissionRate = 10
# ...

if __name__ == '__main__':
N = 1000
emitter = Emitter(N=N)
posx = emitter.particles['posx']
posy = emitter.particles['posy']
posz = emitter.particles['posz']
print emitter.particles[:10]
posx[:] = np.ones((N,),dtype=DTYPE)
posy[:] = np.ones((N,),dtype=DTYPE)*2
posz[:] = np.ones((N,),dtype=DTYPE)*3

import numexpr
posx[:] = numexpr.evaluate('2.*posx + posz - 10.')

print emitter.particles[:10]

For particle simulations that has to be *really fast*, this isn't the
route I'd take for the low-level stuff, though.  All depends on how
fast you want to have something working vs. how fast you need it to
work ;-)

Kurt