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+1 (which we'll call
m[j]), we need to know if:
j+1are both contained within shell
+1are both outside of shell
- One of shell
j+1are outside of
n, and the other within, thus
M[n]+=some fraction of
(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
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.
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 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 js=thisCase 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
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?