Faster implementations of astrometric ssb calculations

[dlakaplan]

This updates the ssb_to_psb_xyz_ICRS and ssb_to_psb_xyz_ECL implementation to be a lot faster while preserving accuracy. It avoids unnecessary coordinate construction/conversion and uses the erfa functions directly.

Specifics:

• In the base Astrometry model, ssb_to_psb_xyz_ICRS explicitly converts coordinates to ICRS but only at POSEPOCH (since we cannot be use what the base system is). Then the rest of the calculations are faster. This will also be used for AstrometryEcliptic models.
• In the AstrometryEquatorial model, ssb_to_psb_xyz_ICRS skips the conversion as it assumes things are already in ICRS. Then the calculations are even faster.
• In the AstrometryEcliptic model, ssb_to_psb_xyz_ECL skips the conversion as it assumes things are already in ECL. Then the calculations are even faster.

[dlakaplan]

There are tests to check accuracy. Finding differences of ~few * 1e-16 compared to the astropy versions. But the speed up is a lot.

ICRS model:
ERFA: 0.01736208299999986
Astropy: 0.9211249159999999
Max difference: 8.326672684688674e-16

ECL model
Accessing ICRS components:
ERFA: 0.7618184169999997
Astropy: 2.318666083
Max difference: 4.996003610813204e-16

Accessing ECL components:
ERFA: 0.019584334000000148
Astropy: 1.797397042
Max difference: 5.551115123125783e-16

So looking at the native components is very fast. But even when converting it's still a factor of 3-4.

[dlakaplan]

Test script:

from astropy import units as u, constants as c
from astropy.time import Time
from astropy.coordinates import SkyCoord
from pint.models import get_model
import io
import timeit
import copy
import numpy as np

parICRS = """PSR              1748-2021E
RAJ       17:48:52.75  1
DECJ      -20:21:29.0  1
PMRA     1000
PMDEC    -500
F0       61.485476554  1
F1         -1.181D-15  1
PEPOCH        53750.000000
POSEPOCH      53750.000000
DM              223.9  1
SOLARN0               0.00
EPHEM               DE421
CLK              UTC(NIST)
UNITS               TDB
TIMEEPH             FB90
T2CMETHOD           TEMPO
CORRECT_TROPOSPHERE N
PLANET_SHAPIRO      N
DILATEFREQ          N
TZRMJD  53801.38605118223
TZRFRQ            1949.609
TZRSITE                  1
"""
print("ICRS model:")
m = get_model(io.StringIO(parICRS))
t0 = m.POSEPOCH.quantity
t0.format = "pulsar_mjd"
t = t0 + np.linspace(0, 20) * u.yr
print(
    f"ERFA: {timeit.timeit('xyz = m.ssb_to_psb_xyz_ICRS(epoch=t)', globals=globals(), number=100)}"
)
print(
    f"Astropy: {timeit.timeit('xyz = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose()', globals=globals(), number=100)}"
)

xyzastropy = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose()
xyznew = m.ssb_to_psb_xyz_ICRS(epoch=t)
print(f"Max difference: {np.abs(xyzastropy-xyznew).max()}")

parECL = """PSR              1748-2021E
ELONG       300.0
ELAT      -55.0
PMELONG     1000
PMELAT    -500
F0       61.485476554  1
F1         -1.181D-15  1
PEPOCH        53750.000000
POSEPOCH      53750.000000
DM              223.9  1
SOLARN0               0.00
EPHEM               DE421
CLK              UTC(NIST)
UNITS               TDB
TIMEEPH             FB90
T2CMETHOD           TEMPO
CORRECT_TROPOSPHERE N
PLANET_SHAPIRO      N
DILATEFREQ          N
TZRMJD  53801.38605118223
TZRFRQ            1949.609
TZRSITE                  1
"""
print("ECL model")
m = get_model(io.StringIO(parECL))
t0 = m.POSEPOCH.quantity
t0.format = "pulsar_mjd"
t = t0 + np.linspace(0, 20) * u.yr
print("Accessing ICRS components:")
print(
    f"ERFA: {timeit.timeit('xyz = m.ssb_to_psb_xyz_ICRS(epoch=t)', globals=globals(), number=100)}"
)
print(
    f"Astropy: {timeit.timeit('xyz = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose()', globals=globals(), number=100)}"
)

xyzastropy = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose()
xyznew = m.ssb_to_psb_xyz_ICRS(epoch=t)
print(f"Max difference: {np.abs(xyzastropy-xyznew).max()}")

print("Accessing ECL components:")
print(
    f"ERFA: {timeit.timeit('xyz = m.ssb_to_psb_xyz_ECL(epoch=t)', globals=globals(), number=100)}"
)
print(
    f"Astropy: {timeit.timeit('xyz = m.coords_as_ECL(epoch=t).cartesian.xyz.transpose()', globals=globals(), number=100)}"
)

xyzastropy = m.coords_as_ECL(epoch=t).cartesian.xyz.transpose()
xyznew = m.ssb_to_psb_xyz_ECL(epoch=t)
print(f"Max difference: {np.abs(xyzastropy-xyznew).max()}")

[dlakaplan]

It will still fall back to the older calculations if it needs a different obliquity.

[dlakaplan]

A number of tests are failing. Some have to relate to tests that compare the analytic parameter derivatives with numerical versions. As far as I can tell, the analytical versions of those are essentially unchanged with this PR, but the numerical versions change appreciably. I don't know if anybody (@scottransom , @paulray ) has any ideas about that part, since they predate me.

[dlakaplan]

So this still isn't working. What I find is that when I compare the new/old methods in general, they are very very close (closer than the linear approximation). But for specific values of the PMRA they differ by a lot more.

For instance, using the files from tests/test_B1855.py I recreate the solar system delay calculation with the new & old methods:

    orig = m.PMRA.value
    dpm = orig * np.linspace(-30, 30, 100)
    ddelay = np.zeros(len(dpm))
    for i in range(len(dpm)):
        m.PMRA.value = orig + dpm[i]
        L_hat = m.ssb_to_psb_xyz_ICRS(epoch=t["tdbld"].astype(np.float64))
        re_dot_L = np.sum(t["ssb_obs_pos"] * L_hat, axis=1)
        delay = -re_dot_L.to_value(pint.ls)

        L_hat_old = m.coords_as_ICRS(epoch=t["tdbld"]).cartesian.xyz.transpose()
        re_dot_L_old = np.sum(t["ssb_obs_pos"] * L_hat_old, axis=1)
        delay_old = -re_dot_L_old.to_value(pint.ls)
        ddelay[i] = np.max(np.abs(delay - delay_old))

This is roughly how the numerical derivatives work. But instead of being close to zero and just changing smoothly I find:

with huge spikes. I don't understand these yet.

[dlakaplan]

the latest tests work fine on a Mac, but fail on ubuntu. not sure why. changing the tolerances

[codecov]

Codecov Report

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

Comparison is base (fb14307) 68.34% compared to head (deb3ec4) 68.39%.
Report is 1 commits behind head on master.

❗ Current head deb3ec4 differs from pull request most recent head 4922800. Consider uploading reports for the commit 4922800 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1646      +/-   ##
==========================================
+ Coverage   68.34%   68.39%   +0.04%     
==========================================
  Files         104      104              
  Lines       24112    24142      +30     
  Branches     4280     4285       +5     
==========================================
+ Hits        16480    16511      +31     
  Misses       6547     6547              
+ Partials     1085     1084       -1     

Files	Coverage Δ
src/pint/models/timing_model.py	83.94% <100.00%> (+0.12%)	⬆️
src/pint/models/astrometry.py	92.72% <95.34%> (-0.01%)	⬇️

... and 4 files with indirect coverage changes

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

[dlakaplan]

So it appears that the pmsafe routine doesn't actually apply the parallax in the way I thought. So it wasn't double-counting. I don't fully understand why, but I am using this to test:

def ssb_delay(m, t, includePX=True):
    toas = t
    delay = np.zeros(len(toas))
    tbl = t.table
    # c selects the non-barycentric TOAs that need actual calculation
    c = np.logical_and.reduce(tbl["ssb_obs_pos"] != 0, axis=1)
    if np.any(c):
        L_hat = m.ssb_to_psb_xyz_ICRS(epoch=tbl["tdbld"][c].astype(np.float64))
        re_dot_L = np.sum(tbl["ssb_obs_pos"][c] * L_hat, axis=1)
        delay[c] = -re_dot_L.to(pint.ls).value
        if m.PX.value != 0.0 and includePX:
            L = (1.0 / m.PX.value) * u.kpc
            # TODO: np.sum currently loses units in some cases...
            re_sqr = (
                np.sum(tbl["ssb_obs_pos"][c] ** 2, axis=1)
                * tbl["ssb_obs_pos"].unit ** 2
            )
            delay[c] += (
                (0.5 * (re_sqr / L) * (1.0 - re_dot_L**2 / re_sqr)).to(pint.ls).value
            )
    return delay

so it will return the delay with or without the explicit parallax term.

If I do this for a model I find that the value for a model with PX=0 matches that for a model with PX>0 if the explicit parallax term is ignored. But that with the explicit parallax included they don't match by an amount proportional to parallax.

[dlakaplan]

That said, if I leave the parallax out of the pmsafe command, then more derivative tests fail.

[dlakaplan]

I don't think that this is doing the parallax wrong or double. What I did was measure the parallax residual amplitude as a function of ELAT and compare with the prediction. (Note that I think the prediction from the HBOPA is wrong: it should be ~cos(ELAT)^2, but it isn't, although papers in the literature like https://ui.adsabs.harvard.edu/abs/2011A%26A...528A.108S/abstract have the cos^2 part). The results are basically identical between the master branch and the new implementation.

[dlakaplan]

P.S. it deviates at high ELAT due to (I think) the finite eccentricity of the Earth's orbit etc, as discussed in the paper cited above.

[dlakaplan]

P.P.S. Michael K has confirmed that there is a typo which is now fixed.

[dlakaplan]

As far as I can tell from this:

arICRS = """PSR              1748-2021E
 RAJ       17:48:52.75  1
 DECJ      -20:21:29.0  1
 PMRA     0
 PMDEC    0
 F0       61.485476554  1
 F1         -1.181D-15  1
 PEPOCH        53750.000000
 POSEPOCH      53750.000000
 DM              223.9  1
 SOLARN0               0.00
 EPHEM               DE421
 CLK              UTC(NIST)
 UNITS               TDB
 TIMEEPH             FB90
 T2CMETHOD           TEMPO
 CORRECT_TROPOSPHERE N
 PLANET_SHAPIRO      N
 DILATEFREQ          N
 TZRMJD  53801.38605118223
 TZRFRQ            1949.609
 TZRSITE                  1
 """

m = get_model(io.StringIO(parICRS), PX=100 * u.mas)
# cc = m.coords_as_ECL()
t = m.POSEPOCH.quantity + np.arange(365) * u.d
starpmout = erfa.pmsafe(
    m.RAJ.quantity.to_value(u.radian),
    m.DECJ.quantity.to_value(u.radian),
    m.PMRA.quantity.to_value(u.radian / u.yr) / np.cos(m.DECJ.quantity).value,
    m.PMDEC.quantity.to_value(u.radian / u.yr),
    m.PX.quantity.to_value(u.arcsec),
    0,
    2400000.5,
    m.POSEPOCH.value,
    2400000.5,
    t.mjd,
)
starpmout_noPX = erfa.pmsafe(
    m.RAJ.quantity.to_value(u.radian),
    m.DECJ.quantity.to_value(u.radian),
    m.PMRA.quantity.to_value(u.radian / u.yr) / np.cos(m.DECJ.quantity).value,
    m.PMDEC.quantity.to_value(u.radian / u.yr),
    0,
    0,
    2400000.5,
    m.POSEPOCH.value,
    2400000.5,
    t.mjd,
)

the pmsafe routine does not actually apply offsets related to parallax. maybe that is only used to convert the positions to a finite distance (and hence when treating a finite RV as well). So this might explain why we are not double-counting.

[dlakaplan]

Yes, I think that's the case from https://github.com/liberfa/erfa/blob/master/src/starpm.c#L130. So the parallax is only used to get a distance (which can change with the RV) but not an astrometric offset.

[abhisrkckl]

OK... I am convinced that this is doing the right thing. I'll merge this if that's OK.

[dlakaplan]

I'm ~99% sure this is good. I think we should merge since this is blocking other things, and if needed can fix later.