Make Astrometry & SolarSystemShapiro faster

CHANGELOG-unreleased.md

@@ -9,6 +9,8 @@ the released changes.
 
 ## Unreleased
 ### Changed
+- Avoided unnecessary creation of `SkyCoord` objects in `AstrometryEquatorial` and `AstrometryEcliptic`.
+- Avoided unnecessary `TOAs` table slices in `SolarSystemShapiro`
 - Allow "CLK UNCORR" in par files (indicates no GPS or BIPM corrections). 
 - Better documentation for `akaike_information_criterion()`
 ### Added

src/pint/models/astrometry.py

@@ -485,7 +485,9 @@ def ssb_to_psb_xyz_ICRS(
         # 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()
+            ra, dec = self.RAJ.quantity, self.DECJ.quantity
+            return self.xyz_from_radec(ra, dec)
+            # return self.coords_as_ICRS(epoch=epoch).cartesian.xyz.transpose()
 
         if isinstance(epoch, Time):
             jd1 = epoch.jd1
@@ -518,6 +520,9 @@ def ssb_to_psb_xyz_ICRS(
             )
         # ra,dec now in radians
         ra, dec = starpmout[0], starpmout[1]
+        return self.xyz_from_radec(ra, dec)
+
+    def xyz_from_radec(self, ra, dec):
         x = np.cos(ra) * np.cos(dec)
         y = np.sin(ra) * np.cos(dec)
         z = np.sin(dec)
@@ -977,7 +982,9 @@ def ssb_to_psb_xyz_ECL(
             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()
+            # return self.coords_as_ECL(epoch=epoch, ecl=ecl).cartesian.xyz.transpose()
+            lon, lat = self.ELONG.quantity, self.ELAT.quantity
+            return self.xyz_from_latlong(lon, lat)
         if isinstance(epoch, Time):
             jd1 = epoch.jd1
             jd2 = epoch.jd2
@@ -1015,6 +1022,9 @@ def ssb_to_psb_xyz_ECL(
             )
         # lon,lat now in radians
         lon, lat = starpmout[0], starpmout[1]
+        return self.xyz_from_latlong(lon, lat)
+
+    def xyz_from_latlong(self, lon, lat):
         x = np.cos(lon) * np.cos(lat)
         y = np.sin(lon) * np.cos(lat)
         z = np.sin(lat)

src/pint/models/solar_system_shapiro.py

@@ -2,7 +2,7 @@
 
 import astropy.constants as const
 import astropy.units as u
-import numpy
+import numpy as np
 from loguru import logger as log
 
 from pint import (
@@ -68,9 +68,9 @@ def ss_obj_shapiro_delay(obj_pos, psr_dir, T_obj):
           psr_dir : unit vector in direction of pulsar
           T_obj : mass of object in seconds (GM/c^3)
         """
-        # TODO: numpy.sum currently loses units in some cases...
-        r = (numpy.sqrt(numpy.sum(obj_pos**2, axis=1))) * obj_pos.unit
-        rcostheta = numpy.sum(obj_pos * psr_dir, axis=1)
+        # TODO: np.sum currently loses units in some cases...
+        r = (np.sqrt(np.sum(obj_pos**2, axis=1))) * obj_pos.unit
+        rcostheta = np.sum(obj_pos * psr_dir, axis=1)
         # This is the 2nd to last term from Eqn 4.6 in Backer &
         # Hellings, ARAA, 1986 with gamma = 1 (as defined by GR).  We
         # have the opposite sign of the cos(theta) term, since our
@@ -79,7 +79,7 @@ def ss_obj_shapiro_delay(obj_pos, psr_dir, T_obj):
         # pulsar (as described after Eqn 4.3 in the paper).
         # See also https://en.wikipedia.org/wiki/Shapiro_time_delay
         # where \Delta t = \frac{2GM}{c^3}\log(1-\vec{R}\cdot\vec{x})
-        return -2.0 * T_obj * numpy.log((r - rcostheta) / const.au).value
+        return -2.0 * T_obj * np.log((r - rcostheta) / const.au).value
 
     def solar_system_shapiro_delay(self, toas, acc_delay=None):
         """
@@ -94,30 +94,32 @@ def solar_system_shapiro_delay(self, toas, acc_delay=None):
         If planets are to be included, TOAs.compute_posvels() must
         have been called with the planets=True argument.
         """
-        # Start out with 0 delay with units of seconds
-        tbl = toas.table
-        delay = numpy.zeros(len(tbl))
-        for key, grp in toas.get_obs_groups():
-            if key.lower() == "barycenter":
-                log.debug("Skipping Shapiro delay for Barycentric TOAs")
-                continue
+
+        # Apply Shapiro delay correction only for non-barycentered TOAs.
+        non_bary_mask = toas.get_obss() != "barycenter"
+        delay = np.zeros(len(toas))
+
+        if np.any(non_bary_mask):
+            tbl_grp = toas.table if np.all(non_bary_mask) else toas.table[non_bary_mask]
+
             psr_dir = self._parent.ssb_to_psb_xyz_ICRS(
-                epoch=tbl[grp]["tdbld"].astype(numpy.float64)
+                epoch=tbl_grp["tdbld"].astype(np.float64)
             )
-            delay[grp] += self.ss_obj_shapiro_delay(
-                tbl[grp]["obs_sun_pos"], psr_dir, self._ss_mass_sec["sun"]
+            delay[non_bary_mask] += self.ss_obj_shapiro_delay(
+                tbl_grp["obs_sun_pos"], psr_dir, self._ss_mass_sec["sun"]
             )
-            try:
-                if self.PLANET_SHAPIRO.value:
+            if self.PLANET_SHAPIRO.value:
+                try:
                     for pl in ("jupiter", "saturn", "venus", "uranus", "neptune"):
-                        delay[grp] += self.ss_obj_shapiro_delay(
-                            tbl[grp][f"obs_{pl}_pos"],
+                        delay[non_bary_mask] += self.ss_obj_shapiro_delay(
+                            tbl_grp[f"obs_{pl}_pos"],
                             psr_dir,
                             self._ss_mass_sec[pl],
                         )
-            except KeyError as e:
-                raise KeyError(
-                    "Planet positions not found when trying to compute Solar System Shapiro delay. "
-                    "Make sure that you include `planets=True` in your `get_TOAs()` call, or use `get_model_and_toas()`."
-                ) from e
+                except KeyError as e:
+                    raise KeyError(
+                        "Planet positions not found when trying to compute Solar System Shapiro delay. "
+                        "Make sure that you include `planets=True` in your `get_TOAs()` call, or use `get_model_and_toas()`."
+                    ) from e
+
         return delay * u.second

src/pint/simulation.py

@@ -3,7 +3,7 @@
 
 from collections import OrderedDict
 from copy import deepcopy
-from typing import Optional, List, Union
+from typing import Optional, List, Tuple, Union
 import pathlib
 
 import astropy.units as u
@@ -566,7 +566,7 @@ def calculate_random_models(
     keep_models: bool = True,
     return_time: bool = False,
     params: str = "all",
-) -> (np.ndarray, Optional[list]):
+) -> Tuple[np.ndarray, Optional[list]]:
     """
     Calculates random models based on the covariance matrix of the `fitter` object.
 
@@ -695,7 +695,7 @@ def _get_freqs_and_times(
     ntoas: int,
     freqs: u.Quantity,
     multi_freqs_in_epoch: bool = True,
-) -> (Union[float, u.Quantity, time.Time], np.ndarray):
+) -> Tuple[Union[float, u.Quantity, time.Time], np.ndarray]:
     freqs = np.atleast_1d(freqs)
     assert (
         len(freqs.shape) == 1 and len(freqs) <= ntoas

tests/test_wavex.py

@@ -249,8 +249,8 @@ def test_multiple_wavex_unit_conversion():
         wxcoses=[2, 3],
         frozens=False,
     )
-    assert getattr(model, f"WXFREQ_0002").value == freqs[0].to(u.d**-1).value
-    assert getattr(model, f"WXFREQ_0003").value == freqs[1].to(u.d**-1).value
+    assert getattr(model, "WXFREQ_0002").value == freqs[0].to(u.d**-1).value
+    assert getattr(model, "WXFREQ_0003").value == freqs[1].to(u.d**-1).value
 
 
 def test_cos_amp_missing():