Repository navigation
slow comoving_distance in astropy.cosmology #3764
Description
My colleague pointed out that the comoving_distance function in astropy.cosmology is a lot slower than equivalent function comdist2 in idlutils by SDSS (https://www.sdss3.org/dr8/software/idlutils.php). For example, the following IDL code reports 0.04 second for 10^5 elements (and 0.35 sec in the case for a 10^7 elements).
z = randomu(1, 100000L)*6
t_begin = systime(/second)
d_l = comdis2(z, 0.3, 0.7, cdtable=cdtable)
t_end = systime(/second)
print, "Time elapsed:", t_end-t_begin
endOn the other hand, an equivalent python code with astropy.cosmology like this takes 23.3 seconds.
from astropy.cosmology import FlatLambdaCDM
import numpy as np
import time
z = np.random.rand(100000)*6.
t_begin = time.time()
cosmo_astropy = FlatLambdaCDM(H0=70., Om0=0.3)
d_astropy = cosmo_astropy.comoving_distance(z)
t_end = time.time()
print("Time elapsed with astropy: %f seconds" % (t_end-t_begin))Instead of astropy, I tried cosmolopy (http://roban.github.io/CosmoloPy/) as follows and it takes 5.1 seconds.
import numpy as np
import time
import cosmolopy.distance as cd
cosmo = {'omega_M_0' : 0.3, 'omega_lambda_0' : 0.7, 'h' : 0.72}
cosmo = cd.set_omega_k_0(cosmo)
z = np.random.rand(100000)*6.
t_begin = time.time()
d_cosmolopy = cd.comoving_distance(z, **cosmo)
t_end = time.time()
print("Time elapsed with cosmolopy: %f seconds" % (t_end-t_begin))I guess it's not expected and there could be a room to be improved.
My computing environment is:
OS X 10.10.3
IDL 8.4
Python 2.7.9 (installed via homebrew)
Astropy 1.0.2
Numpy 1.9.2
Cosmolopy 0.1.104 (laterst version installed by pip)
Activity
astrofrog commentedon May 8, 2015
@aconley @nhmc - could you take a look at this? A factor of >100 seems like a bug, right?
aconley commentedon May 8, 2015
@astrofrog -- not really.
cdist2 computes the distances on a grid and then interpolates. astropy is computing each distance independently. It's a design decision of a sort -- comdis2 is designed to be reused many times for a fixed cosmology, astropy isn't. So if you want to do something like that with astropy you should set up an interpolation object and call that.
However, there astropy isn't quite ideal either because it would be computing the distance to each grid point independently, while comdis2 steps along the grid (actually, in a pretty stupid linear, trap rule way, but still the general idea is what you want to do). We've talked at various points about how one could add something similar to astropy -- give me distances on some grid more efficiently -- but have never done anything about it.
astrofrog commentedon May 8, 2015
At the very least, I wonder if it would be worth showing an example in the docs of how to set up such an interpolation externally, if we don't make that possible to do internally?
eteq commentedon May 8, 2015
@aconley - what about cosmolopy? A quick glance at their driver functions suggest to me it's very similar to astropy, yet 4x slower? (I could be missing some key difference though...)
aconley commentedon May 9, 2015
@astrofrog -- There is actually something in the docs, although not exactly in the most obvious place (see the docs for
z_at_value).@eteq -- That I don't understand. Their integrand is a bit simpler (no photons, etc.), but that shouldn't have that large of an effect.
aconley commentedon May 9, 2015
Ok -- here's one thing worth trying: if you look at
FLRW.comoving_distancethe call toinv_efuncis inside the loop, so Python is looking it up on each call. That isn't free, so we should assign it to a local variable that we pass into quad. Again, it's hard to see that giving a factor of 4...Some of the expense is also coming from wrapping it in a quantity.
nhmc commentedon May 9, 2015
We should really just write a fast version of
comoving_distancealong the lines of the idl code and use that when the length of the input redshift array is large. This will make the code less neat, butastropybeing 2 orders of magnitude slower for a not-uncommon use case is just not acceptable.I will try to do this tomorrow night, but can't make any promises...
aconley commentedon May 9, 2015
So, I've done this sort of thing several times in other pieces of code. Some free advice :)
have to decide how to handle z = 0.
step trap rule integrator. It's pretty trivial to do better in your stepper choice.
I would recommend the
scipy.integrate.odeinterface for stepping along.scipy.interpol.interp1dobject.Be careful with the cubic version of that. It seems prone to a lot of wild flopping
around, more so than any other cubic spline library I've tried.
around an interpol object that knows how to multiply in the extra 1+zs.
aconley commentedon May 9, 2015
@nhmc -- reading your comment more carefully, I think you are actually suggesting modifying
comoving_distanceto cut to a different branch if the number of points is large.I don't think you should do that. It could have some odd consequences, like accuracy loss.
IMHO, I think the thing to do is add a new method to
FLRWthat returns some sort of distance interpolator object with methods likecomoving_distance,luminosity_distance, etc., and tell the user to use that if they want to compute a lot of distances. That way the speed/accuracy tradeoff is much more explicit, and the user can tailor their approach for their problem.Example: say I wanted a bunch of distances near z = 0.4 and a bunch near z = 4.5 (say, comparing two populations). It would be smarter to do this with two interpolators than with one, because the first step out to 4.5 in the second case could be done more intelligently taking advantage of the adaptive integration than if we forced ourselves to plod along tabulating all those uninteresting redshifts between 0.5 and 4 at fixed intervals.
aconley commentedon May 11, 2015
I did some timing tests. Loading the function into a local variable doesn't do much.
However, it does turn out that half of the execution time is spent doing:
We can solve that by breaking out an "inner function" that makes stronger assumptions about the inputs that the integrators can call. I will make a pull for that.
aconley commentedon May 11, 2015
Okay, I made a pull request to show what's involved in breaking the functions out. It's ugly. But it is much faster -- like 10x faster. With that version, astropy is 5x faster than cosmolopy for the same calculation.
The idea is -- if it seems worth it -- to then expand that to the other cosmologies.
Speed up distances with massive neutrinos by 5-7x for 1, 2, 3 flavors
eteq commentedon Sep 24, 2015
@aconley - do you think this should be closed now that #4127 is merged, or is there still more you were planning to do?
aconley commentedon Sep 24, 2015
Yeah, I think it can be closed. Thanks for also closing 3814.
embray commentedon Sep 28, 2015
In wonder how this compares to
comdist2now. But maybe it's not a fair comparison in the first place.aconley commentedon Sep 29, 2015
There's still room for somebody to write a 'grid' convenience class for doing something like what `comdist2' does (i.e.., make a grid and interpolate) . It would probably be useful.
rainwoodman commentedon Mar 23, 2016
For example, adding a 'freeze()' function to cosmology objects, that returns a delegate where all calculations are based on quick, interpolations?
Has there been any updates about the speed issue in general? I was calling WMAP9.age on 100000 points and it took me ~ 10 seconds. This quickly becomes unfeasible approaching billions of points.
pllim commentedon Mar 23, 2016
@rainwoodman , this issue is closed. Would be better to open a new issue for your problem and cross-reference to this one.