Merge pull request #1868 from abhisrkckl/sim-dmgp

Simulate correlated DM noise in wideband TOAs

.github/workflows/ci_test.yml

@@ -115,7 +115,7 @@ jobs:
           shell: bash -el {0}
           run: |
             python -m pip install --upgrade pip
-            python -m pip install tox pytest hypothesis numdifftools pathos setuptools
+            python -m pip install tox pytest hypothesis numdifftools pathos setuptools statsmodels
         - name: Print OS, machine info
           shell: bash -el {0}
           run: |

CHANGELOG-unreleased.md

@@ -10,6 +10,7 @@ the released changes.
 ## Unreleased
 ### Changed
 ### Added
+- Simulate correlated DM noise for wideband TOAs
 - Type hints in `pint.models.timing_model`
 ### Fixed
 - Made `TimingModel.is_binary()` more robust. 

requirements_dev.txt

@@ -39,3 +39,4 @@ loguru
 # click<=8.0.4
 gprof2dot
 py-cpuinfo
+statsmodels

src/pint/models/noise_model.py

@@ -13,6 +13,9 @@
 
 
 class NoiseComponent(Component):
+
+    introduces_dm_errors = False
+
     def __init__(
         self,
     ):
@@ -232,6 +235,7 @@ class ScaleDmError(NoiseComponent):
     category = "scale_dm_error"
 
     introduces_correlated_errors = False
+    introduces_dm_errors = True
 
     def __init__(
         self,
@@ -470,6 +474,7 @@ class PLDMNoise(NoiseComponent):
     category = "pl_DM_noise"
 
     introduces_correlated_errors = True
+    introduces_dm_errors = True
     is_time_correlated = True
 
     def __init__(

src/pint/simulation.py

@@ -10,6 +10,7 @@
 from loguru import logger as log
 from astropy import time
 
+from pint.models.noise_model import NoiseComponent
 from pint.types import time_like, file_like
 import pint.residuals
 import pint.toa
@@ -185,6 +186,7 @@ def update_fake_dms(
     ts: pint.toa.TOAs,
     dm_error: u.Quantity,
     add_noise: bool,
+    add_correlated_noise: bool,
 ) -> pint.toa.TOAs:
     """Update simulated wideband DM information in TOAs.
 
@@ -209,6 +211,20 @@ def update_fake_dms(
     if add_noise:
         dms += scaled_dm_errors.to(pint.dmu) * np.random.randn(len(scaled_dm_errors))
 
+    if add_correlated_noise:
+        dm_noise = np.zeros(len(toas)) * pint.dmu
+        for noise_comp in model.NoiseComponent_list:
+            if (
+                noise_comp.introduces_correlated_errors
+                and noise_comp.introduces_dm_errors
+            ):
+                U = noise_comp.get_noise_basis(toas)
+                b = noise_comp.get_noise_weights(toas)
+                delay = (U @ (b**0.5 * np.random.normal(size=len(b)))) << u.s
+                freqs = model.barycentric_radio_freq(toas)
+                dm_noise += (delay / pint.DMconst * freqs**2).to(pint.dmu)
+        dms += dm_noise
+
     for f, dm in zip(toas.table["flags"], dms):
         f["pp_dm"] = str(dm.to_value(pint.dmu))
 
@@ -338,7 +354,9 @@ def make_fake_toas_uniform(
     )
 
     if wideband:
-        ts = update_fake_dms(model, ts, wideband_dm_error, add_noise)
+        ts = update_fake_dms(
+            model, ts, wideband_dm_error, add_noise, add_correlated_noise
+        )
 
     return make_fake_toas(
         ts,
@@ -466,7 +484,9 @@ def make_fake_toas_fromMJDs(
     )
 
     if wideband:
-        ts = update_fake_dms(model, ts, wideband_dm_error, add_noise)
+        ts = update_fake_dms(
+            model, ts, wideband_dm_error, add_noise, add_correlated_noise
+        )
 
     return make_fake_toas(
         ts,
@@ -517,7 +537,7 @@ def make_fake_toas_fromtim(
 
     if ts.is_wideband():
         dm_errors = ts.get_dm_errors()
-        ts = update_fake_dms(model, ts, dm_errors, add_noise)
+        ts = update_fake_dms(model, ts, dm_errors, add_noise, add_correlated_noise)
 
     return make_fake_toas(
         ts,

tests/test_fake_toas.py

@@ -12,6 +12,7 @@
 from pint.fitter import GLSFitter, DownhillGLSFitter
 from pinttestdata import datadir
 import pytest
+from statsmodels.stats.diagnostic import acorr_ljungbox
 
 
 def roundtrip(toas, model):
@@ -511,3 +512,59 @@ def test_simulate_uniform_multifreq(multifreq):
         multi_freqs_in_epoch=multifreq,
     )
     assert len(t) == ntoas
+
+
+def test_simulate_wideband_dmgp():
+    par = """
+        PSRJ           J0023+0923
+        RAJ             00:23:16.8790858         1  0.00002408141295805134   
+        DECJ           +09:23:23.86936           1  0.00082010713730773120   
+        F0             327.84702062954047136     1  0.00000000000295205483   
+        F1             -1.2278326306812866375e-15 1  3.8219244605614075223e-19
+        PEPOCH         56199.999797564144902       
+        POSEPOCH       56199.999797564144902       
+        DMEPOCH        56200                       
+        DM             14.327978186774068347     1  0.00006751663559857748   
+        BINARY         ELL1
+        PB             0.13879914244858396754    1  0.00000000003514075083   
+        A1             0.034841158415224894973   1  0.00000012173038389247   
+        TASC           56178.804891768506529     1  0.00000007765191894742   
+        EPS1           1.6508830631753595232e-05 1  0.00000477568412215803   
+        EPS2           3.9656838708709247373e-06 1  0.00000458951091435993   
+        CLK            TT(BIPM2015)
+        MODE 1
+        UNITS          TDB
+        TIMEEPH        FB90
+        DILATEFREQ     N
+        PLANET_SHAPIRO N
+        CORRECT_TROPOSPHERE  N
+        EPHEM          DE436
+        TNRedAmp -13.3087
+        TNRedGam 1.5
+        TNRedC 14
+        TNDMAMP -12.2
+        TNDMGAM 3.5
+        TNDMC 15
+    """
+
+    m = get_model(io.StringIO(par))
+    t = pint.simulation.make_fake_toas_uniform(
+        startMJD=54000,
+        endMJD=56000,
+        ntoas=200,
+        model=m,
+        add_noise=True,
+        wideband=True,
+        add_correlated_noise=True,
+    )
+
+    assert t.is_wideband()
+
+    # There is correlated noise in DMs
+    assert (
+        sum(((t.get_dms() - m.total_dm(t)) / m.scaled_dm_uncertainty(t)) ** 2).si.value
+        / len(t)
+        > 10
+    )
+
+    assert all(acorr_ljungbox(t.get_dms() - m.total_dm(t)).lb_pvalue < 1e-5)

tox.ini

@@ -42,6 +42,7 @@ deps =
     numdifftools
     pathos
     setuptools
+    statsmodels
 commands =
     pip freeze
     !cov: pytest --reruns 5
@@ -92,6 +93,7 @@ deps =
     pytest-rerunfailures
     coverage
     hypothesis<=6.72.0
+    statsmodels
 commands =
     pytest --reruns 5