Improve speed of coordinate transformations for arrays of obstime - #10647

After fixing some minor issues with dtypes and quantities, the tests should be passing now.

Starting to do some benchmarks.

@bwinkel do you still have the scripts you used to create the plots from #6068?

Tests for the actual interpolation are missing and it does currently not work due to a shape issue, looking into it.

from astropy.coordinates.astrom_manager import astrom_interpolation_resolution
from astropy.coordinates import SkyCoord, EarthLocation, AltAz
from astropy.time import Time
import astropy.units as u
import numpy as np

import matplotlib.pyplot as plt

from time import perf_counter


t = Time('2020-01-01T20:00:00') + np.random.uniform(0, 24, 100000) * u.hour
location = EarthLocation.of_site('Roque de los Muchachos')
crab = SkyCoord.from_name("Crab Nebula")

altaz = AltAz(obstime=t, location=location)


t0 = perf_counter()
no_interp = crab.transform_to(altaz)
print(f'No Interpolation took {perf_counter() - t0:.4f} s')


seps = {}

for resolution in [10 * u.s, 5 * u.min, 30 * u.min]:
    t0 = perf_counter()
    astrom_interpolation_resolution.set(resolution)
    interp = crab.transform_to(altaz)
    print(f'Interpolation with {resolution.value: 5} {str(resolution.unit):3s} resolution took {perf_counter() - t0:.4f} s')

    seps[resolution.to_value(u.s)] = no_interp.separation(interp)


fig, axs = plt.subplots(len(seps), 1, constrained_layout=True)
for ax, (key, sep) in zip(axs, seps.items()):
    ax.hist(sep.to_value(u.microarcsecond), bins=100)
    ax.set_title(f'{key} s resolution')

    ax.text(
        0.5, 0.95,
        (
            f'mean: {sep.to_value(u.microarcsecond).mean():.4f} µas'
            f'\np95:  {np.percentile(sep.to_value(u.microarcsecond), 95):.4f} µas'
            f'\nmax:  {sep.to_value(u.microarcsecond).max():.4f} µas'
        ),
        va='top',
        ha='center',
        transform=ax.transAxes
    )
    ax.set_xlabel('Angular Separation / µas')

fig.savefig('interpolation_error.png', dpi=300)

Timings:

No Interpolation took 17.1563 s
Interpolation with  10.0 s   resolution took 1.5896 s
Interpolation with   5.0 min resolution took 0.2312 s
Interpolation with  30.0 min resolution took 0.1734 s

@maxnoe - as you will probably have seen from the original issue, I was quite wary of these speed-related work-arounds. That said, clearly you and others need it and looking at it again I think the approach taken, with a generic get_astrom function, is something that can be made to address my concerns, if we change how the ScienceState operates. In particular, I feel quite strongly that a goal in any solution should be to separate the regular and interpolating settings -- basically, I would like the interpolation case to be one that I know will never affect regular use of astropy, and separate enough that it can be maintained separately.

Before going there, it may be good to just make sure that what we implement is the right level of generality, so that users perhaps can change also in other ways how the transformations operate. I have not read back the whole original thread, but is it obviously why one would not override parts of the transformation graph with interpolating routines (and have a ScienceState which essentially does that)? Or is that too high level? Somewhat similarly, I can see more general use for being able to override erfa; would that be sufficient? (I.e., the interpolating case would generally just call the standard routines, but provide overrides for the relevant slow routines to check whether times should be interpolated).

I think it may well be that the get_astrom level is in fact the right one! Right now, though, the actual ScienceState is very limited, in that it just gives an interpolating timescale, one that also very much interleaves the current code with interpolation stuff, thus making the regular case more difficult to follow. Could we instead make it a higher-level ScienceState, where the state is something like a module or class which provides all the astrom functions as attributes or methods (I think that would be better than a long if-else-else sequence, but that is somewhat of an implementation detail); in a way, it would then becomes more similar to the ScienceState used for IERS data, where the IERS tables have different constraints for the precision of the data provided). In this scheme, by default, what is returned just contains the regular code, but it can be replaced by interpolating code (which of course can use the regular module when needed; a class-based scheme would seem easiest for this purpose).

The implementation might look something like

class GetAstrom:
    def apci(self, frame):
        ...

    def apcs(self, frame):
        ...

    def __call__(self, frame, tcode):
        return getattr(self, tcode)(frame)


def GetAstromInterpolate(GetAstrom):
    def __init__(self, interpolation_resolution):
        self.interpolation_resolution = interpolation_resolution

    def apci(self, frame):
        <get interpolated frame>
        result = super().apci(interpolated_frame)
        <interpolate back>
        return final_result

So, internally, whenever needed, code would do something like

coords.get_astrom_methods.set(coords.GetAstromInterpolate(10*u.s)):

Or perhaps with string equivalents, like for the ephemerides,

coords.get_astrom_methods.set('interpolate', 10*u.s):

@mhvk I think your proposal makes a lot of sense! It seems a bit more work than the current solution, but it would indeed be a much more generic solution and leave the default behavior intact. It's been a while, since I looked into this part of the Astropy code, but I'll try to find some time next week to work on this. Perhaps you or @maxnoe would be available for some pair programming session?

After fixing some minor issues with dtypes and quantities, the tests should be passing now.

Starting to do some benchmarks.

@bwinkel do you still have the scripts you used to create the plots from #6068?

One of the plots was posted in the original PR, with the spatial dependency. The other plots are inferior to your histograms, I think, but I could dig that out if necessary.

In particular, I feel quite strongly that a goal in any solution should be to separate the regular and interpolating settings -- basically, I would like the interpolation case to be one that I know will never affect regular use of astropy, and separate enough that it can be maintained separately.

I agree, right now the code is very much entangled and this should be cleanly separated.

would not override parts of the transformation graph with interpolating routines

I started crude approximate transformations (implemented after formulas in the "Explanatory Supplement to the Astronomical Almanac") with a context manager to switch out transformations completely here:

https://github.com/MaxNoe/approximate_coords

But I think the interpolation approach of the astrom is a nice way of not having to replace a full transformation.
Since it is something all transformations, irregardless of how they work, might benefit from.
So having this decoupled from the actual graph is a nice property I think.

The implementation might look something like

I like this proposal a lot better then the imperative structure of the current function, I'll give it a shot.

Perhaps you or @maxnoe would be available for some pair programming session?

Sure!

First I am going to write some basic unit tests and then refactor into the class based approach.

Is codecov failing because the @pytest.mark.remote_data tests are not run? The test do cover the lines codecov complains about and locally report coverage of 100% for erfa_astrom.py

@mhvk Do you have an idea why the coverage is failing?

@maxnoe - I think the coverage is done on a run without remote_data (since with remote_data we just get time-out failures too often). But I think you can have your tests be without that (@pllim removed most of the need for it!), as long as the times used are covered by our base IERS B table. I had only a quick look, but 2020 January should be included. Do the tests fail if you remove remote_data?

@mhvk Ok thanks, seems indeed it is unnecessary. I switched to hard coded coordinates in between without relying on from_name or of_site, so I think for those it was needed before.

@StuartLittlefair Could you re-review? Are all your comments addressed?

@maxnoe sorry for the holiday induced delay! Looks good to me...

@mhvk I adressed most of your comments by simply taking you suggestions but left a few out for now. I have no real problems with following your suggestions everywhere but nitpicked back a bit 🤣

@mhvk Ok, I think that's all comments adressed.

@maxnoe thank you so much for making this happen. Made my day!