Optimize Astrometry

[abhisrkckl]

I am specifically targeting the NGC6440E example dataset in this PR. The benchmarks I am using are (for NGC6440E):

1. Total chi2 computation time
2. Designmatrix computation time

I'll add more benchmarks in the future.

The improvements I have made are as follows.

1. Avoid Quantity.decompose as much as possible (see Quantity.decompose() is slower than Quantity.to() #1641)
2. ssb_to_psb_xyz_ICRS and ssb_to_psb_xyz_ECL functions in the astrometry classes are big performance sinks. I don't fully understand why this is so, only that this has something to do with astropy.coordinates. I have added overrides of ssb_to_psb_xyz_ICRS in AstrometryEquatorial and AstrometryEcliptic that don't use astropy.coordinates, and this seems to take care of this issue.
3. Avoided unnecessary computations in get_d_delay_quantities.
4. update option in Residuals.chi2() to re-compute cached quantities (helpful for profiling).

[codecov]

Codecov Report

Attention: 55 lines in your changes are missing coverage. Please review.

Comparison is base (813acfa) 68.55% compared to head (5e573fc) 68.58%.
Report is 78 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1643      +/-   ##
==========================================
+ Coverage   68.55%   68.58%   +0.03%     
==========================================
  Files         103      103              
  Lines       23516    23562      +46     
  Branches     4102     4113      +11     
==========================================
+ Hits        16122    16161      +39     
- Misses       6378     6380       +2     
- Partials     1016     1021       +5     

Files	Coverage Δ
src/pint/bayesian.py	92.40% <100.00%> (-4.61%)	⬇️
src/pint/models/__init__.py	100.00% <100.00%> (ø)
src/pint/models/wavex.py	87.42% <ø> (ø)
src/pint/models/dispersion_model.py	88.42% <0.00%> (ø)
src/pint/models/solar_system_shapiro.py	94.11% <87.50%> (-0.33%)	⬇️
src/pint/models/astrometry.py	93.33% <95.71%> (+0.60%)	⬆️
src/pint/residuals.py	74.12% <41.66%> (-0.70%)	⬇️
src/pint/models/dmwavex.py	88.55% <88.55%> (ø)
src/pint/utils.py	56.96% <34.28%> (ø)

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[abhisrkckl]

Current performance comparison before and after these changes (updated periodically):

(The benchmarks use example data for NGC6440E.)

Before:
Chisq computation: 14 ms
Design matrix computation: 25 ms

After:
Chisq computation: 7.8 ms
Design matrix computation: 5.3 ms

[dlakaplan]

Is your override function fully accurate? It's basically the linear approximation, but there are cases where higher-order terms are needed (like for very large proper motions and long time spans). Do you have any tests to check the accuracy?

[dlakaplan]

As a quick illustration, here is the linear pm approximation vs. what astropy does:

import numpy as np
from astropy.coordinates import SkyCoord
from astropy import units as u, constants as c
from astropy.time import Time
from matplotlib import pyplot as plt

# astropy example
cc = SkyCoord(
    l=10 * u.degree,
    b=45 * u.degree,
    distance=100 * u.pc,
    pm_l_cosb=34 * u.mas / u.yr,
    pm_b=-117 * u.mas / u.yr,
    frame="galactic",
    obstime=Time("1988-12-18 05:11:23.5"),
)

t = cc.obstime + np.linspace(0, 20) * u.yr
ccnew = cc.apply_space_motion(new_obstime=t)

# linear version
l0, b0 = cc.l, cc.b
pml, pmb = cc.pm_l_cosb, cc.pm_b
t0 = cc.obstime
dt = t - t0
l = l0 + pml * dt / np.cos(b0)
b = b0 + pmb * dt

plt.clf()
plt.plot((t - t0).jd, ccnew.l.value - l.value, label="astropy $l$ - linear")
plt.plot((t - t0).jd, ccnew.b.value - b.value, label="astropy $b$ - linear")
plt.xlabel("$t-t_0$ (d)")
plt.ylabel("Position difference (deg)")
plt.legend()

the differences are small, but get worse near the pole and scale with the proper motions (and are quadratic in time as expected).

[dlakaplan]

If you look at:
https://docs.astropy.org/en/stable/_modules/astropy/coordinates/sky_coordinate.html#SkyCoord.apply_space_motion
you can see how astropy does it. it might be possible to use the same method under the hood:
https://github.com/liberfa/erfa/blob/master/src/starpm.c
which is compiled C code so should be fast. but if you remove/change all of the astropy status checks (which are probably redundant)

[dlakaplan]

So I'm testing just the ERFA function alone, and it seems slower. so maybe it's not that simple.

It's also possible that the differences between the linear and correct versions depend on things like an accurate distance and RV, which in general we don't have. So maybe we don't worry.

[dlakaplan]

Another possibility is that if you only need the cartesian components you could do something like:

cccartesian = cc.cartesian
x0, y0, z0 = cccartesian.x, cccartesian.y, cccartesian.z
xdot0, ydot0, zdot0 = (
    cccartesian.differentials["s"].d_x.to(1 / u.yr),
    cccartesian.differentials["s"].d_y.to(1 / u.yr),
    cccartesian.differentials["s"].d_z.to(1 / u.yr),
)
xnew = x0 + xdot0 * (t - t0)
ynew = y0 + ydot0 * (t - t0)
znew = z0 + zdot0 * (t - t0)

[dlakaplan]

OK, it's not the ERFA function itself that is slow. It's wrapping it into another SkyCoord. So if you just need the bare RA,Dec and or x,y,z you can do:

def erfa_pm_to_cartesian(cc, t):
    icrsrep = cc.icrs.represent_as(
        astropy.coordinates.SphericalRepresentation,
        astropy.coordinates.SphericalDifferential,
    )
    icrsvel = icrsrep.differentials["s"]

    starpm = erfa.pmsafe(
        icrsrep.lon.radian,
        icrsrep.lat.radian,
        icrsvel.d_lon.to_value(u.radian / u.yr),
        icrsvel.d_lat.to_value(u.radian / u.yr),
        0.0,
        0.0,
        cc.obstime.jd1,
        cc.obstime.jd2,
        t.jd1,
        t.jd2,
    )
    ra, dec = u.Quantity(starpm[0], u.radian, copy=False), u.Quantity(
        starpm[1], u.radian, copy=False
    )
    x = np.cos(ra) * np.cos(dec)
    y = np.sin(ra) * np.cos(dec)
    z = np.sin(dec)
    return x, y, z

That is faster even than the linear version (maybe because of unit conversion?) and a lot faster than the astropy version. But I think it agrees exactly with the astropy version.

[dlakaplan]

Comparing:

def linear_pm_to_cartesian(cc, t):
    # linear version
    ra0, dec0 = cc.ra, cc.dec
    pmra, pmdec = cc.pm_ra_cosdec, cc.pm_dec
    t0 = cc.obstime
    dt = t - t0
    ra = ra0 + pmra * dt / np.cos(dec0)
    dec = dec0 + pmdec * dt
    x = np.cos(ra) * np.cos(dec)
    y = np.sin(ra) * np.cos(dec)
    z = np.sin(dec)
    return x, y, z


def erfa_pm_to_cartesian(cc, t):
    icrsrep = cc.icrs.represent_as(
        astropy.coordinates.SphericalRepresentation,
        astropy.coordinates.SphericalDifferential,
    )
    icrsvel = icrsrep.differentials["s"]

    starpm = erfa.pmsafe(
        icrsrep.lon.radian,
        icrsrep.lat.radian,
        icrsvel.d_lon.to_value(u.radian / u.yr),
        icrsvel.d_lat.to_value(u.radian / u.yr),
        0.0,
        0.0,
        cc.obstime.jd1,
        cc.obstime.jd2,
        t.jd1,
        t.jd2,
    )
    ra, dec = u.Quantity(starpm[0], u.radian, copy=False), u.Quantity(
        starpm[1], u.radian, copy=False
    )
    x = np.cos(ra) * np.cos(dec)
    y = np.sin(ra) * np.cos(dec)
    z = np.sin(dec)
    return x, y, z


def full_pm_to_cartesian(cc, t):
    ccnew = cc.apply_space_motion(new_obstime=t)
    x = np.cos(ccnew.ra) * np.cos(ccnew.dec)
    y = np.sin(ccnew.ra) * np.cos(ccnew.dec)
    z = np.sin(ccnew.dec)
    return x, y, z

I find:

Linear: 0.056377207999958046
ERFA: 0.021196374999817635
Full/Astropy: 0.28906691700012743

[abhisrkckl]

Closing this because this only implements the linear approximation to the proper motion. A more rigorous treatment is needed for this.