Speed up distances with massive neutrinos by 6-7x - #3814

Again, I was surprised by just how expensive some numpy functions are that I thought were cheap (like sum), so the solution to making it faster (without Cython) seems to be specialized scalar versions.

The other way to go about this -- which maybe would be better -- would be Cython. But I decided this was simpler for a first attempt.

@aconley - the three nearly-the-same functions seem a bit excessive. Is there really much of a speed gain between the 1-neutrino version and one that simply does a python
for nu in nus: loop? (where nus can contain neutrino species up to three.) Also, I think the python sum function is very fast.

@mhvk -- a straight python loop is about the same speed for 1 massive neutrino, but 10/20% slower for 2/3. Which maybe is close enough -- it's still a big speedup. I would probably prefer the single function approach, but the last pull taught me that other people seem to have a higher tolerance for code ugliness than I do, so I figured I'd start at full ugly for speed. (PS: nus can contain more than 3 neutrinos -- for the fans of sterile neutrinos out there).

As for python sum -- yeah, I thought so too. But that's why we measure! And it turns out we were both just wrong.

To be precise, for a 3 element sum, python sum is 3x slower than just doing arr[0] + arr[1] + arr[2], and numpy sum is 2x slower still (at least on my machine). And for a one element array (so, the Planck cosmologies) the unwrapped loop is 11x faster than python sum, and 30x faster than numpy sum!

@aconley - I think at the few tens of % level, it is not worth making the code way more ugly...

As for sum: the problem is partially that we have np.float rather than python float:

a = [1., 2., 3.]
%timeit a[0] + a[1] + a[2]
# 10000000 loops, best of 3: 132 ns per loop
%timeit sum(a)
# 1000000 loops, best of 3: 194 ns per loop
def loopsum(a):
    r = a[0]
    for e in a[1:]:
        r += e
    return e

%timeit loopsum(a)
# 1000000 loops, best of 3: 347 ns per loop

# Now using an `ndarray`.
a = np.array([1., 2., 3.])
%timeit a[0] + a[1] + a[2]
# 1000000 loops, best of 3: 745 ns per loop
%timeit sum(a) 
# 100000 loops, best of 3: 1.97 µs per loop
%timeit loopsum(a)
# 100000 loops, best of 3: 2.03 µs per loop

And np.sum is much worse, as you say.

Interesting! I'm starting to understand there's a lot more overhead in numpy than I had imagined.

But note that unrolling still seems to win by ~ 50%. And np.floats are what we have.

PS: you shouldn't write loopsum that way. The a[1:] is making a copy of the array, unlike numpy, which uses views.

def loopsum(a):
      r = 0.0
      for e in a:
         r += e
      return r

is about 30% faster.

@mhvk - alas, it appears I messed up while testing the pure python loop, and left part of the calculation out. The 3-headed version is more than 10-20% faster, it's more like 40-70%.
Lesson: make sure your speed hack test still passes all the unit tests!

astropy HEAD: 520ms
astropy HEAD with simple scalar z version: 490ms

3-headed version:
1 massive nu: 85.6ms (note: Planck case, which is likely to become the default cosmology)
2 massive nu: 108ms
3 massive nu: 128ms

python loop:
1 massive nu: 144ms
2 massive nu: 160ms
3 massive nu: 177ms

So, that's a 70% hit in the most common case, which isn't negligible.

Now, the relatively clean loop version is still a good 3.5x faster than HEAD -- but the ugly-ass version is 6x faster.

@mhvk -- ok, you'll like this. If I store a copy of the variable we need as a python list instead of a numpy array, then the pure python loop version is a) faster than the 3-headed version b) only weakly dependent on the number of neutrinos.

So -- ditch numpy, get simpler (as in, only one version) and faster code. Hrm.

@aconley - wow! Are you doing the equivalent of the following?

a = np.array([1., 2., 3.])
%timeit sum(a)
# 100000 loops, best of 3: 1.96 µs per loop
%timeit sum(a.tolist())
# 1000000 loops, best of 3: 354 ns per loop

Not quite... we do have to do some manipulation on the variable (_nu_y; it gets divided by 1+z and raised to various funny powers). So the idea is to also store an instance variable with _nu_y as a list (leaving a ndarray version for when people are actually calling on ndarray zs), and then use a python loop. There's still some ugliness -- you want a special case that is not in the thing being integrated for there being no massive nus -- but it's down to 2 additional private functions instead of 5. And the private functions are easier to understand.

This is making me question everything I thought I knew about numpy...

Btw one should be careful with caching - for the examples @mhvk showed, I get warnings like The slowest run took 17.64 times longer than the fastest. This could mean that an intermediate result is being cached

@astrofrog -- Yeah, I know. That can also reflect something being compiled to bytecode. So you do it multiple times, and the warning goes away.

The timing tests I'm talking about are based on modifying core.py and reinstalling, then checking the distances are still the same, then doing %timeit. If I do the middle step, I no longer get that warning during the timing test.

Apologies for letting this languish for so long; new job, blah, blah. Among other things I was trying to make sure I could legally contribute to astropy...

Anyways, I have now started on a cython version, which will probably make this pull obsolete. It's... a lot faster, although I haven't tried the neutrinos yet.

Replacement cython version is #4127

No need for apologies--it's cool you gave the Cython version a try.

#4127 replaces this, so I'm closing this one.