# Sorting overlapping shells in dynamical simulation - for loops vs array methods

I didn't know whether it was best to start with the general task I'm trying to solve, or the meat of the programming problem. They're both below but the former is rather long and solving the latter may not require it, so feel free to skip at will.

The task to be solved

I'm running a relatively simple dynamical simulation, following the gravitational evolution of a series of spherical shells, between each pair of which there is a fixed mass.

To explain with an example, between shell 3 and 4 there's a mass of 2, when the shells are far apart the density is low but this mass is always contained between them. Other shells may overlap, crossing in and out as the simulation continues, changing the density in at least some of the region between shells 3 and 4, but no mass is transferred.

I am trying to write an efficient function which, for each shell, can find the mass internal (at smaller radii).

To find the contribution to mass internal of shell n (let's call it M[n]), by the mass between shells j and j+1 (which we'll call m[j]), we need to know if:

1. Shells j and j+1 are both contained within shell n, thus M[n]+=m[j]
2. Shells j and j+1 are both outside of shell n, thus M[n]+=0
3. One of shell j or j+1 are outside of n, and the other within, thus M[n]+= some fraction of m[j]

(Let's assume mass is distributed uniformly between shells, meaning the fraction in case c) is just the ratio of volume between the inner shell and shell n, compared to between shells j and j+1.)

The difficulty in solving it

I've written two Python (+ NumPy) solutions to the problem, the first a quick and ugly long for loop, designed just to build the rest of the simulation and to be improved on later. It's neither quick nor beautiful, but it does the job.

I then returned to the problem with the intention of streamlining and speeding it up. The general philosophy I've been operating under is: for loops are slow, array operations are fast.

I carefully worked out how to do the whole thing with no for loops, and mostly operations on integer arrays. And yet to my dismay, it's slower, and seems to scale worse, than the block headed for loop I started with.

My question then is where my understanding of the efficiency of numpy has broken down. Where, in the second solution, are my uses of arrays inefficient, and when does the tradeoff of a simple for loops supersede a complex array treatment?

As you may guess, I'm not a hugely experienced python programmer, with more experience in c++ and IDL, so my mistakes may well be clear as day and easily avoidable.

The solutions

Both solutions are self contained functions that take 2 inputs, an array of radii at which each shell sits, and an array of the fixed masses between each pair of shells. I've been running this with 500 shells (and hence 499 masses), I'd like to increase this by an order of magnitude if possible, but that may be ambitious.

Solution 1: The for loop

For each shell identifies the masses internal (or overlapping) and sums the relevant masses:

def findMass(rads,shellM):
nR=rads.size
massInRad=5+np.zeros(nR)
for i in range(0,nR):
thisRad=rads[i]
innerMass=0
inside,=np.where((rads[:-1]<=thisRad) & (rads[1:]<=thisRad)) #both shell bounds within r
overLR,=np.where((rads[:-1]<thisRad) & (rads[1:]>thisRad)) #lower shell bound within r (returns index below r)
overRL,=np.where((rads[:-1]>thisRad) & (rads[1:]<thisRad)) #upper shell bound within r (returns index above r)

if (inside.size > 0):
innerMass+=np.sum(shellM[inside])
if (overLR.size > 0):
for index in overLR:
innerMass+=shellM[index]*(thisRad**3-rads[index]**3)/(rads[index+1]**3-rads[index]**3)
if (overRL.size > 0):
for index in overRL:
innerMass+=shellM[index]*(thisRad**3-rads[index+1]**3)/(rads[index]**3-rads[index+1]**3)
massInRad[i]+=innerMass
return massInRad


Solution 2: The array method

Using the order of shells identifies which case (internal, external or overlapping) is true between every combination of shells and masses:

def findMass(rads,mass):
def addCase(M,case,whichCase):
thisCase=np.where(case==whichCase)
ns=thisCase[0]
js=thisCase[1]
if (whichCase==3): # case c, mass j contained with shell n
mAdd=mass[js]
if (whichCase==1): # case b, mass j split over shell n, with shell j+1 at or further than shell n
mAdd=mass[js]*(rads[ns]**3-rads[js+1]**3)/(rads[js]**3-rads[js+1]**3)
if (whichCase==-1): # case d, mass j split over shell n, with shell j at or further than shell n
mAdd=mass[js]*(rads[ns]**3-rads[js]**3)/(rads[js+1]**3-rads[js]**3)
# case a, mass j not contained within shell n, requires no action
np.add.at(M,ns,mAdd)
return M

nRads=rads.size
M=5+np.zeros(nRads)
s=np.argsort(np.argsort(rads))

jDiff=np.tile(s,(nRads-1,1)).transpose() - np.tile(s[:-1],(nRads,1))
jSign=np.sign(jDiff+0.1).astype(int)
jPlusDiff=np.tile(s,(nRads-1,1)).transpose() - np.tile(s[1:],(nRads,1))
jPlusSign=np.sign(jPlusDiff+0.1).astype(int)
case=2*jPlusSign + jSign

M=addCase(M,case,3) #case C
M=addCase(M,case,1) #case B
M=addCase(M,case,-1) #case D
return M


Test case For a very simple simulation to run this on:

def runSim(nRads,nSteps):
rads=1+(np.arange(nRads)**2)/nRads
vel=np.zeros(nRads)
mass=1+np.arange(nRads)
lSq=findMass(rads,mass)*rads+1
for step in range(nSteps):
rads=rads+vel
massEnclosed=findMass(rads,mass)
acc=0.001*(lSq-massEnclosed*rads)/np.power(rads,3)
vel=vel+acc


Using

%timeit runSim(500,500)


the first method takes 3.9 seconds whilst the second takes takes 14.5 seconds

I was rather expecting the second to be an order of magnitude faster, not 4 times slower, but can't really see where I'm going wrong?

• Could you maybe post some example data, so we can play around with it? Just enough so the time difference between the two algorithms becomes visible. – Graipher Dec 7 '16 at 10:41