Merge pull request #1802 from abhisrkckl/cmwavex

Fourier series representation of chromatic noise (`CMWaveX`)

CHANGELOG-unreleased.md

@@ -20,7 +20,11 @@ the released changes.
 `pint.erfautils`, `pint.fits_utils`, `pint.logging` and `pint.residuals`
 - Doing `model.par = something` will try to assign to `par.quantity` or `par.value` but will give warning
 - `PLChromNoise` component to model chromatic red noise with a power law spectrum
+- Fourier series representation of chromatic noise (`CMWaveX`)
+- `pint.utils.cmwavex_setup` and `pint.utils.plchromnoise_from_cmwavex` functions
+- More validation for correlated noise components in `TimingModel.validate_component_types()`
 ### Fixed
+- Bug in `DMWaveX.get_indices()` function
 - Explicit type conversion in `woodbury_dot()` function
 - Documentation: Fixed empty descriptions in the timing model components table
 - BIC implementation

docs/examples/rednoise-fit-example.py

@@ -1,5 +1,5 @@
 # %% [markdown]
-# # Red noise and DM noise fitting examples
+# # Red noise, DM noise, and chromatic noise fitting examples
 #
 # This notebook provides an example on how to fit for red noise
 # and DM noise using PINT using simulated datasets.
@@ -22,8 +22,10 @@
 from pint.logging import setup as setup_log
 from pint.fitter import WLSFitter
 from pint.utils import (
+    cmwavex_setup,
     dmwavex_setup,
     find_optimal_nharms,
+    plchromnoise_from_cmwavex,
     wavex_setup,
     plrednoise_from_wavex,
     pldmnoise_from_dmwavex,
@@ -370,3 +372,162 @@
 plt.xlabel("Frequency (Hz)")
 plt.axvline(fyr, color="black", ls="dotted", label="1 yr$^{-1}$")
 plt.legend()
+
+
+# %% [markdown]
+# ## Chromatic noise fitting
+# Let us now do a similar kind of analysis for chromatic noise.
+
+# %%
+par_sim = """
+    PSR           SIM5
+    RAJ           05:00:00     1
+    DECJ          15:00:00     1
+    PEPOCH        55000
+    F0            100          1
+    F1            -1e-15       1 
+    PHOFF         0            1
+    DM            15
+    CM            1.2          1
+    TNCHROMIDX    3.5
+    TNCHROMAMP    -13
+    TNCHROMGAM    3.5
+    TNCHROMC      30
+    TZRMJD        55000
+    TZRFRQ        1400 
+    TZRSITE       gbt
+    UNITS         TDB
+    EPHEM         DE440
+    CLOCK         TT(BIPM2019)
+"""
+
+m = get_model(StringIO(par_sim))
+
+# %%
+# Generate the simulated TOAs.
+ntoas = 2000
+toaerrs = np.random.uniform(0.5, 2.0, ntoas) * u.us
+freqs = np.linspace(500, 1500, 8) * u.MHz
+
+t = make_fake_toas_uniform(
+    startMJD=53001,
+    endMJD=57001,
+    ntoas=ntoas,
+    model=m,
+    freq=freqs,
+    obs="gbt",
+    error=toaerrs,
+    add_noise=True,
+    add_correlated_noise=True,
+    name="fake",
+    include_bipm=True,
+    multi_freqs_in_epoch=True,
+)
+
+# %%
+# Find the optimum number of harmonics by minimizing AIC.
+m1 = deepcopy(m)
+m1.remove_component("PLChromNoise")
+
+m2 = deepcopy(m1)
+
+nharm_opt = m.TNCHROMC.value
+
+# %%
+# Now create a new model with the optimum number of
+# harmonics
+m2 = deepcopy(m1)
+
+Tspan = t.get_mjds().max() - t.get_mjds().min()
+cmwavex_setup(m2, T_span=Tspan, n_freqs=nharm_opt, freeze_params=False)
+
+ftr = WLSFitter(t, m2)
+ftr.fit_toas(maxiter=10)
+m2 = ftr.model
+
+print(m2)
+
+# %% [markdown]
+# ### Estimating the spectral parameters from the `CMWaveX` fit.
+
+# %%
+# Get the Fourier amplitudes and powers and their uncertainties.
+# Note that the `CMWaveX` amplitudes have the units of pc/cm^3/MHz^2.
+# We multiply them by a constant factor to convert them to dimensions
+# of time so that they are consistent with `PLChromNoise`.
+scale = DMconst / 1400**m.TNCHROMIDX.value
+
+idxs = np.array(m2.components["CMWaveX"].get_indices())
+a = np.array(
+    [(scale * m2[f"CMWXSIN_{idx:04d}"].quantity).to_value("s") for idx in idxs]
+)
+da = np.array(
+    [(scale * m2[f"CMWXSIN_{idx:04d}"].uncertainty).to_value("s") for idx in idxs]
+)
+b = np.array(
+    [(scale * m2[f"CMWXCOS_{idx:04d}"].quantity).to_value("s") for idx in idxs]
+)
+db = np.array(
+    [(scale * m2[f"CMWXCOS_{idx:04d}"].uncertainty).to_value("s") for idx in idxs]
+)
+print(len(idxs))
+
+P = (a**2 + b**2) / 2
+dP = ((a * da) ** 2 + (b * db) ** 2) ** 0.5
+
+f0 = (1 / Tspan).to_value(u.Hz)
+fyr = (1 / u.year).to_value(u.Hz)
+
+# %%
+# We can create a `PLChromNoise` model from the `CMWaveX` model.
+# This will estimate the spectral parameters from the `CMWaveX`
+# amplitudes.
+m3 = plchromnoise_from_cmwavex(m2)
+print(m3)
+
+# %%
+# Now let us plot the estimated spectrum with the injected
+# spectrum.
+plt.subplot(211)
+plt.errorbar(
+    idxs * f0,
+    b * 1e6,
+    db * 1e6,
+    ls="",
+    marker="o",
+    label="$\\hat{a}_j$ (CMWXCOS)",
+    color="red",
+)
+plt.errorbar(
+    idxs * f0,
+    a * 1e6,
+    da * 1e6,
+    ls="",
+    marker="o",
+    label="$\\hat{b}_j$ (CMWXSIN)",
+    color="blue",
+)
+plt.axvline(fyr, color="black", ls="dotted")
+plt.axhline(0, color="grey", ls="--")
+plt.ylabel("Fourier coeffs ($\mu$s)")
+plt.xscale("log")
+plt.legend(fontsize=8)
+
+plt.subplot(212)
+plt.errorbar(
+    idxs * f0, P, dP, ls="", marker="o", label="Spectral power (PINT)", color="k"
+)
+P_inj = m.components["PLChromNoise"].get_noise_weights(t)[::2]
+plt.plot(idxs * f0, P_inj, label="Injected Spectrum", color="r")
+P_est = m3.components["PLChromNoise"].get_noise_weights(t)[::2]
+print(len(idxs), len(P_est))
+plt.plot(idxs * f0, P_est, label="Estimated Spectrum", color="b")
+plt.xscale("log")
+plt.yscale("log")
+plt.ylabel("Spectral power (s$^2$)")
+plt.xlabel("Frequency (Hz)")
+plt.axvline(fyr, color="black", ls="dotted", label="1 yr$^{-1}$")
+plt.legend()
+
+
+# %%

src/pint/logging.py

@@ -141,8 +141,7 @@ def __init__(
         never: Optional[List[str]] = None,
         onlyonce_level: str = "INFO",
     ) -> None:
-        r"""
-        Define regexs for messages that will only be seen once.  Use ``\S+`` for a variable that might change.
+        r"""Define regexes for messages that will only be seen once.  Use ``\S+`` for a variable that might change.
         If a message comes through with a new value for that variable, it will be seen.
 
         Make sure to escape other regex commands like ``()``.

src/pint/models/__init__.py

@@ -26,6 +26,7 @@
 from pint.models.binary_ddk import BinaryDDK
 from pint.models.binary_ell1 import BinaryELL1, BinaryELL1H, BinaryELL1k
 from pint.models.chromatic_model import ChromaticCM
+from pint.models.cmwavex import CMWaveX
 from pint.models.dispersion_model import (
     DispersionDM,
     DispersionDMX,

src/pint/models/binary_ddk.py

@@ -42,21 +42,26 @@ def _convert_kom(kom):
 
 
 class BinaryDDK(BinaryDD):
-    """Damour and Deruelle model with kinematics.
+    r"""Damour and Deruelle model with kinematics.
 
     This extends the :class:`pint.models.binary_dd.BinaryDD` model with
     "Shklovskii" and "Kopeikin" terms that account for the finite distance
     of the system from Earth, the finite size of the system, and the
     interaction of these with the proper motion.
 
-    From Kopeikin (1995) this includes :math:`\Delta_{\pi M}` (Equation 17), the mixed annual-orbital parallax term, which changes :math:`a_1` and :math:`\omega`
-    (:meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_a1_parallax` and :meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_omega_parallax`).
+    From Kopeikin (1995) this includes :math:`\Delta_{\pi M}` (Equation 17),
+    the mixed annual-orbital parallax term, which changes :math:`a_1` and :math:`\omega`
+    (:meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_a1_parallax`
+    and :meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_omega_parallax`).
 
-    It does not include :math:`\Delta_{\pi P}`, the pure pulsar orbital parallax term (Equation 14).
+    It does not include :math:`\Delta_{\pi P}`, the pure pulsar orbital parallax term
+    (Equation 14).
 
-    From Kopeikin (1996) this includes apparent changes in :math:`\omega`, :math:`a_1`, and :math:`i` due to the proper motion
-    (:meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_omega_proper_motion`, :meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_a1_proper_motion`,
-    :meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_kin_proper_motion`) (Equations 8, 9, 10).
+    From Kopeikin (1996) this includes apparent changes in :math:`\omega`, :math:`a_1`, and
+    :math:`i` due to the proper motion (:meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_omega_proper_motion`,
+    :meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_a1_proper_motion`,
+    :meth:`~pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel.delta_kin_proper_motion`)
+    (Equations 8, 9, 10).
 
     The actual calculations for this are done in
     :class:`pint.models.stand_alone_psr_binaries.DDK_model.DDKmodel`.
@@ -226,14 +231,14 @@ def validate(self):
             warnings.warn("Using A1DOT with a DDK model is not advised.")
 
     def alternative_solutions(self):
-        """Alternative Kopeikin solutions (potential local minima)
+        r"""Alternative Kopeikin solutions (potential local minima)
 
         There are 4 potential local minima for a DDK model where a1dot is the same
         These are given by where Eqn. 8 in Kopeikin (1996) is equal to the best-fit value.
 
         We first define the symmetry point where a1dot is zero (in equatorial coordinates):
 
-        :math:`KOM_0 = \\tan^{-1} (\mu_{\delta} / \mu_{\\alpha})`
+        :math:`KOM_0 = \tan^{-1} (\mu_{\delta} / \mu_{\alpha})`
 
         The solutions are then: