Vectorized implementation of reflex motion of the Earth correction

[DinoBektesevic]

This is a vectorized solution to the same problem. Here are some timings:

parallax_corrected = correct_parallax_geometrically_vectorized(tbl["ra"][:60], tbl["dec"][:60], tbl["mjd"][:60], earth_location, 40)
CPU times: user 18.9 ms, sys: 0 ns, total: 18.9 ms
Wall time: 18.4 ms

import timeit
timeit.timeit("parallax_corrected = correct_parallax_geometrically_vectorized(tbl['ra'][:60], tbl['dec'][:60], tbl['mjd'][:60], earth_location, 40)", globals=globals(), number=100)
1.6601364687085152

The real gains can be gotten the more SkyCoords we have. For all the pointings in DEEP

parallax_corrected = correct_parallax_geometrically_vectorized(tbl["ra"], tbl["dec"], tbl["mjd"], earth_location, 40)
CPU times: user 1min 41s, sys: 264 ms, total: 1min 42s
Wall time: 1min 41s

len(tbl["ra"])
472047 # !!!

This example takes in ra, dec, obstime and obsgeoloc as arguments because I'm testing from joined DEEP joined collection, so this includes the cost of creating the initial coords object (the current implementation expect a coord object).

The test of the current correct_parallax_geometrically function can't take multiple pointings and times, so they need to be iterated over:

def test():
    res = []
    for ra, dec, obstime in zip(tbl["ra"][:60], tbl["dec"][:60], tbl["mjd"][:60]):
        tt = Time(obstime, format="mjd")
        coord = SkyCoord(ra=ra*u.deg, dec=dec*u.deg)
        res.append(correct_parallax_geometrically(coord, tt, earth_location, 40))
    return res

CPU times: user 294 ms, sys: 14 ms, total: 308 ms
Wall time: 300 ms

which doesn't scale well to many coordinates. Hopefully I'm just not seeing somethign obvious and missusing it. I think it's possible to adapt the function rather easily to take a list of SkyCoords, but there are some corner-cases - it wont' work with different timestamps.

There are no tests etc. in this PR, because I am mostly looking for some guidance on what the interface logic:

• coordinates are given as SkyCoord objects for the current functions
• but timestamps are not given to the SkyCoord object, they are provided separately as Time objects
• and same for earth location

Generally coordinates admit the addition of a timestamp to them, as well as obsgeoloc, I could rework the current implementation in two ways:

• I expect obstime and obsgeoloc attached to every coordinate (most flexible)
• I expect all given coordinates to have the same timestamp and let the user iterate over batches of timestamps

In the first case the change is to go from function calls to coord, time, point_on_earth to coordand then throw an error if it does not have obstime and obsgeoloc attached. Sort out batches per timestamp, pre-fetch earth ephemeris, and return all corrected. This seems very practical because any kinds of coordinates can be given at any timestamp and obsgeoloc.

It's faster than per-coordinate iteration but it'll also be faster for very large numbers of coordinates than the second potential solution.

Second approach vectorizes the current implementation, but just assumes they are all taken at the same time. Then it's on the user to figure out and do the iterations from solution A, but for not-so-very-large-number of coordinates this can potentially be faster since there are less obsgeolocs and obstimes being copied during initial transformation to gcrs. Speedups here are hard to measure without making sure apples-to-apples are being compared.

For large number of coordinates the benefit of first approach is that there's less memory allocations alltogether (one big los_earth_obj is created only and edited in-place. At the low end the execution time is measured in ms, so unless they're run in many loops I don't think it'll ultimately matter whichever one we pick.

[DinoBektesevic]

I've added the invert and fixed the docstring. I added some basic tests too. @ColinOrionChandler if this is good for you then I'll merge.

I did change the order of parameters - this might break some of your scripts, so that I can make the default assumption be geocenter instead of some observatory I think you mentioned that this is useful to you, but it required flipping obsgeoloc and timestamp I think.

[wilsonbb]

This look good to me, but might want to check with someone like Drew or Max that has more familiarity with the original reflex correction approach

[wilsonbb]

Nice idea for testing this

[DinoBektesevic]

I'm merging this, since it's not the default behaviour anywhere so you and Colin can use it from KBMOD instead of copying it elsewhere. We can hash out the details of the interface some other PR. Thanks for all the reviews today!