Merge branch 'master' of https://github.com/nanograv/enterprise

enterprise/pulsar.py

@@ -155,6 +155,14 @@ def to_pickle(self, outdir=None):
         # drop t2pulsar object
         if hasattr(self, "t2pulsar"):
             del self.t2pulsar
+            msg = "t2pulsar object cannot be pickled and has been removed."
+            logger.warning(msg)
+
+        if hasattr(self, "pint_toas"):
+            del self.pint_toas
+            del self.model
+            msg = "pint_toas and model objects cannot be pickled and have been removed."
+            logger.warning(msg)
 
         if outdir is None:
             outdir = os.getcwd()
@@ -278,11 +286,14 @@ def sunssb(self):
 
 
 class PintPulsar(BasePulsar):
-    def __init__(self, toas, model, sort=True, planets=True):
+    def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True):
 
         self._sort = sort
         self.planets = planets
         self.name = model.PSR.value
+        if not drop_pintpsr:
+            self.model = model
+            self.pint_toas = toas
 
         self._toas = np.array(toas.table["tdbld"], dtype="float64") * 86400
         # saving also stoas (e.g., for DMX comparisons)
@@ -522,6 +533,7 @@ def Pulsar(*args, **kwargs):
     planets = kwargs.get("planets", True)
     sort = kwargs.get("sort", True)
     drop_t2pulsar = kwargs.get("drop_t2pulsar", True)
+    drop_pintpsr = kwargs.get("drop_pintpsr", True)
     timing_package = kwargs.get("timing_package", "tempo2")
 
     if pint is not None:
@@ -535,7 +547,7 @@ def Pulsar(*args, **kwargs):
     timfile = [x for x in args if isinstance(x, str) and x.split(".")[-1] in ["tim", "toa"]]
 
     if pint and toas and model:
-        return PintPulsar(toas[0], model[0], sort=sort, planets=planets)
+        return PintPulsar(toas[0], model[0], sort=sort, drop_pintpsr=drop_pintpsr, planets=planets)
     elif t2 and t2pulsar:
         return Tempo2Pulsar(t2pulsar[0], sort=sort, drop_t2pulsar=drop_t2pulsar, planets=planets)
     elif parfile and timfile:
@@ -564,7 +576,7 @@ def Pulsar(*args, **kwargs):
                 relparfile, reltimfile, ephem=ephem, bipm_version=bipm_version, planets=planets
             )
             os.chdir(cwd)
-            return PintPulsar(toas, model, sort=sort, planets=planets)
+            return PintPulsar(toas, model, sort=sort, drop_pintpsr=drop_pintpsr, planets=planets)
 
         elif timing_package.lower() == "tempo2":
 

tests/test_pulsar.py

@@ -159,6 +159,7 @@ def setUpClass(cls):
             datadir + "/B1855+09_NANOGrav_9yv1.gls.par",
             datadir + "/B1855+09_NANOGrav_9yv1.tim",
             ephem="DE430",
+            drop_pintpsr=False,
             timing_package="pint",
         )
 
@@ -171,9 +172,13 @@ def test_dm(self):
         assert hasattr(self.psr, "dm")
 
     def test_planetssb(self):
-        """Place holder for filter_data tests."""
         assert hasattr(self.psr, "planetssb")
 
     def test_sunssb(self):
-        """Place holder for filter_data tests."""
         assert hasattr(self.psr, "sunssb")
+
+    def test_model(self):
+        assert hasattr(self.psr, "model")
+
+    def test_pint_toas(self):
+        assert hasattr(self.psr, "pint_toas")