def test_phase_commands(self): model = pint.models.get_model(parfile) toas = pint.toa.get_TOAs(timfile) # This TOA has PHASE -0.3. Check that assert np.isclose(toas.table[32]["flags"]["phase"], -0.3) # This TOA should have PHASE 0.2 and -padd -0.2 assert np.isclose(toas.table["flags"][9]["padd"], -0.2) assert np.isclose(toas.table["flags"][9]["phase"], 0.2) # The end result should be these residuals if the commands are respected res = Residuals(toas=toas, model=model) assert (res.rms_weighted() - 1602.0293 * u.us) < 0.1 * u.us
class Pulsar: """Wrapper class for a pulsar. Contains the toas, model, residuals, and fitter """ def __init__(self, parfile=None, timfile=None, ephem=None, fitter="GLSFitter"): super(Pulsar, self).__init__() log.info(f"Loading pulsar parfile: {str(parfile)}") if parfile is None or timfile is None: raise ValueError("No valid pulsar model and/or TOAs to load") self.parfile = parfile self.timfile = timfile self.prefit_model = pint.models.get_model(self.parfile) if ephem is not None: log.info( f"Overriding model ephemeris {self.prefit_model.EPHEM.value} with {ephem}" ) self.prefit_model.EPHEM.value = ephem self.all_toas = get_TOAs(self.timfile, model=self.prefit_model, usepickle=True) # Make sure that if we used a model, that any phase jumps from # the parfile have their flags updated in the TOA table if "PhaseJump" in self.prefit_model.components: self.prefit_model.jump_params_to_flags(self.all_toas) # turns pre-existing jump flags in toas.table['flags'] into parameters in parfile self.prefit_model.jump_flags_to_params(self.all_toas) self.selected_toas = copy.deepcopy(self.all_toas) print("The prefit model as a parfile:") print(self.prefit_model.as_parfile()) # adds extra prefix params for fitting self.add_model_params() self.all_toas.print_summary() self.prefit_resids = Residuals(self.selected_toas, self.prefit_model) print( "RMS pre-fit PINT residuals are %.3f us\n" % self.prefit_resids.rms_weighted().to(u.us).value ) # Set of indices from original list that are deleted # We use indices because of the grouping of TOAs by observatory self.deleted = set([]) self.fit_method = fitter self.fitter = None self.fitted = False self.stashed = None # for temporarily stashing some TOAs self.faketoas1 = None # for random models self.faketoas = None # for random models self.use_pulse_numbers = False @property def name(self): return getattr(self.prefit_model, "PSR").value def __getitem__(self, key): try: return getattr(self.prefit_model, key) except AttributeError: log.error(f"Parameter {key} was not found in pulsar model {self.name}") return None def __contains__(self, key): return key in self.prefit_model.params def reset_model(self): self.prefit_model = pint.models.get_model(self.parfile) self.add_model_params() self.postfit_model = None self.postfit_resids = None self.fitted = False self.update_resids() def reset_TOAs(self): self.all_toas = get_TOAs(self.timfile, model=self.prefit_model, usepickle=True) # Make sure that if we used a model, that any phase jumps from # the parfile have their flags updated in the TOA table if "PhaseJump" in self.prefit_model.components: self.prefit_model.jump_params_to_flags(self.all_toas) # turns pre-existing jump flags in toas.table['flags'] into parameters in parfile self.prefit_model.jump_flags_to_params(self.all_toas) self.selected_toas = copy.deepcopy(self.all_toas) self.deleted = set([]) self.stashed = None self.update_resids() def resetAll(self): self.prefit_model = pint.models.get_model(self.parfile) self.postfit_model = None self.postfit_resids = None self.fitted = False self.use_pulse_numbers = False self.reset_TOAs() def _delete_TOAs(self, toa_table): toa_table.group_by("index") del_inds = np.zeros(len(toa_table), dtype=bool) # Set the indices of del_inds to true where the TOA indices are for ii in self.deleted: # There must be a better way di = np.where(toa_table["index"] == ii)[0] if len(di): del_inds[di[0]] |= 1 if del_inds.sum() < len(toa_table): return toa_table[~del_inds].group_by("obs") else: return None def delete_TOAs(self, indices, selected): # note: indices should be a list or an array self.deleted |= set(indices) # update the deleted indices if selected is not None: self.selected_toas.table = self._delete_TOAs(self.selected_toas.table) # Now delete from all_toas self.all_toas.table = self._delete_TOAs(self.all_toas.table) if self.selected_toas.table is None: # all selected were deleted self.selected_toas = copy.deepcopy(self.all_toas) selected = np.zeros(self.selected_toas.ntoas, dtype=bool) else: # Make a new selected list by adding a value if the table # index at that position is not in the new indices to # delete, with a value that is the same as the previous # selected array newselected = [ sel for idx, sel in zip(self.all_toas.table["index"], selected) if idx not in indices ] selected = np.asarray(newselected, dtype=bool) self.selected_toas = self.all_toas[selected] # delete the TOAs from the stashed list also if self.stashed: self.stashed.table = self._delete_TOAs(self.stashed.table) return selected def update_resids(self): # update the pre and post fit residuals using all_toas track_mode = "use_pulse_numbers" if self.use_pulse_numbers else None self.prefit_resids = Residuals( self.all_toas, self.prefit_model, track_mode=track_mode ) if self.fitted: self.postfit_resids = Residuals( self.all_toas, self.postfit_model, track_mode=track_mode ) def orbitalphase(self): """ For a binary pulsar, calculate the orbital phase. Otherwise, return an array of unitless quantities of zeros """ if not self.prefit_model.is_binary: log.warning("This is not a binary pulsar") return u.Quantity(np.zeros(self.all_toas.ntoas)) toas = self.all_toas if self.fitted: phase = self.postfit_model.orbital_phase(toas, anom="mean") else: phase = self.prefit_model.orbital_phase(toas, anom="mean") return phase / (2 * np.pi * u.rad) def dayofyear(self): """ Return the day of the year for all the TOAs of this pulsar """ t = Time(self.all_toas.get_mjds(), format="mjd") year = Time(np.floor(t.decimalyear), format="decimalyear") return np.asarray(t.mjd - year.mjd) << u.day def year(self): """ Return the decimal year for all the TOAs of this pulsar """ t = Time(self.all_toas.get_mjds(), format="mjd") return np.asarray(t.decimalyear) << u.year def add_model_params(self): """This automatically adds the next available unfit prefix parameters to the model so they show up on the GUI """ m = self.prefit_model # Add next spin freq deriv if "Spindown" in m.components: c = m.components["Spindown"] n = len(c.get_prefix_mapping_component("F")) if f"F{n-1}" in m.free_params and not hasattr(m, f"F{n}"): c.add_param(m.F0.new_param(n), setup=True) log.debug(f"Adding F{n} to prefit model") p = getattr(m, f"F{n}") p.quantity = 0.0 * p.units p.frozen = True # Add next orbital freq deriv if "BinaryBT" in m.components: c = m.components["BinaryBT"] n = len(c.get_prefix_mapping_component("FB")) if f"FB{n-1}" in m.free_params and not hasattr(m, f"FB{n}"): c.add_param(m.FB0.new_param(n), setup=True) log.debug(f"Adding FB{n} to prefit model") p = getattr(m, f"FB{n}") p.quantity = 0.0 * p.units p.frozen = True # Add dispersion expansion component if "DispersionDM" in m.components: c = m.components["DispersionDM"] n = len(c.get_prefix_mapping_component("DM")) + 1 # DM1 is always added, but might be unset if n == 2 and m.DM1.value is None: p = m.DM1 p.quantity = 0.0 * p.units p.frozen = True if f"DM{n-1}" in m.free_params and not hasattr(m, f"DM{n}"): c.add_param(m.DM1.new_param(n), setup=True) log.debug(f"Adding DM{n} to prefit model") p = getattr(m, f"DM{n}") p.quantity = 0.0 * p.units p.frozen = True m.setup() # Not sure if this is necessary m.validate() def write_fit_summary(self): """ Summarize fitting results """ if self.fitted: wrms = self.postfit_resids.rms_weighted() print("Post-Fit Chi2: %.8g" % self.postfit_resids.chi2) print("Post-Fit DOF: %8d" % self.postfit_resids.dof) print("Post-Fit Reduced-Chi2: %.8g" % self.postfit_resids.reduced_chi2) print("Post-Fit Weighted RMS: %.8g us" % wrms.to(u.us).value) print("------------------------------------") print( "%19s %24s\t%24s\t%16s %16s %16s" % ( "Parameter", "Pre-Fit", "Post-Fit", "Uncertainty", "Difference", "Diff/Unc", ) ) print("-" * 132) for key in self.prefit_model.free_params: line = "%8s " % key pre = getattr(self.prefit_model, key) post = getattr(self.postfit_model, key) line += "%10s " % ("" if post.units is None else str(post.units)) if post.quantity is not None: line += "%24s\t" % pre.str_quantity(pre.quantity) line += "%24s\t" % post.str_quantity(post.quantity) try: line += "%16.8g " % post.uncertainty.value except: line += "%18s" % "" diff = post.value - pre.value line += "%16.8g " % diff if pre.uncertainty is not None and pre.uncertainty.value != 0.0: line += "%16.8g" % (diff / pre.uncertainty.value) print(line) else: log.warning("Pulsar has not been fitted yet!") def add_phase_wrap(self, selected, phase): """ Add a phase wrap to selected points in the TOAs object Turn on pulse number tracking in the model, if it isn't already :param selected: boolean array to apply to toas, True = selected toa :param phase: phase difference to be added, i.e. -0.5, +2, etc. """ # Check if pulse numbers are in table already, if not, make the column if ( "pulse_number" not in self.all_toas.table.colnames or "pulse_number" not in self.selected_toas.table.colnames ): if self.fitted: self.all_toas.compute_pulse_numbers(self.postfit_model) self.selected_toas.compute_pulse_numbers(self.postfit_model) else: self.all_toas.compute_pulse_numbers(self.prefit_model) self.selected_toas.compute_pulse_numbers(self.prefit_model) if ( "delta_pulse_number" not in self.all_toas.table.colnames or "delta_pulse_number" not in self.selected_toas.table.colnames ): self.all_toas.table["delta_pulse_number"] = np.zeros(self.all_toas.ntoas) self.selected_toas.table["delta_pulse_number"] = np.zeros( self.selected_toas.ntoas ) # add phase wrap and update self.all_toas.table["delta_pulse_number"][selected] += phase self.use_pulse_numbers = True self.update_resids() def add_jump(self, selected): """ jump the toas selected or un-jump them if already jumped :param selected: boolean array to apply to toas, True = selected toa """ # TODO: split into two functions if "PhaseJump" not in self.prefit_model.components: # if no PhaseJump component, add one log.info("PhaseJump component added") a = pint.models.jump.PhaseJump() a.setup() self.prefit_model.add_component(a) retval = self.prefit_model.add_jump_and_flags( self.all_toas.table["flags"][selected] ) if self.fitted: self.postfit_model.add_component(a) return retval # if gets here, has at least one jump param already # and iif it doesn't overlap or cancel, add the param numjumps = self.prefit_model.components["PhaseJump"].get_number_of_jumps() if numjumps == 0: log.warning( "There are no jumps (maskParameter objects) in PhaseJump. Please delete the PhaseJump object and try again. " ) return None # delete jump if perfectly overlaps any existing jump for num in range(1, numjumps + 1): # create boolean array corresponding to TOAs to be jumped toas_jumped = [ "jump" in dict.keys() and str(num) in dict["jump"] for dict in self.all_toas.table["flags"] ] if np.array_equal(toas_jumped, selected): # if current jump exactly matches selected, remove it self.prefit_model.delete_jump_and_flags( self.all_toas.table["flags"], num ) if self.fitted: self.postfit_model.delete_jump_and_flags(None, num) log.info("removed param", f"JUMP{str(num)}") return toas_jumped # if here, then doesn't match anything # add jump flags to selected TOAs at their perspective indices in the TOA tables retval = self.prefit_model.add_jump_and_flags( self.all_toas.table["flags"][selected] ) if ( self.fitted and self.prefit_model.components["PhaseJump"] != self.postfit_model.components["PhaseJump"] ): param = self.prefit_model.components[ "PhaseJump" ].get_jump_param_objects() # array of jump objects self.postfit_model.add_param_from_top( param[-1], "PhaseJump" ) # add last (newest) jump getattr(self.postfit_model, param[-1].name).frozen = False self.postfit_model.components["PhaseJump"].setup() return retval def fit(self, selected, iters=4, compute_random=False): """ Run a fit using the specified fitter """ # Select all the TOAs if none are explicitly set if not any(selected): selected = ~selected if self.fitted: self.prefit_model = self.postfit_model self.prefit_resids = self.postfit_resids self.add_model_params() # Have to change the fitter for each fit since TOAs and models change log.info(f"Using {self.fit_method}") self.fitter = getattr(pint.fitter, self.fit_method)( self.selected_toas, self.prefit_model ) wrms = self.prefit_resids.rms_weighted() print("\n------------------------------------") print(" Pre-Fit Chi2: %.8g" % self.prefit_resids.chi2) print(" Pre-Fit reduced-Chi2: %.8g" % self.prefit_resids.reduced_chi2) print(" Pre-Fit Weighted RMS: %.8g us" % wrms.to(u.us).value) print("------------------------------------") # Do the actual fit and mark things as being fit self.fitter.fit_toas(maxiter=iters) self.fitter.update_model() self.postfit_model = self.fitter.model self.fitted = True # Zero out all of the "delta_pulse_numbers" if they are set self.all_toas.table["delta_pulse_number"] = np.zeros(self.all_toas.ntoas) self.selected_toas.table["delta_pulse_number"] = np.zeros( self.selected_toas.ntoas ) # Re-calculate the pulse numbers here self.all_toas.compute_pulse_numbers(self.postfit_model) self.selected_toas.compute_pulse_numbers(self.postfit_model) # Now compute the residuals using correct pulse numbers self.postfit_resids = Residuals(self.all_toas, self.postfit_model) # Need this since it isn't updated using self.fitter.update_model() self.fitter.model.CHI2.value = self.postfit_resids.chi2 # And print the summary self.write_fit_summary() # Check to see if we should calculate an F-test if ( hasattr(self, "lastfit") and (len(self.postfit_model.free_params) > len(self.lastfit["free_params"])) and (self.lastfit["ntoas"] == self.fitter.toas.ntoas) ): prob = FTest( self.lastfit["chi2"], self.lastfit["dof"], self.postfit_resids.chi2, self.postfit_resids.dof, ) new_params = set(self.postfit_model.free_params) - set( self.lastfit["free_params"] ) print( f"F-test comparing post- to pre-fit models for addition of {new_params}:\n" f" P = {prob:.3g} that the improvement is due to noise." ) # plot the prefit without jumps pm_no_jumps = copy.deepcopy(self.prefit_model) for param in pm_no_jumps.params: if param.startswith("JUMP"): getattr(pm_no_jumps, param).value = 0.0 getattr(pm_no_jumps, param).frozen = True self.prefit_resids_no_jumps = Residuals(self.all_toas, pm_no_jumps) # Store some key params for possible F-testing self.lastfit = { "free_params": self.fitter.model.free_params, "dof": self.postfit_resids.dof, "chi2": self.postfit_resids.chi2, "ntoas": self.fitter.toas.ntoas, } # adds extra prefix params for fitting self.add_model_params() def random_models(self, selected): """Compute and plot random models """ log.info("Computing random models based on parameter covariance matrix.") if [p for p in self.postfit_model.free_params if p.startswith("DM")]: log.warning( "Fitting for DM while using random models can cause strange behavior." ) # These are the currently selected TOAs in the fit sim_sel = copy.deepcopy(self.selected_toas) # These are single TOAs from each cluster of TOAs inds = np.zeros(sim_sel.ntoas, dtype=bool) inds[np.unique(sim_sel.get_clusters(), return_index=True)[1]] |= True sim_sel = sim_sel[inds] # Get the range of MJDs we are using in the fit mjds = sim_sel.get_mjds().value minselMJD, maxselMJD = mjds.min(), mjds.max() extra = 0.1 # Fraction beyond TOAs to plot or calculate random models if self.faketoas1 is None: mjds = self.all_toas.get_mjds().value minallMJD, maxallMJD = mjds.min(), mjds.max() spanMJD = maxallMJD - minallMJD # want roughly 1 per day up to 3 years if spanMJD < 1000: Ntoas = min(400, int(spanMJD)) elif spanMJD < 4000: Ntoas = min(750, int(spanMJD) // 2) else: Ntoas = min(1500, int(spanMJD) // 4) log.debug(f"Generating {Ntoas} fake TOAs for the random models") # By default we will use TOAs from the TopoCenter. This gets done only once. self.faketoas1 = make_fake_toas_uniform( minallMJD - extra * spanMJD, maxallMJD + extra * spanMJD, Ntoas, self.postfit_model, obs="coe", freq=1 * u.THz, # effectively infinite frequency include_bipm=sim_sel.clock_corr_info["include_bipm"], include_gps=sim_sel.clock_corr_info["include_gps"], ) self.faketoas1.compute_pulse_numbers(self.postfit_model) # Combine our TOAs toas = merge_TOAs([sim_sel, self.faketoas1]) zero_residuals(toas, self.postfit_model) toas.table.sort("tdbld") # for plotting # Get a selection array to select the non-fake TOAs refs = np.asarray(toas.get_flag_value("name")[0]) != "fake" # Compute the new random timing models rs = calculate_random_models( self.fitter, toas, Nmodels=15, keep_models=False, return_time=True, ) # Get a selection array for the fake TOAs that covers the fit TOAs (plus extra) mjds = toas.get_mjds().value spanMJD = maxselMJD - minselMJD toplot = np.bitwise_and( mjds > (minselMJD - extra * spanMJD), mjds < (maxselMJD + extra * spanMJD) ) # This is the mean of the reference resids for the selected TOAs if selected.sum(): # shorthand for having some selected ref_mean = self.postfit_resids.time_resids[selected][inds].mean() else: ref_mean = self.postfit_resids.time_resids[inds].mean() # This is the means of the corresponding resids from the random models ran_mean = rs[:, refs].mean(axis=1) # Now adjust each set of random resids so that the ran_mean == ref_mean rs -= ran_mean[:, np.newaxis] rs += ref_mean # And store the key things for plotting self.faketoas = toas self.random_resids = rs
class Pulsar: """Wrapper class for a pulsar. Contains the toas, model, residuals, and fitter """ def __init__(self, parfile=None, timfile=None, ephem=None): super(Pulsar, self).__init__() log.info("STARTING LOADING OF PULSAR %s" % str(parfile)) if parfile is not None and timfile is not None: self.parfile = parfile self.timfile = timfile else: raise ValueError("No valid pulsar to load") self.prefit_model = pint.models.get_model(self.parfile) if ephem is not None: # TODO: EPHEM overwrite message? self.all_toas = get_TOAs(self.timfile, ephem=ephem, planets=True) self.prefit_model.EPHEM.value = ephem elif getattr(self.prefit_model, "EPHEM").value is not None: self.all_toas = get_TOAs( self.timfile, ephem=self.prefit_model.EPHEM.value, planets=True ) else: self.all_toas = get_TOAs(self.timfile, planets=True) # turns pre-existing jump flags in toas.table['flags'] into parameters in parfile # TODO: fix jump_flags_to_params self.prefit_model.jump_flags_to_params(self.all_toas) # adds flags to toas.table for existing jump parameters from .par file if "PhaseJump" in self.prefit_model.components: self.prefit_model.jump_params_to_flags(self.all_toas) self.selected_toas = copy.deepcopy(self.all_toas) print("prefit_model.as_parfile():") print(self.prefit_model.as_parfile()) self.all_toas.print_summary() self.prefit_resids = Residuals(self.selected_toas, self.prefit_model) print( "RMS PINT residuals are %.3f us\n" % self.prefit_resids.rms_weighted().to(u.us).value ) self.fitter = Fitters.WLS self.fitted = False self.use_pulse_numbers = False @property def name(self): return getattr(self.prefit_model, "PSR").value def __getitem__(self, key): try: return getattr(self.prefit_model, key) except AttributeError: log.error( "Parameter %s was not found in pulsar model %s" % (key, self.name) ) return None def __contains__(self, key): return key in self.prefit_model.params def reset_model(self): self.prefit_model = pint.models.get_model(self.parfile) self.postfit_model = None self.postfit_resids = None self.fitted = False self.update_resids() def reset_TOAs(self): if getattr(self.prefit_model, "EPHEM").value is not None: self.all_toas = get_TOAs( self.timfile, ephem=self.prefit_model.EPHEM.value, planets=True ) else: self.all_toas = get_TOAs(self.timfile, planets=True) self.selected_toas = copy.deepcopy(self.all_toas) self.update_resids() def resetAll(self): self.prefit_model = pint.models.get_model(self.parfile) self.postfit_model = None self.postfit_resids = None self.fitted = False self.use_pulse_numbers = False self.reset_TOAs() def update_resids(self): # update the pre and post fit residuals using all_toas if self.use_pulse_numbers: self.prefit_resids = Residuals( self.all_toas, self.prefit_model, track_mode="use_pulse_numbers" ) if self.fitted: self.postfit_resids = Residuals( self.all_toas, self.postfit_model, track_mode="use_pulse_numbers" ) else: self.prefit_resids = Residuals(self.all_toas, self.prefit_model) if self.fitted: self.postfit_resids = Residuals(self.all_toas, self.postfit_model) def orbitalphase(self): """ For a binary pulsar, calculate the orbital phase. Otherwise, return an array of unitless quantities of zeros """ if not self.prefit_model.is_binary: log.warn("This is not a binary pulsar") return u.Quantity(np.zeros(self.selected_toas.ntoas)) toas = self.selected_toas if self.fitted: phase = self.postfit_model.orbital_phase(toas, anom="mean") else: phase = self.prefit_model.orbital_phase(toas, anom="mean") return phase / (2 * np.pi * u.rad) def dayofyear(self): """ Return the day of the year for all the TOAs of this pulsar """ t = Time(self.selected_toas.get_mjds(), format="mjd") year = Time(np.floor(t.decimalyear), format="decimalyear") return (t.mjd - year.mjd) * u.day def year(self): """ Return the decimal year for all the TOAs of this pulsar """ t = Time(self.selected_toas.get_mjds(), format="mjd") return (t.decimalyear) * u.year def write_fit_summary(self): """ Summarize fitting results """ if self.fitted: chi2 = self.postfit_resids.chi2 wrms = self.postfit_resids.rms_weighted() print("Post-Fit Chi2:\t\t%.8g" % chi2) print("Post-Fit Weighted RMS:\t%.8g us" % wrms.to(u.us).value) print( "%19s %24s\t%24s\t%16s %16s %16s" % ( "Parameter", "Pre-Fit", "Post-Fit", "Uncertainty", "Difference", "Diff/Unc", ) ) print("-" * 132) fitparams = [ p for p in self.prefit_model.params if not getattr(self.prefit_model, p).frozen ] for key in fitparams: line = "%8s " % key pre = getattr(self.prefit_model, key) post = getattr(self.postfit_model, key) line += "%10s " % ("" if post.units is None else str(post.units)) if post.quantity is not None: line += "%24s\t" % pre.str_quantity(pre.quantity) line += "%24s\t" % post.str_quantity(post.quantity) try: line += "%16.8g " % post.uncertainty.value except: line += "%18s" % "" try: diff = post.value - pre.value line += "%16.8g " % diff if pre.uncertainty is not None: line += "%16.8g" % (diff / pre.uncertainty.value) except: pass print(line) else: log.warn("Pulsar has not been fitted yet!") def add_phase_wrap(self, selected, phase): """ Add a phase wrap to selected points in the TOAs object Turn on pulse number tracking in the model, if it isn't already :param selected: boolean array to apply to toas, True = selected toa :param phase: phase diffeence to be added, i.e. -0.5, +2, etc. """ # Check if pulse numbers are in table already, if not, make the column if ( "pulse_number" not in self.all_toas.table.colnames or "pulse_number" not in self.selected_toas.table.colnames ): if self.fitted: self.all_toas.compute_pulse_numbers(self.postfit_model) self.selected_toas.compute_pulse_numbers(self.postfit_model) else: self.all_toas.compute_pulse_numbers(self.prefit_model) self.selected_toas.compute_pulse_numbers(self.prefit_model) if ( "delta_pulse_number" not in self.all_toas.table.colnames or "delta_pulse_number" not in self.selected_toas.table.colnames ): self.all_toas.table["delta_pulse_number"] = np.zeros( len(self.all_toas.get_mjds()) ) self.selected_toas.table["delta_pulse_number"] = np.zeros( len(self.selected_toas.get_mjds()) ) # add phase wrap self.all_toas.table["delta_pulse_number"][selected] += phase self.selected_toas.table["delta_pulse_number"] += phase self.use_pulse_numbers = True self.update_resids() def add_jump(self, selected): """ jump the toas selected or unjump them if already jumped :param selected: boolean array to apply to toas, True = selected toa """ # TODO: split into two functions if "PhaseJump" not in self.prefit_model.components: # if no PhaseJump component, add one log.info("PhaseJump component added") a = pint.models.jump.PhaseJump() a.setup() self.prefit_model.add_component(a) retval = self.prefit_model.add_jump_and_flags( self.all_toas.table["flags"][selected] ) if self.fitted: self.postfit_model.add_component(a) return retval # if gets here, has at least one jump param already # if doesnt overlap or cancel, add the param numjumps = self.prefit_model.components["PhaseJump"].get_number_of_jumps() if numjumps == 0: log.warn( "There are no jumps (maskParameter objects) in PhaseJump. Please delete the PhaseJump object and try again. " ) return None # delete jump if perfectly overlaps any existing jump for num in range(1, numjumps + 1): # create boolean array corresponding to TOAs to be jumped toas_jumped = [ "jump" in dict.keys() and str(num) in dict["jump"] for dict in self.all_toas.table["flags"] ] if np.array_equal(toas_jumped, selected): # if current jump exactly matches selected, remove it self.prefit_model.delete_jump_and_flags( self.all_toas.table["flags"], num ) if self.fitted: self.postfit_model.delete_jump_and_flags(None, num) log.info("removed param", "JUMP" + str(num)) return toas_jumped # if here, then doesn't match anything # add jump flags to selected TOAs at their perspective indices in the TOA tables retval = self.prefit_model.add_jump_and_flags( self.all_toas.table["flags"][selected] ) if ( self.fitted and not self.prefit_model.components["PhaseJump"] == self.postfit_model.components["PhaseJump"] ): param = self.prefit_model.components[ "PhaseJump" ].get_jump_param_objects() # array of jump objects self.postfit_model.add_param_from_top( param[-1], "PhaseJump" ) # add last (newest) jump getattr(self.postfit_model, param[-1].name).frozen = False self.postfit_model.components["PhaseJump"].setup() return retval def fit(self, selected, iters=1): """ Run a fit using the specified fitter """ # Select all the TOAs if none are explicitly set if not any(selected): selected = ~selected if self.fitted: self.prefit_model = self.postfit_model self.prefit_resids = self.postfit_resids if self.fitter == Fitters.POWELL: fitter = pint.fitter.PowellFitter(self.selected_toas, self.prefit_model) elif self.fitter == Fitters.WLS: fitter = pint.fitter.WLSFitter(self.selected_toas, self.prefit_model) elif self.fitter == Fitters.GLS: fitter = pint.fitter.GLSFitter(self.selected_toas, self.prefit_model) chi2 = self.prefit_resids.chi2 wrms = self.prefit_resids.rms_weighted() print("Pre-Fit Chi2:\t\t%.8g" % chi2) print("Pre-Fit Weighted RMS:\t%.8g us" % wrms.to(u.us).value) fitter.fit_toas(maxiter=1) self.postfit_model = fitter.model self.postfit_resids = Residuals(self.all_toas, self.postfit_model) self.fitted = True self.write_fit_summary() # TODO: delta_pulse_numbers need some work. They serve both for PHASE and -padd functions from the TOAs # as well as for phase jumps added manually in the GUI. They really should not be zeroed out here because # that will wipe out preexisting values self.all_toas.table["delta_pulse_numbers"] = np.zeros(self.all_toas.ntoas) self.selected_toas.table["delta_pulse_number"] = np.zeros( self.selected_toas.ntoas ) # plot the prefit without jumps pm_no_jumps = copy.deepcopy(self.prefit_model) for param in pm_no_jumps.params: if param.startswith("JUMP"): getattr(pm_no_jumps, param).value = 0.0 getattr(pm_no_jumps, param).frozen = True self.prefit_resids_no_jumps = Residuals(self.selected_toas, pm_no_jumps) f = copy.deepcopy(fitter) no_jumps = [ False if "jump" in dict.keys() else True for dict in f.toas.table["flags"] ] f.toas.select(no_jumps) selectedMJDs = self.selected_toas.get_mjds() if all(no_jumps): q = list(self.all_toas.get_mjds()) index = q.index( [i for i in self.all_toas.get_mjds() if i > selectedMJDs.min()][0] ) rs_mean = ( Residuals(self.all_toas, f.model) .phase_resids[index : index + len(selectedMJDs)] .mean() ) else: rs_mean = self.prefit_resids_no_jumps.phase_resids[no_jumps].mean() # determines how far on either side fake toas go # TODO: hard limit on how far fake toas can go --> can get clkcorr # errors if go before GBT existed, etc. minMJD, maxMJD = selectedMJDs.min(), selectedMJDs.max() spanMJDs = maxMJD - minMJD if spanMJDs < 30 * u.d: redge = ledge = 4 npoints = 400 elif spanMJDs < 90 * u.d: redge = ledge = 2 npoints = 300 elif spanMJDs < 200 * u.d: redge = ledge = 1 npoints = 300 elif spanMJDs < 400 * u.d: redge = ledge = 0.5 npoints = 200 else: redge = ledge = 1.0 npoints = 250 # Check to see if too recent nowish = (Time.now().mjd - 40) * u.d if maxMJD + spanMJDs * redge > nowish: redge = (nowish - maxMJD) / spanMJDs if redge < 0.0: redge = 0.0 f_toas = make_fake_toas_uniform( minMJD - spanMJDs * ledge, maxMJD + spanMJDs * redge, npoints, f.model ) rs = calculate_random_models( f, f_toas, Nmodels=10, keep_models=False, return_time=True ) # subtract the mean residual of each random model from the respective residual set # based ONLY on the mean of the random residuals in the real data range start_index = np.where(abs(f_toas.get_mjds() - minMJD) < 1 * u.d) end_index = np.where(abs(f_toas.get_mjds() - maxMJD) < 1 * u.d) for i in range(len(rs)): # use start_index[0][0] since np.where returns np.array([], dtype), extract index from list in array rs_mean = rs[i][start_index[0][0] : end_index[0][0]].mean() rs[i][:] = [resid - rs_mean for resid in rs[i]] self.random_resids = rs self.fake_toas = f_toas def fake_year(self): """ Function to support plotting of random models on multiple x-axes. Return the decimal year for all the TOAs of this pulsar """ t = Time(self.fake_toas.get_mjds(), format="mjd") return (t.decimalyear) * u.year
class Fitter(object): """ Base class for fitter. The fitting function should be defined as the fit_toas() method. Note that the Fitter object makes a deepcopy of the model, so changes to the model will not be noticed after the Fitter has been instantiated! Use Fitter.model instead. The Fitter also caches a copy of the original model so it can be restored with reset_model() Parameters ---------- toas : a pint TOAs instance The input toas. model : a pint timing model instance The initial timing model for fitting. """ def __init__(self, toas, model, residuals=None): self.toas = toas self.model_init = model if residuals is None: self.resids_init = Residuals(toas=toas, model=model) self.reset_model() else: # residuals were provided, we're just going to use them # probably using GLSFitter to compute a chi-squared self.model = copy.deepcopy(self.model_init) self.resids = residuals self.fitresult = [] self.method = None def reset_model(self): """Reset the current model to the initial model.""" self.model = copy.deepcopy(self.model_init) self.update_resids() self.fitresult = [] def update_resids(self, set_pulse_nums=False): """Update the residuals. Run after updating a model parameter.""" self.resids = Residuals(toas=self.toas, model=self.model, set_pulse_nums=set_pulse_nums) def set_fitparams(self, *params): """Update the "frozen" attribute of model parameters. Ex. fitter.set_fitparams('F0','F1') """ fit_params_name = [] for pn in params: if pn in self.model.params: fit_params_name.append(pn) else: rn = self.model.match_param_aliases(pn) if rn != "": fit_params_name.append(rn) for p in self.model.params: getattr(self.model, p).frozen = p not in fit_params_name def get_allparams(self): """Return a dict of all param names and values.""" return collections.OrderedDict((k, getattr(self.model, k).quantity) for k in self.model.params_ordered) def get_fitparams(self): """Return a dict of fittable param names and quantity.""" return collections.OrderedDict((k, getattr(self.model, k)) for k in self.model.params if not getattr(self.model, k).frozen) def get_fitparams_num(self): """Return a dict of fittable param names and numeric values.""" return collections.OrderedDict((k, getattr(self.model, k).value) for k in self.model.params if not getattr(self.model, k).frozen) def get_fitparams_uncertainty(self): return collections.OrderedDict( (k, getattr(self.model, k).uncertainty_value) for k in self.model.params if not getattr(self.model, k).frozen) def set_params(self, fitp): """Set the model parameters to the value contained in the input dict. Ex. fitter.set_params({'F0':60.1,'F1':-1.3e-15}) """ # In Powell fitter this sometimes fails because after some iterations the values change from # plain float to Quantities. No idea why. if len(fitp.values()) < 1: return if isinstance(list(fitp.values())[0], u.Quantity): for k, v in fitp.items(): getattr(self.model, k).value = v.value else: for k, v in fitp.items(): getattr(self.model, k).value = v def set_param_uncertainties(self, fitp): for k, v in fitp.items(): parunit = getattr(self.model, k).units getattr(self.model, k).uncertainty = v * parunit def get_designmatrix(self): return self.model.designmatrix(toas=self.toas, incfrozen=False, incoffset=True) def minimize_func(self, x, *args): """Wrapper function for the residual class, meant to be passed to scipy.optimize.minimize. The function must take a single list of input values, x, and a second optional tuple of input arguments. It returns a quantity to be minimized (in this case chi^2). """ self.set_params({k: v for k, v in zip(args, x)}) self.update_resids() # Return chi^2 return self.resids.chi2 def fit_toas(self, maxiter=None): raise NotImplementedError def plot(self): """Make residuals plot""" import matplotlib.pyplot as plt from astropy.visualization import quantity_support quantity_support() fig, ax = plt.subplots(figsize=(16, 9)) mjds = self.toas.get_mjds() ax.errorbar(mjds, self.resids.time_resids, yerr=self.toas.get_errors(), fmt="+") ax.set_xlabel("MJD") ax.set_ylabel("Residuals") try: psr = self.model.PSR except: psr = self.model.PSRJ else: psr = "Residuals" ax.set_title(psr) ax.grid(True) plt.show() def get_summary(self, nodmx=False): """Return a human-readable summary of the Fitter results. Parameters ---------- nodmx : bool Set to True to suppress printing DMX parameters in summary """ # Need to check that fit has been done first! if not hasattr(self, "covariance_matrix"): log.warning( "fit_toas() has not been run, so pre-fit and post-fit will be the same!" ) from uncertainties import ufloat import uncertainties.umath as um # First, print fit quality metrics s = "Fitted model using {} method with {} free parameters to {} TOAs\n".format( self.method, len(self.get_fitparams()), self.toas.ntoas) s += "Prefit residuals Wrms = {}, Postfit residuals Wrms = {}\n".format( self.resids_init.rms_weighted(), self.resids.rms_weighted()) s += "Chisq = {:.3f} for {} d.o.f. for reduced Chisq of {:.3f}\n".format( self.resids.chi2, self.resids.dof, self.resids.chi2_reduced) s += "\n" # Next, print the model parameters s += "{:<14s} {:^20s} {:^28s} {}\n".format("PAR", "Prefit", "Postfit", "Units") s += "{:<14s} {:>20s} {:>28s} {}\n".format("=" * 14, "=" * 20, "=" * 28, "=" * 5) for pn in list(self.get_allparams().keys()): if nodmx and pn.startswith("DMX"): continue prefitpar = getattr(self.model_init, pn) par = getattr(self.model, pn) if par.value is not None: if isinstance(par, strParameter): s += "{:14s} {:>20s} {:28s} {}\n".format( pn, prefitpar.value, "", par.units) elif isinstance(par, AngleParameter): # Add special handling here to put uncertainty into arcsec if par.frozen: s += "{:14s} {:>20s} {:>28s} {} \n".format( pn, str(prefitpar.quantity), "", par.units) else: if par.units == u.hourangle: uncertainty_unit = pint.hourangle_second else: uncertainty_unit = u.arcsec s += "{:14s} {:>20s} {:>16s} +/- {:.2g} \n".format( pn, str(prefitpar.quantity), str(par.quantity), par.uncertainty.to(uncertainty_unit), ) else: # Assume a numerical parameter if par.frozen: s += "{:14s} {:20g} {:28s} {} \n".format( pn, prefitpar.value, "", par.units) else: # s += "{:14s} {:20g} {:20g} {:20.2g} {} \n".format( # pn, # prefitpar.value, # par.value, # par.uncertainty.value, # par.units, # ) s += "{:14s} {:20g} {:28SP} {} \n".format( pn, prefitpar.value, ufloat(par.value, par.uncertainty.value), par.units, ) # Now print some useful derived parameters s += "\nDerived Parameters:\n" if hasattr(self.model, "F0"): F0 = self.model.F0.quantity if not self.model.F0.frozen: p, perr = pint.utils.pferrs(F0, self.model.F0.uncertainty) s += "Period = {} +/- {}\n".format(p.to(u.s), perr.to(u.s)) else: s += "Period = {}\n".format((1.0 / F0).to(u.s)) if hasattr(self.model, "F1"): F1 = self.model.F1.quantity if not any([self.model.F1.frozen, self.model.F0.frozen]): p, perr, pd, pderr = pint.utils.pferrs( F0, self.model.F0.uncertainty, F1, self.model.F1.uncertainty) s += "Pdot = {} +/- {}\n".format( pd.to(u.dimensionless_unscaled), pderr.to(u.dimensionless_unscaled)) brakingindex = 3 s += "Characteristic age = {:.4g} (braking index = {})\n".format( pint.utils.pulsar_age(F0, F1, n=brakingindex), brakingindex) s += "Surface magnetic field = {:.3g}\n".format( pint.utils.pulsar_B(F0, F1)) s += "Magnetic field at light cylinder = {:.4g}\n".format( pint.utils.pulsar_B_lightcyl(F0, F1)) I_NS = I = 1.0e45 * u.g * u.cm**2 s += "Spindown Edot = {:.4g} (I={})\n".format( pint.utils.pulsar_edot(F0, F1, I=I_NS), I_NS) if hasattr(self.model, "PX"): if not self.model.PX.frozen: s += "\n" px = ufloat( self.model.PX.quantity.to(u.arcsec).value, self.model.PX.uncertainty.to(u.arcsec).value, ) s += "Parallax distance = {:.3uP} pc\n".format(1.0 / px) # Now binary system derived parameters binary = None for x in self.model.components: if x.startswith("Binary"): binary = x if binary is not None: s += "\n" s += "Binary model {}\n".format(binary) if binary.startswith("BinaryELL1"): if not any([ self.model.EPS1.frozen, self.model.EPS2.frozen, self.model.TASC.frozen, self.model.PB.frozen, ]): eps1 = ufloat( self.model.EPS1.quantity.value, self.model.EPS1.uncertainty.value, ) eps2 = ufloat( self.model.EPS2.quantity.value, self.model.EPS2.uncertainty.value, ) tasc = ufloat( # This is a time in MJD self.model.TASC.quantity.mjd, self.model.TASC.uncertainty.to(u.d).value, ) pb = ufloat( self.model.PB.quantity.to(u.d).value, self.model.PB.uncertainty.to(u.d).value, ) s += "Conversion from ELL1 parameters:\n" ecc = um.sqrt(eps1**2 + eps2**2) s += "ECC = {:P}\n".format(ecc) om = um.atan2(eps1, eps2) * 180.0 / np.pi if om < 0.0: om += 360.0 s += "OM = {:P}\n".format(om) t0 = tasc + pb * om / 360.0 s += "T0 = {:SP}\n".format(t0) s += pint.utils.ELL1_check( self.model.A1.quantity, ecc.nominal_value, self.resids.rms_weighted(), self.toas.ntoas, outstring=True, ) s += "\n" # Masses and inclination if not any([self.model.PB.frozen, self.model.A1.frozen]): pbs = ufloat( self.model.PB.quantity.to(u.s).value, self.model.PB.uncertainty.to(u.s).value, ) a1 = ufloat( self.model.A1.quantity.to(pint.ls).value, self.model.A1.uncertainty.to(pint.ls).value, ) fm = 4.0 * np.pi**2 * a1**3 / (4.925490947e-6 * pbs**2) s += "Mass function = {:SP} Msun\n".format(fm) mcmed = pint.utils.companion_mass( self.model.PB.quantity, self.model.A1.quantity, inc=60.0 * u.deg, mpsr=1.4 * u.solMass, ) mcmin = pint.utils.companion_mass( self.model.PB.quantity, self.model.A1.quantity, inc=90.0 * u.deg, mpsr=1.4 * u.solMass, ) s += "Companion mass min, median (assuming Mpsr = 1.4 Msun) = {:.4f}, {:.4f} Msun\n".format( mcmin, mcmed) if hasattr(self.model, "SINI"): if not self.model.SINI.frozen: si = ufloat( self.model.SINI.quantity.value, self.model.SINI.uncertainty.value, ) s += "From SINI in model:\n" s += " cos(i) = {:SP}\n".format(um.sqrt(1 - si**2)) s += " i = {:SP} deg\n".format( um.asin(si) * 180.0 / np.pi) psrmass = pint.utils.pulsar_mass( self.model.PB.quantity, self.model.A1.quantity, self.model.M2.quantity, np.arcsin(self.model.SINI.quantity), ) s += "Pulsar mass (Shapiro Delay) = {}".format(psrmass) return s def print_summary(self): """Write a summary of the TOAs to stdout.""" print(self.get_summary()) def get_covariance_matrix(self, with_phase=False, pretty_print=False, prec=3): """Show the parameter covariance matrix post-fit. If with_phase, then show and return the phase column as well. If pretty_print, then also pretty-print on stdout the matrix. prec is the precision of the floating point results. """ if hasattr(self, "covariance_matrix"): fps = list(self.get_fitparams().keys()) cm = self.covariance_matrix if with_phase: fps = ["PHASE"] + fps else: cm = cm[1:, 1:] if pretty_print: lens = [max(len(fp) + 2, prec + 8) for fp in fps] maxlen = max(lens) print("\nParameter covariance matrix:") line = "{0:^{width}}".format("", width=maxlen) for fp, ln in zip(fps, lens): line += "{0:^{width}}".format(fp, width=ln) print(line) for ii, fp1 in enumerate(fps): line = "{0:^{width}}".format(fp1, width=maxlen) for jj, (fp2, ln) in enumerate(zip(fps[:ii + 1], lens[:ii + 1])): line += "{0: {width}.{prec}e}".format(cm[ii, jj], width=ln, prec=prec) print(line) print("\n") return cm else: log.error( "You must run .fit_toas() before accessing the covariance matrix" ) raise AttributeError def get_correlation_matrix(self, with_phase=False, pretty_print=False, prec=3): """Show the parameter correlation matrix post-fit. If with_phase, then show and return the phase column as well. If pretty_print, then also pretty-print on stdout the matrix. prec is the precision of the floating point results. """ if hasattr(self, "correlation_matrix"): fps = list(self.get_fitparams().keys()) cm = self.correlation_matrix if with_phase: fps = ["PHASE"] + fps else: cm = cm[1:, 1:] if pretty_print: lens = [max(len(fp) + 2, prec + 4) for fp in fps] maxlen = max(lens) print("\nParameter correlation matrix:") line = "{0:^{width}}".format("", width=maxlen) for fp, ln in zip(fps, lens): line += "{0:^{width}}".format(fp, width=ln) print(line) for ii, fp1 in enumerate(fps): line = "{0:^{width}}".format(fp1, width=maxlen) for jj, (fp2, ln) in enumerate(zip(fps, lens)): line += "{0:^{width}.{prec}f}".format(cm[ii, jj], width=ln, prec=prec) print(line) print("\n") return cm else: log.error( "You must run .fit_toas() before accessing the correlation matrix" ) raise AttributeError
class Pulsar: """Wrapper class for a pulsar. Contains the toas, model, residuals, and fitter """ def __init__(self, parfile=None, timfile=None, ephem=None): super(Pulsar, self).__init__() log.info("STARTING LOADING OF PULSAR %s" % str(parfile)) if parfile is not None and timfile is not None: self.parfile = parfile self.timfile = timfile else: raise ValueError("No valid pulsar to load") self.prefit_model = pint.models.get_model(self.parfile) if ephem is not None: # TODO: EPHEM overwrite message? self.all_toas = get_TOAs(self.timfile, ephem=ephem, planets=True) self.prefit_model.EPHEM.value = ephem elif getattr(self.prefit_model, "EPHEM").value is not None: self.all_toas = get_TOAs(self.timfile, ephem=self.prefit_model.EPHEM.value, planets=True) else: self.all_toas = get_TOAs(self.timfile, planets=True) # turns pre-existing jump flags in toas.table['flags'] into parameters in parfile self.prefit_model.jump_flags_to_params(self.all_toas) self.selected_toas = copy.deepcopy(self.all_toas) print("prefit_model.as_parfile():") print(self.prefit_model.as_parfile()) self.all_toas.print_summary() self.prefit_resids = Residuals(self.selected_toas, self.prefit_model) print("RMS PINT residuals are %.3f us\n" % self.prefit_resids.rms_weighted().to(u.us).value) self.fitter = Fitters.WLS self.fitted = False @property def name(self): return getattr(self.prefit_model, "PSR").value def __getitem__(self, key): try: return getattr(self.prefit_model, key) except AttributeError: log.error("Parameter %s was not found in pulsar model %s" % (key, self.name)) return None def __contains__(self, key): return key in self.prefit_model.params def reset_model(self): self.prefit_model = pint.models.get_model(self.parfile) self.postfit_model = None self.postfit_resids = None self.fitted = False self.update_resids() def reset_TOAs(self): if getattr(self.prefit_model, "EPHEM").value is not None: self.all_toas = get_TOAs(self.timfile, ephem=self.prefit_model.EPHEM.value, planets=True) else: self.all_toas = get_TOAs(self.timfile, planets=True) self.selected_toas = copy.deepcopy(self.all_toas) self.update_resids() def resetAll(self): self.prefit_model = pint.models.get_model(self.parfile) self.postfit_model = None self.postfit_resids = None self.fitted = False self.reset_TOAs() def update_resids(self): # update the pre and post fit residuals using all_toas self.prefit_resids = Residuals(self.all_toas, self.prefit_model) if self.fitted: self.postfit_resids = Residuals(self.all_toas, self.postfit_model) def orbitalphase(self): """ For a binary pulsar, calculate the orbital phase. Otherwise, return an array of unitless quantities of zeros """ if not self.prefit_model.is_binary: log.warn("This is not a binary pulsar") return u.Quantity(np.zeros(self.selected_toas.ntoas)) toas = self.selected_toas if self.fitted: phase = self.postfit_model.orbital_phase(toas, anom="mean") else: phase = self.prefit_model.orbital_phase(toas, anom="mean") return phase / (2 * np.pi * u.rad) def dayofyear(self): """ Return the day of the year for all the TOAs of this pulsar """ t = Time(self.selected_toas.get_mjds(), format="mjd") year = Time(np.floor(t.decimalyear), format="decimalyear") return (t.mjd - year.mjd) * u.day def year(self): """ Return the decimal year for all the TOAs of this pulsar """ t = Time(self.selected_toas.get_mjds(), format="mjd") return (t.decimalyear) * u.year def write_fit_summary(self): """ Summarize fitting results """ if self.fitted: chi2 = self.postfit_resids.chi2 wrms = self.postfit_resids.rms_weighted() print("Post-Fit Chi2:\t\t%.8g us^2" % chi2) print("Post-Fit Weighted RMS:\t%.8g us" % wrms.to(u.us).value) print("%19s %24s\t%24s\t%16s %16s %16s" % ( "Parameter", "Pre-Fit", "Post-Fit", "Uncertainty", "Difference", "Diff/Unc", )) print("-" * 132) fitparams = [ p for p in self.prefit_model.params if not getattr(self.prefit_model, p).frozen ] for key in fitparams: line = "%8s " % key pre = getattr(self.prefit_model, key) post = getattr(self.postfit_model, key) line += "%10s " % ("" if post.units is None else str(post.units)) if post.quantity is not None: line += "%24s\t" % pre.print_quantity(pre.quantity) line += "%24s\t" % post.print_quantity(post.quantity) try: line += "%16.8g " % post.uncertainty.value except: line += "%18s" % "" try: diff = post.value - pre.value line += "%16.8g " % diff if pre.uncertainty is not None: line += "%16.8g" % (diff / pre.uncertainty.value) except: pass print(line) else: log.warn("Pulsar has not been fitted yet!") def add_phase_wrap(self, selected, phase): """ Add a phase wrap to selected points in the TOAs object Turn on pulse number tracking in the model, if it isn't already :param selected: boolean array to apply to toas, True = selected toa :param phase: phase diffeence to be added, i.e. -0.5, +2, etc. """ # Check if pulse numbers are in table already, if not, make the column if ("pn" not in self.all_toas.table.colnames or "pn" not in self.selected_toas.table.colnames): self.all_toas.compute_pulse_numbers(self.prefit_model) self.selected_toas.compute_pulse_numbers(self.prefit_model) if ("delta_pulse_number" not in self.all_toas.table.colnames or "delta_pulse_number" not in self.selected_toas.table.colnames): self.all_toas.table["delta_pulse_number"] = np.zeros( len(self.all_toas.get_mjds())) self.selected_toas.table["delta_pulse_number"] = np.zeros( len(self.selected_toas.get_mjds())) # add phase wrap self.all_toas.table["delta_pulse_number"][selected] += phase self.selected_toas.table["delta_pulse_number"] += phase self.update_resids() def add_jump(self, selected): """ jump the toas selected or unjump them if already jumped :param selected: boolean array to apply to toas, True = selected toa """ # TODO: split into two functions if "PhaseJump" not in self.prefit_model.components: # if no PhaseJump component, add one log.info("PhaseJump component added") a = pint.models.jump.PhaseJump() a.setup() self.prefit_model.add_component(a) self.prefit_model.remove_param("JUMP1") param = pint.models.parameter.maskParameter( name="JUMP", index=1, key="-gui_jump", key_value=1, value=0.0, units="second", ) self.prefit_model.add_param_from_top(param, "PhaseJump") getattr(self.prefit_model, param.name).frozen = False self.prefit_model.components[ "PhaseJump"]._parent = self.prefit_model if self.fitted: self.postfit_model.add_component(a) for dict1, dict2 in zip( self.all_toas.table["flags"][selected], self.selected_toas.table["flags"], ): dict1["gui_jump"] = 1 dict1["jump"] = 1 dict2["gui_jump"] = 1 dict2["jump"] = 1 return param.name # if gets here, has at least one jump param already # if doesnt overlap or cancel, add the param jump_nums = [ int(dict["jump"]) if "jump" in dict.keys() else np.nan for dict in self.all_toas.table["flags"] ] numjumps = self.prefit_model.components[ "PhaseJump"].get_number_of_jumps() if numjumps == 0: log.warn( "There are no jumps (maskParameter objects) in PhaseJump. Please delete the PhaseJump object and try again. " ) return None # if only par file jumps in PhaseJump object if np.isnan(np.nanmax(jump_nums)): # for every jump, set appropriate flag for TOAs it jumps for jump_par in self.prefit_model.components[ "PhaseJump"].get_jump_param_objects(): # find TOAs jump applies to mask = jump_par.select_toa_mask(self.all_toas) # apply to dictionaries for future use for dict in self.all_toas.table["flags"][mask]: dict["jump"] = jump_par.index jump_nums = [ int(dict["jump"]) if "jump" in dict.keys() else np.nan for dict in self.all_toas.table["flags"] ] for num in range(1, numjumps + 1): num = int(num) jump_select = [num == jump_num for jump_num in jump_nums] if np.array_equal(jump_select, selected): # if current jump exactly matches selected, remove it self.prefit_model.remove_param("JUMP" + str(num)) if self.fitted: self.postfit_model.remove_param("JUMP" + str(num)) for dict1, dict2 in zip( self.all_toas.table["flags"][selected], self.selected_toas.table["flags"], ): if "jump" in dict1.keys() and dict1["jump"] == num: del dict1["jump"] if "gui_jump" in dict1.keys(): del dict1["gui_jump"] if "jump" in dict2.keys() and dict2["jump"] == num: del dict2["jump"] if "gui_jump" in dict2.keys(): del dict2["gui_jump"] nums_subset = range(num + 1, numjumps + 1) for n in nums_subset: # iterate through jump params and rename them so that they are always in numerical order starting with JUMP1 n = int(n) param = getattr(self.prefit_model.components["PhaseJump"], "JUMP" + str(n)) for dict in self.all_toas.table["flags"]: if "jump" in dict.keys() and dict["jump"] == n: dict["jump"] = n - 1 if "gui_jump" in dict.keys(): dict["gui_jump"] = n - 1 param.key_value = n - 1 newpar = param.new_param(index=(n - 1), copy_all=True) self.prefit_model.add_param_from_top(newpar, "PhaseJump") self.prefit_model.remove_param(param.name) if self.fitted: self.postfit_model.add_param_from_top( newpar, "PhaseJump") self.postfit_model.remove_param(param.name) if "JUMP1" not in self.prefit_model.params: # remove PhaseJump component if no jump params comp_list = getattr(self.prefit_model, "PhaseComponent_list") for item in comp_list: if isinstance(item, pint.models.jump.PhaseJump): self.prefit_model.remove_component(item) if self.fitted: comp_list = getattr(self.postfit_model, "PhaseComponent_list") for item in comp_list: if isinstance(item, pint.models.jump.PhaseJump): self.postfit_model.remove_component(item) else: self.prefit_model.components["PhaseJump"].setup() if self.fitted: self.postfit_model.components["PhaseJump"].setup() log.info("removed param", "JUMP" + str(num)) return jump_select elif True in [a and b for a, b in zip(jump_select, selected)]: # if current jump overlaps selected, raise and error and end log.warn( "The selected toa(s) overlap an existing jump. Remove all interfering jumps before attempting to jump these toas." ) return None # if here, then doesn't overlap or match anything for dict1, dict2 in zip(self.all_toas.table["flags"][selected], self.selected_toas.table["flags"]): dict1["jump"] = numjumps + 1 dict1["gui_jump"] = numjumps + 1 dict2["jump"] = numjumps + 1 dict2["gui_jump"] = numjumps + 1 param = pint.models.parameter.maskParameter( name="JUMP", index=numjumps + 1, key="-gui_jump", key_value=numjumps + 1, value=0.0, units="second", aliases=["JUMP"], ) self.prefit_model.add_param_from_top(param, "PhaseJump") getattr(self.prefit_model, param.name).frozen = False self.prefit_model.components["PhaseJump"].setup() if (self.fitted and not self.prefit_model.components["PhaseJump"] == self.postfit_model.components["PhaseJump"]): self.postfit_model.add_param_from_top(param, "PhaseJump") getattr(self.postfit_model, param.name).frozen = False self.postfit_model.components["PhaseJump"].setup() return param.name def fit(self, selected, iters=1): """ Run a fit using the specified fitter """ # Select all the TOAs if none are explicitly set if not any(selected): selected = ~selected """JUMP check, TODO: put in fitter?""" if "PhaseJump" in self.prefit_model.components: # Modifies jump flags. If attempted fit (selected) # A) contains only jumps, don't do the fit and return an error # B) excludes a jump, turn that jump off # C) partially contains a jump, redefine that jump only with the overlap fit_jumps = [] for param in self.prefit_model.params: if getattr(self.prefit_model, param).frozen == False and param.startswith("JUMP"): fit_jumps.append(int(param[4:])) numjumps = self.prefit_model.components[ "PhaseJump"].get_number_of_jumps() if numjumps == 0: log.warn( "There are no jumps (maskParameter objects) in PhaseJump. Please delete the PhaseJump object and try again. " ) return None # boolean array to determine if all selected toas are jumped jumps = [ True if "jump" in dict.keys() and dict["jump"] in fit_jumps else False for dict in self.all_toas.table["flags"][selected] ] # check if par file jumps in PhaseJump object if not any(jumps): # for every jump, set appropriate flag for TOAs it jumps for jump_par in self.prefit_model.components[ "PhaseJump"].get_jump_param_objects(): # find TOAs jump applies to mask = jump_par.select_toa_mask(self.all_toas) # apply to dictionaries for future use for dict in self.all_toas.table["flags"][mask]: dict["jump"] = jump_par.index jumps = [ True if "jump" in dict.keys() and dict["jump"] in fit_jumps else False for dict in self.all_toas.table["flags"][selected] ] if all(jumps): log.warn( "toas being fit must not all be jumped. Remove or uncheck at least one jump in the selected toas before fitting." ) return None # numerical array of selected jump flags sel_jump_nums = [ dict["jump"] if "jump" in dict.keys() else np.nan for dict in self.all_toas.table["flags"][selected] ] # numerical array of all jump flags full_jump_nums = [ dict["jump"] if "jump" in dict.keys() else np.nan for dict in self.all_toas.table["flags"] ] for num in range(1, numjumps + 1): num = int(num) if num not in sel_jump_nums: getattr(self.prefit_model, "JUMP" + str(num)).frozen = True continue jump_select = [num == jump_num for jump_num in full_jump_nums] overlap = [a and b for a, b in zip(jump_select, selected)] # remove the jump flags for that num for dict in self.all_toas.table["flags"]: if "jump" in dict.keys() and dict["jump"] == num: del dict["jump"] # re-add the jump using overlap as 'selected' for dict in self.all_toas.table["flags"][overlap]: dict["jump"] = num if self.fitted: self.prefit_model = self.postfit_model self.prefit_resids = self.postfit_resids if self.fitter == Fitters.POWELL: fitter = pint.fitter.PowellFitter(self.selected_toas, self.prefit_model) elif self.fitter == Fitters.WLS: fitter = pint.fitter.WLSFitter(self.selected_toas, self.prefit_model) elif self.fitter == Fitters.GLS: fitter = pint.fitter.GLSFitter(self.selected_toas, self.prefit_model) chi2 = self.prefit_resids.chi2 wrms = self.prefit_resids.rms_weighted() print("Pre-Fit Chi2:\t\t%.8g us^2" % chi2) print("Pre-Fit Weighted RMS:\t%.8g us" % wrms.to(u.us).value) fitter.fit_toas(maxiter=1) self.postfit_model = fitter.model self.postfit_resids = Residuals(self.all_toas, self.postfit_model) self.fitted = True self.write_fit_summary() # TODO: delta_pulse_numbers need some work. They serve both for PHASE and -padd functions from the TOAs # as well as for phase jumps added manually in the GUI. They really should not be zeroed out here because # that will wipe out preexisting values self.all_toas.table["delta_pulse_numbers"] = np.zeros( self.all_toas.ntoas) self.selected_toas.table["delta_pulse_number"] = np.zeros( self.selected_toas.ntoas) # plot the prefit without jumps pm_no_jumps = copy.deepcopy(self.postfit_model) for param in pm_no_jumps.params: if param.startswith("JUMP"): getattr(pm_no_jumps, param).value = 0.0 getattr(pm_no_jumps, param).frozen = True self.prefit_resids_no_jumps = Residuals(self.selected_toas, pm_no_jumps) f = copy.deepcopy(fitter) no_jumps = [ False if "jump" in dict.keys() else True for dict in f.toas.table["flags"] ] f.toas.select(no_jumps) selectedMJDs = self.selected_toas.get_mjds() if all(no_jumps): q = list(self.all_toas.get_mjds()) index = q.index([ i for i in self.all_toas.get_mjds() if i > selectedMJDs.min() ][0]) rs_mean = (Residuals( self.all_toas, f.model).phase_resids[index:index + len(selectedMJDs)].mean()) else: rs_mean = self.prefit_resids_no_jumps.phase_resids[no_jumps].mean() # determines how far on either side fake toas go # TODO: hard limit on how far fake toas can go --> can get clkcorr # errors if go before GBT existed, etc. minMJD, maxMJD = selectedMJDs.min(), selectedMJDs.max() spanMJDs = maxMJD - minMJD if spanMJDs < 30 * u.d: redge = ledge = 4 npoints = 400 elif spanMJDs < 90 * u.d: redge = ledge = 2 npoints = 300 elif spanMJDs < 200 * u.d: redge = ledge = 1 npoints = 300 elif spanMJDs < 400 * u.d: redge = ledge = 0.5 npoints = 200 else: redge = ledge = 1.0 npoints = 250 # Check to see if too recent nowish = (Time.now().mjd - 40) * u.d if maxMJD + spanMJDs * redge > nowish: redge = (nowish - maxMJD) / spanMJDs if redge < 0.0: redge = 0.0 f_toas, rs, mrands = random_models( f, rs_mean=rs_mean, redge_multiplier=redge, ledge_multiplier=ledge, npoints=npoints, iter=10, ) self.random_resids = rs self.fake_toas = f_toas def fake_year(self): """ Function to support plotting of random models on multiple x-axes. Return the decimal year for all the TOAs of this pulsar """ t = Time(self.fake_toas.get_mjds(), format="mjd") return (t.decimalyear) * u.year