Adding PintPulsar SSB Vectors

enterprise/pulsar.py

@@ -7,6 +7,7 @@
 import os
 
 import astropy.units as u
+import astropy.constants as const
 import numpy as np
 from ephem import Ecliptic, Equatorial
 
@@ -322,7 +323,8 @@ def __init__(self, toas, model, sort=True, planets=True):
         self._pdist = self._get_pdist()
         self._raj, self._decj = self._get_radec(model)
         self._pos = self._get_pos()
-        self._planetssb = self._get_planetssb()
+        self._planetssb = self._get_planetssb(toas, model)
+        self._sunssb = self._get_sunssb(toas, model)
 
         # TODO: pos_t not currently implemented
         self._pos_t = np.zeros((len(self._toas), 3))
@@ -359,11 +361,41 @@ def _get_radec(self, model):
             elong, elat = model.ELONG.value, model.ELAT.value
             return self._get_radec_from_ecliptic(elong * d2r, elat * d2r)
 
-    def _get_planetssb(self):
-        return np.zeros((len(self._toas), 9, 6))
+    def _get_ssb_lsec(self, toas, obs_planet):
+        """Get the planet to SSB vector in lightseconds from Pint table"""
+        vec = toas.table[obs_planet] + toas.table["ssb_obs_pos"]
+        return (vec / const.c).to("s").value
 
-    def _get_sunssb(self):
-        return np.zeros((len(self._toas), 6))
+    def _get_planetssb(self, toas, model):
+        planetssb = None
+        if self.planets:
+            planetssb = np.zeros((len(self._toas), 9, 6))
+            # planetssb[:, 0, :] = self.t2pulsar.mercury_ssb
+            # planetssb[:, 1, :] = self.t2pulsar.venus_ssb
+            planetssb[:, 2, :3] = self._get_ssb_lsec(toas, "obs_earth_pos")
+            # planetssb[:, 3, :] = self.t2pulsar.mars_ssb
+            planetssb[:, 4, :3] = self._get_ssb_lsec(toas, "obs_jupiter_pos")
+            planetssb[:, 5, :3] = self._get_ssb_lsec(toas, "obs_saturn_pos")
+            planetssb[:, 6, :3] = self._get_ssb_lsec(toas, "obs_uranus_pos")
+            planetssb[:, 7, :3] = self._get_ssb_lsec(toas, "obs_neptune_pos")
+            # planetssb[:, 8, :] = self.t2pulsar.pluto_ssb
+
+            # if hasattr(model, "ELAT") and hasattr(model, "ELONG"):
+            #     for ii in range(9):
+            #         planetssb[:, ii, :3] = utils.ecl2eq_vec(planetssb[:, ii, :3])
+            #         # planetssb[:, ii, 3:] = utils.ecl2eq_vec(planetssb[:, ii, 3:])
+        return planetssb
+
+    def _get_sunssb(self, toas, model):
+        sunssb = None
+        if self.planets:
+            sunssb = np.zeros((len(self._toas), 6))
+            sunssb[:, :3] = self._get_ssb_lsec(toas, "obs_sun_pos")
+
+            # if hasattr(model, "ELAT") and hasattr(model, "ELONG"):
+            #     sunssb[:, :3] = utils.ecl2eq_vec(sunssb[:, :3])
+            # #     sunssb[:, 3:] = utils.ecl2eq_vec(sunssb[:, 3:])
+        return sunssb
 
 
 class Tempo2Pulsar(BasePulsar):
@@ -529,7 +561,9 @@ def Pulsar(*args, **kwargs):
         if timing_package.lower() == "pint":
             if (clk is not None) and (bipm_version is None):
                 bipm_version = clk.split("(")[1][:-1]
-            model, toas = get_model_and_toas(relparfile, reltimfile, ephem=ephem, bipm_version=bipm_version)
+            model, toas = get_model_and_toas(
+                relparfile, reltimfile, ephem=ephem, bipm_version=bipm_version, planets=planets
+            )
             os.chdir(cwd)
             return PintPulsar(toas, model, sort=sort, planets=planets)
 

tests/test_pulsar.py

@@ -151,13 +151,15 @@ def setUpClass(cls):
     # exclude tests pending implementation of .stoas, .dm, .dmx in PintPulsar
 
     def test_stoas(self):
-        pass
+        assert hasattr(self.psr, "stoas")
 
     def test_dm(self):
-        pass
+        assert hasattr(self.psr, "dm")
 
     def test_planetssb(self):
-        pass
+        """Place holder for filter_data tests."""
+        assert hasattr(self.psr, "planetssb")
 
     def test_sunssb(self):
-        pass
+        """Place holder for filter_data tests."""
+        assert hasattr(self.psr, "sunssb")