Evaluating N orbits at N times in monte carlo simulations

[mahollands]

Hi,

I'm doing some Monte Carlo simulations for orbits within the MW. The time of flight is a free parameter in my model. So for each orbit there is a corresponding time value to evaluate at. However, using the Orbit.SkyCoord method evaluates all orbits at all times, and so returns a SkyCoord object with dimensions N^2. Of course, with few orbits/timestamps, I can just use SkyCoord.diagonal(), but for larger numbers of samples, the SkyCoord takes a long time to evaluate and uses a lot of memory. I have tried doing a list comp over the orbits and time array, which solves the O(N^2) memory issue, but this is clunky and also very slow.

Is there a better way to do this, or could one be added (e.g. have a broadcast_t argument in Orbit.SkyCoord)?

[jobovy]

I don't think there's a better way than a list comprehension to do this currently. I don't think it's that clunky and I'm surprised it's that slow. Is it slow to just list-comprehend o.R(time), that is, is the slow part the conversion to a SkyCoord or just calling the orbit? If the former, you could first get the Galactocentric coordinates at all times and do the SkyCoord conversion yourself on the arrays. Could you elaborate on what you mean you mean by adding a broadcast_t argument? I can't find anything relevant about this. Thanks!

[mahollands]

Thanks for the reply. I have checked and you are correct that it is the SkyCoord part that is slow. List comping o.R(time) takes 9 seconds for my ~16000 pairs of orbits and times, while the SkyCoord list comp takes 2.5 minutes. Of course, I can survive waiting a few minutes, but it is getting frustrating waiting on this part of the code each time I rerun, when the orbit integrations themselves only take a few seconds.

To elaborate on the clunkinesss, it would be okay if I could just do c = SkyCoord([o_i.SkyCoord(t_i) for o_i, t_i, in zip(o, t)]), but for some reason that drops the differential information. So after the first list comp I'm having to extract each component with another 6 list comps.

By adding a broadcast_t argument (or whatever name seems most appropriate), I meant as a suggestion, adding an argument to o.SkyCoord (and to o.R etc) that is by default False, but if True checks t has the same shape as o, and if so maps the elements of t to each orbit in o, returning a SkyCoord with the same shape as o. I.e. something that looks like normal numpy broadcasting rules.

In summary, my code currently looks like

coord = [o_i.SkyCoord(t_i).galactocentric for o_i, t_i in zip(o, t)]
coord = SkyCoord(
    x = Quantity([c.x for c in coord]),
    y = Quantity([c.y for c in coord]),
    z = Quantity([c.z for c in coord]),
    v_x = Quantity([c.v_x for c in coord]),
    v_y = Quantity([c.v_y for c in coord]),
    v_z = Quantity([c.v_z for c in coord]),
    frame = 'galactocentric',
)

and is slow, whereas my ideal would be to write something like
coord = o.SkyCoord(t, broadcast_t=True).galactocentric
to do the same thing, but without the overhead of instantiating SkyCoord thousands of times.

[jobovy]

Okay, got it! Yes, I agree that such a broadcast_t argument could be useful, also in other contexts, so perhaps you should open an issue about this and give a self-contained example along the lines of your current example that could then serve as a test of the implementation. I'm not sure whether I will be able to do this soon though.

In the meantime, I think you can avoid the thousands of SkyCoord instantiations as follows:

coord = SkyCoord(
    x = [-o_i.x(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    y = [o_i.y(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    z = [o_i.z(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_x = [-o_i.vx(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_y = [o_i.vy(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_z = [o_i.vz(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    frame = 'galactocentric',
)

I've made sure to convert to astropy's right-handed frame (galpy uses a left-handed frame). I tested this by just evaluating the orbit of the Sun at 100 times:

from galpy.potential import MWPotential2014
from galpy.orbit import Orbit
from astropy.coordinates import SkyCoord
from astropy.units import Quantity
import numpy as np
t = np.linspace(0.,10.,101)
o = Orbit([Orbit() for t_i in t])
o.integrate(t,MWPotential2014)

then

%%timeit
coord = SkyCoord(
    x = [-o_i.x(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    y = [o_i.y(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    z = [o_i.z(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_x = [-o_i.vx(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_y = [o_i.vy(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_z = [o_i.vz(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    frame = 'galactocentric',
)
# 41.6 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

while yours

%%timeit
coord = [o_i.SkyCoord(t_i).galactocentric for o_i, t_i in zip(o, t)]
coord = SkyCoord(
    x = Quantity([c.x for c in coord]),
    y = Quantity([c.y for c in coord]),
    z = Quantity([c.z for c in coord]),
    v_x = Quantity([c.v_x for c in coord]),
    v_y = Quantity([c.v_y for c in coord]),
    v_z = Quantity([c.v_z for c in coord]),
    frame = 'galactocentric',
)
# 501 ms ± 2.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

so my implementation is substantially faster and hopefully fast enough for you. I don't think that broadcast_t would provide much more of a speed-up, although likely a bit.

In testing this, I also noticed that your code is probably not working as you intend, because doing SkyCoord.galactocentric in astropy does not propagate the Galactocentric frame specification and the Galactocentric coordinates that you get are therefore not the correct ones (because galpy returns an ICRS SkyCoord. For example, compare

o[0].SkyCoord().galcen_v_sun

with

o[0].SkyCoord().galactocentric.galcen_v_sun

The latter is just the default from astropy, while the former is the correct one from the galpy orbit. My implementation is correct, but if you want to go to the ICRS frame, you'd have to adjust to

from astropy.coordinates import CartesianDifferential
from astropy import units
coord = SkyCoord(
    x = [-o_i.x(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    y = [o_i.y(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    z = [o_i.z(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_x = [-o_i.vx(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_y = [o_i.vy(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    v_z = [o_i.vz(t_i,quantity=True) for o_i, t_i in zip(o, t)],
    frame = 'galactocentric',
    galcen_distance = numpy.sqrt(o._ro**2.0 + o._zo**2.0) * units.kpc,
    z_sun = o._zo * units.kpc,
    galcen_v_sun = CartesianDifferential(np.array([-o._solarmotion[0],o._solarmotion[1] + o._vo,o._solarmotion[2]])* units.km / units.s)
)

To get the correct answer using your implementation, do

from astropy.coordinates import CartesianDifferential, Galactocentric
galcen_v_sun = CartesianDifferential(np.array([-o._solarmotion[0],o._solarmotion[1] + o._vo,o._solarmotion[2]])* units.km / units.s)
coord = [o_i.SkyCoord(t_i).transform_to(Galactocentric(galcen_distance=np.sqrt(o._ro**2.0 + o._zo**2.0) * units.kpc,z_sun=o._zo * units.kpc,galcen_v_sun=galcen_v_sun))for o_i, t_i in zip(o, t)]
coord = SkyCoord(
    x = Quantity([c.x for c in coord]),
    y = Quantity([c.y for c in coord]),
    z = Quantity([c.z for c in coord]),
    v_x = Quantity([c.v_x for c in coord]),
    v_y = Quantity([c.v_y for c in coord]),
    v_z = Quantity([c.v_z for c in coord]),
    frame = 'galactocentric',
)

[mahollands]

Thanks for taking the time to look into this, and your thorough reply! Your suggested alternative to generate a SkyCoord works much better. Personally, I've already found two specific scenarios where a broadcasting option would be useful: MC sampling the orbits of a single object where time of flight is a free parameter, and simulating an ensemble of objects that each have their own time of flight. I'll raise an issue as soon when I get a free moment myself.

Regarding correctness, I think I'm okay, since I'm initialising Orbit with a SkyCoord in the galactocentric frame. So the integration is implicitly using the astropy values for ro, vo etc. I've done a quick test and found that after the orbit integration, o[0].SkyCoord().galcen_v_sun vs o[0].SkyCoord().galactocentric.galcen_v_sun give identical values in my case. As a double check I have the performed the process in reverse, intialising with the intermediate SkyCoord, and arrived back at the original location with the same velocity components to very high precision. I recall running into the problem of mixing units from astropy/galpy when I first started out, but it looks like everything is internally consistent now.

[jobovy]

Okay, great! With the test you describe, it seems like things are working well. However, it is good to check carefully, because currently when you provide a SkyCoord, galpy will start from the ICRS representation and convert that internally to Galactocentric coordinates. It would probably be better to just grab the provided Galactocentric coordinates when given a SkyCoord that is in Galactocentric representation already...