Merge pull request #1534 from abhisrkckl/pintk-wb

Plot wideband DMs and DM residuals in `pintk`

CHANGELOG-unreleased.md

@@ -43,6 +43,7 @@ the released changes.
 - `ScaleToaError.d_toasigma_d_EFAC` and `ScaleToaError.d_toasigma_d_EQUAD` methods.
 - Separate `.fullname` for all observatories
 - `Residuals.calc_whitened_resids()` method
+- Plot wideband DM measurements, wideband DM residuals, and wideband DM errors in `pintk`. (Disabled for narrowband data.)
 - Optionally generate multi-frequency TOAs in an epoch using `make_fake_toas_uniform` and `make_fake_toas_fromMJDs`
 - Documentation: Example notebook for simulations
 ### Fixed

src/pint/pintk/plk.py

@@ -8,7 +8,7 @@
 from astropy.time import Time
 import astropy.units as u
 import matplotlib as mpl
-from matplotlib import figure
+import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
 
@@ -23,6 +23,8 @@
 import pint.logging
 from loguru import logger as log
 
+from pint.residuals import WidebandDMResiduals
+
 
 try:
     from matplotlib.backends.backend_tkagg import NavigationToolbar2Tk
@@ -50,6 +52,9 @@
     "frequency": r"Observing Frequency (MHz)",
     "TOA error": r"TOA uncertainty ($\mu$s)",
     "rounded MJD": r"MJD",
+    "WB DM": "Wideband DM (pc/cm3)",
+    "WB DM res": "Wideband DM residual (pc/cm3)",
+    "WB DM err": "Wideband DM error (pc/cm3)",
     "elongation": r"Solar Elongation (deg)",
 }
 
@@ -475,6 +480,9 @@ def __init__(self, master=None, **kwargs):
         self.xvar = tk.StringVar()
         self.yvar = tk.StringVar()
 
+        # This will be set in PlkWidget.setPulsar and PlkWidget.update methods.
+        self.wideband = False
+
         self.initPlkXYChoice()
 
     def initPlkXYChoice(self):
@@ -531,7 +539,7 @@ def setChoice(self, xid="mjd", yid="pre-fit"):
                 else self.master.psr.prefit_model
             )
             if choice == "elongation" and not any(
-                isinstance(x, Astrometry) for x in model.components
+                isinstance(x, Astrometry) for x in model.components.values()
             ):
                 self.xbuttons[ii].configure(state="disabled")
                 self.ybuttons[ii].configure(state="disabled")
@@ -544,6 +552,12 @@ def setChoice(self, xid="mjd", yid="pre-fit"):
             ):
                 self.xbuttons[ii].configure(state="disabled")
                 self.ybuttons[ii].configure(state="disabled")
+            if (
+                choice in ["WB DM", "WB DM res", "WB DM err"]
+                and not self.master.psr.all_toas.is_wideband()
+            ):
+                self.xbuttons[ii].configure(state="disabled")
+                self.ybuttons[ii].configure(state="disabled")
 
     def setCallbacks(self, updatePlot):
         """
@@ -718,7 +732,7 @@ def initPlk(self):
         self.colorModeWidget = PlkColorModeBoxes(master=self)
 
         self.plkDpi = 100
-        self.plkFig = mpl.figure.Figure(dpi=self.plkDpi)
+        self.plkFig = plt.Figure(dpi=self.plkDpi)
         self.plkCanvas = FigureCanvasTkAgg(self.plkFig, self)
         self.plkCanvas.mpl_connect("button_press_event", self.canvasClickEvent)
         self.plkCanvas.mpl_connect("button_release_event", self.canvasReleaseEvent)
@@ -773,6 +787,7 @@ def update(self):
             self.randomboxWidget.addRandomCheckbox(self)
             self.colorModeWidget.addColorModeCheckbox(self.color_modes)
             self.fitterWidget.updateFitterChoices(self.psr.all_toas.wideband)
+            self.xyChoiceWidget.wideband = self.psr.all_toas.wideband
             self.xyChoiceWidget.setChoice()
             self.updatePlot(keepAxes=True)
             self.plkToolbar.update()
@@ -797,6 +812,7 @@ def setPulsar(self, psr, updates):
 
         self.fitboxesWidget.setCallbacks(self.fitboxChecked)
         self.colorModeWidget.setCallbacks(self.updateGraphColors)
+        self.xyChoiceWidget.wideband = self.psr.all_toas.wideband
         self.xyChoiceWidget.setCallbacks(self.updatePlot)
         self.actionsWidget.setCallbacks(
             self.fit, self.reset, self.writePar, self.writeTim, self.revert
@@ -919,7 +935,7 @@ def writePar(self, format="pint"):
                         f"Pulsar has not been fitted! Saving pre-fit parfile to {filename} in {format} format"
                     )
 
-        except:
+        except Exception:
             if filename in [(), ""]:
                 print("Write Par cancelled.")
             else:
@@ -938,7 +954,7 @@ def writeTim(self, format="tempo2"):
             log.info(f"Choose output file {filename}")
             self.psr.all_toas.write_TOA_file(filename, format=format)
             log.info(f"Wrote TOAs to {filename} with format {format}")
-        except:
+        except Exception:
             if filename in [(), ""]:
                 print("Write Tim cancelled.")
             else:
@@ -1250,6 +1266,37 @@ def psr_data_from_label(self, label):
         elif label == "rounded MJD":
             data = np.floor(self.psr.all_toas.get_mjds() + 0.5 * u.d)
             error = self.psr.all_toas.get_errors().to(u.d)
+        elif label == "WB DM":
+            if self.psr.all_toas.wideband:
+                data = self.psr.all_toas.get_dms().to(pint.dmu)
+                error = self.psr.all_toas.get_dm_errors().to(pint.dmu)
+            else:
+                log.warning("Cannot plot WB DMs for NB TOAs.")
+                data = None
+                error = None
+        elif label == "WB DM res":
+            if self.psr.all_toas.wideband:
+                if self.psr.fitter is not None:
+                    data = self.psr.fitter.resids.dm.calc_resids().to(pint.dmu)
+                else:
+                    data = (
+                        WidebandDMResiduals(self.psr.all_toas, self.psr.prefit_model)
+                        .calc_resids()
+                        .to(pint.dmu)
+                    )
+                error = self.psr.all_toas.get_dm_errors().to(pint.dmu)
+            else:
+                log.warning("Cannot plot WB DM resids for NB TOAs.")
+                data = None
+                error = None
+        elif label == "WB DM err":
+            if self.psr.all_toas.wideband:
+                data = self.psr.all_toas.get_dm_errors().to(pint.dmu)
+                error = None
+            else:
+                log.warning("Cannot plot WB DM errors for NB TOAs.")
+                data = None
+                error = None
         elif label == "elongation":
             data = np.degrees(
                 self.psr.prefit_model.sun_angle(self.psr.all_toas, also_distance=False)

src/pint/pintk/pulsar.py

@@ -1,4 +1,4 @@
-"""A wrapper around pulsar functions for pintk to use.
+"""A wrapper around pulsar functions for `pintk` to use.
 
 This object will be shared between widgets in the main frame
 and will contain the pre/post fit model, toas,
@@ -37,6 +37,9 @@
     "frequency",
     "TOA error",
     "rounded MJD",
+    "WB DM",
+    "WB DM res",
+    "WB DM err",
     "elongation",
 ]
 
@@ -324,7 +327,7 @@ def write_fit_summary(self):
                     line += "%24s\t" % post.str_quantity(post.quantity)
                     try:
                         line += "%16.8g  " % post.uncertainty.value
-                    except:
+                    except Exception:
                         line += "%18s" % ""
                     diff = post.value - pre.value
                     line += "%16.8g  " % diff

src/pint/residuals.py

@@ -93,7 +93,7 @@ def __new__(
             except KeyError as e:
                 raise ValueError(
                     f"'{residual_type}' is not a PINT supported residual. Currently supported data types are {list(residual_map.keys())}"
-                )
+                ) from e
 
         return super().__new__(cls)
 
@@ -370,11 +370,15 @@ def calc_phase_resids(
             # and delta_pulse_numbers (from PHASE lines or adding phase jumps in GUI)
             i = pulse_num.copy()
             f = np.zeros_like(pulse_num)
-            if np.any(np.isnan(pulse_num)):
+
+            c = np.isnan(pulse_num)
+            if np.any(c):
+                # i[c] = 0
                 raise ValueError("Pulse numbers are missing on some TOAs")
             residualphase = modelphase - Phase(i, f)
             # This converts from a Phase object to a np.float128
             full = residualphase.int + residualphase.frac
+
         elif self.track_mode == "nearest":
             # Compute model phase
             modelphase = self.model.phase(self.toas) + delta_pulse_numbers