2008/7/10 Dan Lussier <dtlussier@gmail.com>:

> I am relatively new to numpy and am having trouble with the speed of
> a specific array based calculation that I'm trying to do.
>
> What I'm trying to do is to calculate the total total potential
> energy and coordination number of each atom within a relatively large
> simulation.  Each atom is at a position (x,y,z) given by a row in a
> large array (approximately 1e6 by 3) and presently I have no
> information about its nearest neighbours so each its position must be
> checked against all others before cutting the list down prior to
> calculating the energy.

The way you've implemented this goes as the square of the number of
atoms. This is going to absolutely kill performance, and you can spend
weeks trying to optimize the code for a factor of two speedup. I would
look hard for algorithms that do this in less than O(n^2).

This problem of finding the neighbours of an object has seen
substantial research, and there are a variety of well-established
techniques covering many possible situations. My knowledge of it is
far from up-to-date, but since you have only a three-dimensional
problem, you could try a three-dimensional grid (if your atoms aren't
too clumpy) or octrees (if they are somewhat clumpy); kd-trees are
probably overkill (since they're designed for higher-dimensional
problems).

> My current numpy code is below and in it I have tried to push as much
> of the work for this computation into compiled C (ideally) of numpy.
> However it is still taking a very long time at approx 1 second per
> row.  At this speed even some simple speed ups like weave  won't
> really help.
>
> Are there any other numpy ways that it would be possible to calculate
> this, or should I start looking at going to C/C++?
>
> I am also left wondering if there is something wrong with my
> installation of numpy (i.e. not making use of lapack/ATLAS).  Other
> than that if it is relevant - I am running 32 bit x86 Centos 5.1
> linux on a dual Xeon 3.0 GHz Dell tower with 4 GB of memory.

Unfortunately, implementing most of the algorithms in the literature
within numpy is going to be fairly cumbersome. But I think there are
better algorithms you could try:

* Put all your atoms in a grid. Ideally the cells would be about the
size of your cutoff radius, so that for each atom you need to pull out
all atoms in at most eight cells for checking against the cutoff. I
think this can be done fairly efficiently in numpy.

* Sort the atoms along a coordinate and use searchsorted to pull out
those for which that coordinate is within the cutoff. This should get
you down to reasonably modest lists fairly quickly.

There is of course a tradeoff between preprocessing time and lookup time.

We seem to get quite a few posts from people wanting some kind of
spatial data structure (whether they know it or not). Would it make
sense to come up with some kind of compiled spatial data structure to
include in scipy?