Simulate correlated noise

CHANGELOG-unreleased.md

@@ -20,6 +20,7 @@ the released changes.
 - Added radial velocity methods for binary models
 - Support for wideband data in `pint.bayesian` (no correlated noise).
 - Added `DMWaveX` model (Fourier representation of DM noise)
+- Simulate correlated noise using `pint.simulate` functions and the `zima` script.
 ### Fixed
 - Wave model `validate()` can correctly use PEPOCH to assign WAVEEPOCH parameter
 - Fixed RTD by specifying theme explicitly.

src/pint/data/examples/B1855+09_NANOGrav_9yv1.gls.par

@@ -453,7 +453,7 @@ FD3  1.07526915D-04  1      2.50177766D-05
 SOLARN0               0.00
 EPHEM               DE421
 ECL                 IERS2003
-CLK                 TT(BIPM)
+CLK                 TT(BIPM2019)
 UNITS               TDB
 TIMEEPH             FB90
 #T2CMETHOD           TEMPO

src/pint/scripts/zima.py

@@ -65,6 +65,12 @@ def main(argv=None):
         default=False,
         help="Actually add in random noise, or just populate the column",
     )
+    parser.add_argument(
+        "--addcorrnoise",
+        action="store_true",
+        default=False,
+        help="Add in a correlated noise realization if it's present in the model",
+    )
     parser.add_argument(
         "--wideband",
         action="store_true",
@@ -127,13 +133,17 @@ def main(argv=None):
             freq=np.atleast_1d(args.freq) * u.MHz,
             fuzz=args.fuzzdays * u.d,
             add_noise=args.addnoise,
+            add_correlated_noise=args.addcorrnoise,
             wideband=args.wideband,
             wideband_dm_error=args.dmerror * pint.dmu,
         )
     else:
         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,
+            add_correlated_noise=args.addcorrnoise,
         )
 
     # Write TOAs to a file

src/pint/simulation.py

@@ -128,7 +128,7 @@ def get_fake_toa_clock_versions(model, include_bipm=False, include_gps=True):
     }
 
 
-def make_fake_toas(ts, model, add_noise=False, name="fake"):
+def make_fake_toas(ts, model, add_noise=False, add_correlated_noise=False, name="fake"):
     """Make toas from an array of times
 
     Can include alternating frequencies if fed an array of frequencies,
@@ -142,6 +142,8 @@ def make_fake_toas(ts, model, add_noise=False, name="fake"):
         current model
     add_noise : bool, optional
         Add noise to the TOAs (otherwise `error` just populates the column)
+    add_correlated_noise : bool, optional
+        Add correlated noise to the TOAs if it's present in the timing mode.
     name : str, optional
         Name for the TOAs (goes into the flags)
 
@@ -156,6 +158,13 @@ def make_fake_toas(ts, model, add_noise=False, name="fake"):
     """
     tsim = deepcopy(ts)
     zero_residuals(tsim, model)
+
+    if add_correlated_noise:
+        U = model.noise_model_designmatrix(tsim)
+        b = model.noise_model_basis_weight(tsim)
+        corrn = (U @ (b**0.5 * np.random.randn(len(b)))) << u.s
+        tsim.adjust_TOAs(time.TimeDelta(corrn))
+
     if add_noise:
         # this function will include EFAC and EQUAD
         err = model.scaled_toa_uncertainty(tsim) * np.random.normal(size=len(tsim))
@@ -198,6 +207,7 @@ def make_fake_toas_uniform(
     obs="GBT",
     error=1 * u.us,
     add_noise=False,
+    add_correlated_noise=False,
     wideband=False,
     wideband_dm_error=1e-4 * pint.dmu,
     name="fake",
@@ -229,6 +239,8 @@ 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)
+    add_correlated_noise : bool, optional
+        Add correlated noise to the TOAs if it's present in the timing mode.
     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
@@ -301,7 +313,13 @@ def make_fake_toas_uniform(
     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)
+    return make_fake_toas(
+        ts,
+        model=model,
+        add_noise=add_noise,
+        add_correlated_noise=add_correlated_noise,
+        name=name,
+    )
 
 
 def make_fake_toas_fromMJDs(
@@ -311,6 +329,7 @@ def make_fake_toas_fromMJDs(
     obs="GBT",
     error=1 * u.us,
     add_noise=False,
+    add_correlated_noise=False,
     wideband=False,
     wideband_dm_error=1e-4 * pint.dmu,
     name="fake",
@@ -336,6 +355,8 @@ 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)
+    add_correlated_noise : bool, optional
+        Add correlated noise to the TOAs if it's present in the timing mode.
     wideband : astropy.units.Quantity, optional
         Whether to include wideband DM values with each TOA; default is
         not to include any DM information
@@ -396,10 +417,18 @@ def make_fake_toas_fromMJDs(
     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)
+    return make_fake_toas(
+        ts,
+        model=model,
+        add_noise=add_noise,
+        add_correlated_noise=add_correlated_noise,
+        name=name,
+    )
 
 
-def make_fake_toas_fromtim(timfile, model, add_noise=False, name="fake"):
+def make_fake_toas_fromtim(
+    timfile, model, add_noise=False, add_correlated_noise=False, name="fake"
+):
     """Make fake toas with the same times as an input tim file
 
     Can include alternating frequencies if fed an array of frequencies,
@@ -413,6 +442,8 @@ def make_fake_toas_fromtim(timfile, model, add_noise=False, name="fake"):
         current model
     add_noise : bool, optional
         Add noise to the TOAs (otherwise `error` just populates the column)
+    add_correlated_noise : bool, optional
+        Add correlated noise to the TOAs if it's present in the timing mode.
     name : str, optional
         Name for the TOAs (goes into the flags)
 
@@ -432,7 +463,13 @@ def make_fake_toas_fromtim(timfile, model, add_noise=False, name="fake"):
         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)
+    return make_fake_toas(
+        input_ts,
+        model=model,
+        add_noise=add_noise,
+        add_correlated_noise=add_correlated_noise,
+        name=name,
+    )
 
 
 def calculate_random_models(

tests/.gitignore

@@ -0,0 +1,2 @@
+*.ipynb
+chains/*

tests/datafile/B1855+09_NANOGrav_9yv1.gls.par

@@ -453,7 +453,7 @@ FD3  1.07526915D-04  1      2.50177766D-05
 SOLARN0               0.00
 EPHEM               DE421
 ECL                 IERS2003
-CLK                 TT(BIPM)
+CLK                 TT(BIPM2019)
 UNITS               TDB
 TIMEEPH             FB90
 #T2CMETHOD           TEMPO

tests/test_zima.py

@@ -92,3 +92,11 @@ def test_zima_fuzzdays(tmp_path):
         lines = sys.stdout.getvalue()
     finally:
         sys.stdout = saved_stdout
+
+
+def test_zima_corrnoise(tmp_path):
+    parfile = datadir / "B1855+09_NANOGrav_9yv1.gls.par"
+    output_timfile = tmp_path / "fake_testzima.tim"
+    cmd = f"--addnoise --addcorrnoise {parfile} {output_timfile}"
+    zima.main(cmd.split())
+    assert os.path.isfile(output_timfile)