diff --git a/enterprise/pulsar.py b/enterprise/pulsar.py index 176df342..18a25cd0 100644 --- a/enterprise/pulsar.py +++ b/enterprise/pulsar.py @@ -23,9 +23,7 @@ try: import libstempo as t2 except ImportError: - logger.warning( - "libstempo not installed. Will use PINT instead." - ) # pragma: no cover + logger.warning("libstempo not installed. Will use PINT instead.") # pragma: no cover t2 = None try: @@ -34,9 +32,7 @@ from pint.residuals import Residuals as resids from pint.toa import TOAs except ImportError: - logger.warning( - "PINT not installed. Will use libstempo instead." - ) # pragma: no cover + logger.warning("PINT not installed. Will use libstempo instead.") # pragma: no cover pint = None if pint is None and t2 is None: @@ -136,9 +132,7 @@ def filter_data(self, start_time=None, end_time=None): if start_time is None and end_time is None: mask = np.ones(self._toas.shape, dtype=bool) else: - mask = np.logical_and( - self._toas >= start_time * 86400, self._toas <= end_time * 86400 - ) + mask = np.logical_and(self._toas >= start_time * 86400, self._toas <= end_time * 86400) self._toas = self._toas[mask] self._toaerrs = self._toaerrs[mask] @@ -244,11 +238,7 @@ def dmx(self): def flags(self): """Return a dictionary of tim-file flags.""" - flagnames = ( - self._flags.dtype.names - if isinstance(self._flags, np.ndarray) - else self._flags.keys() - ) + flagnames = self._flags.dtype.names if isinstance(self._flags, np.ndarray) else self._flags.keys() return {flag: self._flags[flag][self._isort] for flag in flagnames} @@ -271,25 +261,16 @@ def backend_flags(self): """ # collect flag names - flagnames = ( - self._flags.dtype.names - if isinstance(self._flags, np.ndarray) - else list(self._flags.keys()) - ) + flagnames = self._flags.dtype.names if isinstance(self._flags, np.ndarray) else list(self._flags.keys()) # allocate array with widest dtype - ret = np.zeros( - len(self._toas), dtype=max([self._flags[name].dtype for name in flagnames]) - ) + ret = np.zeros(len(self._toas), dtype=max([self._flags[name].dtype for name in flagnames])) # go through the flags in reverse order of preference # setting or replacing values for each TOA if "fe" in flagnames and "be" in flagnames: - ret[:] = [ - (a + "_" + b if (a and b) else "") - for a, b in zip(self._flags["fe"], self._flags["be"]) - ] + ret[:] = [(a + "_" + b if (a and b) else "") for a, b in zip(self._flags["fe"], self._flags["be"])] for flag in ["f", "i", "sys", "g", "group"]: if flag in flagnames: @@ -345,23 +326,17 @@ def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): # these are TDB but not barycentered # self._toas = np.array(toas.table["tdbld"], dtype="float64") * 86400 - self._toas = ( - np.array(model.get_barycentric_toas(toas).value, dtype="float64") * 86400 - ) + self._toas = np.array(model.get_barycentric_toas(toas).value, dtype="float64") * 86400 # saving also stoas (e.g., for DMX comparisons) 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._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._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 - ] + 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) @@ -393,9 +368,7 @@ def __init__(self, toas, model, sort=True, drop_pintpsr=True, planets=True): self._sunssb = self._get_sunssb(toas, model) which_astrometry = ( - "AstrometryEquatorial" - if "AstrometryEquatorial" in model.components - else "AstrometryEcliptic" + "AstrometryEquatorial" if "AstrometryEquatorial" in model.components else "AstrometryEcliptic" ) self._pos_t = ( @@ -415,9 +388,7 @@ def _set_dm(self, model): dmx = { par: { "DMX": float(model[par].value), - "DMXerr": None - if model[par].uncertainty_value is None - else float(model[par].uncertainty_value), + "DMXerr": None if model[par].uncertainty_value is None else float(model[par].uncertainty_value), "DMXR1": float(model[par[:3] + "R1" + par[3:]].value), "DMXR2": float(model[par[:3] + "R2" + par[3:]].value), "fit": par in pars, @@ -511,9 +482,7 @@ def __init__(self, t2pulsar, sort=True, drop_t2pulsar=True, planets=True): flags[key] = t2pulsar.flagvals(key) # 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()] - ) + 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 @@ -527,9 +496,7 @@ def __init__(self, t2pulsar, sort=True, drop_t2pulsar=True, planets=True): self._set_dm(t2pulsar) self._pos_t = t2pulsar.psrPos.copy() - if "ELONG" and "ELAT" in np.concatenate( - (t2pulsar.pars(which="fit"), t2pulsar.pars(which="set")) - ): + if "ELONG" and "ELAT" in np.concatenate((t2pulsar.pars(which="fit"), t2pulsar.pars(which="set"))): self._pos_t = utils.ecl2eq_vec(self._pos_t) self.sort_data() @@ -560,9 +527,7 @@ def _set_dm(self, t2pulsar): self._dmx = dmx def _get_radec(self, t2pulsar): - if "RAJ" in np.concatenate( - (t2pulsar.pars(which="fit"), t2pulsar.pars(which="set")) - ): + if "RAJ" in np.concatenate((t2pulsar.pars(which="fit"), t2pulsar.pars(which="set"))): return (np.double(t2pulsar["RAJ"].val), np.double(t2pulsar["DECJ"].val)) else: @@ -589,9 +554,7 @@ def _get_planetssb(self, t2pulsar): planetssb[:, 7, :] = self.t2pulsar.neptune_ssb planetssb[:, 8, :] = self.t2pulsar.pluto_ssb - if "ELONG" and "ELAT" in np.concatenate( - (t2pulsar.pars(), t2pulsar.pars(which="set")) - ): + if "ELONG" and "ELAT" in np.concatenate((t2pulsar.pars(), t2pulsar.pars(which="set"))): for ii in range(9): planetssb[:, ii, :3] = utils.ecl2eq_vec(planetssb[:, ii, :3]) planetssb[:, ii, 3:] = utils.ecl2eq_vec(planetssb[:, ii, 3:]) @@ -607,9 +570,7 @@ def _get_sunssb(self, t2pulsar): sunssb = np.zeros((len(self._toas), 6)) sunssb[:, :] = self.t2pulsar.sun_ssb - if "ELONG" and "ELAT" in np.concatenate( - (t2pulsar.pars(), t2pulsar.pars(which="set")) - ): + if "ELONG" and "ELAT" in np.concatenate((t2pulsar.pars(), t2pulsar.pars(which="set"))): sunssb[:, :3] = utils.ecl2eq_vec(sunssb[:, :3]) sunssb[:, 3:] = utils.ecl2eq_vec(sunssb[:, 3:]) return sunssb @@ -666,24 +627,16 @@ def Pulsar(*args, **kwargs): t2pulsar = [x for x in args if isinstance(x, t2.tempopulsar)] parfile = [x for x in args if isinstance(x, str) and x.split(".")[-1] == "par"] - timfile = [ - x for x in args if isinstance(x, str) and x.split(".")[-1] in ["tim", "toa"] - ] + 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, drop_pintpsr=drop_pintpsr, 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 - ) + return Tempo2Pulsar(t2pulsar[0], sort=sort, drop_t2pulsar=drop_t2pulsar, planets=planets) elif parfile and timfile: # Check whether the two files exist if not os.path.isfile(parfile[0]) or not os.path.isfile(timfile[0]): - msg = "Cannot find parfile {0} or timfile {1}!".format( - parfile[0], timfile[0] - ) + msg = "Cannot find parfile {0} or timfile {1}!".format(parfile[0], timfile[0]) raise IOError(msg) # Obtain the directory name of the timfile, and change to it @@ -710,19 +663,13 @@ def Pulsar(*args, **kwargs): planets=planets, ) os.chdir(cwd) - return PintPulsar( - toas, model, sort=sort, drop_pintpsr=drop_pintpsr, planets=planets - ) + 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 - ) + t2pulsar = t2.tempopulsar(relparfile, reltimfile, maxobs=maxobs, ephem=ephem, clk=clk) os.chdir(cwd) - return Tempo2Pulsar( - t2pulsar, sort=sort, drop_t2pulsar=drop_t2pulsar, planets=planets - ) + return Tempo2Pulsar(t2pulsar, sort=sort, drop_t2pulsar=drop_t2pulsar, planets=planets) raise ValueError("Unknown arguments {}".format(args))