Merge pull request #1429 from abhisrkckl/wbsim

Simulate wideband TOAs using `zima`

CHANGELOG.md

@@ -20,6 +20,7 @@ and this project, at least loosely, adheres to [Semantic Versioning](https://sem
 - H.E.S.S. telescope to the list of known observatories
 - Documentation: making TOAs from a list of times added to HowTo
 - Clock correction for LEAP
+- Wideband TOA simulation feature in `pint.simulation` and `zima`
 ### Fixed
 - Broken notebooks CI test
 - BIPM correction for simulated TOAs

src/pint/residuals.py

@@ -437,8 +437,7 @@ def calc_chi2(self, full_cov=False):
                 # return ((self.time_resids.to(u.s) / self.toas.get_errors().to(u.s)).value**2.0).sum()
 
                 # This the fastest way, but highly depend on the assumption of time_resids and
-                # error units.
-                # insure only a pure number returned
+                # error units. Ensure only a pure number is returned.
                 try:
                     return ((self.time_resids / toa_errors.to(u.s)) ** 2.0).sum().value
                 except ValueError:

src/pint/scripts/zima.py

@@ -9,6 +9,7 @@
 
 pint.logging.setup(level=pint.logging.script_level)
 
+import pint
 import pint.fitter
 import pint.models
 import pint.simulation
@@ -64,6 +65,18 @@ def main(argv=None):
         default=False,
         help="Actually add in random noise, or just populate the column",
     )
+    parser.add_argument(
+        "--wideband",
+        action="store_true",
+        default=False,
+        help="Add DM information to simulated TOAs. Generates wideband toas.",
+    )
+    parser.add_argument(
+        "--dmerror",
+        help="Random error to apply to simulated DM measurements (dmu)",
+        type=float,
+        default=1e-4,
+    )
     parser.add_argument(
         "--fuzzdays",
         help="Standard deviation of 'fuzz' distribution (jd)",
@@ -74,7 +87,7 @@ def main(argv=None):
         "--plot", help="Plot residuals", action="store_true", default=False
     )
     parser.add_argument(
-        "--format", help="The format of out put .tim file.", default="TEMPO2"
+        "--format", help="The format of output .tim file.", default="TEMPO2"
     )
     parser.add_argument(
         "--log-level",
@@ -114,33 +127,37 @@ def main(argv=None):
             freq=np.atleast_1d(args.freq) * u.MHz,
             fuzz=args.fuzzdays * u.d,
             add_noise=args.addnoise,
+            wideband=args.wideband,
+            wideband_dm_error=args.dmerror * pint.dmu,
         )
     else:
-        log.info("Reading initial TOAs from {0}".format(args.inputtim))
+        log.info(f"Reading initial TOAs from {args.inputtim}")
         ts = pint.simulation.make_fake_toas_fromtim(
-            args.inputtim,
-            model=m,
-            add_noise=args.addnoise,
+            args.inputtim, model=m, add_noise=args.addnoise
         )
 
     # Write TOAs to a file
     ts.write_TOA_file(args.timfile, name="fake", format=out_format)
 
     if args.plot:
-        # This should be a very boring plot with all residuals flat at 0.0!
-        import matplotlib.pyplot as plt
-        from astropy.visualization import quantity_support
-
-        quantity_support()
-
-        r = pint.residuals.Residuals(ts, m)
-        plt.errorbar(
-            ts.get_mjds(),
-            r.calc_time_resids(calctype="taylor").to(u.us),
-            yerr=ts.get_errors().to(u.us),
-            fmt=".",
-        )
-        plt.xlabel("MJD")
-        plt.ylabel("Residual (us)")
-        plt.grid(True)
-        plt.show()
+        plot_simulated_toas(ts, m)
+
+
+def plot_simulated_toas(ts, m):
+    # This should be a very boring plot with all residuals flat at 0.0!
+    import matplotlib.pyplot as plt
+    from astropy.visualization import quantity_support
+
+    quantity_support()
+
+    r = pint.residuals.Residuals(ts, m)
+    plt.errorbar(
+        ts.get_mjds(),
+        r.calc_time_resids(calctype="taylor").to(u.us),
+        yerr=ts.get_errors().to(u.us),
+        fmt=".",
+    )
+    plt.xlabel("MJD")
+    plt.ylabel("Residual (us)")
+    plt.grid(True)
+    plt.show()

src/pint/simulation.py

@@ -169,6 +169,26 @@ def make_fake_toas(ts, model, add_noise=False, name="fake"):
     return tsim
 
 
+def update_fake_dms(model, ts, dm_error, add_noise):
+    """Update simulated wideband DM information in TOAs."""
+    toas = deepcopy(ts)
+
+    dm_errors = dm_error * np.ones(len(toas))
+
+    for f, dme in zip(toas.table["flags"], dm_errors):
+        f["pp_dme"] = str(dme.to_value(pint.dmu))
+
+    scaled_dm_errors = model.scaled_dm_uncertainty(toas)
+    dms = model.total_dm(toas)
+    if add_noise:
+        dms += scaled_dm_errors.to(pint.dmu) * np.random.randn(len(scaled_dm_errors))
+
+    for f, dm in zip(toas.table["flags"], dms):
+        f["pp_dm"] = str(dm.to_value(pint.dmu))
+
+    return toas
+
+
 def make_fake_toas_uniform(
     startMJD,
     endMJD,
@@ -179,8 +199,8 @@ def make_fake_toas_uniform(
     obs="GBT",
     error=1 * u.us,
     add_noise=False,
-    dm=None,
-    dm_error=1e-4 * pint.dmu,
+    wideband=False,
+    wideband_dm_error=1e-4 * pint.dmu,
     name="fake",
     include_bipm=False,
     include_gps=True,
@@ -210,9 +230,11 @@ def make_fake_toas_uniform(
         uncertainty to attach to each TOA
     add_noise : bool, optional
         Add noise to the TOAs (otherwise `error` just populates the column)
-    dm : astropy.units.Quantity, optional
-        DM value to include with each TOA; default is
-        not to include any DM information
+    wideband : bool, optional
+        Whether to include wideband DM information with each TOA; default is
+        not to include any wideband DM information. If True, the DM associated
+        with each TOA will be computed using the model, and the `-ppdm` and
+        `-ppdme` flags will be set.
     dm_error : astropy.units.Quantity
         uncertainty to attach to each DM measurement
     name : str, optional
@@ -231,7 +253,12 @@ def make_fake_toas_uniform(
 
     Notes
     -----
-    `add_noise` respects any ``EFAC`` or ``EQUAD`` present in the `model`
+    1. `add_noise` respects any ``EFAC`` or ``EQUAD`` present in the `model`
+    2. When `wideband` is set, wideband DM measurement noise will be included
+       only if `add_noise` is set. Otherwise, the `-pp_dme` flags will be set
+       without adding the measurement noise to the simulated DM values.
+    3. The simulated DM measurement noise respects ``DMEFAC`` and ``DMEQUAD``
+       values in the `model`.
 
     See Also
     --------
@@ -262,12 +289,11 @@ def make_fake_toas_uniform(
         include_gps=clk_version["include_gps"],
         planets=model["PLANET_SHAPIRO"].value,
     )
-
     ts.table["error"] = error
-    if dm is not None:
-        for f in ts.table["flags"]:
-            f["pp_dm"] = str(dm.to_value(pint.dmu))
-            f["pp_dme"] = str(dm_error.to_value(pint.dmu))
+
+    if wideband:
+        ts = update_fake_dms(model, ts, wideband_dm_error, add_noise)
+
     return make_fake_toas(ts, model=model, add_noise=add_noise, name=name)
 
 
@@ -278,8 +304,8 @@ def make_fake_toas_fromMJDs(
     obs="GBT",
     error=1 * u.us,
     add_noise=False,
-    dm=None,
-    dm_error=1e-4 * pint.dmu,
+    wideband=False,
+    wideband_dm_error=1e-4 * pint.dmu,
     name="fake",
     include_bipm=False,
     include_gps=True,
@@ -303,10 +329,10 @@ def make_fake_toas_fromMJDs(
         uncertainty to attach to each TOA
     add_noise : bool, optional
         Add noise to the TOAs (otherwise `error` just populates the column)
-    dm : astropy.units.Quantity, optional
-        DM value to include with each TOA; default is
+    wideband : astropy.units.Quantity, optional
+        Whether to include wideband DM values with each TOA; default is
         not to include any DM information
-    dm_error : astropy.units.Quantity
+    wideband_dm_error : astropy.units.Quantity
         uncertainty to attach to each DM measurement
     name : str, optional
         Name for the TOAs (goes into the flags)
@@ -352,10 +378,10 @@ def make_fake_toas_fromMJDs(
         planets=model["PLANET_SHAPIRO"].value,
     )
     ts.table["error"] = error
-    if dm is not None:
-        for f in ts.table["flags"]:
-            f["pp_dm"] = str(dm.to_value(pint.dmu))
-            f["pp_dme"] = str(dm_error.to_value(pint.dmu))
+
+    if wideband:
+        ts = update_fake_dms(model, ts, wideband_dm_error, add_noise)
+
     return make_fake_toas(ts, model=model, add_noise=add_noise, name=name)
 
 
@@ -387,6 +413,11 @@ def make_fake_toas_fromtim(timfile, model, add_noise=False, name="fake"):
     :func:`make_fake_toas`
     """
     input_ts = pint.toa.get_TOAs(timfile)
+
+    if input_ts.is_wideband():
+        dm_errors = input_ts.get_dm_errors()
+        ts = update_fake_dms(model, ts, dm_errors, add_noise)
+
     return make_fake_toas(input_ts, model=model, add_noise=add_noise, name=name)
 
 
@@ -454,12 +485,12 @@ def calculate_random_models(
     freqs = np.zeros((Nmodels, Nmjd), dtype=np.float128) * u.Hz
 
     cov_matrix = fitter.parameter_covariance_matrix
-    # this is a list of the parameter names in the order they appear in the coviarance matrix
+    # this is a list of the parameter names in the order they appear in the covariance matrix
     param_names = cov_matrix.get_label_names(axis=0)
     # this is a dictionary with the parameter values, but it might not be in the same order
     # and it leaves out the Offset parameter
     param_values = fitter.model.get_params_dict("free", "value")
-    mean_vector = np.array([param_values[x] for x in param_names if not x == "Offset"])
+    mean_vector = np.array([param_values[x] for x in param_names if x != "Offset"])
     if params == "all":
         # remove the first column and row (absolute phase)
         if param_names[0] == "Offset":
@@ -512,7 +543,4 @@ def calculate_random_models(
     if return_time:
         dphase /= freqs
 
-    if keep_models:
-        return dphase, random_models
-    else:
-        return dphase
+    return (dphase, random_models) if keep_models else dphase

tests/test_downhill_fitter.py

@@ -80,7 +80,7 @@ def model_eccentric_toas_wb():
                 model_eccentric,
                 freq=1000 * u.MHz,
                 obs="@",
-                dm=10 * u.pc / u.cm**3,
+                wideband=True,
             ),
             make_fake_toas_uniform(
                 57000,
@@ -89,7 +89,7 @@ def model_eccentric_toas_wb():
                 model_eccentric,
                 freq=2000 * u.MHz,
                 obs="@",
-                dm=10 * u.pc / u.cm**3,
+                wideband=True,
             ),
         ]
     )
@@ -180,6 +180,7 @@ def test_wls_two_step(model_eccentric_toas):
     with pytest.raises(pint.fitter.MaxiterReached):
         f.fit_toas(maxiter=2)
     assert not f.converged
+
     f2 = pint.fitter.DownhillWLSFitter(toas, model_wrong)
     f2.model.free_params = ["ECC"]
     with pytest.raises(pint.fitter.MaxiterReached):
@@ -226,7 +227,7 @@ def test_wb_two_step(model_eccentric_toas_wb, full_cov):
         f2.fit_toas(maxiter=1, full_cov=full_cov)
     with pytest.raises(pint.fitter.MaxiterReached):
         f2.fit_toas(maxiter=1, full_cov=full_cov)
-    # FIXME: The full_cov version differs at the 1e-10 level fror some reason, is it a failure really?
+    # FIXME: The full_cov version differs at the 1e-10 level for some reason, is it a failure really?
     assert np.abs(f.model.ECC.value - f2.model.ECC.value) < 1e-9
 
 

tests/test_fake_toas.py

@@ -274,3 +274,19 @@ def test_fake_DMfit():
 
     assert np.isclose(DMs.std(), f.model.DM.uncertainty, rtol=0.5)
     assert np.abs(DMs.mean() - f.model.DM.quantity) < 5 * f.model.DM.uncertainty
+
+
+def test_fake_wb_toas():
+    parfile = pint.config.examplefile("NGC6440E.par")
+    m = get_model(parfile)
+
+    tfu = pint.simulation.make_fake_toas_uniform(
+        startMJD=50000,
+        endMJD=51000,
+        ntoas=100,
+        model=m,
+        wideband=True,
+        add_noise=True,
+    )
+    assert len(tfu) == 100
+    assert all("pp_dm" in f and "pp_dme" in f for f in tfu.get_flags())

tests/test_fitter.py

@@ -106,25 +106,25 @@ def test_ftest_nb():
         frozen=False,
     )
     ft = f.ftest(F2, "Spindown", remove=False)
-    assert isinstance(ft["ft"], float) or isinstance(ft["ft"], bool)
+    assert isinstance(ft["ft"], (float, bool))
     # Test return the full output
     Ftest_dict = f.ftest(F2, "Spindown", remove=False, full_output=True)
-    assert isinstance(Ftest_dict["ft"], float) or isinstance(Ftest_dict["ft"], bool)
+    assert isinstance(Ftest_dict["ft"], (float, bool))
     # Test removing parameter
     F1 = param.prefixParameter(
         parameter_type="float", name="F1", value=0.0, units=u.Hz / u.s, frozen=False
     )
     ft = f.ftest(F1, "Spindown", remove=True)
-    assert isinstance(ft["ft"], float) or isinstance(ft["ft"], bool)
+    assert isinstance(ft["ft"], (float, bool))
     Ftest_dict = f.ftest(F1, "Spindown", remove=True, full_output=True)
-    assert isinstance(Ftest_dict["ft"], float) or isinstance(Ftest_dict["ft"], bool)
+    assert isinstance(Ftest_dict["ft"], (float, bool))
 
 
 def test_ftest_wb():
     """Test for wideband fitter class F-test."""
     wb_m = tm.get_model(os.path.join(datadir, "J0023+0923_ell1_simple.par"))
     wb_t = simulation.make_fake_toas_uniform(
-        56000.0, 56001.0, 10, wb_m, freq=1400.0 * u.MHz, obs="GBT", dm=wb_m.DM.quantity
+        56000.0, 56001.0, 10, wb_m, freq=1400.0 * u.MHz, obs="GBT", wideband=True
     )
     wb_f = fitter.WidebandTOAFitter(wb_t, wb_m)
     wb_f.fit_toas()

tests/test_fitter_compare.py

@@ -38,7 +38,7 @@ def wls():
 def wb():
     m = get_model(join(datadir, "NGC6440E.par"))
     t = make_fake_toas_uniform(
-        55000, 58000, 20, model=m, freq=1400 * u.MHz, dm=10 * pint.dmu
+        55000, 58000, 20, model=m, freq=1400 * u.MHz, wideband=True
     )
 
     wb = WidebandTOAFitter(t, m)
@@ -107,7 +107,7 @@ def m_t():
         )
     )
     toas = make_fake_toas_uniform(
-        55000, 57000, 20, model=model, add_noise=True, dm=10 * u.pc / u.cm**3
+        55000, 57000, 20, model=model, add_noise=True, wideband=True
     )
     return model, toas