From ca9a5a97add6be92a5005599e6f338bdafb3a553 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Wed, 27 Sep 2023 11:32:20 +0200 Subject: [PATCH 1/8] Modified pulsar.py to allow for additional pulsar types from external packages --- enterprise/pulsar.py | 129 ++++++++++++++++++++++++++++++++----------- 1 file changed, 98 insertions(+), 31 deletions(-) diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index b8f23cd5..7929182c 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -2,10 +2,12 @@ """Class containing pulsar data from timing package [tempo2/PINT]. """ +import contextlib import json import logging import os import pickle +from io import StringIO import astropy.constants as const import astropy.units as u @@ -149,9 +151,18 @@ def filter_data(self, start_time=None, end_time=None): self.sort_data() + def drop_not_picklable(self): + """Drop all attributes that cannot be pickled. + + Derived classes should implement this if they have + any such attributes. + """ + pass + def to_pickle(self, outdir=None): """Save object to pickle file.""" + self.drop_not_picklable() # drop t2pulsar object if hasattr(self, "t2pulsar"): del self.t2pulsar @@ -318,7 +329,11 @@ def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): if not drop_pintpsr: self.model = model + self.parfile = model.as_parfile() self.pint_toas = toas + with StringIO() as tim: + toas.write_TOA_file(tim) + self.timfile = tim.getvalue() # these are TDB but not barycentered # self._toas = np.array(toas.table["tdbld"], dtype="float64") * 86400 @@ -327,20 +342,17 @@ def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): self._stoas = np.array(toas.get_mjds().value, dtype="float64") * 86400 self._residuals = np.array(resids(toas, model).time_resids.to(u.s), dtype="float64") self._toaerrs = np.array(toas.get_errors().to(u.s), dtype="float64") - self._designmatrix = model.designmatrix(toas)[0] + self._designmatrix, self.fitpars, self.designmatrix_units = model.designmatrix(toas) self._ssbfreqs = np.array(model.barycentric_radio_freq(toas), dtype="float64") self._telescope = np.array(toas.get_obss()) - # fitted parameters - self.fitpars = ["Offset"] + [par for par in model.params if not getattr(model, par).frozen] - # gather DM/DMX information if available self._set_dm(model) # set parameters - spars = [par for par in model.params] - self.setpars = [sp for sp in spars if sp not in self.fitpars] + self.setpars = [sp for sp in model.params if sp not in self.fitpars] + # FIXME: this can be done more cleanly using PINT self._flags = {} for ii, obsflags in enumerate(toas.get_flags()): for jj, flag in enumerate(obsflags): @@ -351,6 +363,7 @@ def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): # convert flags to arrays # TODO probably better way to do this + # -- PINT always stores flags as strings for key, val in self._flags.items(): if isinstance(val[0], u.quantity.Quantity): self._flags[key] = np.array([v.value for v in val]) @@ -371,6 +384,21 @@ def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): self.sort_data() + def drop_pintpsr(self): + with contextlib.suppress(NameError): + del self.model + del self.parfile + del self.pint_toas + del self.timfile + + def drop_not_picklable(self): + with contextlib.suppress(AttributeError): + del self.model + del self.pint_toas + logger.warning("pint_toas and model objects cannot be pickled and have been removed.") + + return super().drop_not_picklable() + def _set_dm(self, model): pars = [par for par in model.params if not getattr(model, par).frozen] @@ -447,7 +475,15 @@ def _get_sunssb(self, toas, model): class Tempo2Pulsar(BasePulsar): - def __init__(self, t2pulsar, sort=True, drop_t2pulsar=True, planets=True): + def __init__( + self, + t2pulsar, + sort=True, + drop_t2pulsar=True, + planets=True, + par_name=None, + tim_name=None, + ): self._sort = sort self.t2pulsar = t2pulsar @@ -496,6 +532,11 @@ def __init__(self, t2pulsar, sort=True, drop_t2pulsar=True, planets=True): if drop_t2pulsar: del self.t2pulsar + else: + if par_name is not None and os.path.exists(par_name): + self.parfile = open(par_name).read() + if tim_name is not None and os.path.exists(tim_name): + self.timfile = open(tim_name).read() # gather DM/DMX information if available def _set_dm(self, t2pulsar): @@ -557,7 +598,7 @@ def _get_sunssb(self, t2pulsar): sunssb = None if self.planets: # for ii in range(1, 10): - # tag = 'DMASSPLANET' + str(ii) + # tag = 'DMASSPLANET' + str(ii)@pytest.mark.skipif(t2 is None, reason="TEMPO2/libstempo not available") # self.t2pulsar[tag].val = 0.0 self.t2pulsar.formbats() sunssb = np.zeros((len(self._toas), 6)) @@ -574,6 +615,12 @@ def _get_sunssb(self, t2pulsar): # then replace them with pickleable objects that can be inflated # to numpy arrays with SharedMemory storage + def drop_not_picklable(self): + with contextlib.suppress(AttributeError): + del self.t2pulsar + logger.warning("t2pulsar object cannot be pickled and has been removed.") + return super().drop_not_picklable() + _todeflate = ["_designmatrix", "_planetssb", "_sunssb", "_flags"] _deflated = "pristine" @@ -610,7 +657,9 @@ def Pulsar(*args, **kwargs): 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") + timing_package = kwargs.get("timing_package", None) + if timing_package is not None: + timing_package = timing_package.lower() if pint is not None: toas = [x for x in args if isinstance(x, TOAs)] @@ -638,28 +687,46 @@ def Pulsar(*args, **kwargs): reltimfile = timfiletup[-1] relparfile = os.path.relpath(parfile[0], dirname) + if timing_package is None: + if t2 is not None: + timing_package = "tempo2" + elif pint is not None: + timing_package = "pint" + else: + raise ValueError("No timing package available with which to load a pulsar") + # get current directory cwd = os.getcwd() - - # Change directory to the base directory of the tim-file to deal with - # INCLUDE statements in the tim-file - os.chdir(dirname) - - 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, planets=planets - ) - os.chdir(cwd) - return PintPulsar(toas, model, sort=sort, drop_pintpsr=drop_pintpsr, planets=planets) - - elif timing_package.lower() == "tempo2": - - # hack to set maxobs - maxobs = get_maxobs(reltimfile) + 100 - t2pulsar = t2.tempopulsar(relparfile, reltimfile, maxobs=maxobs, ephem=ephem, clk=clk) + try: + # Change directory to the base directory of the tim-file to deal with + # INCLUDE statements in the tim-file + os.chdir(dirname) + if timing_package.lower == "tempo2": + if t2 is None: + raise ValueError("tempo2 requested but tempo2 is not available") + # hack to set maxobs + maxobs = get_maxobs(reltimfile) + 100 + t2pulsar = t2.tempopulsar(relparfile, reltimfile, maxobs=maxobs, ephem=ephem, clk=clk) + return Tempo2Pulsar( + t2pulsar, + sort=sort, + drop_t2pulsar=drop_t2pulsar, + planets=planets, + par_name=relparfile, + tim_name=reltimfile, + ) + elif timing_package.lower() == "pint": + if pint is None: + raise ValueError("PINT requested but PINT is not available") + 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, planets=planets + ) + os.chdir(cwd) + return PintPulsar(toas, model, sort=sort, drop_pintpsr=drop_pintpsr, planets=planets) + else: + raise ValueError(f"Unknown timing package {timing_package}") + finally: os.chdir(cwd) - return Tempo2Pulsar(t2pulsar, sort=sort, drop_t2pulsar=drop_t2pulsar, planets=planets) - - raise ValueError("Unknown arguments {}".format(args)) + raise ValueError("Pulsar (par/tim) not specified in {args} or {kwargs}") From 6703327454403e5a9eb2fa82fff26bec83759feb Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Wed, 27 Sep 2023 14:54:55 +0200 Subject: [PATCH 2/8] Moved drop t2pulsar code --- enterprise/pulsar.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index 7929182c..f8589db4 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -163,17 +163,6 @@ def to_pickle(self, outdir=None): """Save object to pickle file.""" self.drop_not_picklable() - # 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() @@ -316,7 +305,7 @@ def sunssb(self): @property def telescope(self): - """Return telescope vector at all timestamps""" + """Return telescope name at all timestamps""" return self._telescope[self._isort] From 860e2b6223b1caa51001524a6d9ec5d58b3abd17 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Thu, 28 Sep 2023 13:41:28 +0200 Subject: [PATCH 3/8] Linting --- enterprise/pulsar.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index f8589db4..9217e5c4 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -311,7 +311,6 @@ def telescope(self): class PintPulsar(BasePulsar): def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): - self._sort = sort self.planets = planets self.name = model.PSR.value @@ -473,7 +472,6 @@ def __init__( par_name=None, tim_name=None, ): - self._sort = sort self.t2pulsar = t2pulsar self.planets = planets From 863a3887dff3c3fbc54a12a1b1d549504653b8a0 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Thu, 9 Nov 2023 14:17:09 +0100 Subject: [PATCH 4/8] Fixed typo --- enterprise/pulsar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index 9217e5c4..489a94eb 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -688,7 +688,7 @@ def Pulsar(*args, **kwargs): # Change directory to the base directory of the tim-file to deal with # INCLUDE statements in the tim-file os.chdir(dirname) - if timing_package.lower == "tempo2": + if timing_package.lower() == "tempo2": if t2 is None: raise ValueError("tempo2 requested but tempo2 is not available") # hack to set maxobs From 74ec30f989774ce7deecb872232559829f8bca81 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Thu, 9 Nov 2023 16:29:26 +0100 Subject: [PATCH 5/8] Added the necessary unit tests --- enterprise/pulsar.py | 14 +++++++++----- tests/test_pulsar.py | 43 +++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index 489a94eb..4a79c037 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -518,13 +518,17 @@ def __init__( self.sort_data() if drop_t2pulsar: - del self.t2pulsar + self.drop_tempopsr() else: if par_name is not None and os.path.exists(par_name): self.parfile = open(par_name).read() if tim_name is not None and os.path.exists(tim_name): self.timfile = open(tim_name).read() + def drop_tempopsr(self): + with contextlib.suppress(NameError): + del self.t2pulsar + # gather DM/DMX information if available def _set_dm(self, t2pulsar): pars = t2pulsar.pars(which="set") @@ -677,9 +681,9 @@ def Pulsar(*args, **kwargs): if timing_package is None: if t2 is not None: timing_package = "tempo2" - elif pint is not None: + elif pint is not None: # pragma: no cover timing_package = "pint" - else: + else: # pragma: no cover raise ValueError("No timing package available with which to load a pulsar") # get current directory @@ -689,7 +693,7 @@ def Pulsar(*args, **kwargs): # INCLUDE statements in the tim-file os.chdir(dirname) if timing_package.lower() == "tempo2": - if t2 is None: + if t2 is None: # pragma: no cover raise ValueError("tempo2 requested but tempo2 is not available") # hack to set maxobs maxobs = get_maxobs(reltimfile) + 100 @@ -703,7 +707,7 @@ def Pulsar(*args, **kwargs): tim_name=reltimfile, ) elif timing_package.lower() == "pint": - if pint is None: + if pint is None: # pragma: no cover raise ValueError("PINT requested but PINT is not available") if (clk is not None) and (bipm_version is None): bipm_version = clk.split("(")[1][:-1] diff --git a/tests/test_pulsar.py b/tests/test_pulsar.py index 6ab617e3..9fb1d0fc 100644 --- a/tests/test_pulsar.py +++ b/tests/test_pulsar.py @@ -25,18 +25,34 @@ from pint.models import get_model_and_toas +class TestTimingPackageExceptions(unittest.TestCase): + def test_unkown_timing_package(self): + # initialize Pulsar class + with self.assertRaises(ValueError): + self.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", timing_package='foobar') + + def test_clk_but_no_bipm(self): + self.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", clk='TT(BIPM2020)', timing_package='pint') + class TestPulsar(unittest.TestCase): @classmethod def setUpClass(cls): """Setup the Pulsar object.""" # initialize Pulsar class - cls.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim") + cls.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=True) + cls.psr_nodrop = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=False) @classmethod def tearDownClass(cls): shutil.rmtree("pickle_dir", ignore_errors=True) + def test_droppsr(self): + self.psr_nodrop.drop_tempopsr() + + with self.assertRaises(AttributeError): + _ = self.psr.t2pulsar + def test_residuals(self): """Check Residual shape.""" @@ -195,6 +211,14 @@ def setUpClass(cls): # initialize Pulsar class cls.psr = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", + datadir + "/B1855+09_NANOGrav_9yv1.tim", + ephem="DE430", + drop_pintpsr=True, + timing_package="pint", + ) + + cls.psr_nodrop = Pulsar( datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", ephem="DE430", @@ -202,6 +226,21 @@ def setUpClass(cls): timing_package="pint", ) + def test_droppsr(self): + self.psr_nodrop.drop_pintpsr() + + with self.assertRaises(AttributeError): + _ = self.psr_nodrop.model + + with self.assertRaises(AttributeError): + _ = self.psr_nodrop.parfile + + with self.assertRaises(AttributeError): + _ = self.psr_nodrop.pint_toas + + with self.assertRaises(AttributeError): + _ = self.psr_nodrop.timfile + def test_deflate_inflate(self): pass @@ -225,7 +264,7 @@ def test_no_planet(self): model, toas = get_model_and_toas( datadir + "/J0030+0451_NANOGrav_9yv1.gls.par", datadir + "/J0030+0451_NANOGrav_9yv1.tim", planets=False ) - Pulsar(model, toas, planets=True) + Pulsar(model, toas, planets=True, drop_pintpsr=False) msg = "obs_earth_pos is not in toas.table.colnames. Either " msg += "`planet` flag is not True in `toas` or further Pint " msg += "development to add additional planets is needed." From 8b7565068f14c929947cc7244c553909914fce63 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Thu, 9 Nov 2023 16:37:12 +0100 Subject: [PATCH 6/8] Linting --- enterprise/pulsar.py | 6 +++--- tests/test_pulsar.py | 22 ++++++++++++++++++---- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index 4a79c037..1af00e18 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -681,9 +681,9 @@ def Pulsar(*args, **kwargs): if timing_package is None: if t2 is not None: timing_package = "tempo2" - elif pint is not None: # pragma: no cover + elif pint is not None: # pragma: no cover timing_package = "pint" - else: # pragma: no cover + else: # pragma: no cover raise ValueError("No timing package available with which to load a pulsar") # get current directory @@ -707,7 +707,7 @@ def Pulsar(*args, **kwargs): tim_name=reltimfile, ) elif timing_package.lower() == "pint": - if pint is None: # pragma: no cover + if pint is None: # pragma: no cover raise ValueError("PINT requested but PINT is not available") if (clk is not None) and (bipm_version is None): bipm_version = clk.split("(")[1][:-1] diff --git a/tests/test_pulsar.py b/tests/test_pulsar.py index 9fb1d0fc..ce8cde49 100644 --- a/tests/test_pulsar.py +++ b/tests/test_pulsar.py @@ -29,10 +29,20 @@ class TestTimingPackageExceptions(unittest.TestCase): def test_unkown_timing_package(self): # initialize Pulsar class with self.assertRaises(ValueError): - self.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", timing_package='foobar') + self.psr = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", + datadir + "/B1855+09_NANOGrav_9yv1.tim", + timing_package="foobar", + ) def test_clk_but_no_bipm(self): - self.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", clk='TT(BIPM2020)', timing_package='pint') + self.psr = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", + datadir + "/B1855+09_NANOGrav_9yv1.tim", + clk="TT(BIPM2020)", + timing_package="pint", + ) + class TestPulsar(unittest.TestCase): @classmethod @@ -40,8 +50,12 @@ def setUpClass(cls): """Setup the Pulsar object.""" # initialize Pulsar class - cls.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=True) - cls.psr_nodrop = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=False) + cls.psr = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=True + ) + cls.psr_nodrop = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=False + ) @classmethod def tearDownClass(cls): From 4b94571dc3aaca4730909bd807736054b45b7fb9 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Thu, 9 Nov 2023 17:28:04 +0100 Subject: [PATCH 7/8] Added extra unit test --- tests/test_pulsar.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_pulsar.py b/tests/test_pulsar.py index ce8cde49..25bfc022 100644 --- a/tests/test_pulsar.py +++ b/tests/test_pulsar.py @@ -255,6 +255,15 @@ def test_droppsr(self): with self.assertRaises(AttributeError): _ = self.psr_nodrop.timfile + def test_drop_not_picklable(self): + self.psr_nodrop.drop_not_picklable() + + with self.assertRaises(AttributeError): + _ = self.psr_nodrop.model + + with self.assertRaises(AttributeError): + _ = self.psr_nodrop.pint_toas + def test_deflate_inflate(self): pass From 87813ee2dd43fabfbcb2a7c4b4d20d4914c019d1 Mon Sep 17 00:00:00 2001 From: Rutger van Haasteren Date: Thu, 9 Nov 2023 18:52:00 +0100 Subject: [PATCH 8/8] Updated the test_pulsar.py --- tests/test_pulsar.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/test_pulsar.py b/tests/test_pulsar.py index 25bfc022..6b490070 100644 --- a/tests/test_pulsar.py +++ b/tests/test_pulsar.py @@ -53,15 +53,16 @@ def setUpClass(cls): cls.psr = Pulsar( datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=True ) - cls.psr_nodrop = Pulsar( - datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=False - ) @classmethod def tearDownClass(cls): shutil.rmtree("pickle_dir", ignore_errors=True) def test_droppsr(self): + self.psr_nodrop = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", drop_t2pulsar=False + ) + self.psr_nodrop.drop_tempopsr() with self.assertRaises(AttributeError): @@ -232,7 +233,8 @@ def setUpClass(cls): timing_package="pint", ) - cls.psr_nodrop = Pulsar( + def test_droppsr(self): + self.psr_nodrop = Pulsar( datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim", ephem="DE430", @@ -240,7 +242,6 @@ def setUpClass(cls): timing_package="pint", ) - def test_droppsr(self): self.psr_nodrop.drop_pintpsr() with self.assertRaises(AttributeError): @@ -256,6 +257,14 @@ def test_droppsr(self): _ = self.psr_nodrop.timfile def test_drop_not_picklable(self): + self.psr_nodrop = Pulsar( + datadir + "/B1855+09_NANOGrav_9yv1.gls.par", + datadir + "/B1855+09_NANOGrav_9yv1.tim", + ephem="DE430", + drop_pintpsr=False, + timing_package="pint", + ) + self.psr_nodrop.drop_not_picklable() with self.assertRaises(AttributeError):