Speed up distance calculations by ~10-20x - #3767

Would changing all isiterable() tests to a try-except structure also work? If so, that would be a cleaner solution.

Short answer: No.

That's what isiterable is -- a try-except. I guess those are expensive... so it isn't asarray, that's not being called ever in the quad loop.

Note that getting rid of the isiterables is only about a factor of 2-3x speedup. Having individual functions that avoid calculating things they don't need to is another 2-3x speedup.

To provide a slightly more nuanced option: a big chunk of the expense is that np.sqrt and np.exp are much more expensive than the math versions on scalars. Maybe they are doing a wrap, unwrap. Also, the massive neutrino calculation is very expensive, as is testing conditions inside the inner function. So the important things are:

• Have a scalar only version of the integrand that doesn't need to test for being a scalar and can call the math versions of functions.
• Have a branch outside the integrand to test if we need to compute neutrino stuff.

So we could only have two "special" versions of the scalar integrand: with massive neutrinos, and without, and could eliminate the no-radiation version. That's only like 5-10% slower than the fast version. My thinking was once I'm committed to having two ugly internal versions of inv_efunc, might as well squeeze that extra bit and have three.

About isiterable: there's no overhead for code executed in the try part of a try-except statement, but a large overhead associated with catching an exception. isiterable throws an exception every time it runs on a scalar, which is slow - on my machine it's ~7 times slower for a scalar compared to an array. So if we use a try-except that tries to do an integral assuming the redshift is a scalar, then that will avoid the expensive exception in isiterable.

I see your point about the individual functions to avoid unnecessary calculations, though. I guess we can't avoid making special versions of integrands for that.

Wouldn't the computation that fails on an array just be the "special" functions that use math.sqrt? That wouldn't really avoid any of the ugliness... and avoiding doing if/then tests inside those functions makes a pretty big difference too. So it seems kinda like the same thing in a different order, unless you have something cleverer in mind.

@aconley - Speed-ups for common cases are a good idea, and I think that by factoring out some of the boilerplate code, this could be made to look a lot less ugly (see inline comments).

In case you want to profile the code, you can see if https://jiffyclub.github.io/snakeviz/ works for you.

@pllim - I was just using cProfile, which I think told me everything I needed to know.

@mhvk -- let's not kid ourselves. This is ugly as hell either way.

The idea here is not to yet put this in it's final form, just to see if people are perhaps willing to pay this price.

If it is 10x improvement, I think speed outweighs ugliness. Pretty codes that no one uses are useless.

I should point out: the Planck cosmologies, which I imagine will become the default fairly soon, are still very slow. That's because they include massive neutrinos, and that turns out to be quite expensive. So the next effort -- even if this one is rejected -- should be to try to figure out how to make that faster.

I completely agree a 10x speedup is worth any ugliness.

Note the current version breaks backwards compatibility for _tfunc and _xfunc, since they no longer accept array input. In theory this is ok because the leading underscore is supposed to tell people "don't use this, it's an internal method which might change in the future". But in practice people are probably using them anyway, and will get annoyed when it's changed. This could be avoided by making new functions (for example) _xfunc_scalar and then having _xfunc wrap it with the extra isiterable check (or the try/except example I gave in the in-line comment). I strongly prefer doing this to avoid breaking backwards compatibility.

Hmm... I don't like people using _tfunc, _xfunc, and I don't think people should have much of a right to be annoyed at us changing them.

Would it be acceptable to create promoted non-scalar versions xfunc and tfunc, where that creates an implicit promise not to mess with them except at major version number boundaries?

@nhmc -- the recent update returns _tfunc and _xfunc to their previous behavior, adds new public versions (tfunc and xfunc), and throws deprecation warnings on the underscore versions. After all, we do provide public access to the distance integrand through efunc and inv_efunc.

You don't want to use np.vectorize (or, rather, the new wrapper) to go between them -- and between the scalar and ndarray versions in general -- because it can be much slower than just computing on the array. vectorize is just a for loop, and therefore doesn't take full advantage of numpy's capabilities. We use it in the integrals because quad already kills the ndarray advantage, and because the individual integrations are pretty expensive already.

They definitely have no right to be annoyed, but that doesn't stop them from getting annoyed and tweeting 'astropy broke my code!'

I left some inline comments, but everything else looks fine to me.

Hmm... something went screwy with Travis. It doesn't seem to have scipy.integrate installed.

@aconley - remember that only some of the travis builds have optional dependencies (including scipy) installed - you need to make sure the tests are skipped if scipy is not installed.

@astrofrog -- yeah, I knew that. I just didn't think I had added any tests that required scipy, and it turns out I had in a slightly subtle fashion. Anyways, finally should be ready to go.

I also agree this is enough of a speedup to justify the ugliness, especially given the importance of this to plenty of modern astro.

I'm happy with most of it, but two relatively straightforward items:

• I'm fine with the publicizing of xfunc and tfunc, but I think they should get more descriptive names. E.g., xfunc -> abs_distance_integrand tfunc -> lookback_time_integrand (or lb_time_integ if that seems to awkward, although this is still mostly internal so I tend towards explicit names if it's not going to be typed alot). To complement this, it might make sense to replace efunc with hubble_scaling (although it should probably call efunc to maintain backwards compatibility for the foreseeable future).
• It might make sense to link to this issue or even write a brief note somewhere in the class docstrings about the key speedups made here. It might help possible future contributors (or future cosmologies...) have an easier time of it if they want to take similar roads.

And just to make sure we're all on the same page, you'd rather address any possible optimizations for the massive neutrinos in a later separate PR, right @aconley?

@eteq -- definitely. I looked at the neutrinos briefly, and it seems like it will be harder, so it makes sense to do this in stages.

I don't think we should rename efunc -- people may actually use that. And the kind of people looking for it will not be surprised by the name. The E(z) notation is incredibly common.

Awesome. This looks great and everything is passing, so I'll go ahead and merge. Thanks @aconley !

@eteq -- There is now a first attempt at the neutrinos are slow problem: #3814