diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index 8611eef0..20a8933e 100644 --- a/enterprise/pulsar.py +++ b/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) diff --git a/tests/test_pulsar.py b/tests/test_pulsar.py index d4128d90..95cad51e 100644 --- a/tests/test_pulsar.py +++ b/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")