Merge pull request #1646 from dlakaplan/fasterastrom

Faster implementations of astrometric ssb calculations

CHANGELOG-unreleased.md

@@ -12,6 +12,8 @@ the released changes.
 - `WAVE` parameters can be added to a `Wave` model with `add_wave_component()` in `wave.py` 
 - Moved design matrix normalization code from `pint.fitter` to the new `pint.utils.normalize_designmatrix()` function.
 - Made `Residuals` independent of `GLSFitter` (GLS chi2 is now computed using the new function `Residuals._calc_gls_chi2()`).
+- `ssb_to_psb_ICRS` implementation is now a lot faster (uses erfa functions directly)
+- `ssb_to_psb_ECL` implementation within `AstrometryEcliptic` model is now a lot faster (uses erfa functions directly)
 - Upgraded versioneer for compatibility with Python 3.12
 - Creation of `Fitter` objects will fail if there are free unfittable parameters in the timing model.
 - Only fittable parameters will be listed as check boxes in the `plk` interface.

src/pint/models/astrometry.py

@@ -11,7 +11,7 @@
 
 from loguru import logger as log
 
-from erfa import ErfaWarning
+from erfa import ErfaWarning, pmsafe
 from pint import ls
 from pint.models.parameter import (
     AngleParameter,
@@ -23,7 +23,6 @@
 from pint.pulsar_ecliptic import OBL, PulsarEcliptic
 from pint.utils import add_dummy_distance, remove_dummy_distance
 
-
 astropy_version = sys.modules["astropy"].__version__
 mas_yr = u.mas / u.yr
 
@@ -61,8 +60,20 @@ def ssb_to_psb_xyz_ICRS(self, epoch=None):
         """Returns unit vector(s) from SSB to pulsar system barycenter under ICRS.
 
         If epochs (MJD) are given, proper motion is included in the calculation.
+
+        Parameters
+        ----------
+        epoch : float or astropy.time.Time, optional
+
+        Returns
+        -------
+        np.ndarray :
+            (len(epoch), 3) array of unit vectors
         """
         # TODO: would it be better for this to return a 6-vector (pos, vel)?
+
+        # this is somewhat slow, since it repeatedly created different SkyCoord Objects
+        # but for consistency only change the method in the subclasses below
         return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose()
 
     def ssb_to_psb_xyz_ECL(self, epoch=None, ecl=None):
@@ -75,6 +86,11 @@ def ssb_to_psb_xyz_ECL(self, epoch=None, ecl=None):
         epoch : float, optional
         ecl : str, optional
             Obliquity (IERS2010 by default)
+
+        Returns
+        -------
+        np.ndarray :
+            (len(epoch), 3) array of unit vectors
         """
         # TODO: would it be better for this to return a 6-vector (pos, vel)?
         return self.coords_as_ECL(epoch=epoch, ecl=ecl).cartesian.xyz.transpose()
@@ -371,7 +387,7 @@ def coords_as_ECL(self, epoch=None, ecl=None):
         If `ecl` is left unspecified, the global default IERS2010 will be used.
         """
         if ecl is None:
-            log.info("ECL not specified; using IERS2010.")
+            log.debug("ECL not specified; using IERS2010.")
             ecl = "IERS2010"
 
         pos_icrs = self.get_psr_coords(epoch=epoch)
@@ -390,6 +406,63 @@ def get_params_as_ICRS(self):
             "PMDEC": self.PMDEC.quantity,
         }
 
+    def ssb_to_psb_xyz_ICRS(self, epoch=None):
+        """Returns unit vector(s) from SSB to pulsar system barycenter under ICRS.
+
+        If epochs (MJD) are given, proper motion is included in the calculation.
+
+        Parameters
+        ----------
+        epoch : float or astropy.time.Time, optional
+
+        Returns
+        -------
+        np.ndarray :
+            (len(epoch), 3) array of unit vectors
+        """
+        # TODO: would it be better for this to return a 6-vector (pos, vel)?
+
+        # this was somewhat slow, since it repeatedly created different SkyCoord Objects
+        # return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose()
+
+        # Instead look at what https://docs.astropy.org/en/stable/_modules/astropy/coordinates/sky_coordinate.html#SkyCoord.apply_space_motion
+        # does, which is to use https://github.com/liberfa/erfa/blob/master/src/starpm.c
+        # and then just use the relevant pieces of that
+        if epoch is None or (self.PMRA.quantity == 0 and self.PMDEC.quantity == 0):
+            return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose()
+
+        if isinstance(epoch, Time):
+            jd1 = epoch.jd1
+            jd2 = epoch.jd2
+        else:
+            # assume MJD
+            jd1 = 2400000.5
+            jd2 = epoch
+        # compared to the general case above we can assume that the coordinates are ICRS
+        # so just access those components
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", ErfaWarning)
+            # note that starpm wants mu_alpha not mu_alpha * cos(delta)
+            starpmout = pmsafe(
+                self.RAJ.quantity.to_value(u.radian),
+                self.DECJ.quantity.to_value(u.radian),
+                self.PMRA.quantity.to_value(u.radian / u.yr)
+                / np.cos(self.DECJ.quantity).value,
+                self.PMDEC.quantity.to_value(u.radian / u.yr),
+                self.PX.quantity.to_value(u.arcsec),
+                0.0,
+                self.POSEPOCH.quantity.jd1,
+                self.POSEPOCH.quantity.jd2,
+                jd1,
+                jd2,
+            )
+        # ra,dec now in radians
+        ra, dec = starpmout[0], starpmout[1]
+        x = np.cos(ra) * np.cos(dec)
+        y = np.sin(ra) * np.cos(dec)
+        z = np.sin(dec)
+        return u.Quantity([x, y, z]).T
+
     def d_delay_astrometry_d_RAJ(self, toas, param="", acc_delay=None):
         """Calculate the derivative wrt RAJ
 
@@ -754,6 +827,78 @@ def coords_as_ECL(self, epoch=None, ecl=None):
             pos_ecl = pos_ecl.transform_to(PulsarEcliptic(ecl=ecl))
         return pos_ecl
 
+    def ssb_to_psb_xyz_ECL(self, epoch=None, ecl=None):
+        """Returns unit vector(s) from SSB to pulsar system barycenter under ECL.
+
+        If epochs (MJD) are given, proper motion is included in the calculation.
+
+        Parameters
+        ----------
+        epoch : float or astropy.time.Time, optional
+        ecl : str, optional
+            Obliquity (IERS2010 by default)
+
+        Returns
+        -------
+        np.ndarray :
+            (len(epoch), 3) array of unit vectors
+        """
+        # TODO: would it be better for this to return a 6-vector (pos, vel)?
+
+        # this was somewhat slow, since it repeatedly created different SkyCoord Objects
+        # return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose()
+
+        # Instead look at what https://docs.astropy.org/en/stable/_modules/astropy/coordinates/sky_coordinate.html#SkyCoord.apply_space_motion
+        # does, which is to use https://github.com/liberfa/erfa/blob/master/src/starpm.c
+        # and then just use the relevant pieces of that
+
+        # but we need to check that the obliquity is the same
+        if ecl is not None and ecl != self.ECL.quantity:
+            return super().ssb_to_psb_xyz_ECL(epoch=epoch, ecl=ecl)
+
+        if ecl is None:
+            log.debug("ECL not specified; using IERS2010.")
+            ecl = "IERS2010"
+        if epoch is None or (self.PMELONG.value == 0 and self.PMELAT.value == 0):
+            return self.coords_as_ECL(epoch=epoch, ecl=ecl).cartesian.xyz.transpose()
+        if isinstance(epoch, Time):
+            jd1 = epoch.jd1
+            jd2 = epoch.jd2
+        else:
+            jd1 = 2400000.5
+            jd2 = epoch
+        # compared to the general case above we can assume that the coordinates are ECL
+        # so just access those components
+        lon = self.ELONG.quantity.to_value(u.radian)
+        lat = self.ELAT.quantity.to_value(u.radian)
+        pm_lon = (
+            self.PMELONG.quantity.to_value(u.radian / u.yr)
+            / np.cos(self.ELAT.quantity).value
+        )
+        pm_lat = self.PMELAT.quantity.to_value(u.radian / u.yr)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", ErfaWarning)
+            # note that pmsafe wants mu_lon not mu_lon * cos(lat)
+            starpmout = pmsafe(
+                lon,
+                lat,
+                pm_lon,
+                pm_lat,
+                self.PX.quantity.to_value(u.arcsec),
+                0.0,
+                self.POSEPOCH.quantity.jd1,
+                self.POSEPOCH.quantity.jd2,
+                jd1,
+                jd2,
+            )
+        # lon,lat now in radians
+        lon, lat = starpmout[0], starpmout[1]
+        x = np.cos(lon) * np.cos(lat)
+        y = np.sin(lon) * np.cos(lat)
+        z = np.sin(lat)
+        return u.Quantity([x, y, z]).T
+
     def get_d_delay_quantities_ecliptical(self, toas):
         """Calculate values needed for many d_delay_d_param functions."""
         # TODO: Move all these calculations in a separate class for elegance

src/pint/models/timing_model.py

@@ -1898,7 +1898,7 @@ def d_phase_d_param_num(self, toas, param, step=1e-2):
         res_f = -phase_f[:, 0] + phase_f[:, 1]
         result = (res_i + res_f) / (2.0 * h * unit)
         # shift value back to the original value
-        par.quantity = ori_value
+        par.value = ori_value
         return result
 
     def d_delay_d_param_num(self, toas, param, step=1e-2):

tests/test_B1855.py

@@ -1,7 +1,6 @@
 """Various tests to assess the performance of the B1855+09."""
 import logging
 import os
-import pytest
 
 import astropy.units as u
 import numpy as np

tests/test_astrometry.py

@@ -81,3 +81,123 @@ def test_pm_acquires_posepoch_equatorial():
 def test_pm_one_set_other_not():
     m = get_model(StringIO("\n".join([par_basic_equatorial, "PMRA 7"])))
     assert m.POSEPOCH.quantity == m.PEPOCH.quantity
+
+
+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
+"""
+
+
+ntests = 200
+
+
+@pytest.mark.parametrize(
+    ("ra,dec,pmra,pmdec"),
+    list(
+        zip(
+            np.random.uniform(0, 360, ntests),
+            np.random.uniform(-90, 90, ntests),
+            np.random.normal(0, scale=10, size=ntests),
+            np.random.normal(0, scale=10, size=ntests),
+        )
+    ),
+)
+def test_ssb_accuracy_ICRS(ra, dec, pmra, pmdec):
+    m = get_model(StringIO(parICRS), RAJ=ra, DECJ=dec, PMRA=pmra, PMDEC=pmdec)
+    t0 = m.POSEPOCH.quantity
+    t0.format = "pulsar_mjd"
+    t = t0 + np.linspace(0, 20) * u.yr
+    xyzastropy = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose()
+    xyznew = m.ssb_to_psb_xyz_ICRS(epoch=t)
+    assert np.allclose(xyznew, xyzastropy)
+
+
+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
+"""
+
+
+@pytest.mark.parametrize(
+    ("elong,elat,pmelong,pmelat"),
+    list(
+        zip(
+            np.random.uniform(0, 360, ntests),
+            np.random.uniform(-90, 90, ntests),
+            np.random.normal(0, scale=10, size=ntests),
+            np.random.normal(0, scale=10, size=ntests),
+        )
+    ),
+)
+def test_ssb_accuracy_ICRS_ECLmodel(elong, elat, pmelong, pmelat):
+    m = get_model(
+        StringIO(parECL), ELONG=elong, ELAT=elat, PMELONG=pmelong, PMELAT=pmelat
+    )
+    t0 = m.POSEPOCH.quantity
+    t0.format = "pulsar_mjd"
+    t = t0 + np.linspace(0, 20) * u.yr
+    xyzastropy = m.coords_as_ICRS(epoch=t).cartesian.xyz.transpose()
+    xyznew = m.ssb_to_psb_xyz_ICRS(epoch=t)
+    assert np.allclose(xyznew, xyzastropy)
+
+
+@pytest.mark.parametrize(
+    ("elong,elat,pmelong,pmelat"),
+    list(
+        zip(
+            np.random.uniform(0, 360, ntests),
+            np.random.uniform(-90, 90, ntests),
+            np.random.normal(0, scale=10, size=ntests),
+            np.random.normal(0, scale=10, size=ntests),
+        )
+    ),
+)
+def test_ssb_accuracy_ECL_ECLmodel(elong, elat, pmelong, pmelat):
+    m = get_model(
+        StringIO(parECL), ELONG=elong, ELAT=elat, PMELONG=pmelong, PMELAT=pmelat
+    )
+    t0 = m.POSEPOCH.quantity
+    t0.format = "pulsar_mjd"
+    t = t0 + np.linspace(0, 20) * u.yr
+    xyzastropy = m.coords_as_ECL(epoch=t).cartesian.xyz.transpose()
+    xyznew = m.ssb_to_psb_xyz_ECL(epoch=t)
+    assert np.allclose(xyznew, xyzastropy)

tests/test_derivative_utils.py

@@ -68,10 +68,8 @@ def get_derivative_params(model):
             h = 2e-1
         elif p in ["FD2"]:
             h = 3e-2
-        elif p in ["F1", "M2", "PMELONG", "PMELAT"]:
+        elif p in ["F1", "M2", "PMELONG", "PMELAT", "PMRA", "PMDEC"]:
             h = 2
-        elif p in ["PMDEC"]:
-            h = 3e-2
         elif p in ["PBDOT"]:
             h = 200
         elif p in ["FB1"]:

tests/test_modelconversions.py

@@ -155,7 +155,10 @@ def test_ICRS_to_ECL_uncertainties():
     m2 = fit_ECL.model.as_ICRS()
 
     for p in ("RAJ", "DECJ", "PMRA", "PMDEC"):
-        assert np.isclose(m1.__getitem__(p).value, m2.__getitem__(p).value)
+        tol = 1e-12 if p in ["RAJ", "DECJ"] else 1e-3
+        assert np.isclose(
+            m1.__getitem__(p).value, m2.__getitem__(p).value, atol=tol
+        ), f"Paramter {p} with values {m1.__getitem__(p).value}, {m2.__getitem__(p).value} was too far apart (difference={m1.__getitem__(p).value-m2.__getitem__(p).value})"
         # do a generous test here since the uncertainties could change
         assert np.isclose(
             m1.__getitem__(p).uncertainty, m2.__getitem__(p).uncertainty, rtol=0.5