From 5cbd67da4c60e3a262326e51bd6aa9220c658f7d Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 10:50:50 +0200
Subject: [PATCH 01/22] add cmwavex

---
 src/pint/models/cmwavex.py | 387 +++++++++++++++++++++++++++++++++++++
 src/pint/models/dmwavex.py |   4 +-
 2 files changed, 389 insertions(+), 2 deletions(-)
 create mode 100644 src/pint/models/cmwavex.py

diff --git a/src/pint/models/cmwavex.py b/src/pint/models/cmwavex.py
new file mode 100644
index 000000000..465700565
--- /dev/null
+++ b/src/pint/models/cmwavex.py
@@ -0,0 +1,387 @@
+"""Chromatic variations expressed as a sum of sinusoids."""
+
+import astropy.units as u
+import numpy as np
+from loguru import logger as log
+from warnings import warn
+
+from pint.models.parameter import MJDParameter, prefixParameter
+from pint.models.timing_model import MissingParameter
+from pint.models.chromatic_model import Chromatic, cmu
+from pint import DMconst
+
+
+class CMWaveX(Chromatic):
+    """
+    Fourier representation of chromatic variations.
+
+    Used for decomposition of chromatic noise into a series of sine/cosine components with the amplitudes as fitted parameters.
+
+    Parameters supported:
+
+    .. paramtable::
+        :class: pint.models.cmwavex.CMWaveX
+
+    To set up a CMWaveX model, users can use the `pint.utils` function `cmwavex_setup()` with either a list of frequencies or a choice
+    of harmonics of a base frequency determined by 2 * pi /Timespan
+    """
+
+    register = True
+    category = "cmwavex"
+
+    def __init__(self):
+        super().__init__()
+        self.add_param(
+            MJDParameter(
+                name="CMWXEPOCH",
+                description="Reference epoch for Fourier representation of chromatic noise",
+                time_scale="tdb",
+                tcb2tdb_scale_factor=u.Quantity(1),
+            )
+        )
+        self.add_cmwavex_component(0.1, index=1, cmwxsin=0, cmwxcos=0, frozen=False)
+        self.set_special_params(["CMWXFREQ_0001", "CMWXSIN_0001", "CMWXCOS_0001"])
+        self.cm_value_funcs += [self.cmwavex_cm]
+        self.delay_funcs_component += [self.cmwavex_delay]
+
+    def add_cmwavex_component(
+        self, cmwxfreq, index=None, cmwxsin=0, cmwxcos=0, frozen=True
+    ):
+        """
+        Add CMWaveX component
+
+        Parameters
+        ----------
+
+        cmwxfreq : float or astropy.quantity.Quantity
+            Base frequency for CMWaveX component
+        index : int, None
+            Interger label for CMWaveX component. If None, will increment largest used index by 1.
+        cmwxsin : float or astropy.quantity.Quantity
+            Sine amplitude for CMWaveX component
+        cmwxcos : float or astropy.quantity.Quantity
+            Cosine amplitude for CMWaveX component
+        frozen : iterable of bool or bool
+            Indicates whether CMWaveX parameters will be fit
+
+        Returns
+        -------
+
+        index : int
+            Index that has been assigned to new CMWaveX component
+        """
+
+        #### If index is None, increment the current max CMWaveX index by 1. Increment using CMWXFREQ
+        if index is None:
+            dct = self.get_prefix_mapping_component("CMWXFREQ_")
+            index = np.max(list(dct.keys())) + 1
+        i = f"{int(index):04d}"
+
+        if int(index) in self.get_prefix_mapping_component("CMWXFREQ_"):
+            raise ValueError(
+                f"Index '{index}' is already in use in this model. Please choose another"
+            )
+
+        if isinstance(cmwxsin, u.quantity.Quantity):
+            cmwxsin = cmwxsin.to_value(cmu)
+        if isinstance(cmwxcos, u.quantity.Quantity):
+            cmwxcos = cmwxcos.to_value(cmu)
+        if isinstance(cmwxfreq, u.quantity.Quantity):
+            cmwxfreq = cmwxfreq.to_value(1 / u.d)
+        self.add_param(
+            prefixParameter(
+                name=f"CMWXFREQ_{i}",
+                description="Component frequency for Fourier representation of chromatic noise",
+                units="1/d",
+                value=cmwxfreq,
+                parameter_type="float",
+                tcb2tdb_scale_factor=u.Quantity(1),
+            )
+        )
+        self.add_param(
+            prefixParameter(
+                name=f"CMWXSIN_{i}",
+                description="Sine amplitudes for Fourier representation of chromatic noise",
+                units=cmu,
+                value=cmwxsin,
+                frozen=frozen,
+                parameter_type="float",
+                tcb2tdb_scale_factor=DMconst,
+            )
+        )
+        self.add_param(
+            prefixParameter(
+                name=f"CMWXCOS_{i}",
+                description="Cosine amplitudes for Fourier representation of chromatic noise",
+                units=cmu,
+                value=cmwxcos,
+                frozen=frozen,
+                parameter_type="float",
+                tcb2tdb_scale_factor=DMconst,
+            )
+        )
+        self.setup()
+        self.validate()
+        return index
+
+    def add_cmwavex_components(
+        self, cmwxfreqs, indices=None, cmwxsins=0, cmwxcoses=0, frozens=True
+    ):
+        """
+        Add CMWaveX components with specified base frequencies
+
+        Parameters
+        ----------
+
+        cmwxfreqs : iterable of float or astropy.quantity.Quantity
+            Base frequencies for CMWaveX components
+        indices : iterable of int, None
+            Interger labels for CMWaveX components. If None, will increment largest used index by 1.
+        cmwxsins : iterable of float or astropy.quantity.Quantity
+            Sine amplitudes for CMWaveX components
+        cmwxcoses : iterable of float or astropy.quantity.Quantity
+            Cosine amplitudes for CMWaveX components
+        frozens : iterable of bool or bool
+            Indicates whether sine and cosine amplitudes of CMwavex components will be fit
+
+        Returns
+        -------
+
+        indices : list
+            Indices that have been assigned to new CMWaveX components
+        """
+
+        if indices is None:
+            indices = [None] * len(cmwxfreqs)
+        cmwxsins = np.atleast_1d(cmwxsins)
+        cmwxcoses = np.atleast_1d(cmwxcoses)
+        if len(cmwxsins) == 1:
+            cmwxsins = np.repeat(cmwxsins, len(cmwxfreqs))
+        if len(cmwxcoses) == 1:
+            cmwxcoses = np.repeat(cmwxcoses, len(cmwxfreqs))
+        if len(cmwxsins) != len(cmwxfreqs):
+            raise ValueError(
+                f"Number of base frequencies {len(cmwxfreqs)} doesn't match number of sine ampltudes {len(cmwxsins)}"
+            )
+        if len(cmwxcoses) != len(cmwxfreqs):
+            raise ValueError(
+                f"Number of base frequencies {len(cmwxfreqs)} doesn't match number of cosine ampltudes {len(cmwxcoses)}"
+            )
+        frozens = np.atleast_1d(frozens)
+        if len(frozens) == 1:
+            frozens = np.repeat(frozens, len(cmwxfreqs))
+        if len(frozens) != len(cmwxfreqs):
+            raise ValueError(
+                "Number of base frequencies must match number of frozen values"
+            )
+        #### If indices is None, increment the current max CMWaveX index by 1. Increment using CMWXFREQ
+        dct = self.get_prefix_mapping_component("CMWXFREQ_")
+        last_index = np.max(list(dct.keys()))
+        added_indices = []
+        for cmwxfreq, index, cmwxsin, cmwxcos, frozen in zip(
+            cmwxfreqs, indices, cmwxsins, cmwxcoses, frozens
+        ):
+            if index is None:
+                index = last_index + 1
+                last_index += 1
+            elif index in list(dct.keys()):
+                raise ValueError(
+                    f"Attempting to insert CMWXFREQ_{index:04d} but it already exists"
+                )
+            added_indices.append(index)
+            i = f"{int(index):04d}"
+
+            if int(index) in dct:
+                raise ValueError(
+                    f"Index '{index}' is already in use in this model. Please choose another"
+                )
+            if isinstance(cmwxfreq, u.quantity.Quantity):
+                cmwxfreq = cmwxfreq.to_value(u.d**-1)
+            if isinstance(cmwxsin, u.quantity.Quantity):
+                cmwxsin = cmwxsin.to_value(cmu)
+            if isinstance(cmwxcos, u.quantity.Quantity):
+                cmwxcos = cmwxcos.to_value(cmu)
+            log.trace(f"Adding CMWXSIN_{i} and CMWXCOS_{i} at frequency CMWXFREQ_{i}")
+            self.add_param(
+                prefixParameter(
+                    name=f"CMWXFREQ_{i}",
+                    description="Component frequency for Fourier representation of chromatic noise",
+                    units="1/d",
+                    value=cmwxfreq,
+                    parameter_type="float",
+                    tcb2tdb_scale_factor=u.Quantity(1),
+                )
+            )
+            self.add_param(
+                prefixParameter(
+                    name=f"CMWXSIN_{i}",
+                    description="Sine amplitude for Fourier representation of chromatic noise",
+                    units=cmu,
+                    value=cmwxsin,
+                    parameter_type="float",
+                    frozen=frozen,
+                    tcb2tdb_scale_factor=DMconst,
+                )
+            )
+            self.add_param(
+                prefixParameter(
+                    name=f"CMWXCOS_{i}",
+                    description="Cosine amplitude for Fourier representation of chromatic noise",
+                    units=cmu,
+                    value=cmwxcos,
+                    parameter_type="float",
+                    frozen=frozen,
+                    tcb2tdb_scale_factor=DMconst,
+                )
+            )
+        self.setup()
+        self.validate()
+        return added_indices
+
+    def remove_cmwavex_component(self, index):
+        """
+        Remove all CMWaveX components associated with a given index or list of indices
+
+        Parameters
+        ----------
+        index : float, int, list, np.ndarray
+            Number or list/array of numbers corresponding to CMWaveX indices to be removed from model.
+        """
+
+        if isinstance(index, (int, float, np.int64)):
+            indices = [index]
+        elif isinstance(index, (list, set, np.ndarray)):
+            indices = index
+        else:
+            raise TypeError(
+                f"index most be a float, int, set, list, or array - not {type(index)}"
+            )
+        for index in indices:
+            index_rf = f"{int(index):04d}"
+            for prefix in ["CMWXFREQ_", "CMWXSIN_", "CMWXCOS_"]:
+                self.remove_param(prefix + index_rf)
+        self.validate()
+
+    def get_indices(self):
+        """
+        Returns an array of intergers corresponding to CMWaveX component parameters using CMWXFREQs
+
+        Returns
+        -------
+        inds : np.ndarray
+        Array of CMWaveX indices in model.
+        """
+        inds = [int(p.split("_")[-1]) for p in self.params if "CMWXFREQ_" in p]
+        return np.array(inds)
+
+    # Initialize setup
+    def setup(self):
+        super().setup()
+        # Get CMWaveX mapping and register CMWXSIN and CMWXCOS derivatives
+        for prefix_par in self.get_params_of_type("prefixParameter"):
+            if prefix_par.startswith("CMWXSIN_"):
+                self.register_deriv_funcs(self.d_delay_d_cmparam, prefix_par)
+                self.register_cm_deriv_funcs(self.d_cm_d_CMWXSIN, prefix_par)
+            if prefix_par.startswith("CMWXCOS_"):
+                self.register_deriv_funcs(self.d_delay_d_cmparam, prefix_par)
+                self.register_cm_deriv_funcs(self.d_cm_d_CMWXCOS, prefix_par)
+            self.cmwavex_freqs = list(
+                self.get_prefix_mapping_component("CMWXFREQ_").keys()
+            )
+            self.num_cmwavex_freqs = len(self.cmwavex_freqs)
+
+    def validate(self):
+        # Validate all the CMWaveX parameters
+        super().validate()
+        self.setup()
+        CMWXFREQ_mapping = self.get_prefix_mapping_component("CMWXFREQ_")
+        CMWXSIN_mapping = self.get_prefix_mapping_component("CMWXSIN_")
+        CMWXCOS_mapping = self.get_prefix_mapping_component("CMWXCOS_")
+        if CMWXFREQ_mapping.keys() != CMWXSIN_mapping.keys():
+            raise ValueError(
+                "CMWXFREQ_ parameters do not match CMWXSIN_ parameters."
+                "Please check your prefixed parameters"
+            )
+        if CMWXFREQ_mapping.keys() != CMWXCOS_mapping.keys():
+            raise ValueError(
+                "CMWXFREQ_ parameters do not match CMWXCOS_ parameters."
+                "Please check your prefixed parameters"
+            )
+        # if len(CMWXFREQ_mapping.keys()) != len(CMWXSIN_mapping.keys()):
+        #     raise ValueError(
+        #         "The number of CMWXFREQ_ parameters do not match the number of CMWXSIN_ parameters."
+        #         "Please check your prefixed parameters"
+        #     )
+        # if len(CMWXFREQ_mapping.keys()) != len(CMWXCOS_mapping.keys()):
+        #     raise ValueError(
+        #         "The number of CMWXFREQ_ parameters do not match the number of CMWXCOS_ parameters."
+        #         "Please check your prefixed parameters"
+        #     )
+        if CMWXSIN_mapping.keys() != CMWXCOS_mapping.keys():
+            raise ValueError(
+                "CMWXSIN_ parameters do not match CMWXCOS_ parameters."
+                "Please check your prefixed parameters"
+            )
+        if len(CMWXSIN_mapping.keys()) != len(CMWXCOS_mapping.keys()):
+            raise ValueError(
+                "The number of CMWXSIN_ and CMWXCOS_ parameters do not match"
+                "Please check your prefixed parameters"
+            )
+        wfreqs = np.zeros(len(CMWXFREQ_mapping))
+        for j, index in enumerate(CMWXFREQ_mapping):
+            if (getattr(self, f"CMWXFREQ_{index:04d}").value == 0) or (
+                getattr(self, f"CMWXFREQ_{index:04d}").quantity is None
+            ):
+                raise ValueError(
+                    f"CMWXFREQ_{index:04d} is zero or None. Please check your prefixed parameters"
+                )
+            if getattr(self, f"CMWXFREQ_{index:04d}").value < 0.0:
+                warn(f"Frequency CMWXFREQ_{index:04d} is negative")
+            wfreqs[j] = getattr(self, f"CMWXFREQ_{index:04d}").value
+        wfreqs.sort()
+        # if np.any(np.diff(wfreqs) <= (1.0 / (2.0 * 364.25))):
+        #     warn("Frequency resolution is greater than 1/yr")
+        if self.CMWXEPOCH.value is None and self._parent is not None:
+            if self._parent.PEPOCH.value is None:
+                raise MissingParameter(
+                    "CMWXEPOCH or PEPOCH are required if CMWaveX is being used"
+                )
+            else:
+                self.CMWXEPOCH.quantity = self._parent.PEPOCH.quantity
+
+    def validate_toas(self, toas):
+        return super().validate_toas(toas)
+
+    def cmwavex_cm(self, toas):
+        total_cm = np.zeros(toas.ntoas) * cmu
+        cmwave_freqs = self.get_prefix_mapping_component("CMWXFREQ_")
+        cmwave_sins = self.get_prefix_mapping_component("CMWXSIN_")
+        cmwave_cos = self.get_prefix_mapping_component("CMWXCOS_")
+
+        base_phase = toas.table["tdbld"].data * u.d - self.CMWXEPOCH.value * u.d
+        for idx, param in cmwave_freqs.items():
+            freq = getattr(self, param).quantity
+            cmwxsin = getattr(self, cmwave_sins[idx]).quantity
+            cmwxcos = getattr(self, cmwave_cos[idx]).quantity
+            arg = 2.0 * np.pi * freq * base_phase
+            total_cm += cmwxsin * np.sin(arg.value) + cmwxcos * np.cos(arg.value)
+        return total_cm
+
+    def cmwavex_delay(self, toas, acc_delay=None):
+        return self.chromatic_type_delay(toas)
+
+    def d_cm_d_CMWXSIN(self, toas, param, acc_delay=None):
+        par = getattr(self, param)
+        freq = getattr(self, f"CMWXFREQ_{int(par.index):04d}").quantity
+        base_phase = toas.table["tdbld"].data * u.d - self.CMWXEPOCH.value * u.d
+        arg = 2.0 * np.pi * freq * base_phase
+        deriv = np.sin(arg.value)
+        return deriv * cmu / par.units
+
+    def d_cm_d_CMWXCOS(self, toas, param, acc_delay=None):
+        par = getattr(self, param)
+        freq = getattr(self, f"CMWXFREQ_{int(par.index):04d}").quantity
+        base_phase = toas.table["tdbld"].data * u.d - self.CMWXEPOCH.value * u.d
+        arg = 2.0 * np.pi * freq * base_phase
+        deriv = np.cos(arg.value)
+        return deriv * cmu / par.units
diff --git a/src/pint/models/dmwavex.py b/src/pint/models/dmwavex.py
index 258eb3556..368862391 100644
--- a/src/pint/models/dmwavex.py
+++ b/src/pint/models/dmwavex.py
@@ -271,7 +271,7 @@ def get_indices(self):
         inds : np.ndarray
         Array of DMWaveX indices in model.
         """
-        inds = [int(p.split("_")[-1]) for p in self.params if "WXFREQ_" in p]
+        inds = [int(p.split("_")[-1]) for p in self.params if "DMWXFREQ_" in p]
         return np.array(inds)
 
     # Initialize setup
@@ -299,7 +299,7 @@ def validate(self):
         DMWXCOS_mapping = self.get_prefix_mapping_component("DMWXCOS_")
         if DMWXFREQ_mapping.keys() != DMWXSIN_mapping.keys():
             raise ValueError(
-                "WXFREQ_ parameters do not match DMWXSIN_ parameters."
+                "DMWXFREQ_ parameters do not match DMWXSIN_ parameters."
                 "Please check your prefixed parameters"
             )
         if DMWXFREQ_mapping.keys() != DMWXCOS_mapping.keys():

From e9b0d3553d25965680b7f7ccb3b17bbad2b6234c Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 12:35:59 +0200
Subject: [PATCH 02/22] init

---
 src/pint/models/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/pint/models/__init__.py b/src/pint/models/__init__.py
index ab7820027..4c77ce9e1 100644
--- a/src/pint/models/__init__.py
+++ b/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,

From 557ab55c4b19c821a22afef44862232707a1578f Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 12:38:33 +0200
Subject: [PATCH 03/22] test_dmwavex

---
 tests/test_dmwavex.py | 45 ++++++++++++++++++++++++++++---------------
 1 file changed, 30 insertions(+), 15 deletions(-)

diff --git a/tests/test_dmwavex.py b/tests/test_dmwavex.py
index 0b37877aa..5fdaf0cba 100644
--- a/tests/test_dmwavex.py
+++ b/tests/test_dmwavex.py
@@ -11,23 +11,23 @@
 import pytest
 import astropy.units as u
 
-par = """
-    PSR              B1937+21
-    LAMBDA   301.9732445337270
-    BETA      42.2967523367957
-    PMLAMBDA           -0.0175
-    PMBETA             -0.3971
-    PX                  0.1515
-    POSEPOCH        55321.0000
-    F0    641.9282333345536244  1  0.0000000000000132
-    F1     -4.330899370129D-14  1  2.149749089617D-22
-    PEPOCH        55321.000000
-    DM               71.016633
-    UNITS                  TDB
-"""
-
 
 def test_dmwavex():
+    par = """
+        PSR              B1937+21
+        LAMBDA   301.9732445337270
+        BETA      42.2967523367957
+        PMLAMBDA           -0.0175
+        PMBETA             -0.3971
+        PX                  0.1515
+        POSEPOCH        55321.0000
+        F0    641.9282333345536244  1  0.0000000000000132
+        F1     -4.330899370129D-14  1  2.149749089617D-22
+        PEPOCH        55321.000000
+        DM               71.016633
+        UNITS                  TDB
+    """
+
     m = get_model(StringIO(par))
 
     with pytest.raises(ValueError):
@@ -111,6 +111,21 @@ def test_dmwavex_badpar():
 
 
 def test_add_dmwavex():
+    par = """
+        PSR              B1937+21
+        LAMBDA   301.9732445337270
+        BETA      42.2967523367957
+        PMLAMBDA           -0.0175
+        PMBETA             -0.3971
+        PX                  0.1515
+        POSEPOCH        55321.0000
+        F0    641.9282333345536244  1  0.0000000000000132
+        F1     -4.330899370129D-14  1  2.149749089617D-22
+        PEPOCH        55321.000000
+        DM               71.016633
+        UNITS                  TDB
+    """
+
     m = get_model(StringIO(par))
     idxs = dmwavex_setup(m, 3600, n_freqs=5)
 

From fd87992f99227d2fa2dfac29b26bba854aaa3b36 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 12:41:42 +0200
Subject: [PATCH 04/22] cmwavex_setup

---
 src/pint/utils.py | 102 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 98 insertions(+), 4 deletions(-)

diff --git a/src/pint/utils.py b/src/pint/utils.py
index 6d59055f1..bb80b37ee 100644
--- a/src/pint/utils.py
+++ b/src/pint/utils.py
@@ -1570,7 +1570,7 @@ def dmwavex_setup(
     model : pint.models.timing_model.TimingModel
     T_span : float, astropy.quantity.Quantity
         Time span used to calculate nyquist frequency when using freqs
-        Time span used to calculate WaveX frequencies when using n_freqs
+        Time span used to calculate DMWaveX frequencies when using n_freqs
         Usually to be set as the length of the timing baseline the model is being used for
     freqs : iterable of float or astropy.quantity.Quantity, None
         User inputed base frequencies
@@ -1643,6 +1643,100 @@ def dmwavex_setup(
     return model.components["DMWaveX"].get_indices()
 
 
+def cmwavex_setup(
+    model: "pint.models.TimingModel",
+    T_span: Union[float, u.Quantity],
+    freqs: Optional[Iterable[Union[float, u.Quantity]]] = None,
+    n_freqs: Optional[int] = None,
+    freeze_params: bool = False,
+) -> List[int]:
+    """
+    Set-up a CMWaveX model based on either an array of user-provided frequencies or the wave number
+    frequency calculation. Sine and Cosine amplitudes are initially set to zero
+
+    User specifies T_span and either freqs or n_freqs. This function assumes that the timing model does not already
+    have any CMWaveX components. See add_cmwavex_component() or add_cmwavex_components() to add components
+    to an existing CMWaveX model.
+
+    Parameters
+    ----------
+
+    model : pint.models.timing_model.TimingModel
+    T_span : float, astropy.quantity.Quantity
+        Time span used to calculate nyquist frequency when using freqs
+        Time span used to calculate CMWaveX frequencies when using n_freqs
+        Usually to be set as the length of the timing baseline the model is being used for
+    freqs : iterable of float or astropy.quantity.Quantity, None
+        User inputed base frequencies
+    n_freqs : int, None
+        Number of wave frequencies to calculate using the equation: freq_n = 2 * pi * n / T_span
+        Where n is the wave number, and T_span is the total time span of the toas in the fitter object
+    freeze_params : bool, optional
+        Whether the new parameters should be frozen
+
+    Returns
+    -------
+
+    indices : list
+            Indices that have been assigned to new WaveX components
+    """
+    from pint.models.cmwavex import CMWaveX
+
+    if (freqs is None) and (n_freqs is None):
+        raise ValueError(
+            "CMWaveX component base frequencies are not specified. "
+            "Please input either freqs or n_freqs"
+        )
+
+    if (freqs is not None) and (n_freqs is not None):
+        raise ValueError(
+            "Both freqs and n_freqs are specified. Only one or the other should be used"
+        )
+
+    if n_freqs is not None and n_freqs <= 0:
+        raise ValueError("Must use a non-zero number of wave frequencies")
+
+    model.add_component(CMWaveX())
+    if isinstance(T_span, u.quantity.Quantity):
+        T_span.to(u.d)
+    else:
+        T_span *= u.d
+
+    nyqist_freq = 1.0 / (2.0 * T_span)
+    if freqs is not None:
+        if isinstance(freqs, u.quantity.Quantity):
+            freqs.to(u.d**-1)
+        else:
+            freqs *= u.d**-1
+        if len(freqs) == 1:
+            model.CMWXFREQ_0001.quantity = freqs
+        else:
+            freqs = np.array(freqs)
+            freqs.sort()
+            if min(np.diff(freqs)) < nyqist_freq:
+                warnings.warn(
+                    "CMWaveX frequency spacing is finer than frequency resolution of data"
+                )
+            model.CMWXFREQ_0001.quantity = freqs[0]
+            model.components["CMWaveX"].add_cmwavex_components(freqs[1:])
+
+    if n_freqs is not None:
+        if n_freqs == 1:
+            wave_freq = 1 / T_span
+            model.CMWXFREQ_0001.quantity = wave_freq
+        else:
+            wave_numbers = np.arange(1, n_freqs + 1)
+            wave_freqs = wave_numbers / T_span
+            model.CMWXFREQ_0001.quantity = wave_freqs[0]
+            model.components["CMWaveX"].add_cmwavex_components(wave_freqs[1:])
+
+    for p in model.params:
+        if p.startswith("CMWXSIN") or p.startswith("CMWXCOS"):
+            model[p].frozen = freeze_params
+
+    return model.components["CMWaveX"].get_indices()
+
+
 def _translate_wave_freqs(om: Union[float, u.Quantity], k: int) -> u.Quantity:
     """
     Use Wave model WAVEOM parameter to calculate a WaveX WXFREQ_ frequency parameter for wave number k
@@ -2396,12 +2490,12 @@ def info_string(
         else:
             s += f"{os.linesep}Comment: {comment}"
 
-    if (prefix_string is not None) and (len(prefix_string) > 0):
+    if prefix_string is not None and prefix_string != "":
         s = os.linesep.join([prefix_string + x for x in s.splitlines()])
     return s
 
 
-def list_parameters(class_: Type = None) -> List[Dict[str, Union[str, List]]]:
+def list_parameters(class_: Optional[Type] = None) -> List[Dict[str, Union[str, List]]]:
     """List parameters understood by PINT.
 
     Parameters
@@ -3039,7 +3133,7 @@ def _get_wx2pl_lnlike(
     from pint.models.noise_model import powerlaw
     from pint import DMconst
 
-    assert component_name in ["WaveX", "DMWaveX"]
+    assert component_name in {"WaveX", "DMWaveX"}
     prefix = "WX" if component_name == "WaveX" else "DMWX"
 
     idxs = np.array(model.components[component_name].get_indices())

From 86a83b97742059a2b070f5130335cd1169df469b Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 12:41:59 +0200
Subject: [PATCH 05/22] test_cmwavex

---
 tests/test_cmwavex.py | 177 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 177 insertions(+)
 create mode 100644 tests/test_cmwavex.py

diff --git a/tests/test_cmwavex.py b/tests/test_cmwavex.py
new file mode 100644
index 000000000..d0f2652fb
--- /dev/null
+++ b/tests/test_cmwavex.py
@@ -0,0 +1,177 @@
+from io import StringIO
+
+import numpy as np
+
+from pint.models import get_model
+from pint.fitter import Fitter
+from pint.simulation import make_fake_toas_uniform
+from pint.utils import cmwavex_setup
+from pint.models.chromatic_model import cmu
+
+import pytest
+import astropy.units as u
+
+
+def test_cmwavex():
+    par = """`
+        PSR              B1937+21
+        LAMBDA   301.9732445337270
+        BETA      42.2967523367957
+        PMLAMBDA           -0.0175
+        PMBETA             -0.3971
+        PX                  0.1515
+        POSEPOCH        55321.0000
+        F0    641.9282333345536244  1  0.0000000000000132
+        F1     -4.330899370129D-14  1  2.149749089617D-22
+        PEPOCH        55321.000000
+        DM               71.016633
+        CM                     0.1
+        TNCHROMIDX               4
+        UNITS                  TDB
+    """
+    m = get_model(StringIO(par))
+
+    with pytest.raises(ValueError):
+        idxs = cmwavex_setup(m, 3600)
+
+    idxs = cmwavex_setup(m, 3600, n_freqs=5)
+
+    assert "CMWaveX" in m.components
+    assert m.num_cmwavex_freqs == len(idxs)
+
+    m.components["CMWaveX"].remove_cmwavex_component(5)
+    assert m.num_cmwavex_freqs == len(idxs) - 1
+
+    t = make_fake_toas_uniform(54000, 56000, 200, m, add_noise=True)
+
+    ftr = Fitter.auto(t, m)
+    ftr.fit_toas()
+
+    assert ftr.resids.reduced_chi2 < 2
+
+
+def test_cmwavex_badpar():
+    with pytest.raises(ValueError):
+        par = """
+            PSR              B1937+21
+            LAMBDA   301.9732445337270
+            BETA      42.2967523367957
+            PMLAMBDA           -0.0175
+            PMBETA             -0.3971
+            PX                  0.1515
+            POSEPOCH        55321.0000
+            F0    641.9282333345536244  1  0.0000000000000132
+            F1     -4.330899370129D-14  1  2.149749089617D-22
+            PEPOCH        55321.000000
+            DM               71.016633
+            CM                     0.1
+            TNCHROMIDX               4
+            UNITS                  TDB
+            CMWXFREQ_0001         0.01
+            CMWXSIN_0001             0
+            CMWXSIN_0002             0
+        """
+        get_model(StringIO(par))
+
+    with pytest.raises(ValueError):
+        par = """
+            PSR              B1937+21
+            LAMBDA   301.9732445337270
+            BETA      42.2967523367957
+            PMLAMBDA           -0.0175
+            PMBETA             -0.3971
+            PX                  0.1515
+            POSEPOCH        55321.0000
+            F0    641.9282333345536244  1  0.0000000000000132
+            F1     -4.330899370129D-14  1  2.149749089617D-22
+            PEPOCH        55321.000000
+            DM               71.016633
+            CM                     0.1
+            TNCHROMIDX               4
+            UNITS                  TDB
+            CMWXFREQ_0001         0.01
+            CMWXCOS_0001             0
+            CMWXCOS_0002             0
+        """
+        get_model(StringIO(par))
+
+    with pytest.raises(ValueError):
+        par = """
+            PSR              B1937+21
+            LAMBDA   301.9732445337270
+            BETA      42.2967523367957
+            PMLAMBDA           -0.0175
+            PMBETA             -0.3971
+            PX                  0.1515
+            POSEPOCH        55321.0000
+            F0    641.9282333345536244  1  0.0000000000000132
+            F1     -4.330899370129D-14  1  2.149749089617D-22
+            PEPOCH        55321.000000
+            DM               71.016633
+            CM                     0.1
+            TNCHROMIDX               4
+            UNITS                  TDB
+            CMWXFREQ_0001         0.00
+            CMWXCOS_0001             0
+        """
+        get_model(StringIO(par))
+
+
+def test_add_cmwavex():
+    par = """
+        PSR              B1937+21
+        LAMBDA   301.9732445337270
+        BETA      42.2967523367957
+        PMLAMBDA           -0.0175
+        PMBETA             -0.3971
+        PX                  0.1515
+        POSEPOCH        55321.0000
+        F0    641.9282333345536244  1  0.0000000000000132
+        F1     -4.330899370129D-14  1  2.149749089617D-22
+        PEPOCH        55321.000000
+        DM               71.016633
+        CM                     0.1
+        TNCHROMIDX               4
+        UNITS                  TDB
+    """
+    m = get_model(StringIO(par))
+    idxs = cmwavex_setup(m, 3600, n_freqs=5)
+
+    with pytest.raises(ValueError):
+        m.components["CMWaveX"].add_cmwavex_component(1, index=5, cmwxsin=0, cmwxcos=0)
+
+    m.components["CMWaveX"].add_cmwavex_component(1, index=6, cmwxsin=0, cmwxcos=0)
+    assert m.num_cmwavex_freqs == len(idxs) + 1
+
+    m.components["CMWaveX"].add_cmwavex_component(
+        1 / u.day, index=7, cmwxsin=0 * cmu, cmwxcos=0 * cmu
+    )
+    assert m.num_cmwavex_freqs == len(idxs) + 2
+
+    m.components["CMWaveX"].add_cmwavex_component(2 / u.day)
+    assert m.num_cmwavex_freqs == len(idxs) + 3
+
+    m.components["CMWaveX"].add_cmwavex_components(
+        np.array([3]) / u.day,
+        cmwxsins=np.array([0]) * cmu,
+        cmwxcoses=np.array([0]) * cmu,
+    )
+    assert m.num_cmwavex_freqs == len(idxs) + 4
+
+    with pytest.raises(ValueError):
+        m.components["CMWaveX"].add_cmwavex_components(
+            [2 / u.day, 3 / u.day], cmwxsins=[0, 0], cmwxcoses=[0, 0, 0]
+        )
+
+    with pytest.raises(ValueError):
+        m.components["CMWaveX"].add_cmwavex_components(
+            [2 / u.day, 3 / u.day], cmwxsins=[0, 0, 0], cmwxcoses=[0, 0]
+        )
+
+    with pytest.raises(ValueError):
+        m.components["CMWaveX"].add_cmwavex_components(
+            [2 / u.day, 3 / u.day],
+            cmwxsins=[0, 0],
+            cmwxcoses=[0, 0],
+            frozens=[False, False, False],
+        )

From 19c35214ad9df346fd6e6009d4e3e9b516e99025 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 13:07:37 +0200
Subject: [PATCH 06/22] CHANGELOG

---
 CHANGELOG-unreleased.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md
index 2899b99f8..8e7d1fc3a 100644
--- a/CHANGELOG-unreleased.md
+++ b/CHANGELOG-unreleased.md
@@ -10,5 +10,8 @@ the released changes.
 ## Unreleased
 ### Changed
 ### Added
+- Fourier series representation of chromatic noise (`CMWaveX`)
+- `pint.utils.cmwavex_setup` function
 ### Fixed
+- Bug in `DMWaveX.get_indices()` function
 ### Removed

From 5ddf707ba7fa20e2b1b3e895af9e1d0757bb1137 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 13:17:58 +0200
Subject: [PATCH 07/22] validation for correlated noise components

---
 src/pint/models/timing_model.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py
index 52c7cd6f1..469b26db5 100644
--- a/src/pint/models/timing_model.py
+++ b/src/pint/models/timing_model.py
@@ -501,23 +501,36 @@ def num_components_of_type(type):
             num_components_of_type(SolarWindDispersionBase) <= 1
         ), "Model can have at most one solar wind dispersion component."
 
-        from pint.models.dispersion_model import DispersionDMX
+        from pint.models.dispersion_model import DispersionDM, DispersionDMX
+        from pint.models.chromatic_model import ChromaticCM
         from pint.models.wave import Wave
         from pint.models.wavex import WaveX
         from pint.models.dmwavex import DMWaveX
+        from pint.models.cmwavex import CMWaveX
         from pint.models.noise_model import PLRedNoise, PLDMNoise
+        from pint.models.ifunc import IFunc
 
         if num_components_of_type((DispersionDMX, PLDMNoise, DMWaveX)) > 1:
             log.warning(
                 "DispersionDMX, PLDMNoise, and DMWaveX cannot be used together. "
                 "They are ways of modelling the same effect."
             )
-        if num_components_of_type((Wave, WaveX, PLRedNoise)) > 1:
+        if num_components_of_type((Wave, WaveX, PLRedNoise, IFunc)) > 1:
             log.warning(
                 "Wave, WaveX, and PLRedNoise cannot be used together. "
                 "They are ways of modelling the same effect."
             )
 
+        if num_components_of_type((PLDMNoise, DMWaveX)) == 1:
+            assert (
+                num_components_of_type(DispersionDM) == 1
+            ), "PLDMNoise / DMWaveX component cannot be used without the DispersionDM component."
+
+        if num_components_of_type((CMWaveX)) == 1:
+            assert (
+                num_components_of_type(ChromaticCM) == 1
+            ), "PLChromNoise / CMWaveX component cannot be used without the ChromaticCM component."
+
     # def __str__(self):
     #    result = ""
     #    comps = self.components

From 75a5cf564acc8537bc4c2306491274c4a900bba1 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 13:19:27 +0200
Subject: [PATCH 08/22] changelog#

---
 CHANGELOG-unreleased.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md
index 8e7d1fc3a..7b6f00ddd 100644
--- a/CHANGELOG-unreleased.md
+++ b/CHANGELOG-unreleased.md
@@ -12,6 +12,7 @@ the released changes.
 ### Added
 - Fourier series representation of chromatic noise (`CMWaveX`)
 - `pint.utils.cmwavex_setup` function
+- More validation for correlated noise components in `TimingModel`
 ### Fixed
 - Bug in `DMWaveX.get_indices()` function
 ### Removed

From 9c97955e716045975fc92f204c583a436dba6d1b Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 13:29:39 +0200
Subject: [PATCH 09/22] sourcery

---
 src/pint/models/timing_model.py | 249 +++++++++++++-------------------
 1 file changed, 100 insertions(+), 149 deletions(-)

diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py
index 469b26db5..2160de9b2 100644
--- a/src/pint/models/timing_model.py
+++ b/src/pint/models/timing_model.py
@@ -2708,16 +2708,12 @@ def use_aliases(self, reset_to_default=True, alias_translation=None):
             if reset_to_default:
                 po.use_alias = None
             if alias_translation is not None:
-                if hasattr(po, "origin_name"):
-                    try:
-                        po.use_alias = alias_translation[po.origin_name]
-                    except KeyError:
-                        pass
-                else:
-                    try:
-                        po.use_alias = alias_translation[p]
-                    except KeyError:
-                        pass
+                with contextlib.suppress(KeyError):
+                    po.use_alias = (
+                        alias_translation[po.origin_name]
+                        if hasattr(po, "origin_name")
+                        else alias_translation[p]
+                    )
 
     def as_parfile(
         self,
@@ -2746,7 +2742,7 @@ def as_parfile(
              Parfile output format. PINT outputs in 'tempo', 'tempo2' and 'pint'
              formats. The defaul format is `pint`.
         """
-        if not format.lower() in _parfile_formats:
+        if format.lower() not in _parfile_formats:
             raise ValueError(f"parfile format must be one of {_parfile_formats}")
 
         self.validate()
@@ -2768,33 +2764,30 @@ def as_parfile(
                 continue
             result_begin += getattr(self, p).as_parfile_line(format=format)
         for cat in start_order:
-            if cat in list(cates_comp.keys()):
-                # print("Starting: %s" % cat)
-                cp = cates_comp[cat]
-                for cpp in cp:
-                    result_begin += cpp.print_par(format=format)
-                printed_cate.append(cat)
-            else:
+            if cat not in list(cates_comp.keys()):
                 continue
 
+            # print("Starting: %s" % cat)
+            cp = cates_comp[cat]
+            for cpp in cp:
+                result_begin += cpp.print_par(format=format)
+            printed_cate.append(cat)
         for cat in last_order:
-            if cat in list(cates_comp.keys()):
-                # print("Ending: %s" % cat)
-                cp = cates_comp[cat]
-                for cpp in cp:
-                    result_end += cpp.print_par(format=format)
-                printed_cate.append(cat)
-            else:
+            if cat not in list(cates_comp.keys()):
                 continue
 
+            # print("Ending: %s" % cat)
+            cp = cates_comp[cat]
+            for cpp in cp:
+                result_end += cpp.print_par(format=format)
+            printed_cate.append(cat)
         for cat in list(cates_comp.keys()):
             if cat in printed_cate:
                 continue
-            else:
-                cp = cates_comp[cat]
-                for cpp in cp:
-                    result_middle += cpp.print_par(format=format)
-                printed_cate.append(cat)
+            cp = cates_comp[cat]
+            for cpp in cp:
+                result_middle += cpp.print_par(format=format)
+            printed_cate.append(cat)
 
         return result_begin + result_middle + result_end
 
@@ -2934,8 +2927,7 @@ def __len__(self):
         return len(self.params)
 
     def __iter__(self):
-        for p in self.params:
-            yield p
+        yield from self.params
 
     def as_ECL(self, epoch=None, ecl="IERS2010"):
         """Return TimingModel in PulsarEcliptic frame.
@@ -3237,9 +3229,7 @@ def get_derived_params(self, rms=None, ntoas=None, returndict=False):
                     )
                     s += f"Pulsar mass (Shapiro Delay) = {psrmass}"
                     outdict["Mp (Msun)"] = psrmass
-        if not returndict:
-            return s
-        return s, outdict
+        return (s, outdict) if returndict else s
 
 
 class ModelMeta(abc.ABCMeta):
@@ -3253,8 +3243,8 @@ class ModelMeta(abc.ABCMeta):
     """
 
     def __init__(cls, name, bases, dct):
-        regname = "component_types"
         if "register" in dct and cls.register:
+            regname = "component_types"
             getattr(cls, regname)[name] = cls
         super().__init__(name, bases, dct)
 
@@ -3334,10 +3324,10 @@ def param_prefixs(self):
         for p in self.params:
             par = getattr(self, p)
             if par.is_prefix:
-                if par.prefix not in prefixs.keys():
-                    prefixs[par.prefix] = [p]
-                else:
+                if par.prefix in prefixs:
                     prefixs[par.prefix].append(p)
+                else:
+                    prefixs[par.prefix] = [p]
         return prefixs
 
     @property_exists
@@ -3407,10 +3397,7 @@ def add_param(self, param, deriv_func=None, setup=False):
             exist_par = getattr(self, param.name)
             if exist_par.value is not None:
                 raise ValueError(
-                    "Tried to add a second parameter called {}. "
-                    "Old value: {} New value: {}".format(
-                        param.name, getattr(self, param.name), param
-                    )
+                    f"Tried to add a second parameter called {param.name}. Old value: {getattr(self, param.name)} New value: {param}"
                 )
             else:
                 setattr(self, param.name, param)
@@ -3433,10 +3420,7 @@ def remove_param(self, param):
         param : str or pint.models.Parameter
             The parameter to remove.
         """
-        if isinstance(param, str):
-            param_name = param
-        else:
-            param_name = param.name
+        param_name = param if isinstance(param, str) else param.name
         if param_name not in self.params:
             raise ValueError(
                 f"Tried to remove parameter {param_name} but it is not listed: {self.params}"
@@ -3473,10 +3457,7 @@ def get_params_of_type(self, param_type):
             par = getattr(self, p)
             par_type = type(par).__name__
             par_prefix = par_type[:-9]
-            if (
-                param_type.upper() == par_type.upper()
-                or param_type.upper() == par_prefix.upper()
-            ):
+            if param_type.upper() in [par_type.upper(), par_prefix.upper()]:
                 result.append(par.name)
         return result
 
@@ -3521,37 +3502,35 @@ def match_param_aliases(self, alias):
         # Split the alias prefix, see if it is a perfix alias
         try:
             prefix, idx_str, idx = split_prefixed_name(alias)
-        except PrefixError:  # Not a prefixed name
-            if pname is not None:
-                par = getattr(self, pname)
-                if par.is_prefix:
-                    raise UnknownParameter(
-                        f"Prefix {alias} maps to mulitple parameters"
-                        ". Please specify the index as well."
-                    )
-            else:
+        except PrefixError as e:  # Not a prefixed name
+            if pname is None:
                 # Not a prefix, not an alias
-                raise UnknownParameter(f"Unknown parameter name or alias {alias}")
-        # When the alias is a prefixed name but not in the parameter list yet
-        if pname is None:
-            prefix_pname = self.aliases_map.get(prefix, None)
-            if prefix_pname:
-                par = getattr(self, prefix_pname)
-                if par.is_prefix:
-                    raise UnknownParameter(
-                        f"Found a similar prefixed parameter '{prefix_pname}'"
-                        f" But parameter {par.prefix}{idx} need to be added"
-                        f" to the model."
-                    )
-                else:
-                    raise UnknownParameter(
-                        f"{par} is not a prefixed parameter, howere the input"
-                        f" {alias} has index with it."
-                    )
+                raise UnknownParameter(
+                    f"Unknown parameter name or alias {alias}"
+                ) from e
+            par = getattr(self, pname)
+            if par.is_prefix:
+                raise UnknownParameter(
+                    f"Prefix {alias} maps to mulitple parameters"
+                    ". Please specify the index as well."
+                ) from e
+        if pname is not None:
+            return pname
+        if prefix_pname := self.aliases_map.get(prefix, None):
+            par = getattr(self, prefix_pname)
+            if par.is_prefix:
+                raise UnknownParameter(
+                    f"Found a similar prefixed parameter '{prefix_pname}'"
+                    f" But parameter {par.prefix}{idx} need to be added"
+                    f" to the model."
+                )
             else:
-                raise UnknownParameter(f"Unknown parameter name or alias {alias}")
+                raise UnknownParameter(
+                    f"{par} is not a prefixed parameter, howere the input"
+                    f" {alias} has index with it."
+                )
         else:
-            return pname
+            raise UnknownParameter(f"Unknown parameter name or alias {alias}")
 
     def register_deriv_funcs(self, func, param):
         """Register the derivative function in to the deriv_func dictionaries.
@@ -3568,15 +3547,10 @@ def register_deriv_funcs(self, func, param):
 
         if pn not in list(self.deriv_funcs.keys()):
             self.deriv_funcs[pn] = [func]
+        elif func in self.deriv_funcs[pn]:
+            return
         else:
-            # TODO:
-            # Runing setup() mulitple times can lead to adding derivative
-            # function multiple times. This prevent it from happening now. But
-            # in the future, we should think a better way to do so.
-            if func in self.deriv_funcs[pn]:
-                return
-            else:
-                self.deriv_funcs[pn] += [func]
+            self.deriv_funcs[pn] += [func]
 
     def is_in_parfile(self, para_dict):
         """Check if this subclass included in parfile.
@@ -3594,11 +3568,7 @@ def is_in_parfile(self, para_dict):
 
         """
         if self.component_special_params:
-            for p in self.component_special_params:
-                if p in para_dict:
-                    return True
-            return False
-
+            return any(p in para_dict for p in self.component_special_params)
         pNames_inpar = list(para_dict.keys())
         pNames_inModel = self.params
 
@@ -3606,22 +3576,14 @@ def is_in_parfile(self, para_dict):
         # should go in them.
         # For solar system Shapiro delay component
         if hasattr(self, "PLANET_SHAPIRO"):
-            if "NO_SS_SHAPIRO" in pNames_inpar:
-                return False
-            else:
-                return True
-
+            return "NO_SS_SHAPIRO" not in pNames_inpar
         try:
             bmn = getattr(self, "binary_model_name")
         except AttributeError:
             # This isn't a binary model, keep looking
             pass
         else:
-            if "BINARY" in para_dict:
-                return bmn == para_dict["BINARY"][0]
-            else:
-                return False
-
+            return bmn == para_dict["BINARY"][0] if "BINARY" in para_dict else False
         # Compare the componets parameter names with par file parameters
         compr = list(set(pNames_inpar).intersection(pNames_inModel))
 
@@ -3654,10 +3616,9 @@ def print_par(self, format="pint"):
         -------
         str : formatted line for par file
         """
-        result = ""
-        for p in self.params:
-            result += getattr(self, p).as_parfile_line(format=format)
-        return result
+        return "".join(
+            getattr(self, p).as_parfile_line(format=format) for p in self.params
+        )
 
 
 class DelayComponent(Component):
@@ -3797,7 +3758,7 @@ def _param_unit_map(self):
         units = {}
         for k, cp in self.components.items():
             for p in cp.params:
-                if p in units.keys() and units[p] != getattr(cp, p).units:
+                if p in units and units[p] != getattr(cp, p).units:
                     raise TimingModelError(
                         f"Units of parameter '{p}' in component '{cp}' ({getattr(cp, p).units}) do not match those of existing parameter ({units[p]})"
                     )
@@ -3815,11 +3776,9 @@ def repeatable_param(self):
             for p in cp.params:
                 par = getattr(cp, p)
                 if par.repeatable:
-                    repeatable.append(p)
-                    repeatable.append(par._parfile_name)
+                    repeatable.extend((p, par._parfile_name))
                     # also add the aliases to the repeatable param
-                    for als in par.aliases:
-                        repeatable.append(als)
+                    repeatable.extend(iter(par.aliases))
         return set(repeatable)
 
     @lazyproperty
@@ -3849,10 +3808,7 @@ def component_category_map(self):
             The mapping from components to its categore. The key is the component
             name and the value is the component's category name.
         """
-        cp_ca = {}
-        for k, cp in self.components.items():
-            cp_ca[k] = cp.category
-        return cp_ca
+        return {k: cp.category for k, cp in self.components.items()}
 
     @lazyproperty
     def component_unique_params(self):
@@ -3894,38 +3850,35 @@ def search_binary_components(self, system_name):
             model.
         """
         all_systems = self.category_component_map["pulsar_system"]
-        # Search the system name first
         if system_name in all_systems:
             return self.components[system_name]
-        else:  # search for the pulsar system aliases
-            for cp_name in all_systems:
-                if system_name == self.components[cp_name].binary_model_name:
-                    return self.components[cp_name]
-
-            if system_name == "BTX":
-                raise UnknownBinaryModel(
-                    "`BINARY  BTX` is not supported bt PINT. Use "
-                    "`BINARY  BT` instead. It supports both orbital "
-                    "period (PB, PBDOT) and orbital frequency (FB0, ...) "
-                    "parametrizations."
-                )
-            elif system_name == "DDFWHE":
-                raise UnknownBinaryModel(
-                    "`BINARY  DDFWHE` is not supported, but the same model "
-                    "is available as `BINARY  DDH`."
-                )
-            elif system_name in ["MSS", "EH", "H88", "DDT", "BT1P", "BT2P"]:
-                # Binary model list taken from
-                # https://tempo.sourceforge.net/ref_man_sections/binary.txt
-                raise UnknownBinaryModel(
-                    f"`The binary model {system_name} is not yet implemented."
-                )
+        for cp_name in all_systems:
+            if system_name == self.components[cp_name].binary_model_name:
+                return self.components[cp_name]
 
+        if system_name == "BTX":
+            raise UnknownBinaryModel(
+                "`BINARY  BTX` is not supported bt PINT. Use "
+                "`BINARY  BT` instead. It supports both orbital "
+                "period (PB, PBDOT) and orbital frequency (FB0, ...) "
+                "parametrizations."
+            )
+        elif system_name == "DDFWHE":
+            raise UnknownBinaryModel(
+                "`BINARY  DDFWHE` is not supported, but the same model "
+                "is available as `BINARY  DDH`."
+            )
+        elif system_name in ["MSS", "EH", "H88", "DDT", "BT1P", "BT2P"]:
+            # Binary model list taken from
+            # https://tempo.sourceforge.net/ref_man_sections/binary.txt
             raise UnknownBinaryModel(
-                f"Pulsar system/Binary model component"
-                f" {system_name} is not provided."
+                f"`The binary model {system_name} is not yet implemented."
             )
 
+        raise UnknownBinaryModel(
+            f"Pulsar system/Binary model component" f" {system_name} is not provided."
+        )
+
     def alias_to_pint_param(self, alias):
         """Translate a alias to a PINT parameter name.
 
@@ -3989,14 +3942,11 @@ def alias_to_pint_param(self, alias):
                 # count length of idx_str and dectect leading zeros
                 # TODO fix the case for searching `DMX`
                 num_lzero = len(idx_str) - len(str(idx))
-                if num_lzero > 0:  # Has leading zero
-                    fmt = len(idx_str)
-                else:
-                    fmt = 0
+                fmt = len(idx_str) if num_lzero > 0 else 0
                 first_init_par = None
                 # Handle the case of start index from 0 and 1
                 for start_idx in [0, 1]:
-                    first_init_par_alias = prefix + f"{start_idx:0{fmt}}"
+                    first_init_par_alias = f"{prefix}{start_idx:0{fmt}}"
                     first_init_par = self._param_alias_map.get(
                         first_init_par_alias, None
                     )
@@ -4006,13 +3956,14 @@ def alias_to_pint_param(self, alias):
                         break
             except PrefixError:
                 pint_par = None
-
         else:
             first_init_par = pint_par
+
         if pint_par is None:
             raise UnknownParameter(
-                "Can not find matching PINT parameter for '{}'".format(alias)
+                f"Can not find matching PINT parameter for '{alias}'"
             )
+
         return pint_par, first_init_par
 
     def param_to_unit(self, name):
@@ -4072,7 +4023,7 @@ def __init__(self, module, param, msg=None):
         self.msg = msg
 
     def __str__(self):
-        result = self.module + "." + self.param
+        result = f"{self.module}.{self.param}"
         if self.msg is not None:
             result += "\n  " + self.msg
         return result

From 30c8f4bcad02de5f050fc90dddfde89b2741115f Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 11 Jul 2024 13:37:38 +0200
Subject: [PATCH 10/22] validation

---
 src/pint/models/timing_model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/pint/models/timing_model.py b/src/pint/models/timing_model.py
index 2160de9b2..d381a91bc 100644
--- a/src/pint/models/timing_model.py
+++ b/src/pint/models/timing_model.py
@@ -507,7 +507,7 @@ def num_components_of_type(type):
         from pint.models.wavex import WaveX
         from pint.models.dmwavex import DMWaveX
         from pint.models.cmwavex import CMWaveX
-        from pint.models.noise_model import PLRedNoise, PLDMNoise
+        from pint.models.noise_model import PLRedNoise, PLDMNoise, PLChromNoise
         from pint.models.ifunc import IFunc
 
         if num_components_of_type((DispersionDMX, PLDMNoise, DMWaveX)) > 1:
@@ -526,7 +526,7 @@ def num_components_of_type(type):
                 num_components_of_type(DispersionDM) == 1
             ), "PLDMNoise / DMWaveX component cannot be used without the DispersionDM component."
 
-        if num_components_of_type((CMWaveX)) == 1:
+        if num_components_of_type((PLChromNoise, CMWaveX)) == 1:
             assert (
                 num_components_of_type(ChromaticCM) == 1
             ), "PLChromNoise / CMWaveX component cannot be used without the ChromaticCM component."

From fa6a9d6a0dfd793bcfec8d7aed07b233c482b8a3 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Wed, 7 Aug 2024 13:54:40 +0200
Subject: [PATCH 11/22] --

---
 src/pint/logging.py           |  2 +-
 src/pint/models/binary_ddk.py | 29 +++++++++++++++++------------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/pint/logging.py b/src/pint/logging.py
index e87a039bd..f7de1a5aa 100644
--- a/src/pint/logging.py
+++ b/src/pint/logging.py
@@ -129,7 +129,7 @@ class LogFilter:
 
     def __init__(self, onlyonce=None, never=None, onlyonce_level="INFO"):
         """
-        Define regexs for messages that will only be seen once.  Use ``\S+`` for a variable that might change.
+        Define regexs 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 ``()``.
diff --git a/src/pint/models/binary_ddk.py b/src/pint/models/binary_ddk.py
index 8a31f6209..452774ead 100644
--- a/src/pint/models/binary_ddk.py
+++ b/src/pint/models/binary_ddk.py
@@ -49,14 +49,19 @@ class BinaryDDK(BinaryDD):
     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`.
@@ -67,7 +72,7 @@ class BinaryDDK(BinaryDD):
        KIN
             the inclination angle: :math:`i`
        KOM
-            the longitude of the ascending node, Kopeikin (1995) Eq 9: :math:`\Omega`
+            the longitude of the ascending node, Kopeikin (1995) Eq 9: :math:`\\Omega`
        K96
             flag for Kopeikin binary model proper motion correction
 
@@ -233,19 +238,19 @@ def alternative_solutions(self):
 
         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:
 
         :math:`(KIN, KOM)`
 
-        :math:`(KIN, 2KOM_0 - KOM - 180^{\circ})`
+        :math:`(KIN, 2KOM_0 - KOM - 180^{\\circ})`
 
-        :math:`(180^{\circ}-KIN, KOM+180^{\circ})`
+        :math:`(180^{\\circ}-KIN, KOM+180^{\\circ})`
 
-        :math:`(180^{\circ}-KIN, 2KOM_0 - KOM)`
+        :math:`(180^{\\circ}-KIN, 2KOM_0 - KOM)`
 
-        All values will be between 0 and :math:`360^{\circ}`.
+        All values will be between 0 and :math:`360^{\\circ}`.
 
         Returns
         -------

From 4b931c3072fb2449bc606e4b619fa70c803d9213 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Wed, 7 Aug 2024 15:40:19 +0200
Subject: [PATCH 12/22] plchromnoise_from_cmwavex

---
 src/pint/utils.py | 75 +++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 69 insertions(+), 6 deletions(-)

diff --git a/src/pint/utils.py b/src/pint/utils.py
index 3ef7a7783..0327e79c1 100644
--- a/src/pint/utils.py
+++ b/src/pint/utils.py
@@ -3132,8 +3132,9 @@ def _get_wx2pl_lnlike(
     from pint.models.noise_model import powerlaw
     from pint import DMconst
 
-    assert component_name in {"WaveX", "DMWaveX"}
-    prefix = "WX" if component_name == "WaveX" else "DMWX"
+    assert component_name in {"WaveX", "DMWaveX", "CMWaveX"}
+    prefix_dict = {"WaveX": "WX", "DMWaveX": "DMWX", "CMWaveX": "CMWX"}
+    prefix = prefix_dict[component_name]
 
     idxs = np.array(model.components[component_name].get_indices())
 
@@ -3145,7 +3146,7 @@ def _get_wx2pl_lnlike(
 
     assert np.allclose(
         np.diff(np.diff(fs)), 0
-    ), "[DM]WaveX frequencies must be uniformly spaced."
+    ), "WaveX/DMWaveX/CMWaveX frequencies must be uniformly spaced for this conversion to work."
 
     if ignore_fyr:
         year_mask = np.abs(((fs - fyr) / f0)) > 0.5
@@ -3156,7 +3157,15 @@ def _get_wx2pl_lnlike(
         )
         f0 = np.min(fs)
 
-    scaling_factor = 1 if component_name == "WaveX" else DMconst / (1400 * u.MHz) ** 2
+    scaling_factor = (
+        1
+        if component_name == "WaveX"
+        else (
+            DMconst / (1400 * u.MHz) ** 2
+            if component_name == "DMWaveX"
+            else DMconst / 1400**model.TNCHROMIDX.value
+        )
+    )
 
     a = np.array(
         [
@@ -3185,14 +3194,14 @@ def _get_wx2pl_lnlike(
 
     def powl_model(params: Tuple[float, float]) -> float:
         """Get the powerlaw spectrum for the WaveX frequencies for a given
-        set of parameters. This calls the powerlaw function used by `PLRedNoise`/`PLDMNoise`.
+        set of parameters. This calls the powerlaw function used by `PLRedNoise`/`PLDMNoise`/`PLChromNoise`.
         """
         gamma, log10_A = params
         return (powerlaw(fs, A=10**log10_A, gamma=gamma) * f0) ** 0.5
 
     def mlnlike(params: Tuple[float, ...]) -> float:
         """Negative of the likelihood function that acts on the
-        `[DM]WaveX` amplitudes."""
+        `[DM/CM]WaveX` amplitudes."""
         sigma = powl_model(params)
         return 0.5 * np.sum(
             (a**2 / (sigma**2 + da**2))
@@ -3308,6 +3317,60 @@ def pldmnoise_from_dmwavex(
     return model1
 
 
+def plchromnoise_from_cmwavex(
+    model: "pint.models.TimingModel", ignore_fyr: bool = False
+) -> "pint.models.TimingModel":
+    """Convert a `CMWaveX` representation of red noise to a `PLChromNoise`
+    representation. This is done by minimizing a likelihood function
+    that acts on the `CMWaveX` amplitudes over the powerlaw spectral
+    parameters.
+
+    Parameters
+    ----------
+    model: pint.models.timing_model.TimingModel
+        The timing model with a `CMWaveX` component.
+
+    Returns
+    -------
+    pint.models.timing_model.TimingModel
+        The timing model with a converted `PLChromNoise` component.
+    """
+    from pint.models.noise_model import PLChromNoise
+
+    mlnlike = _get_wx2pl_lnlike(model, "CMWaveX", ignore_fyr=ignore_fyr)
+
+    result = minimize(mlnlike, [4, -13], method="Nelder-Mead")
+    if not result.success:
+        raise ValueError("Log-likelihood maximization failed to converge.")
+
+    gamma_val, log10_A_val = result.x
+
+    hess = Hessian(mlnlike)
+
+    H = hess((gamma_val, log10_A_val))
+    assert np.all(np.linalg.eigvals(H) > 0), "The Hessian is not positive definite!"
+
+    Hinv = np.linalg.pinv(H)
+    assert np.all(
+        np.linalg.eigvals(Hinv) > 0
+    ), "The inverse Hessian is not positive definite!"
+
+    gamma_err, log10_A_err = np.sqrt(np.diag(Hinv))
+
+    tndmc = len(model.components["CMWaveX"].get_indices())
+
+    model1 = deepcopy(model)
+    model1.remove_component("CMWaveX")
+    model1.add_component(PLChromNoise())
+    model1.TNCHROMAMP.value = log10_A_val
+    model1.TNCHROMGAM.value = gamma_val
+    model1.TNCHROMC.value = tndmc
+    model1.TNCHROMAMP.uncertainty_value = log10_A_err
+    model1.TNCHROMGAM.uncertainty_value = gamma_err
+
+    return model1
+
+
 def find_optimal_nharms(
     model: "pint.models.TimingModel",
     toas: "pint.toa.TOAs",

From b39395aeadd75f10c5073c234cf71a2b45bc885f Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Wed, 7 Aug 2024 15:46:39 +0200
Subject: [PATCH 13/22] tests

---
 tests/test_wx2pl.py | 76 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/tests/test_wx2pl.py b/tests/test_wx2pl.py
index 788a238dc..8f426f2e3 100644
--- a/tests/test_wx2pl.py
+++ b/tests/test_wx2pl.py
@@ -3,8 +3,10 @@
 from pint.simulation import make_fake_toas_uniform
 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,
@@ -107,6 +109,54 @@ def data_dmwx():
     return m, t
 
 
+@pytest.fixture
+def data_cmwx():
+    par_sim_cmwx = """
+            PSR           SIM3
+            RAJ           05:00:00     1
+            DECJ          15:00:00     1
+            PEPOCH        55000
+            F0            100          1
+            F1            -1e-15       1 
+            PHOFF         0            1
+            DM            15           1
+            TNCHROMIDX    4
+            CM            10
+            TNCHROMAMP    -13
+            TNCHROMGAM    3.5
+            TNCHROMC      10
+            TZRMJD        55000
+            TZRFRQ        1400 
+            TZRSITE       gbt
+            UNITS         TDB
+            EPHEM         DE440
+            CLOCK         TT(BIPM2019)
+        """
+
+    m = get_model(StringIO(par_sim_cmwx))
+
+    ntoas = 200
+    toaerrs = np.random.uniform(0.5, 2.0, ntoas) * u.us
+    freqs = np.linspace(500, 1500, 4) * u.MHz
+
+    t = make_fake_toas_uniform(
+        startMJD=54001,
+        endMJD=56001,
+        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,
+    )
+
+    return m, t
+
+
 def test_wx2pl(data_wx):
     m, t = data_wx
 
@@ -147,6 +197,32 @@ def test_dmwx2pldm(data_dmwx):
     assert abs(m.TNDMGAM.value - m2.TNDMGAM.value) / m2.TNDMGAM.uncertainty_value < 5
 
 
+def test_cmwx2pldm(data_cmwx):
+    m, t = data_cmwx
+
+    m1 = deepcopy(m)
+    m1.remove_component("PLChromNoise")
+
+    Tspan = t.get_mjds().max() - t.get_mjds().min()
+    cmwavex_setup(m1, Tspan, n_freqs=int(m.TNCHROMC.value), freeze_params=False)
+
+    ftr = WLSFitter(t, m1)
+    ftr.fit_toas(maxiter=10)
+    m1 = ftr.model
+
+    m2 = plchromnoise_from_cmwavex(m1)
+
+    assert "PLChromNoise" in m2.components
+    assert (
+        abs(m.TNCHROMAMP.value - m2.TNCHROMAMP.value) / m2.TNCHROMAMP.uncertainty_value
+        < 5
+    )
+    assert (
+        abs(m.TNCHROMGAM.value - m2.TNCHROMGAM.value) / m2.TNCHROMGAM.uncertainty_value
+        < 5
+    )
+
+
 def test_find_optimal_nharms_wx(data_wx):
     m, t = data_wx
 

From e5f67837b5344c8ec7585cb7fb4c9f564cc1ccd6 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Wed, 7 Aug 2024 15:47:19 +0200
Subject: [PATCH 14/22] CHANGELOG

---
 CHANGELOG-unreleased.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md
index 08e7ebd6f..d201b1d22 100644
--- a/CHANGELOG-unreleased.md
+++ b/CHANGELOG-unreleased.md
@@ -16,7 +16,7 @@ the released changes.
 - 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` function
+- `pint.utils.cmwavex_setup` and `pint.utils.plchromnoise_from_cmwavex` functions
 - More validation for correlated noise components in `TimingModel`
 ### Fixed
 - Bug in `DMWaveX.get_indices()` function

From e9e1771e3aa4a17368650bfe4b6486e712209551 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Thu, 8 Aug 2024 12:08:52 +0200
Subject: [PATCH 15/22] float

---
 src/pint/utils.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/pint/utils.py b/src/pint/utils.py
index 0327e79c1..758307558 100644
--- a/src/pint/utils.py
+++ b/src/pint/utils.py
@@ -3203,11 +3203,13 @@ def mlnlike(params: Tuple[float, ...]) -> float:
         """Negative of the likelihood function that acts on the
         `[DM/CM]WaveX` amplitudes."""
         sigma = powl_model(params)
-        return 0.5 * np.sum(
-            (a**2 / (sigma**2 + da**2))
-            + (b**2 / (sigma**2 + db**2))
-            + np.log(sigma**2 + da**2)
-            + np.log(sigma**2 + db**2)
+        return 0.5 * float(
+            np.sum(
+                (a**2 / (sigma**2 + da**2))
+                + (b**2 / (sigma**2 + db**2))
+                + np.log(sigma**2 + da**2)
+                + np.log(sigma**2 + db**2)
+            )
         )
 
     return mlnlike

From 38e3d30cbcefbafd8fa7390d3e3ebf6c7926bec1 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Mon, 26 Aug 2024 10:07:42 +0200
Subject: [PATCH 16/22] docstring

---
 src/pint/models/binary_ddk.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/pint/models/binary_ddk.py b/src/pint/models/binary_ddk.py
index 452774ead..e6b250a72 100644
--- a/src/pint/models/binary_ddk.py
+++ b/src/pint/models/binary_ddk.py
@@ -42,22 +42,22 @@ 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`
+    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
+    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
+    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`)
@@ -72,7 +72,7 @@ class BinaryDDK(BinaryDD):
        KIN
             the inclination angle: :math:`i`
        KOM
-            the longitude of the ascending node, Kopeikin (1995) Eq 9: :math:`\\Omega`
+            the longitude of the ascending node, Kopeikin (1995) Eq 9: :math:`\Omega`
        K96
             flag for Kopeikin binary model proper motion correction
 

From aaad96026055402a884648f214f4fd223701b7e3 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Mon, 26 Aug 2024 10:08:46 +0200
Subject: [PATCH 17/22] docs

---
 src/pint/logging.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/pint/logging.py b/src/pint/logging.py
index f92fb4bbf..203b70025 100644
--- a/src/pint/logging.py
+++ b/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 ``()``.

From a30e234806ce65b7ba49cde86ac1db389bb3ad1f Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Mon, 26 Aug 2024 10:10:01 +0200
Subject: [PATCH 18/22] docs

---
 src/pint/models/binary_ddk.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/pint/models/binary_ddk.py b/src/pint/models/binary_ddk.py
index e6b250a72..37e31856c 100644
--- a/src/pint/models/binary_ddk.py
+++ b/src/pint/models/binary_ddk.py
@@ -231,26 +231,26 @@ 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:
 
         :math:`(KIN, KOM)`
 
-        :math:`(KIN, 2KOM_0 - KOM - 180^{\\circ})`
+        :math:`(KIN, 2KOM_0 - KOM - 180^{\circ})`
 
-        :math:`(180^{\\circ}-KIN, KOM+180^{\\circ})`
+        :math:`(180^{\circ}-KIN, KOM+180^{\circ})`
 
-        :math:`(180^{\\circ}-KIN, 2KOM_0 - KOM)`
+        :math:`(180^{\circ}-KIN, 2KOM_0 - KOM)`
 
-        All values will be between 0 and :math:`360^{\\circ}`.
+        All values will be between 0 and :math:`360^{\circ}`.
 
         Returns
         -------

From aa9472012b75369c178d3bbe75b708cd1a2cb1c1 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Mon, 26 Aug 2024 10:23:28 +0200
Subject: [PATCH 19/22] test_cmwavex

---
 tests/test_cmwavex.py | 81 ++++++++++++++++++++-----------------------
 1 file changed, 37 insertions(+), 44 deletions(-)

diff --git a/tests/test_cmwavex.py b/tests/test_cmwavex.py
index d0f2652fb..c496cce48 100644
--- a/tests/test_cmwavex.py
+++ b/tests/test_cmwavex.py
@@ -12,7 +12,8 @@
 import astropy.units as u
 
 
-def test_cmwavex():
+@pytest.fixture()
+def model():
     par = """`
         PSR              B1937+21
         LAMBDA   301.9732445337270
@@ -29,22 +30,30 @@ def test_cmwavex():
         TNCHROMIDX               4
         UNITS                  TDB
     """
-    m = get_model(StringIO(par))
+    return get_model(StringIO(par))
 
-    with pytest.raises(ValueError):
-        idxs = cmwavex_setup(m, 3600)
 
-    idxs = cmwavex_setup(m, 3600, n_freqs=5)
+def test_cmwavex_setup(model):
+    idxs = cmwavex_setup(model, 3600, n_freqs=5)
+
+    assert "CMWaveX" in model.components
+    assert model.num_cmwavex_freqs == len(idxs)
 
-    assert "CMWaveX" in m.components
-    assert m.num_cmwavex_freqs == len(idxs)
+    model.components["CMWaveX"].remove_cmwavex_component(5)
+    assert model.num_cmwavex_freqs == len(idxs) - 1
 
-    m.components["CMWaveX"].remove_cmwavex_component(5)
-    assert m.num_cmwavex_freqs == len(idxs) - 1
 
-    t = make_fake_toas_uniform(54000, 56000, 200, m, add_noise=True)
+def test_cwavex_setup_error(model):
+    with pytest.raises(ValueError):
+        cmwavex_setup(model, 3600)
+
 
-    ftr = Fitter.auto(t, m)
+def test_fit_cmwavex(model):
+    cmwavex_setup(model, 3600, n_freqs=5)
+
+    t = make_fake_toas_uniform(54000, 56000, 200, model, add_noise=True)
+
+    ftr = Fitter.auto(t, model)
     ftr.fit_toas()
 
     assert ftr.resids.reduced_chi2 < 2
@@ -117,59 +126,43 @@ def test_cmwavex_badpar():
         get_model(StringIO(par))
 
 
-def test_add_cmwavex():
-    par = """
-        PSR              B1937+21
-        LAMBDA   301.9732445337270
-        BETA      42.2967523367957
-        PMLAMBDA           -0.0175
-        PMBETA             -0.3971
-        PX                  0.1515
-        POSEPOCH        55321.0000
-        F0    641.9282333345536244  1  0.0000000000000132
-        F1     -4.330899370129D-14  1  2.149749089617D-22
-        PEPOCH        55321.000000
-        DM               71.016633
-        CM                     0.1
-        TNCHROMIDX               4
-        UNITS                  TDB
-    """
-    m = get_model(StringIO(par))
-    idxs = cmwavex_setup(m, 3600, n_freqs=5)
-
-    with pytest.raises(ValueError):
-        m.components["CMWaveX"].add_cmwavex_component(1, index=5, cmwxsin=0, cmwxcos=0)
+def test_add_cmwavex(model):
+    idxs = cmwavex_setup(model, 3600, n_freqs=5)
 
-    m.components["CMWaveX"].add_cmwavex_component(1, index=6, cmwxsin=0, cmwxcos=0)
-    assert m.num_cmwavex_freqs == len(idxs) + 1
+    model.components["CMWaveX"].add_cmwavex_component(1, index=6, cmwxsin=0, cmwxcos=0)
+    assert model.num_cmwavex_freqs == len(idxs) + 1
 
-    m.components["CMWaveX"].add_cmwavex_component(
+    model.components["CMWaveX"].add_cmwavex_component(
         1 / u.day, index=7, cmwxsin=0 * cmu, cmwxcos=0 * cmu
     )
-    assert m.num_cmwavex_freqs == len(idxs) + 2
+    assert model.num_cmwavex_freqs == len(idxs) + 2
 
-    m.components["CMWaveX"].add_cmwavex_component(2 / u.day)
-    assert m.num_cmwavex_freqs == len(idxs) + 3
+    model.components["CMWaveX"].add_cmwavex_component(2 / u.day)
+    assert model.num_cmwavex_freqs == len(idxs) + 3
 
-    m.components["CMWaveX"].add_cmwavex_components(
+    model.components["CMWaveX"].add_cmwavex_components(
         np.array([3]) / u.day,
         cmwxsins=np.array([0]) * cmu,
         cmwxcoses=np.array([0]) * cmu,
     )
-    assert m.num_cmwavex_freqs == len(idxs) + 4
+    assert model.num_cmwavex_freqs == len(idxs) + 4
+
+
+def test_add_cmwavex_errors(model):
+    idxs = cmwavex_setup(model, 3600, n_freqs=5)
 
     with pytest.raises(ValueError):
-        m.components["CMWaveX"].add_cmwavex_components(
+        model.components["CMWaveX"].add_cmwavex_components(
             [2 / u.day, 3 / u.day], cmwxsins=[0, 0], cmwxcoses=[0, 0, 0]
         )
 
     with pytest.raises(ValueError):
-        m.components["CMWaveX"].add_cmwavex_components(
+        model.components["CMWaveX"].add_cmwavex_components(
             [2 / u.day, 3 / u.day], cmwxsins=[0, 0, 0], cmwxcoses=[0, 0]
         )
 
     with pytest.raises(ValueError):
-        m.components["CMWaveX"].add_cmwavex_components(
+        model.components["CMWaveX"].add_cmwavex_components(
             [2 / u.day, 3 / u.day],
             cmwxsins=[0, 0],
             cmwxcoses=[0, 0],

From bc480d7612ccc28325d6da490099c5f4f811f86f Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Mon, 26 Aug 2024 10:29:18 +0200
Subject: [PATCH 20/22] CHANGELOG

---
 CHANGELOG-unreleased.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md
index f612c8797..e6ff78b11 100644
--- a/CHANGELOG-unreleased.md
+++ b/CHANGELOG-unreleased.md
@@ -21,7 +21,7 @@ the released changes.
 - `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`
+- 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

From fbd662b0748f857c46602f42f80ce57a34dea3d8 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Tue, 27 Aug 2024 16:00:04 +0200
Subject: [PATCH 21/22] cmwavex example

---
 docs/examples/rednoise-fit-example.py | 161 ++++++++++++++++++++++++++
 1 file changed, 161 insertions(+)

diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py
index ce8bac1c4..78c546a7d 100644
--- a/docs/examples/rednoise-fit-example.py
+++ b/docs/examples/rednoise-fit-example.py
@@ -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()
+
+
+# %%

From 3a02f65e12025c804d5b9019d5a7d8f26952bae0 Mon Sep 17 00:00:00 2001
From: Abhimanyu Susobhanan <abhisrkckl@gmail.com>
Date: Tue, 27 Aug 2024 16:18:26 +0200
Subject: [PATCH 22/22] title

---
 docs/examples/rednoise-fit-example.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/examples/rednoise-fit-example.py b/docs/examples/rednoise-fit-example.py
index 78c546a7d..f0ab81058 100644
--- a/docs/examples/rednoise-fit-example.py
+++ b/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.