Add MockPulsar class, for easy use of Mock data with Enterprise

enterprise/pulsar.py

@@ -13,7 +13,7 @@
 from io import StringIO
 
 import numpy as np
-from ephem import Ecliptic, Equatorial
+from ephem import Ecliptic, Equatorial, J2000
 from astropy.time import Time
 
 import enterprise
@@ -673,6 +673,88 @@
             psr._deflated = "destroyed"
 
 
+class MockPulsar(BasePulsar):
+    """Class to allow mock pulsars to be used with Enterprise"""
+
+    _noastropy_warning_issued = False  # Class variable
+
+    def __init__(
+        self,
+        obs_times_mjd,
+        elong=None,
+        elat=None,
+        raj=None,
+        decj=None,
+        ssbfreqs=1440.0,
+        residuals=None,
+        toaerrs=1e-6,
+        sort=True,
+        flags={},
+        telescope="GBT",
+        spindown_order=2,
+        inc_astrometry=True,
+    ):
+        if inc_astrometry and not hasattr(const, "c"):  # pragma: no cover
+            # We requested astromery parameters, but there's no astropy
+            if not MockPulsar._noastropy_warning_issued:
+                msg = "WARNING: Astropy not installed but user requested "
+                msg += "astrometry timing model in MockPulsar. "
+                msg += "Switching off astrometry in all instances of MockPulsar."
+                logger.warning(msg)
+
+                MockPulsar._noastropy_warning_issued = True
+
+            inc_astrometry = False
+
+        self.name = utils.get_psrname_from_pos(elong=elong, elat=elat, raj=raj, decj=decj)
+
+        if elong is not None and elat is not None:
+            ec = Ecliptic(elong * np.pi / 180.0, elat * np.pi / 180.0)
+            eq = Equatorial(ec, epoch=J2000)
+            raj, decj = np.double(eq.ra), np.double(eq.dec)
+
+        self._raj, self._decj = raj, decj
+
+        self._sort = sort
+        self.planets = False
+
+        self._toas = np.double(obs_times_mjd) * 86400
+        self._stoas = np.double(obs_times_mjd) * 86400
+        self._residuals = residuals if residuals is not None else np.zeros_like(obs_times_mjd)
+        self._toaerrs = np.ones_like(obs_times_mjd) * toaerrs
+        self._posepoch = np.mean(self._toas)
+
+        if inc_astrometry:
+            self._designmatrix, self.fitpars = utils.create_astrometry_spin_timing_model(
+                self._toas, self._raj, self._decj, self._posepoch, spindown_order=spindown_order
+            )
+        else:
+            self._designmatrix, self.fitpars = utils.create_spindown_timing_model(self._toas, order=spindown_order)
+        self._ssbfreqs = np.ones_like(self._toas) * ssbfreqs / 1e6
+        self._telescope = telescope
+
+        # set parameters
+        self.setpars = [fp for fp in self.fitpars]
+
+        # new-style storage of flags as a numpy record array (previously, psr._flags = flags)
+        self._flags = np.zeros(len(self._toas), dtype=[(key, val.dtype) for key, val in flags.items()])
+        for key, val in flags.items():
+            self._flags[key] = val
+
+        self._pdist = self._get_pdist()
+
+        self._pos = self._get_pos()
+        self._planetssb = None
+        self._sunssb = None
+
+        self._pos_t = np.tile(self._pos, (len(self._toas), 1))
+
+        self.sort_data()
+
+    def set_residuals(self, residuals):
+        self._residuals = residuals
+
+
 class FeatherPulsar:
     columns = ["toas", "stoas", "toaerrs", "residuals", "freqs", "backend_flags", "telescope"]
     vector_columns = ["Mmat", "sunssb", "pos_t"]
@@ -866,4 +948,6 @@
                 raise ValueError(f"Unknown timing package {timing_package}")
         finally:
             os.chdir(cwd)
-    raise ValueError("Pulsar (par/tim) not specified in {args} or {kwargs}")
+            return PintPulsar(toas, model, sort=sort, drop_pintpsr=drop_pintpsr, planets=planets)
+
+    raise ValueError("Unknown arguments {}".format(args))