Simulate correlated noise

CHANGELOG-unreleased.md

@@ -21,6 +21,7 @@ the released changes.
 - Support for wideband data in `pint.bayesian` (no correlated noise).
 - Added `DMWaveX` model (Fourier representation of DM noise)
 - Piecewise orbital model (`BinaryBTPiecewise`)
+- Simulate correlated noise using `pint.simulation` (also available via 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.normal(size=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/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

@@ -8,8 +8,10 @@
 
 import matplotlib
 import pint.scripts.zima as zima
-from pint.models import get_model_and_toas
+from pint.models import get_model_and_toas, get_model
+from pint.simulation import make_fake_toas_uniform
 from pint.residuals import Residuals
+from pint.fitter import DownhillGLSFitter
 
 
 @pytest.mark.parametrize("addnoise", ["", "--addnoise"])
@@ -92,3 +94,41 @@ def test_zima_fuzzdays(tmp_path):
         lines = sys.stdout.getvalue()
     finally:
         sys.stdout = saved_stdout
+
+
+def test_simulate_corrnoise(tmp_path):
+    parfile = datadir / "B1855+09_NANOGrav_9yv1.gls.par"
+
+    m = get_model(parfile)
+
+    # Simulated TOAs won't have the correct flags for some of these to work.
+    m.remove_component("ScaleToaError")
+    m.remove_component("EcorrNoise")
+    m.remove_component("DispersionDMX")
+    m.remove_component("PhaseJump")
+    m.remove_component("FD")
+    m.PLANET_SHAPIRO.value = False
+
+    t = make_fake_toas_uniform(
+        m.START.value,
+        m.FINISH.value,
+        1000,
+        m,
+        add_noise=True,
+        add_correlated_noise=True,
+    )
+
+    # Check if the created TOAs can be whitened using
+    # the original timing model. This won't work if the
+    # noise is not realized correctly.
+    ftr = DownhillGLSFitter(t, m)
+    ftr.fit_toas()
+    rc = sum(ftr.resids.noise_resids.values())
+    r = ftr.resids.time_resids
+    rw = r - rc
+    sigma = ftr.resids.get_data_error()
+
+    # This should be independent and standard-normal distributed.
+    x = (rw / sigma).to_value("")
+    assert np.isclose(np.std(x), 1, atol=0.2)
+    assert np.isclose(np.mean(x), 0, atol=0.01)