From caef4a30f1b8061681137fb1d9edddf8ef21d26c Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Mon, 24 Feb 2025 16:08:47 +0100
Subject: [PATCH 01/29] Initial FFT-based method. No sparse TNT yet

---
 enterprise/signals/gp_priors.py  |   9 +++
 enterprise/signals/gp_signals.py |  49 ++++++++++++++
 enterprise/signals/utils.py      | 113 +++++++++++++++++++++++++++++++
 3 files changed, 171 insertions(+)

diff --git a/enterprise/signals/gp_priors.py b/enterprise/signals/gp_priors.py
index e515552e..c787782f 100644
--- a/enterprise/signals/gp_priors.py
+++ b/enterprise/signals/gp_priors.py
@@ -67,6 +67,15 @@ def t_process_adapt(f, log10_A=-15, gamma=4.33, alphas_adapt=None, nfreq=None):
     return powerlaw(f, log10_A=log10_A, gamma=gamma) * alpha_model
 
 
+@function
+def powerlaw_flat_tail(f, log10_A=-16, gamma=5, log10_kappa=-7, components=2):
+    """Powerlaw with a flat tail (similar to broken powerlaw)"""
+    df = np.diff(np.concatenate((np.array([0]), f[::components])))
+    pl = (10**log10_A) ** 2 / 12.0 / np.pi**2 * const.fyr ** (gamma - 3) * f ** (-gamma) * np.repeat(df, components)
+    flat = 10 ** (2*log10_kappa)
+    return np.maximum(pl, flat)
+
+
 def InvGammaPrior(value, alpha=1, gamma=1):
     """Prior function for InvGamma parameters."""
     return scipy.stats.invgamma.pdf(value, alpha, scale=gamma)
diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index a504a410..44bbac8a 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -216,6 +216,55 @@ class FourierBasisGP(BaseClass):
     return FourierBasisGP
 
 
+def FFTBasisGP(
+    spectrum,
+    coefficients=False,
+    combine=True,
+    components=20,
+    selection=Selection(selections.no_selection),
+    oversample=3,
+    cutoff=1,
+    Tspan=None,
+    name="red_noise",
+):
+    """Convenience function to return a BasisGP class with a
+    coarse time basis basis."""
+
+    basis = utils.create_fft_time_basis(
+        nmodes=components,
+        Tspan=Tspan
+    )
+    BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
+
+    class FFTBasisGP(BaseClass):
+        signal_type = "basis"
+        signal_name = "red noise"
+        signal_id = name
+
+        if coefficients:
+            pass
+
+        else:
+            def get_phi(self, params):
+                """Over-load constructing Phi to deal with the FFT"""
+
+                self._construct_basis(params)
+
+                for key, slc in self._slices.items():
+                    t_knots = self._labels[key]
+                    freqs, zeros = utils.knots_to_freqs(t_knots, oversample=oversample, cutoff=cutoff)
+                    psd = np.concatenate([
+                        self._prior[key](freqs, params=params, components=1),
+                        zeros
+                    ])
+
+                    phislc = utils.psd2cov(t_knots, freqs, psd)
+                    self._phi = self._phi.set(phislc, slc)
+                return self._phi
+
+    return FFTBasisGP
+
+
 def get_timing_model_basis(use_svd=False, normed=True, idx_exclude=None):
     if use_svd:
         if normed is not True:
diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index e1f32430..dc23d68b 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -841,6 +841,119 @@ def linear_interp_basis(toas, dt=30 * 86400):
     return M[:, idx], x[idx]
 
 
+@function
+def create_fft_time_basis(
+    toas, 
+    nmodes=30,
+    Tspan=None,
+    start_time=None
+):
+    """
+    Construct coarse time-domain design matrix from eq 11 of Chrisostomi et al., 2025
+    :param toas: vector of time series in seconds
+    :param nmodes: number of fourier coefficients to use
+    :param Tspan: option to some other Tspan
+    :param start_time: option to set some other start epoch of basis
+
+    :return B: coarse time-domain design matrix
+    :return t_coarse: timestamps of coarse time grid
+    """
+    if start_time is None:
+        start_time = np.min(toas)
+    else:
+        if start_time > np.min(toas):
+            raise ValueError('Coarse time basis start must be earlier than earliest TOA.')
+
+    if Tspan is None:
+        Tspan = np.max(toas) - np.min(toas)
+
+    t_fine = toas
+    t_coarse = np.linspace(start_time, start_time + Tspan, nmodes)
+    dt_coarse = t_coarse[1] - t_coarse[0]
+
+    idx = np.arange(len(t_fine))
+    idy = np.searchsorted(t_coarse, t_fine)
+    idy[idy == 0] = 1
+
+    Bmat = np.zeros((len(t_fine), len(t_coarse)), 'd')
+
+    Bmat[idx, idy] = (t_fine - t_coarse[idy - 1]) / dt_coarse
+    Bmat[idx, idy - 1] = (t_coarse[idy] - t_fine) / dt_coarse
+
+    return Bmat, t_coarse
+
+
+def psd2cov(
+        t_knots,
+        freqs,
+        psd,
+    ):
+    """
+    Convert a power spectral density function, defined by (freqs, psd), to a covariance matrix
+
+    :param t_knots: Timestamps of the coarse time grid
+    :param freqs: frequencies of the PSD
+    :param psd: values of the PSD at frequencies freqs
+
+    :return covmat: Covariance matrix at coarse time grid
+    """
+    df = freqs[1] - freqs[0]
+
+    # Toeplitz helper (symmetric matrix defined by its first column)
+    def toeplitz(c):
+        c = np.asarray(c)
+        n = len(c)
+        i = np.arange(n).reshape(-1, 1)
+        j = np.arange(n).reshape(1, -1)
+        return c[np.abs(i - j)]
+
+    def covmat(*args):
+        # Build the full symmetric PSD by mirroring (skip the first and last duplicate).
+        fullpsd = np.concatenate([psd, psd[-2:0:-1]])
+
+        # Compute the inverse FFT. The 'norm' parameter 'backward' corresponds to NumPy's default.
+        Cfreq = np.fft.ifft(fullpsd, norm='backward')
+        # Scale to get the covariance at lag τ.
+        Ctau = Cfreq.real * len(fullpsd) * df / 2
+
+        # Construct and return the Toeplitz covariance matrix.
+        return toeplitz(Ctau[:len(t_knots)])
+
+    return covmat()
+
+
+def knots_to_freqs(t_knots, oversample=3, cutoff=1):
+    """
+    Convert knots of coarse time grid to frequencies
+
+    :param t_knots: Timestamps of the coarse time grid
+    :param oversample: amount by which to over-sample the frequency grid
+    :param cutoff: frequency 1 / (cutoff * T) at which to do
+                   low-frequency cut-off of the PSD
+
+    :return covmat: Covariance matrix at coarse time grid
+    """
+    nmodes = len(t_knots)
+    Tspan = np.max(t_knots) - np.min(t_knots)
+
+    if nmodes % 2 == 0:
+        raise ValueError('psd2cov number of nmodes must be odd.')
+
+    n_freqs = int((nmodes - 1) / 2 * oversample + 1)
+    fmax = (nmodes - 1) / Tspan / 2
+    freqs = np.linspace(0, fmax, n_freqs)
+    df = 1 / Tspan / oversample
+
+    if cutoff is not None:
+        i_cutoff = int(np.ceil(oversample / cutoff))
+        fs, zs = np.array(freqs[i_cutoff:]), np.zeros(i_cutoff)
+    else:
+        fs = np.array(freqs)
+        zs = np.array([], dtype=fs.dtype)
+
+    return fs, zs
+
+
 # overlap reduction functions
 
 

From d57be32a28644591564ac3132dee985192071eec Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 26 Feb 2025 08:02:31 +0000
Subject: [PATCH 02/29] Linting

---
 enterprise/signals/gp_bases.py   |  38 ++++++++++++
 enterprise/signals/gp_signals.py | 101 +++++++++++++++++++++++++++----
 enterprise/signals/utils.py      |  68 +++------------------
 3 files changed, 138 insertions(+), 69 deletions(-)

diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py
index f924be12..63dd73d7 100644
--- a/enterprise/signals/gp_bases.py
+++ b/enterprise/signals/gp_bases.py
@@ -12,6 +12,7 @@
 
 __all__ = [
     "createfourierdesignmatrix_red",
+    "create_fft_time_basis",
     "createfourierdesignmatrix_dm",
     "createfourierdesignmatrix_dm_tn",
     "createfourierdesignmatrix_env",
@@ -91,6 +92,43 @@ def createfourierdesignmatrix_red(
     return F, Ffreqs
 
 
+@function
+def create_fft_time_basis(toas, nmodes=30, Tspan=None, start_time=None):
+    """
+    Construct coarse time-domain design matrix from eq 11 of Chrisostomi et al., 2025
+    :param toas: vector of time series in seconds
+    :param nmodes: number of fourier coefficients to use
+    :param Tspan: option to some other Tspan
+    :param start_time: option to set some other start epoch of basis
+
+    :return B: coarse time-domain design matrix
+    :return t_coarse: timestamps of coarse time grid
+    """
+    if start_time is None:
+        start_time = np.min(toas)
+    else:
+        if start_time > np.min(toas):
+            raise ValueError("Coarse time basis start must be earlier than earliest TOA.")
+
+    if Tspan is None:
+        Tspan = np.max(toas) - np.min(toas)
+
+    t_fine = toas
+    t_coarse = np.linspace(start_time, start_time + Tspan, nmodes)
+    dt_coarse = t_coarse[1] - t_coarse[0]
+
+    idx = np.arange(len(t_fine))
+    idy = np.searchsorted(t_coarse, t_fine)
+    idy = np.clip(idy, 1, nmodes - 1)
+
+    Bmat = np.zeros((len(t_fine), len(t_coarse)), "d")
+
+    Bmat[idx, idy] = (t_fine - t_coarse[idy - 1]) / dt_coarse
+    Bmat[idx, idy - 1] = (t_coarse[idy] - t_fine) / dt_coarse
+
+    return Bmat, t_coarse
+
+
 @function
 def createfourierdesignmatrix_dm(
     toas, freqs, nmodes=30, Tspan=None, pshift=False, fref=1400, logf=False, fmin=None, fmax=None, modes=None
diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 44bbac8a..93cfd9d4 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -225,15 +225,13 @@ def FFTBasisGP(
     oversample=3,
     cutoff=1,
     Tspan=None,
+    start_time=None,
     name="red_noise",
 ):
     """Convenience function to return a BasisGP class with a
     coarse time basis basis."""
 
-    basis = utils.create_fft_time_basis(
-        nmodes=components,
-        Tspan=Tspan
-    )
+    basis = utils.create_fft_time_basis(nmodes=components, Tspan=Tspan, start_time=start_time)
     BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
 
     class FFTBasisGP(BaseClass):
@@ -242,9 +240,10 @@ class FFTBasisGP(BaseClass):
         signal_id = name
 
         if coefficients:
-            pass
+            raise NotImplementedError("Coefficients not supported for FFTBasisGP")
 
         else:
+
             def get_phi(self, params):
                 """Over-load constructing Phi to deal with the FFT"""
 
@@ -253,10 +252,7 @@ def get_phi(self, params):
                 for key, slc in self._slices.items():
                     t_knots = self._labels[key]
                     freqs, zeros = utils.knots_to_freqs(t_knots, oversample=oversample, cutoff=cutoff)
-                    psd = np.concatenate([
-                        self._prior[key](freqs, params=params, components=1),
-                        zeros
-                    ])
+                    psd = np.concatenate([self._prior[key](freqs, params=params, components=1), zeros])
 
                     phislc = utils.psd2cov(t_knots, freqs, psd)
                     self._phi = self._phi.set(phislc, slc)
@@ -475,7 +471,6 @@ def FourierBasisCommonGP(
     pshift=False,
     pseed=None,
 ):
-
     if coefficients and Tspan is None:
         raise ValueError(
             "With coefficients=True, FourierBasisCommonGP " + "requires that you specify Tspan explicitly."
@@ -503,7 +498,7 @@ def __init__(self, psr):
         # since this function has side-effects, it can only be cached
         # with limit=1, so it will run again if called with params different
         # than the last time
-        @signal_base.cache_call("basis_params", 1)
+        @signal_base.cache_call("basis_params", limit=1)
         def _construct_basis(self, params={}):
             span = Tspan if Tspan is not None else max(FourierBasisCommonGP._Tmax) - min(FourierBasisCommonGP._Tmin)
             self._basis, self._labels = self._bases(params=params, Tspan=span)
@@ -511,6 +506,90 @@ def _construct_basis(self, params={}):
     return FourierBasisCommonGP
 
 
+def FFTBasisCommonGP(
+    spectrum,
+    orf,
+    coefficients=False,
+    combine=True,
+    components=20,
+    Tspan=None,
+    start_time=None,
+    cutoff=1,
+    oversample=3,
+    name="common_fft",
+):
+    if coefficients and Tspan is None:
+        raise ValueError(
+            "With coefficients=True, FourierBasisCommonGP " + "requires that you specify Tspan explicitly."
+        )
+
+    basis = utils.create_fft_time_basis(nmodes=components, Tspan=Tspan, start_time=start_time)
+    BaseClass = BasisCommonGP(spectrum, basis, orf, coefficients=coefficients, combine=combine, name=name)
+
+    class FFTBasisCommonGP(BaseClass):
+        signal_type = "common basis"
+        signal_name = "common red noise"
+        signal_id = name
+
+        _Tmin, _Tmax = [], []
+
+        def __init__(self, psr):
+            super(FFTBasisCommonGP, self).__init__(psr)
+
+            if Tspan is None:
+                FFTBasisCommonGP._Tmin.append(psr.toas.min())
+                FFTBasisCommonGP._Tmax.append(psr.toas.max())
+
+        # since this function has side-effects, it can only be cached
+        # with limit=1, so it will run again if called with params different
+        # than the last time
+        @signal_base.cache_call("basis_params", limit=1)
+        def _construct_basis(self, params={}):
+            span = Tspan if Tspan is not None else max(FFTBasisCommonGP._Tmax) - min(FFTBasisCommonGP._Tmin)
+            start = start_time if start_time is not None else min(FFTBasisCommonGP._Tmin)
+            self._basis, self._labels = self._bases(params=params, Tspan=span, start_time=start)
+
+            self._t_knots = self._labels
+            freqs, zeros = utils.knots_to_freqs(self._t_knots, oversample=oversample, cutoff=cutoff)
+            self._freqs = freqs
+            self._zeros = zeros
+
+        if coefficients:
+            raise NotImplementedError("Coefficients not supported for FFTBasisCommonGP")
+
+        else:
+
+            @signal_base.cache_call("basis_params", limit=1)
+            def get_phi_matrix(self, params):
+                """
+                Compute and cache the time-domain covariance ('phi') for *this* signal's basis.
+                """
+                psd = np.concatenate([FFTBasisCommonGP._prior(self._freqs, params=params, components=1), self._zeros])
+                return utils.psd2cov(self._t_knots, self._freqs, psd)
+
+            def get_phi(self, params):
+                """Over-load constructing Phi to deal with the FFT"""
+                self._construct_basis(params)
+                phi = self.get_phi_matrix(params)
+
+                orf = FFTBasisCommonGP._orf(self._psrpos, self._psrpos, params=params)
+
+                return orf * phi
+
+            @classmethod
+            def get_phicross(cls, signal1, signal2, params):
+                """Use Phi from signal1, ORF from signal1 vs signal2"""
+
+                phi1 = signal1.get_phi_matrix(params)
+                # phi2 = signal2.get_phi_matrix(params)
+
+                orf = FFTBasisCommonGP._orf(signal1._psrpos, signal2._psrpos, params=params)
+
+                return phi1 * orf
+
+    return FFTBasisCommonGP
+
+
 # for simplicity, we currently do not handle Tspan automatically
 def FourierBasisCommonGP_ephem(spectrum, components, Tspan, name="ephem_gp"):
     basis = utils.createfourierdesignmatrix_ephem(nmodes=components, Tspan=Tspan)
diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index dc23d68b..f535e938 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -20,6 +20,7 @@
 from enterprise import signals as sigs  # noqa: F401
 from enterprise.signals.gp_bases import (  # noqa: F401
     createfourierdesignmatrix_red,
+    create_fft_time_basis,
     createfourierdesignmatrix_dm,
     createfourierdesignmatrix_dm_tn,
     createfourierdesignmatrix_env,
@@ -841,65 +842,21 @@ def linear_interp_basis(toas, dt=30 * 86400):
     return M[:, idx], x[idx]
 
 
-@function
-def create_fft_time_basis(
-    toas, 
-    nmodes=30,
-    Tspan=None,
-    start_time=None
-):
-    """
-    Construct coarse time-domain design matrix from eq 11 of Chrisostomi et al., 2025
-    :param toas: vector of time series in seconds
-    :param nmodes: number of fourier coefficients to use
-    :param Tspan: option to some other Tspan
-    :param start_time: option to set some other start epoch of basis
-
-    :return B: coarse time-domain design matrix
-    :return t_coarse: timestamps of coarse time grid
-    """
-    if start_time is None:
-        start_time = np.min(toas)
-    else:
-        if start_time > np.min(toas):
-            raise ValueError('Coarse time basis start must be earlier than earliest TOA.')
-
-    if Tspan is None:
-        Tspan = np.max(toas) - np.min(toas)
-
-    t_fine = toas
-    t_coarse = np.linspace(start_time, start_time + Tspan, nmodes)
-    dt_coarse = t_coarse[1] - t_coarse[0]
-
-    idx = np.arange(len(t_fine))
-    idy = np.searchsorted(t_coarse, t_fine)
-    idy[idy == 0] = 1
-
-    Bmat = np.zeros((len(t_fine), len(t_coarse)), 'd')
-
-    Bmat[idx, idy] = (t_fine - t_coarse[idy - 1]) / dt_coarse
-    Bmat[idx, idy - 1] = (t_coarse[idy] - t_fine) / dt_coarse
-
-    return Bmat, t_coarse
-
-
 def psd2cov(
-        t_knots,
-        freqs,
-        psd,
-    ):
+    t_knots,
+    freqs,
+    psd,
+):
     """
     Convert a power spectral density function, defined by (freqs, psd), to a covariance matrix
 
     :param t_knots: Timestamps of the coarse time grid
     :param freqs: frequencies of the PSD
-    :param psd: values of the PSD at frequencies freqs
+    :param psd: values of the PSD at frequencies freqs (assumes *delta_f in psd)
 
     :return covmat: Covariance matrix at coarse time grid
     """
-    df = freqs[1] - freqs[0]
 
-    # Toeplitz helper (symmetric matrix defined by its first column)
     def toeplitz(c):
         c = np.asarray(c)
         n = len(c)
@@ -908,16 +865,12 @@ def toeplitz(c):
         return c[np.abs(i - j)]
 
     def covmat(*args):
-        # Build the full symmetric PSD by mirroring (skip the first and last duplicate).
         fullpsd = np.concatenate([psd, psd[-2:0:-1]])
 
-        # Compute the inverse FFT. The 'norm' parameter 'backward' corresponds to NumPy's default.
-        Cfreq = np.fft.ifft(fullpsd, norm='backward')
-        # Scale to get the covariance at lag τ.
-        Ctau = Cfreq.real * len(fullpsd) * df / 2
+        Cfreq = np.fft.ifft(fullpsd, norm="backward")
+        Ctau = Cfreq.real * len(fullpsd) / 2
 
-        # Construct and return the Toeplitz covariance matrix.
-        return toeplitz(Ctau[:len(t_knots)])
+        return toeplitz(Ctau[: len(t_knots)])
 
     return covmat()
 
@@ -937,12 +890,11 @@ def knots_to_freqs(t_knots, oversample=3, cutoff=1):
     Tspan = np.max(t_knots) - np.min(t_knots)
 
     if nmodes % 2 == 0:
-        raise ValueError('psd2cov number of nmodes must be odd.')
+        raise ValueError("psd2cov number of nmodes must be odd.")
 
     n_freqs = int((nmodes - 1) / 2 * oversample + 1)
     fmax = (nmodes - 1) / Tspan / 2
     freqs = np.linspace(0, fmax, n_freqs)
-    df = 1 / Tspan / oversample
 
     if cutoff is not None:
         i_cutoff = int(np.ceil(oversample / cutoff))

From 6f7a92e61a750f32afc64aeed3274c536ae3e130 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 26 Feb 2025 08:51:29 +0000
Subject: [PATCH 03/29] Change some comments

---
 enterprise/signals/gp_signals.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 93cfd9d4..18592d4c 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -229,7 +229,7 @@ def FFTBasisGP(
     name="red_noise",
 ):
     """Convenience function to return a BasisGP class with a
-    coarse time basis basis."""
+    coarse time basis."""
 
     basis = utils.create_fft_time_basis(nmodes=components, Tspan=Tspan, start_time=start_time)
     BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
@@ -245,8 +245,6 @@ class FFTBasisGP(BaseClass):
         else:
 
             def get_phi(self, params):
-                """Over-load constructing Phi to deal with the FFT"""
-
                 self._construct_basis(params)
 
                 for key, slc in self._slices.items():

From 6c43e1dcb6b107e5aca667380e23b8f2081ab99a Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 26 Feb 2025 12:28:06 +0000
Subject: [PATCH 04/29] Removed caching for common FFT signals

---
 enterprise/signals/gp_priors.py  | 2 +-
 enterprise/signals/gp_signals.py | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/enterprise/signals/gp_priors.py b/enterprise/signals/gp_priors.py
index c787782f..4d1a42ec 100644
--- a/enterprise/signals/gp_priors.py
+++ b/enterprise/signals/gp_priors.py
@@ -72,7 +72,7 @@ def powerlaw_flat_tail(f, log10_A=-16, gamma=5, log10_kappa=-7, components=2):
     """Powerlaw with a flat tail (similar to broken powerlaw)"""
     df = np.diff(np.concatenate((np.array([0]), f[::components])))
     pl = (10**log10_A) ** 2 / 12.0 / np.pi**2 * const.fyr ** (gamma - 3) * f ** (-gamma) * np.repeat(df, components)
-    flat = 10 ** (2*log10_kappa)
+    flat = 10 ** (2 * log10_kappa)
     return np.maximum(pl, flat)
 
 
diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 18592d4c..a966fc19 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -556,8 +556,7 @@ def _construct_basis(self, params={}):
             raise NotImplementedError("Coefficients not supported for FFTBasisCommonGP")
 
         else:
-
-            @signal_base.cache_call("basis_params", limit=1)
+            # @signal_base.cache_call("basis_params", limit=1)
             def get_phi_matrix(self, params):
                 """
                 Compute and cache the time-domain covariance ('phi') for *this* signal's basis.

From 9cd98990a490f77d13b83a3625d5bbd7e5fdc203 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 26 Feb 2025 18:42:24 +0000
Subject: [PATCH 05/29] Tweaks to interpolation basis

---
 enterprise/signals/gp_bases.py   | 5 ++++-
 enterprise/signals/gp_signals.py | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py
index 63dd73d7..ae6c13b8 100644
--- a/enterprise/signals/gp_bases.py
+++ b/enterprise/signals/gp_bases.py
@@ -111,7 +111,10 @@ def create_fft_time_basis(toas, nmodes=30, Tspan=None, start_time=None):
             raise ValueError("Coarse time basis start must be earlier than earliest TOA.")
 
     if Tspan is None:
-        Tspan = np.max(toas) - np.min(toas)
+        Tspan = np.max(toas) - start_time
+    else:
+        if start_time + Tspan < np.max(toas):
+            raise ValueError("Coarse time basis end must be later than latest TOA.")
 
     t_fine = toas
     t_coarse = np.linspace(start_time, start_time + Tspan, nmodes)
diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index a966fc19..07282371 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -543,8 +543,8 @@ def __init__(self, psr):
         # than the last time
         @signal_base.cache_call("basis_params", limit=1)
         def _construct_basis(self, params={}):
-            span = Tspan if Tspan is not None else max(FFTBasisCommonGP._Tmax) - min(FFTBasisCommonGP._Tmin)
             start = start_time if start_time is not None else min(FFTBasisCommonGP._Tmin)
+            span = Tspan if Tspan is not None else max(FFTBasisCommonGP._Tmax) - start
             self._basis, self._labels = self._bases(params=params, Tspan=span, start_time=start)
 
             self._t_knots = self._labels

From 2baf8ebfed16889d8155c28a16b329de2fd6335b Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 10:18:33 +0000
Subject: [PATCH 06/29] FIXED: bug in FFT signals due to Enterprise prior
 functions df calculation

---
 enterprise/signals/gp_signals.py | 80 ++++++++++++++++++++++++++------
 enterprise/signals/utils.py      | 17 +------
 2 files changed, 69 insertions(+), 28 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 07282371..e628d287 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -221,9 +221,11 @@ def FFTBasisGP(
     coefficients=False,
     combine=True,
     components=20,
+    knots=None,
     selection=Selection(selections.no_selection),
     oversample=3,
-    cutoff=1,
+    cutoff=None,
+    cutbins=1,
     Tspan=None,
     start_time=None,
     name="red_noise",
@@ -231,7 +233,18 @@ def FFTBasisGP(
     """Convenience function to return a BasisGP class with a
     coarse time basis."""
 
-    basis = utils.create_fft_time_basis(nmodes=components, Tspan=Tspan, start_time=start_time)
+    if knots is None:
+        knots = 2 * components + 1
+
+    elif knots is not None and knots % 2 == 0:
+        raise ValueError("Knots needs to be an odd number")
+
+    if cutoff is not None:
+        #:param cutoff: frequency 1 / (cutoff * T) at which to do
+        #            low-frequency cut-off of the PSD
+        cutbins = int(np.ceil(oversample / cutoff))
+
+    basis = utils.create_fft_time_basis(nmodes=knots, Tspan=Tspan, start_time=start_time)
     BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
 
     class FFTBasisGP(BaseClass):
@@ -249,11 +262,25 @@ def get_phi(self, params):
 
                 for key, slc in self._slices.items():
                     t_knots = self._labels[key]
-                    freqs, zeros = utils.knots_to_freqs(t_knots, oversample=oversample, cutoff=cutoff)
-                    psd = np.concatenate([self._prior[key](freqs, params=params, components=1), zeros])
+
+                    freqs = utils.knots_to_freqs(t_knots, oversample=oversample)
+
+                    # Hack, because Enterprise adds in f=0 and then calculates df,
+                    # meaning we cannot simply start freqs from 0. Thus, we use
+                    # a modified frequency spacing, such that:
+                    # [0, f1, 0, f1, f2, f3] => [df, -df, df, df, df]
+                    if cutbins == 0:
+                        freqs_prior = np.concatenate([[freqs[1]], freqs])
+                        psd_prior = self._prior[key](freqs_prior, params=params, components=1)
+                        psd = np.concatenate([[-psd_prior[1]], psd_prior[2:]])
+
+                    else:
+                        psd_prior = self._prior[key](freqs[1:], params=params, components=1)
+                        psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
 
                     phislc = utils.psd2cov(t_knots, freqs, psd)
                     self._phi = self._phi.set(phislc, slc)
+
                 return self._phi
 
     return FFTBasisGP
@@ -510,18 +537,31 @@ def FFTBasisCommonGP(
     coefficients=False,
     combine=True,
     components=20,
+    knots=None,
     Tspan=None,
     start_time=None,
-    cutoff=1,
+    cutoff=None,
+    cutbins=1,
     oversample=3,
     name="common_fft",
 ):
-    if coefficients and Tspan is None:
+    if coefficients and (Tspan is None or start_time is None):
         raise ValueError(
-            "With coefficients=True, FourierBasisCommonGP " + "requires that you specify Tspan explicitly."
+            "With coefficients=True, FFTBasisCommonGP " + "requires that you specify Tspan/start_time explicitly."
         )
 
-    basis = utils.create_fft_time_basis(nmodes=components, Tspan=Tspan, start_time=start_time)
+    if knots is None:
+        knots = 2 * components + 1
+
+    elif knots is not None and knots % 2 == 0:
+        raise ValueError("Knots needs to be an odd number")
+
+    if cutoff is not None:
+        #:param cutoff: frequency 1 / (cutoff * T) at which to do
+        #            low-frequency cut-off of the PSD
+        cutbins = int(np.ceil(oversample / cutoff))
+
+    basis = utils.create_fft_time_basis(nmodes=knots, Tspan=Tspan, start_time=start_time)
     BaseClass = BasisCommonGP(spectrum, basis, orf, coefficients=coefficients, combine=combine, name=name)
 
     class FFTBasisCommonGP(BaseClass):
@@ -534,8 +574,10 @@ class FFTBasisCommonGP(BaseClass):
         def __init__(self, psr):
             super(FFTBasisCommonGP, self).__init__(psr)
 
-            if Tspan is None:
+            if start_time is None:
                 FFTBasisCommonGP._Tmin.append(psr.toas.min())
+
+            if Tspan is None:
                 FFTBasisCommonGP._Tmax.append(psr.toas.max())
 
         # since this function has side-effects, it can only be cached
@@ -548,9 +590,8 @@ def _construct_basis(self, params={}):
             self._basis, self._labels = self._bases(params=params, Tspan=span, start_time=start)
 
             self._t_knots = self._labels
-            freqs, zeros = utils.knots_to_freqs(self._t_knots, oversample=oversample, cutoff=cutoff)
+            freqs = utils.knots_to_freqs(self._t_knots, oversample=oversample)
             self._freqs = freqs
-            self._zeros = zeros
 
         if coefficients:
             raise NotImplementedError("Coefficients not supported for FFTBasisCommonGP")
@@ -561,8 +602,21 @@ def get_phi_matrix(self, params):
                 """
                 Compute and cache the time-domain covariance ('phi') for *this* signal's basis.
                 """
-                psd = np.concatenate([FFTBasisCommonGP._prior(self._freqs, params=params, components=1), self._zeros])
-                return utils.psd2cov(self._t_knots, self._freqs, psd)
+
+                # Hack, because Enterprise adds in f=0 and then calculates df,
+                # meaning we cannot simply start freqs from 0. Thus, we use
+                # a modified frequency spacing, such that:
+                # [0, f1, 0, f1, f2, f3] => [df, -df, df, df, df]
+                if cutbins == 0:
+                    freqs_prior = np.concatenate([[self._freqs[1]], self._freqs])
+                    psd_prior = FFTBasisCommonGP._prior(freqs_prior, params=params, components=1)
+                    psd = np.concatenate([[-psd_prior[1]], psd_prior[2:]])
+
+                else:
+                    psd_prior = FFTBasisCommonGP._prior(self._freqs[1:], params=params, components=1)
+                    psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
+
+                return utils.psd2cov(self._t_knots, psd)
 
             def get_phi(self, params):
                 """Over-load constructing Phi to deal with the FFT"""
diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index f535e938..ebfe855a 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -844,14 +844,12 @@ def linear_interp_basis(toas, dt=30 * 86400):
 
 def psd2cov(
     t_knots,
-    freqs,
     psd,
 ):
     """
     Convert a power spectral density function, defined by (freqs, psd), to a covariance matrix
 
     :param t_knots: Timestamps of the coarse time grid
-    :param freqs: frequencies of the PSD
     :param psd: values of the PSD at frequencies freqs (assumes *delta_f in psd)
 
     :return covmat: Covariance matrix at coarse time grid
@@ -875,14 +873,12 @@ def covmat(*args):
     return covmat()
 
 
-def knots_to_freqs(t_knots, oversample=3, cutoff=1):
+def knots_to_freqs(t_knots, oversample=3):
     """
     Convert knots of coarse time grid to frequencies
 
     :param t_knots: Timestamps of the coarse time grid
     :param oversample: amount by which to over-sample the frequency grid
-    :param cutoff: frequency 1 / (cutoff * T) at which to do
-                   low-frequency cut-off of the PSD
 
     :return covmat: Covariance matrix at coarse time grid
     """
@@ -894,16 +890,7 @@ def knots_to_freqs(t_knots, oversample=3, cutoff=1):
 
     n_freqs = int((nmodes - 1) / 2 * oversample + 1)
     fmax = (nmodes - 1) / Tspan / 2
-    freqs = np.linspace(0, fmax, n_freqs)
-
-    if cutoff is not None:
-        i_cutoff = int(np.ceil(oversample / cutoff))
-        fs, zs = np.array(freqs[i_cutoff:]), np.zeros(i_cutoff)
-    else:
-        fs = np.array(freqs)
-        zs = np.array([], dtype=fs.dtype)
-
-    return fs, zs
+    return np.linspace(0, fmax, n_freqs)
 
 
 # overlap reduction functions

From 01c3655835147452ee9ee5ee4ab308b9de611469 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 11:05:04 +0000
Subject: [PATCH 07/29] Block comment should start with '# '

---
 enterprise/signals/gp_signals.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index e628d287..2e88dbd7 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -240,8 +240,8 @@ def FFTBasisGP(
         raise ValueError("Knots needs to be an odd number")
 
     if cutoff is not None:
-        #:param cutoff: frequency 1 / (cutoff * T) at which to do
-        #            low-frequency cut-off of the PSD
+        # :param cutoff: frequency 1 / (cutoff * T) at which to do
+        #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
     basis = utils.create_fft_time_basis(nmodes=knots, Tspan=Tspan, start_time=start_time)
@@ -557,8 +557,8 @@ def FFTBasisCommonGP(
         raise ValueError("Knots needs to be an odd number")
 
     if cutoff is not None:
-        #:param cutoff: frequency 1 / (cutoff * T) at which to do
-        #            low-frequency cut-off of the PSD
+        # :param cutoff: frequency 1 / (cutoff * T) at which to do
+        #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
     basis = utils.create_fft_time_basis(nmodes=knots, Tspan=Tspan, start_time=start_time)

From 357f1e9e44ad843aa2444db877bc401086f794b2 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 15:22:21 +0000
Subject: [PATCH 08/29] Changed knots -> nknots, and added tests for FFT
 signals

---
 enterprise/signals/gp_bases.py   |  8 +--
 enterprise/signals/gp_signals.py | 22 ++++----
 tests/test_gp_signals.py         | 87 +++++++++++++++++++++++++++++++-
 3 files changed, 100 insertions(+), 17 deletions(-)

diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py
index ae6c13b8..02cc20b1 100644
--- a/enterprise/signals/gp_bases.py
+++ b/enterprise/signals/gp_bases.py
@@ -93,11 +93,11 @@ def createfourierdesignmatrix_red(
 
 
 @function
-def create_fft_time_basis(toas, nmodes=30, Tspan=None, start_time=None):
+def create_fft_time_basis(toas, nknots=30, Tspan=None, start_time=None):
     """
     Construct coarse time-domain design matrix from eq 11 of Chrisostomi et al., 2025
     :param toas: vector of time series in seconds
-    :param nmodes: number of fourier coefficients to use
+    :param nknots: number of coarse time samples to use (knots)
     :param Tspan: option to some other Tspan
     :param start_time: option to set some other start epoch of basis
 
@@ -117,12 +117,12 @@ def create_fft_time_basis(toas, nmodes=30, Tspan=None, start_time=None):
             raise ValueError("Coarse time basis end must be later than latest TOA.")
 
     t_fine = toas
-    t_coarse = np.linspace(start_time, start_time + Tspan, nmodes)
+    t_coarse = np.linspace(start_time, start_time + Tspan, nknots)
     dt_coarse = t_coarse[1] - t_coarse[0]
 
     idx = np.arange(len(t_fine))
     idy = np.searchsorted(t_coarse, t_fine)
-    idy = np.clip(idy, 1, nmodes - 1)
+    idy = np.clip(idy, 1, nknots - 1)
 
     Bmat = np.zeros((len(t_fine), len(t_coarse)), "d")
 
diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 2e88dbd7..15830f21 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -221,7 +221,7 @@ def FFTBasisGP(
     coefficients=False,
     combine=True,
     components=20,
-    knots=None,
+    nknots=None,
     selection=Selection(selections.no_selection),
     oversample=3,
     cutoff=None,
@@ -233,10 +233,10 @@ def FFTBasisGP(
     """Convenience function to return a BasisGP class with a
     coarse time basis."""
 
-    if knots is None:
-        knots = 2 * components + 1
+    if nknots is None:
+        nknots = 2 * components + 1
 
-    elif knots is not None and knots % 2 == 0:
+    elif nknots is not None and nknots % 2 == 0:
         raise ValueError("Knots needs to be an odd number")
 
     if cutoff is not None:
@@ -244,7 +244,7 @@ def FFTBasisGP(
         #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
-    basis = utils.create_fft_time_basis(nmodes=knots, Tspan=Tspan, start_time=start_time)
+    basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time)
     BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
 
     class FFTBasisGP(BaseClass):
@@ -278,7 +278,7 @@ def get_phi(self, params):
                         psd_prior = self._prior[key](freqs[1:], params=params, components=1)
                         psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
 
-                    phislc = utils.psd2cov(t_knots, freqs, psd)
+                    phislc = utils.psd2cov(t_knots, psd)
                     self._phi = self._phi.set(phislc, slc)
 
                 return self._phi
@@ -537,7 +537,7 @@ def FFTBasisCommonGP(
     coefficients=False,
     combine=True,
     components=20,
-    knots=None,
+    nknots=None,
     Tspan=None,
     start_time=None,
     cutoff=None,
@@ -550,10 +550,10 @@ def FFTBasisCommonGP(
             "With coefficients=True, FFTBasisCommonGP " + "requires that you specify Tspan/start_time explicitly."
         )
 
-    if knots is None:
-        knots = 2 * components + 1
+    if nknots is None:
+        nknots = 2 * components + 1
 
-    elif knots is not None and knots % 2 == 0:
+    elif nknots is not None and nknots % 2 == 0:
         raise ValueError("Knots needs to be an odd number")
 
     if cutoff is not None:
@@ -561,7 +561,7 @@ def FFTBasisCommonGP(
         #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
-    basis = utils.create_fft_time_basis(nmodes=knots, Tspan=Tspan, start_time=start_time)
+    basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time)
     BaseClass = BasisCommonGP(spectrum, basis, orf, coefficients=coefficients, combine=combine, name=name)
 
     class FFTBasisCommonGP(BaseClass):
diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 24197488..2ba683c5 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -15,7 +15,7 @@
 import scipy.linalg as sl
 
 from enterprise.pulsar import Pulsar
-from enterprise.signals import gp_signals, parameter, selections, signal_base, utils
+from enterprise.signals import gp_signals, parameter, selections, signal_base, utils, white_signals
 from enterprise.signals.selections import Selection
 from tests.enterprise_test_data import datadir
 from tests.enterprise_test_data import LIBSTEMPO_INSTALLED, PINT_INSTALLED
@@ -28,7 +28,6 @@ def create_quant_matrix(toas, dt=1):
     # return value slightly different than 1 to get around ECORR columns
     return U * 1.0000001, avetoas
 
-
 @signal_base.function
 def se_kernel(etoas, log10_sigma=-7, log10_lam=np.log10(30 * 86400)):
     tm = np.abs(etoas[None, :] - etoas[:, None])
@@ -36,6 +35,18 @@ def se_kernel(etoas, log10_sigma=-7, log10_lam=np.log10(30 * 86400)):
     return 10 ** (2 * log10_sigma) * np.exp(-(tm**2) / 2 / 10 ** (2 * log10_lam)) + d
 
 
+@signal_base.function
+def psd_matern32(f, l=365 * 86400.0, log10_sigma_sqr=-14, components=2):
+    df = np.diff(np.concatenate((np.array([0]), f[::components])))
+    return (
+        (10**log10_sigma_sqr) * 24 * np.sqrt(3) * l / (3 + (2 * np.pi * f * l) ** 2) ** 2 * np.repeat(df, components)
+    )
+
+
+def matern32_kernel(tau, l=365 * 86400.0, log10_sigma_sqr=-14):
+    return (10**log10_sigma_sqr) * (1 + np.sqrt(3) * tau / l) * np.exp(-np.sqrt(3) * tau / l)
+
+
 class TestGPSignals(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -355,6 +366,78 @@ def test_fourier_red_noise_backend(self):
         msg = "F matrix shape incorrect"
         assert rnm.get_basis(params).shape == F.shape, msg
 
+    def test_fft_red_noise(self):
+        """Test the FFT implementation of red noise signals"""
+        # set up signal parameter
+        mpsd = psd_matern32(
+            l=parameter.Uniform(365 * 86400.0, 3650 * 86400.0), log10_sigma_sqr=parameter.Uniform(-17, -9)
+        )
+        rn_cb0 = gp_signals.FFTBasisGP(spectrum=mpsd, components=15, oversample=3, cutbins=0)
+        rn_cb1 = gp_signals.FFTBasisGP(spectrum=mpsd, nknots=31, oversample=3, cutoff=3)
+        rnm0 = rn_cb0(self.psr)
+        rnm1 = rn_cb1(self.psr)
+
+        # parameter values
+        l, log10_sigma_sqr = 1.5 * 365 * 86400.0, -14.0
+        params = {"B1855+09_red_noise_l": l, "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr}
+
+        # basis matrix test
+        start_time = np.min(self.psr.toas)
+        Tspan = np.max(self.psr.toas) - start_time
+        B, tc = utils.create_fft_time_basis(self.psr.toas, nknots=31)
+        B1, _ = utils.create_fft_time_basis(self.psr.toas, nknots=31, Tspan=Tspan, start_time=start_time)
+
+        msg = "B matrix incorrect for GP FFT signal."
+        assert np.allclose(B, rnm0.get_basis(params)), msg
+        assert np.allclose(B1, rnm1.get_basis(params)), msg
+        assert np.allclose(np.sum(B, axis=1), np.ones(B.shape[0])), msg
+
+        # spectrum test
+        tau = np.abs(tc[:, None] - tc[None, :])
+        phi_K = matern32_kernel(tau, l, log10_sigma_sqr)
+        phi_E = rnm0.get_phi(params)
+
+        msg = "Prior incorrect for GP FFT signal."
+        assert np.allclose(phi_K, phi_E), msg
+
+        # spectrum test with low-frequency cut-off
+        freqs = utils.knots_to_freqs(tc, oversample=3)
+        psd = psd_matern32(freqs[1:], l=l, log10_sigma_sqr=log10_sigma_sqr, components=1)
+        psd = np.concatenate([[0.0], psd])
+        phi_K = utils.psd2cov(tc, psd)
+        phi_E = rnm1.get_phi(params)
+
+        msg = f"Prior incorrect for GP FFT signal: {phi_K[:3,:3], phi_E[:3,:3]}"
+        assert np.allclose(phi_K, phi_E), msg
+
+    def test_fft_common(self):
+        """Test the FFT implementation of common red noise signals"""
+        # set up signal parameters
+        log10_A, gamma = -14.5, 4.33
+        params = {"B1855+09_red_noise_log10_A": log10_A, "B1855+09_red_noise_gamma": gamma}
+        pl = utils.powerlaw(log10_A=parameter.Uniform(-18, -12), gamma=parameter.Uniform(1, 7))
+        orf = utils.hd_orf()
+
+        # set up the basis and the model
+        start_time = np.min(self.psr.toas)
+        Tspan = np.max(self.psr.toas) - start_time
+        mn = white_signals.MeasurementNoise(efac=parameter.Constant(1.0), selection=Selection(selections.no_selection))
+        crn = gp_signals.FFTBasisCommonGP(
+            pl, orf, nknots=31, name="gw", oversample=3, cutoff=3, Tspan=Tspan, start_time=start_time
+        )
+        model = mn + crn
+        pta = signal_base.PTA([model(psr) for psr in [self.psr, self.psr]])
+
+        # test the prior matrices, including ORF
+        phi_full = pta.get_phi(params)
+        phi_1 = phi_full[:31, :31]
+        phi_12 = phi_full[31:, :31]
+        phi_2 = phi_full[31:, 31:]
+
+        msg = f"Common mode not equal {phi_full.shape}"
+        assert np.allclose(phi_1, phi_2), msg
+        assert np.allclose(0.5 * phi_1, phi_12), msg
+
     def test_red_noise_add(self):
         """Test that red noise addition only returns independent columns."""
         # set up signals

From bb071b424bd24b4b30b08ae923c155f4c01fce02 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 15:40:04 +0000
Subject: [PATCH 09/29] Linting test_gp_signals

---
 tests/test_gp_signals.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 2ba683c5..8eefce0b 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -28,6 +28,7 @@ def create_quant_matrix(toas, dt=1):
     # return value slightly different than 1 to get around ECORR columns
     return U * 1.0000001, avetoas
 
+
 @signal_base.function
 def se_kernel(etoas, log10_sigma=-7, log10_lam=np.log10(30 * 86400)):
     tm = np.abs(etoas[None, :] - etoas[:, None])
@@ -135,7 +136,6 @@ def test_ecorr_backend(self):
         assert ecm.get_basis(params).shape == U.shape, msg
 
     def test_kernel(self):
-
         log10_sigma = parameter.Uniform(-10, -5)
         log10_lam = parameter.Uniform(np.log10(86400), np.log10(1500 * 86400))
         basis = create_quant_matrix(dt=7 * 86400)
@@ -467,7 +467,6 @@ def test_red_noise_add(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
-
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn
@@ -486,9 +485,7 @@ def test_red_noise_add(self):
                 F = F1 if nf1 > nf2 else F2
                 phi[: 2 * nf1] = p1
                 phi[: 2 * nf2] += p2
-                F[
-                    :,
-                ]  # noqa: E231
+                F[:,]  # noqa: E231
             else:
                 phi = np.concatenate((p1, p2))
                 F = np.hstack((F1, F2))
@@ -544,7 +541,6 @@ def test_red_noise_add_backend(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
-
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1, selection=selection)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn

From c2927c547ab1db5ceb1aadcbbddd2e94d9b42891 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:07:16 +0000
Subject: [PATCH 10/29] Removed linting of un-changed code in
 test_gp_signals.py

---
 tests/test_gp_signals.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 8eefce0b..15c6e11c 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -136,6 +136,7 @@ def test_ecorr_backend(self):
         assert ecm.get_basis(params).shape == U.shape, msg
 
     def test_kernel(self):
+
         log10_sigma = parameter.Uniform(-10, -5)
         log10_lam = parameter.Uniform(np.log10(86400), np.log10(1500 * 86400))
         basis = create_quant_matrix(dt=7 * 86400)
@@ -467,6 +468,7 @@ def test_red_noise_add(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
+
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn
@@ -485,7 +487,9 @@ def test_red_noise_add(self):
                 F = F1 if nf1 > nf2 else F2
                 phi[: 2 * nf1] = p1
                 phi[: 2 * nf2] += p2
-                F[:,]  # noqa: E231
+                F[
+                    :,
+                ]  # noqa: E231
             else:
                 phi = np.concatenate((p1, p2))
                 F = np.hstack((F1, F2))
@@ -541,6 +545,7 @@ def test_red_noise_add_backend(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
+
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1, selection=selection)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn

From 7e79acb71a922b775624572fb5078e6555e06a31 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:08:16 +0000
Subject: [PATCH 11/29] Removed fft tests to see whether linting goes now

---
 tests/test_gp_signals.py | 72 ----------------------------------------
 1 file changed, 72 deletions(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 15c6e11c..64a01635 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -367,78 +367,6 @@ def test_fourier_red_noise_backend(self):
         msg = "F matrix shape incorrect"
         assert rnm.get_basis(params).shape == F.shape, msg
 
-    def test_fft_red_noise(self):
-        """Test the FFT implementation of red noise signals"""
-        # set up signal parameter
-        mpsd = psd_matern32(
-            l=parameter.Uniform(365 * 86400.0, 3650 * 86400.0), log10_sigma_sqr=parameter.Uniform(-17, -9)
-        )
-        rn_cb0 = gp_signals.FFTBasisGP(spectrum=mpsd, components=15, oversample=3, cutbins=0)
-        rn_cb1 = gp_signals.FFTBasisGP(spectrum=mpsd, nknots=31, oversample=3, cutoff=3)
-        rnm0 = rn_cb0(self.psr)
-        rnm1 = rn_cb1(self.psr)
-
-        # parameter values
-        l, log10_sigma_sqr = 1.5 * 365 * 86400.0, -14.0
-        params = {"B1855+09_red_noise_l": l, "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr}
-
-        # basis matrix test
-        start_time = np.min(self.psr.toas)
-        Tspan = np.max(self.psr.toas) - start_time
-        B, tc = utils.create_fft_time_basis(self.psr.toas, nknots=31)
-        B1, _ = utils.create_fft_time_basis(self.psr.toas, nknots=31, Tspan=Tspan, start_time=start_time)
-
-        msg = "B matrix incorrect for GP FFT signal."
-        assert np.allclose(B, rnm0.get_basis(params)), msg
-        assert np.allclose(B1, rnm1.get_basis(params)), msg
-        assert np.allclose(np.sum(B, axis=1), np.ones(B.shape[0])), msg
-
-        # spectrum test
-        tau = np.abs(tc[:, None] - tc[None, :])
-        phi_K = matern32_kernel(tau, l, log10_sigma_sqr)
-        phi_E = rnm0.get_phi(params)
-
-        msg = "Prior incorrect for GP FFT signal."
-        assert np.allclose(phi_K, phi_E), msg
-
-        # spectrum test with low-frequency cut-off
-        freqs = utils.knots_to_freqs(tc, oversample=3)
-        psd = psd_matern32(freqs[1:], l=l, log10_sigma_sqr=log10_sigma_sqr, components=1)
-        psd = np.concatenate([[0.0], psd])
-        phi_K = utils.psd2cov(tc, psd)
-        phi_E = rnm1.get_phi(params)
-
-        msg = f"Prior incorrect for GP FFT signal: {phi_K[:3,:3], phi_E[:3,:3]}"
-        assert np.allclose(phi_K, phi_E), msg
-
-    def test_fft_common(self):
-        """Test the FFT implementation of common red noise signals"""
-        # set up signal parameters
-        log10_A, gamma = -14.5, 4.33
-        params = {"B1855+09_red_noise_log10_A": log10_A, "B1855+09_red_noise_gamma": gamma}
-        pl = utils.powerlaw(log10_A=parameter.Uniform(-18, -12), gamma=parameter.Uniform(1, 7))
-        orf = utils.hd_orf()
-
-        # set up the basis and the model
-        start_time = np.min(self.psr.toas)
-        Tspan = np.max(self.psr.toas) - start_time
-        mn = white_signals.MeasurementNoise(efac=parameter.Constant(1.0), selection=Selection(selections.no_selection))
-        crn = gp_signals.FFTBasisCommonGP(
-            pl, orf, nknots=31, name="gw", oversample=3, cutoff=3, Tspan=Tspan, start_time=start_time
-        )
-        model = mn + crn
-        pta = signal_base.PTA([model(psr) for psr in [self.psr, self.psr]])
-
-        # test the prior matrices, including ORF
-        phi_full = pta.get_phi(params)
-        phi_1 = phi_full[:31, :31]
-        phi_12 = phi_full[31:, :31]
-        phi_2 = phi_full[31:, 31:]
-
-        msg = f"Common mode not equal {phi_full.shape}"
-        assert np.allclose(phi_1, phi_2), msg
-        assert np.allclose(0.5 * phi_1, phi_12), msg
-
     def test_red_noise_add(self):
         """Test that red noise addition only returns independent columns."""
         # set up signals

From 5074199752955969d6253841d47f5c05b2ea2d86 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:13:20 +0000
Subject: [PATCH 12/29] Added FFT red noise test again

---
 tests/test_gp_signals.py | 51 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 4 deletions(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 64a01635..b04d2c0a 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -37,15 +37,15 @@ def se_kernel(etoas, log10_sigma=-7, log10_lam=np.log10(30 * 86400)):
 
 
 @signal_base.function
-def psd_matern32(f, l=365 * 86400.0, log10_sigma_sqr=-14, components=2):
+def psd_matern32(f, length_scale=365 * 86400.0, log10_sigma_sqr=-14, components=2):
     df = np.diff(np.concatenate((np.array([0]), f[::components])))
     return (
-        (10**log10_sigma_sqr) * 24 * np.sqrt(3) * l / (3 + (2 * np.pi * f * l) ** 2) ** 2 * np.repeat(df, components)
+        (10**log10_sigma_sqr) * 24 * np.sqrt(3) * length_scale / (3 + (2 * np.pi * f * length_scale) ** 2) ** 2 * np.repeat(df, components)
     )
 
 
-def matern32_kernel(tau, l=365 * 86400.0, log10_sigma_sqr=-14):
-    return (10**log10_sigma_sqr) * (1 + np.sqrt(3) * tau / l) * np.exp(-np.sqrt(3) * tau / l)
+def matern32_kernel(tau, length_scale=365 * 86400.0, log10_sigma_sqr=-14):
+    return (10**log10_sigma_sqr) * (1 + np.sqrt(3) * tau / length_scale) * np.exp(-np.sqrt(3) * tau / length_scale)
 
 
 class TestGPSignals(unittest.TestCase):
@@ -367,6 +367,49 @@ def test_fourier_red_noise_backend(self):
         msg = "F matrix shape incorrect"
         assert rnm.get_basis(params).shape == F.shape, msg
 
+    def test_fft_red_noise(self):
+        """Test the FFT implementation of red noise signals"""
+        # set up signal parameter
+        mpsd = psd_matern32(
+            length_scale=parameter.Uniform(365 * 86400.0, 3650 * 86400.0), log10_sigma_sqr=parameter.Uniform(-17, -9)
+        )
+        rn_cb0 = gp_signals.FFTBasisGP(spectrum=mpsd, components=15, oversample=3, cutbins=0)
+        rn_cb1 = gp_signals.FFTBasisGP(spectrum=mpsd, nknots=31, oversample=3, cutoff=3)
+        rnm0 = rn_cb0(self.psr)
+        rnm1 = rn_cb1(self.psr)
+
+        # parameter values
+        length_scale, log10_sigma_sqr = 1.5 * 365 * 86400.0, -14.0
+        params = {"B1855+09_red_noise_length_scale": length_scale, "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr}
+
+        # basis matrix test
+        start_time = np.min(self.psr.toas)
+        Tspan = np.max(self.psr.toas) - start_time
+        B, tc = utils.create_fft_time_basis(self.psr.toas, nknots=31)
+        B1, _ = utils.create_fft_time_basis(self.psr.toas, nknots=31, Tspan=Tspan, start_time=start_time)
+
+        msg = "B matrix incorrect for GP FFT signal."
+        assert np.allclose(B, rnm0.get_basis(params)), msg
+        assert np.allclose(B1, rnm1.get_basis(params)), msg
+        assert np.allclose(np.sum(B, axis=1), np.ones(B.shape[0])), msg
+
+        # spectrum test
+        tau = np.abs(tc[:, None] - tc[None, :])
+        phi_K = matern32_kernel(tau, length_scale, log10_sigma_sqr)
+        phi_E = rnm0.get_phi(params)
+
+        msg = "Prior incorrect for GP FFT signal."
+        assert np.allclose(phi_K, phi_E), msg
+
+        # spectrum test with low-frequency cut-off
+        freqs = utils.knots_to_freqs(tc, oversample=3)
+        psd = psd_matern32(freqs[1:], length_scale=length_scale, log10_sigma_sqr=log10_sigma_sqr, components=1)
+        psd = np.concatenate([[0.0], psd])
+        phi_K = utils.psd2cov(tc, psd)
+        phi_E = rnm1.get_phi(params)
+
+        msg = f"Prior incorrect for GP FFT signal: {phi_K[:3,:3], phi_E[:3,:3]}"
+        assert np.allclose(phi_K, phi_E), msg
     def test_red_noise_add(self):
         """Test that red noise addition only returns independent columns."""
         # set up signals

From 094c67f8c3a54e05d5a1b3460b2f668455f91d38 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:17:25 +0000
Subject: [PATCH 13/29] Switched FFT test function, see if this one lints

---
 tests/test_gp_signals.py | 61 +++++++++++++++-------------------------
 1 file changed, 23 insertions(+), 38 deletions(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index b04d2c0a..7df6c199 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -367,49 +367,34 @@ def test_fourier_red_noise_backend(self):
         msg = "F matrix shape incorrect"
         assert rnm.get_basis(params).shape == F.shape, msg
 
-    def test_fft_red_noise(self):
-        """Test the FFT implementation of red noise signals"""
-        # set up signal parameter
-        mpsd = psd_matern32(
-            length_scale=parameter.Uniform(365 * 86400.0, 3650 * 86400.0), log10_sigma_sqr=parameter.Uniform(-17, -9)
-        )
-        rn_cb0 = gp_signals.FFTBasisGP(spectrum=mpsd, components=15, oversample=3, cutbins=0)
-        rn_cb1 = gp_signals.FFTBasisGP(spectrum=mpsd, nknots=31, oversample=3, cutoff=3)
-        rnm0 = rn_cb0(self.psr)
-        rnm1 = rn_cb1(self.psr)
-
-        # parameter values
-        length_scale, log10_sigma_sqr = 1.5 * 365 * 86400.0, -14.0
-        params = {"B1855+09_red_noise_length_scale": length_scale, "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr}
+    def test_fft_common(self):
+        """Test the FFT implementation of common red noise signals"""
+        # set up signal parameters
+        log10_A, gamma = -14.5, 4.33
+        params = {"B1855+09_red_noise_log10_A": log10_A, "B1855+09_red_noise_gamma": gamma}
+        pl = utils.powerlaw(log10_A=parameter.Uniform(-18, -12), gamma=parameter.Uniform(1, 7))
+        orf = utils.hd_orf()
 
-        # basis matrix test
+        # set up the basis and the model
         start_time = np.min(self.psr.toas)
         Tspan = np.max(self.psr.toas) - start_time
-        B, tc = utils.create_fft_time_basis(self.psr.toas, nknots=31)
-        B1, _ = utils.create_fft_time_basis(self.psr.toas, nknots=31, Tspan=Tspan, start_time=start_time)
+        mn = white_signals.MeasurementNoise(efac=parameter.Constant(1.0), selection=Selection(selections.no_selection))
+        crn = gp_signals.FFTBasisCommonGP(
+            pl, orf, nknots=31, name="gw", oversample=3, cutoff=3, Tspan=Tspan, start_time=start_time
+        )
+        model = mn + crn
+        pta = signal_base.PTA([model(psr) for psr in [self.psr, self.psr]])
 
-        msg = "B matrix incorrect for GP FFT signal."
-        assert np.allclose(B, rnm0.get_basis(params)), msg
-        assert np.allclose(B1, rnm1.get_basis(params)), msg
-        assert np.allclose(np.sum(B, axis=1), np.ones(B.shape[0])), msg
+        # test the prior matrices, including ORF
+        phi_full = pta.get_phi(params)
+        phi_1 = phi_full[:31, :31]
+        phi_12 = phi_full[31:, :31]
+        phi_2 = phi_full[31:, 31:]
+
+        msg = f"Common mode not equal {phi_full.shape}"
+        assert np.allclose(phi_1, phi_2), msg
+        assert np.allclose(0.5 * phi_1, phi_12), msg
 
-        # spectrum test
-        tau = np.abs(tc[:, None] - tc[None, :])
-        phi_K = matern32_kernel(tau, length_scale, log10_sigma_sqr)
-        phi_E = rnm0.get_phi(params)
-
-        msg = "Prior incorrect for GP FFT signal."
-        assert np.allclose(phi_K, phi_E), msg
-
-        # spectrum test with low-frequency cut-off
-        freqs = utils.knots_to_freqs(tc, oversample=3)
-        psd = psd_matern32(freqs[1:], length_scale=length_scale, log10_sigma_sqr=log10_sigma_sqr, components=1)
-        psd = np.concatenate([[0.0], psd])
-        phi_K = utils.psd2cov(tc, psd)
-        phi_E = rnm1.get_phi(params)
-
-        msg = f"Prior incorrect for GP FFT signal: {phi_K[:3,:3], phi_E[:3,:3]}"
-        assert np.allclose(phi_K, phi_E), msg
     def test_red_noise_add(self):
         """Test that red noise addition only returns independent columns."""
         # set up signals

From 9e1d1a3bf37c3d795d18a0ed4442acc6e76f967c Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:25:32 +0000
Subject: [PATCH 14/29] Added FFT red noise and common tests

---
 tests/test_gp_signals.py | 51 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 7df6c199..63a607ae 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -40,7 +40,12 @@ def se_kernel(etoas, log10_sigma=-7, log10_lam=np.log10(30 * 86400)):
 def psd_matern32(f, length_scale=365 * 86400.0, log10_sigma_sqr=-14, components=2):
     df = np.diff(np.concatenate((np.array([0]), f[::components])))
     return (
-        (10**log10_sigma_sqr) * 24 * np.sqrt(3) * length_scale / (3 + (2 * np.pi * f * length_scale) ** 2) ** 2 * np.repeat(df, components)
+        (10**log10_sigma_sqr)
+        * 24
+        * np.sqrt(3)
+        * length_scale
+        / (3 + (2 * np.pi * f * length_scale) ** 2) ** 2
+        * np.repeat(df, components)
     )
 
 
@@ -367,6 +372,50 @@ def test_fourier_red_noise_backend(self):
         msg = "F matrix shape incorrect"
         assert rnm.get_basis(params).shape == F.shape, msg
 
+    def test_fft_red_noise(self):
+        """Test the FFT implementation of red noise signals"""
+        # set up signal parameter
+        mpsd = psd_matern32(
+            length_scale=parameter.Uniform(365 * 86400.0, 3650 * 86400.0), log10_sigma_sqr=parameter.Uniform(-17, -9)
+        )
+        rn_cb0 = gp_signals.FFTBasisGP(spectrum=mpsd, components=15, oversample=3, cutbins=0)
+        rn_cb1 = gp_signals.FFTBasisGP(spectrum=mpsd, nknots=31, oversample=3, cutoff=3)
+        rnm0 = rn_cb0(self.psr)
+        rnm1 = rn_cb1(self.psr)
+
+        # parameter values
+        length_scale, log10_sigma_sqr = 1.5 * 365 * 86400.0, -14.0
+        params = {"B1855+09_red_noise_length_scale": length_scale, "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr}
+
+        # basis matrix test
+        start_time = np.min(self.psr.toas)
+        Tspan = np.max(self.psr.toas) - start_time
+        B, tc = utils.create_fft_time_basis(self.psr.toas, nknots=31)
+        B1, _ = utils.create_fft_time_basis(self.psr.toas, nknots=31, Tspan=Tspan, start_time=start_time)
+
+        msg = "B matrix incorrect for GP FFT signal."
+        assert np.allclose(B, rnm0.get_basis(params)), msg
+        assert np.allclose(B1, rnm1.get_basis(params)), msg
+        assert np.allclose(np.sum(B, axis=1), np.ones(B.shape[0])), msg
+
+        # spectrum test
+        tau = np.abs(tc[:, None] - tc[None, :])
+        phi_K = matern32_kernel(tau, length_scale, log10_sigma_sqr)
+        phi_E = rnm0.get_phi(params)
+
+        msg = "Prior incorrect for GP FFT signal."
+        assert np.allclose(phi_K, phi_E), msg
+
+        # spectrum test with low-frequency cut-off
+        freqs = utils.knots_to_freqs(tc, oversample=3)
+        psd = psd_matern32(freqs[1:], length_scale=length_scale, log10_sigma_sqr=log10_sigma_sqr, components=1)
+        psd = np.concatenate([[0.0], psd])
+        phi_K = utils.psd2cov(tc, psd)
+        phi_E = rnm1.get_phi(params)
+
+        msg = f"Prior incorrect for GP FFT signal: {phi_K[:3,:3], phi_E[:3,:3]}"
+        assert np.allclose(phi_K, phi_E), msg
+
     def test_fft_common(self):
         """Test the FFT implementation of common red noise signals"""
         # set up signal parameters

From 0649d1f086b13e00a929582abefe1f1735c1cead Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:31:06 +0000
Subject: [PATCH 15/29] Used different black config

---
 tests/test_gp_signals.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 63a607ae..61e5ba35 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -141,7 +141,6 @@ def test_ecorr_backend(self):
         assert ecm.get_basis(params).shape == U.shape, msg
 
     def test_kernel(self):
-
         log10_sigma = parameter.Uniform(-10, -5)
         log10_lam = parameter.Uniform(np.log10(86400), np.log10(1500 * 86400))
         basis = create_quant_matrix(dt=7 * 86400)
@@ -385,7 +384,10 @@ def test_fft_red_noise(self):
 
         # parameter values
         length_scale, log10_sigma_sqr = 1.5 * 365 * 86400.0, -14.0
-        params = {"B1855+09_red_noise_length_scale": length_scale, "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr}
+        params = {
+            "B1855+09_red_noise_length_scale": length_scale,
+            "B1855+09_red_noise_log10_sigma_sqr": log10_sigma_sqr,
+        }
 
         # basis matrix test
         start_time = np.min(self.psr.toas)
@@ -473,7 +475,6 @@ def test_red_noise_add(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
-
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn
@@ -492,9 +493,7 @@ def test_red_noise_add(self):
                 F = F1 if nf1 > nf2 else F2
                 phi[: 2 * nf1] = p1
                 phi[: 2 * nf2] += p2
-                F[
-                    :,
-                ]  # noqa: E231
+                F[:,]  # noqa: E231
             else:
                 phi = np.concatenate((p1, p2))
                 F = np.hstack((F1, F2))
@@ -550,7 +549,6 @@ def test_red_noise_add_backend(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
-
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1, selection=selection)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn

From 0afe73b655876543e745b1f5a39eeb0d9dad1f65 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:37:38 +0000
Subject: [PATCH 16/29] Yet again, linting

---
 tests/test_gp_signals.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index 61e5ba35..e82fc77d 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -141,6 +141,7 @@ def test_ecorr_backend(self):
         assert ecm.get_basis(params).shape == U.shape, msg
 
     def test_kernel(self):
+
         log10_sigma = parameter.Uniform(-10, -5)
         log10_lam = parameter.Uniform(np.log10(86400), np.log10(1500 * 86400))
         basis = create_quant_matrix(dt=7 * 86400)
@@ -475,6 +476,7 @@ def test_red_noise_add(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
+
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn
@@ -493,7 +495,9 @@ def test_red_noise_add(self):
                 F = F1 if nf1 > nf2 else F2
                 phi[: 2 * nf1] = p1
                 phi[: 2 * nf2] += p2
-                F[:,]  # noqa: E231
+                F[
+                    :,
+                ]  # noqa: E231
             else:
                 phi = np.concatenate((p1, p2))
                 F = np.hstack((F1, F2))
@@ -549,6 +553,7 @@ def test_red_noise_add_backend(self):
         ]
 
         for nf1, nf2, T1, T2 in tpars:
+
             rn = gp_signals.FourierBasisGP(spectrum=pl, components=nf1, Tspan=T1, selection=selection)
             crn = gp_signals.FourierBasisGP(spectrum=cpl, components=nf2, Tspan=T2)
             s = rn + crn

From 5d435aba36fdb41b700ca7ab1ba92c8c058d7c8f Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 17:42:31 +0000
Subject: [PATCH 17/29] Removed debugging test statements in test_gp_signals.py

---
 tests/test_gp_signals.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_gp_signals.py b/tests/test_gp_signals.py
index e82fc77d..43956ae3 100644
--- a/tests/test_gp_signals.py
+++ b/tests/test_gp_signals.py
@@ -416,7 +416,7 @@ def test_fft_red_noise(self):
         phi_K = utils.psd2cov(tc, psd)
         phi_E = rnm1.get_phi(params)
 
-        msg = f"Prior incorrect for GP FFT signal: {phi_K[:3,:3], phi_E[:3,:3]}"
+        msg = f"Prior incorrect for GP FFT signal."
         assert np.allclose(phi_K, phi_E), msg
 
     def test_fft_common(self):
@@ -443,7 +443,7 @@ def test_fft_common(self):
         phi_12 = phi_full[31:, :31]
         phi_2 = phi_full[31:, 31:]
 
-        msg = f"Common mode not equal {phi_full.shape}"
+        msg = f"Common mode FFT Prior not equal between pulsars."
         assert np.allclose(phi_1, phi_2), msg
         assert np.allclose(0.5 * phi_1, phi_12), msg
 

From 085f12855d13a04bb1fc3560970320a0719b35ca Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 21:30:33 +0000
Subject: [PATCH 18/29] Added DMV and Chromatic noise FFT interpolation GP
 bases

---
 enterprise/signals/gp_bases.py | 57 ++++++++++++++++++++++++++++++++++
 enterprise/signals/utils.py    |  2 ++
 tests/test_utils.py            | 24 ++++++++++++++
 3 files changed, 83 insertions(+)

diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py
index 02cc20b1..a14ca9f0 100644
--- a/enterprise/signals/gp_bases.py
+++ b/enterprise/signals/gp_bases.py
@@ -14,11 +14,13 @@
     "createfourierdesignmatrix_red",
     "create_fft_time_basis",
     "createfourierdesignmatrix_dm",
+    "create_fft_time_basis_dm",
     "createfourierdesignmatrix_dm_tn",
     "createfourierdesignmatrix_env",
     "createfourierdesignmatrix_ephem",
     "createfourierdesignmatrix_eph",
     "createfourierdesignmatrix_chromatic",
+    "create_fft_time_basis_chromatic",
     "createfourierdesignmatrix_general",
 ]
 
@@ -168,6 +170,33 @@ def createfourierdesignmatrix_dm(
     return F * Dm[:, None], Ffreqs
 
 
+@function
+def create_fft_time_basis_dm(toas, freqs, nknots=30, Tspan=None, start_time=None, fref=1400):
+    """
+    Construct DM-variation linear interpolation design matrix. Current
+    normalization expresses DM signal as a deviation [seconds]
+    at fref [MHz]
+
+    :param toas: vector of time series in seconds
+    :param freqs: radio frequencies of observations [MHz]
+    :param nknots: number of coarse time samples to use (knots)
+    :param Tspan: option to some other Tspan
+    :param start_time: option to set some other start epoch of basis
+    :param fref: reference frequency [MHz]
+
+    :return B: coarse time-domain design matrix
+    :return t_coarse: timestamps of coarse time grid
+    """
+
+    # get base course time-domain matrix and times
+    Bmat, t_coarse = create_fft_time_basis(toas, nknots=nknots, Tspan=Tspan, start_time=start_time)
+
+    # compute the DM-variation vectors
+    Dm = (fref / freqs) ** 2
+
+    return Bmat * Dm[:, None], t_coarse
+
+
 @function
 def createfourierdesignmatrix_dm_tn(
     toas, freqs, nmodes=30, Tspan=None, pshift=False, fref=1400, logf=False, fmin=None, fmax=None, idx=2, modes=None
@@ -333,6 +362,34 @@ def createfourierdesignmatrix_chromatic(
     return F * Dm[:, None], Ffreqs
 
 
+@function
+def create_fft_time_basis_chromatic(toas, freqs, nknots=30, Tspan=None, start_time=None, fref=1400, idx=4):
+    """
+    Construct scattering linear interpolation design matrix. Current
+    normalization expresses DM signal as a deviation [seconds]
+    at fref [MHz]
+
+    :param toas: vector of time series in seconds
+    :param freqs: radio frequencies of observations [MHz]
+    :param nknots: number of coarse time samples to use (knots)
+    :param Tspan: option to some other Tspan
+    :param start_time: option to set some other start epoch of basis
+    :param fref: reference frequency [MHz]
+    :param idx: Index of chromatic effects
+
+    :return B: coarse time-domain design matrix
+    :return t_coarse: timestamps of coarse time grid
+    """
+
+    # get base course time-domain matrix and times
+    Bmat, t_coarse = create_fft_time_basis(toas, nknots=nknots, Tspan=Tspan, start_time=start_time)
+
+    # compute the DM-variation vectors
+    Dm = (fref / freqs) ** idx
+
+    return Bmat * Dm[:, None], t_coarse
+
+
 @function
 def createfourierdesignmatrix_general(
     toas,
diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index ebfe855a..e2308e52 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -22,11 +22,13 @@
     createfourierdesignmatrix_red,
     create_fft_time_basis,
     createfourierdesignmatrix_dm,
+    create_fft_time_basis_dm,
     createfourierdesignmatrix_dm_tn,
     createfourierdesignmatrix_env,
     createfourierdesignmatrix_ephem,
     createfourierdesignmatrix_eph,
     createfourierdesignmatrix_chromatic,
+    create_fft_time_basis_chromatic,
     createfourierdesignmatrix_general,
 )
 from enterprise.signals.gp_priors import powerlaw, turnover  # noqa: F401
diff --git a/tests/test_utils.py b/tests/test_utils.py
index f85f1a17..9b3ec31f 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -35,6 +35,12 @@ def setUpClass(cls):
 
         cls.Fdm, _ = utils.createfourierdesignmatrix_dm(cls.psr.toas, freqs=cls.psr.freqs, nmodes=30)
 
+        cls.B, _ = utils.create_fft_time_basis(cls.psr.toas, nknots=30)
+
+        cls.Bdm, _ = utils.create_fft_time_basis_dm(cls.psr.toas, freqs=cls.psr.freqs, nknots=30)
+
+        cls.Bchr, _ = utils.create_fft_time_basis_chromatic(cls.psr.toas, freqs=cls.psr.freqs, nknots=30)
+
         cls.Feph, cls.feph = utils.createfourierdesignmatrix_ephem(cls.psr.toas, cls.psr.pos, nmodes=30)
 
         cls.Mm = utils.create_stabletimingdesignmatrix(cls.psr.Mmat)
@@ -51,12 +57,30 @@ def test_createfourierdesignmatrix_red(self, nf=30):
         msg = "Fourier design matrix shape incorrect"
         assert self.F.shape == (4005, 2 * nf), msg
 
+    def test_create_fft_time_basis(self, nk=30):
+        """Check FFT interpolation design matrix shape."""
+
+        msg = "FFT interpolation design matrix shape incorrect"
+        assert self.B.shape == (4005, nk), msg
+
     def test_createfourierdesignmatrix_dm(self, nf=30):
         """Check DM-variation Fourier design matrix shape."""
 
         msg = "DM-variation Fourier design matrix shape incorrect"
         assert self.Fdm.shape == (4005, 2 * nf), msg
 
+    def test_create_fft_time_basis_dm(self, nk=30):
+        """Check FFT interpolation design matrix shape."""
+
+        msg = "DM-variation FFT interpolation design matrix shape incorrect"
+        assert self.Bdm.shape == (4005, nk), msg
+
+    def test_create_fft_time_basis_chromatic(self, nk=30):
+        """Check FFT interpolation design matrix shape."""
+
+        msg = "DM-variation FFT interpolation design matrix shape incorrect"
+        assert self.Bchr.shape == (4005, nk), msg
+
     def test_createfourierdesignmatrix_ephem(self, nf=30):
         """Check x-axis ephemeris Fourier design matrix shape."""
 

From 4d8435679e5c53cd99cdc01d94e4805a75802c1f Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 27 Feb 2025 21:37:28 +0000
Subject: [PATCH 19/29] Added option 'basis' to FFTBasisGP for DM-variations
 and chromatic noise

---
 enterprise/signals/gp_signals.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 15830f21..698e9acb 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -218,6 +218,7 @@ class FourierBasisGP(BaseClass):
 
 def FFTBasisGP(
     spectrum,
+    basis=None,
     coefficients=False,
     combine=True,
     components=20,
@@ -244,7 +245,9 @@ def FFTBasisGP(
         #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
-    basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time)
+    if basis is None:
+        basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time)
+
     BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
 
     class FFTBasisGP(BaseClass):

From c0b4b15dfde3c359031106599c5af962529b0d98 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Mon, 3 Mar 2025 10:30:26 +0100
Subject: [PATCH 20/29] Added FFT documentation

---
 docs/fft.rst   | 103 +++++++++++++++++++++++++++++++++++++++++++++++++
 docs/index.rst |   1 +
 2 files changed, 104 insertions(+)
 create mode 100644 docs/fft.rst

diff --git a/docs/fft.rst b/docs/fft.rst
new file mode 100644
index 00000000..a257f7f6
--- /dev/null
+++ b/docs/fft.rst
@@ -0,0 +1,103 @@
+
+.. module:: enterprise
+   :noindex:
+
+.. note:: This tutorial was generated from a Jupyter notebook that can be
+          downloaded `here <_static/notebooks/mdc.ipynb>`_.
+
+.. _mdc:
+
+Red noise modeling
+=======================
+
+In the beginning of Enterprise red noise modeling, the red noise prior was
+always modeled as a diagonal matrix, meaning that the Fourier coefficients
+were assumed to be uncorrelated. This model was introduced by Lentati et al.
+(2013), and explained by van Haasteren and Vallisneri (2014). In practice this
+has been a good-enough approximation, but it is not exact.
+
+As of early 2025 we now have a more realistic implementation of red noise
+priors that include correlations between the basis functions. The `FFT`
+method as it is called is a rank-reduced time-domain implementation, meaning
+it does not rely on Fourier modes, but on regularly sampled coarse grained
+time samples. Here we briefly explain how to use it.
+
+
+
+Red noise modeling
+-------------------
+
+The traditional old-style way of modeling was done like:
+
+.. code:: python
+
+    rn_pl = powerlaw(log10_A=rn_log10_A, gamma=rn_gamma)
+    rn_phi = gp_signals.FourierBasisGP(spectrum=rn_pl, components=n_components, Tspan=Tspan)
+
+For the FFT time-domain model, one would do:
+
+.. code:: python
+
+    rn_pl = powerlaw(log10_A=rn_log10_A, gamma=rn_gamma)
+    rn_fft = gp_signals.FFTBasisGP(spectrum=rn_pl, components=n_components, oversample=3, cutoff=3)
+
+The same spectral function can be used. Free spectrum is NOT supported yet.
+Instead of `components`, we can also pass `knots=`, where it is understood that
+`knots=2*n_components+1`. This is because `components` actually means
+frequencies.  In the time-domain, the number of `knots` is the number of
+`modes+1`, because we cannot just omit the DC term.
+
+The `oversample` parameter determines how densely the PSD is sampled in
+frequencies. With `oversample=1` we would use frequencies at spacing of
+`df=1/T`.  With `oversample=3` (the default), the frequency spacing is
+`df=1/(3T)`. Note that this is a way to numerically approximate the
+Wiener-Khinchin integral. With oversample sufficiently large, the FFT is an
+excellent approximation of the analytical integral. For powerlaw signals,
+`oversample=3` seems a very reasonable choice.
+
+The `cutoff` parameter is used to specify below which frequency `fcut = 1 /
+(cutoff*Tspan)` we set the PSD equal to zero. Note that this parameterization
+(which is also in Discovery) is a bit ambiguous, as fcut may not correspond to
+an actual bin of the FFT: especially if oversample is not a high number this
+can cause a mismatch. In case of a mismatch, `fcut` is rounded up to the next
+oversampled-FFT frequency bin. Instead of `cutoff`, the parameter `cutbins` can
+also be used (this overrides cutoff). With cutbins the low-frequency cutoff is
+set at: `fcut = cutbins / (oversample * Tspan)`, and its interpretation is less
+ambiguous: it is the number of bins of the over-sampled PSD of the FFT that is
+being zeroed out.
+
+Common signals
+--------------
+
+For common signals, instead of:
+
+.. code:: python
+
+    gw_pl = powerlaw(log10_A=gw_log10_A, gamma=gw_gamma)
+    orf = utils.hd_orf()
+    crn_phi = gp_signals.FourierBasisCommonGP(gw_pl, orf, components=20, name='gw', Tspan=Tspan)
+
+
+one would do:
+
+.. code:: python
+
+    gw_pl = powerlaw(log10_A=gw_log10_A, gamma=gw_gamma)
+    orf = utils.hd_orf()
+    crn_fft = gp_signals.FFTBasisCommonGP(gw_pl, orf, components=20, name='gw', Tspan=Tspan, start_time=start_time)
+
+Chromatic signals
+-----------------
+
+DM-variations and Chromatic noise can be similarly set up:
+
+.. code:: python
+
+    nknots = 81
+    dm_basis = utils.create_fft_time_basis_dm(nknots=nknots)
+    dm_pl = powerlaw(log10_A=dm_log10_A, gamma=dm_gamma)
+    dm_fft = gp_signals.FFTBasisGP(dm_pl, basis=dm_basis, nknots=nknots, name='dmgp')
+
+    chrom_basis = utils.create_fft_time_basis_chromatic(nknots=nknots, idx=chrom_idx)
+    chrom_pl = powerlaw(log10_A=chrom_log10_A, gamma=chrom_gamma)
+    chrom_fft = gp_signals.FFTBasisGP(chrom_pl, basis=chrom_basis, nknots=nknots, name='chromgp')
diff --git a/docs/index.rst b/docs/index.rst
index f7d95823..7ff78a17 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -42,6 +42,7 @@ searches, and timing model analysis.
 
    mdc
    nano9
+   fft
 
 .. toctree::
    :maxdepth: 2

From 78e9797b74d22b6ccf12047f8065d74185f3e36f Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Mon, 3 Mar 2025 17:05:33 +0100
Subject: [PATCH 21/29] Updated documentation on FFT models

---
 docs/fft.rst | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/docs/fft.rst b/docs/fft.rst
index a257f7f6..5df11f87 100644
--- a/docs/fft.rst
+++ b/docs/fft.rst
@@ -101,3 +101,20 @@ DM-variations and Chromatic noise can be similarly set up:
     chrom_basis = utils.create_fft_time_basis_chromatic(nknots=nknots, idx=chrom_idx)
     chrom_pl = powerlaw(log10_A=chrom_log10_A, gamma=chrom_gamma)
     chrom_fft = gp_signals.FFTBasisGP(chrom_pl, basis=chrom_basis, nknots=nknots, name='chromgp')
+
+Subtleties
+----------
+
+Enterprise allows one to combine basis functions when they are the same. This
+is especially useful when analyzing common signals which have the same basis as
+a single-pulsar signal, such as one would have with red noise and a correlated
+GWB. This can be done with the `combine=True` option in `FFTBasisGP` and
+`FFTBasisCommonGP`. Default is `combine=True`. The subtlety is that modern PTA
+datasets typically have large gaps, which causes some of the time-domain basis
+functions to basically be all zeros. Therefore, some basis functions that you
+would not expect to be identical will be combined.
+
+The above is not a bug. Combining such bases and the corresponding Phi matrix
+does not matter, because the basis is zero, and that part of the signal has no
+bearing on the data or the model. However, when doing signal reconstruction,
+such as with `la_forge`, make sure to set `combine=False`.

From 245355e034c102185dd1d553e93df7bd563a2e89 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Mon, 3 Mar 2025 17:37:46 +0100
Subject: [PATCH 22/29] Updated FFT red signals docs to mention
 (...Tspan=Tspan, start_time=start_time)

---
 docs/fft.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/fft.rst b/docs/fft.rst
index 5df11f87..b2c880fa 100644
--- a/docs/fft.rst
+++ b/docs/fft.rst
@@ -39,7 +39,7 @@ For the FFT time-domain model, one would do:
 .. code:: python
 
     rn_pl = powerlaw(log10_A=rn_log10_A, gamma=rn_gamma)
-    rn_fft = gp_signals.FFTBasisGP(spectrum=rn_pl, components=n_components, oversample=3, cutoff=3)
+    rn_fft = gp_signals.FFTBasisGP(spectrum=rn_pl, components=n_components, oversample=3, cutoff=3, Tspan=Tspan, start_time=start_time)
 
 The same spectral function can be used. Free spectrum is NOT supported yet.
 Instead of `components`, we can also pass `knots=`, where it is understood that

From 7aaacbbde08ab019b52cda72bd9cf58b3bc0f0dc Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 12 Mar 2025 17:53:43 +0100
Subject: [PATCH 23/29] Updated docstring of knots_to_freqs

---
 enterprise/signals/utils.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index e2308e52..9d628ea6 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -853,6 +853,7 @@ def psd2cov(
 
     :param t_knots: Timestamps of the coarse time grid
     :param psd: values of the PSD at frequencies freqs (assumes *delta_f in psd)
+                so psd is assumed to be in units of [s^2]
 
     :return covmat: Covariance matrix at coarse time grid
     """
@@ -882,13 +883,14 @@ def knots_to_freqs(t_knots, oversample=3):
     :param t_knots: Timestamps of the coarse time grid
     :param oversample: amount by which to over-sample the frequency grid
 
-    :return covmat: Covariance matrix at coarse time grid
+    :return freqs: Frequencies, regularly sampled with
+                   delta-f = 1/(oversample*T), fmax=1/(2*delta_t_knots)
     """
     nmodes = len(t_knots)
     Tspan = np.max(t_knots) - np.min(t_knots)
 
     if nmodes % 2 == 0:
-        raise ValueError("psd2cov number of nmodes must be odd.")
+        raise ValueError("len(t_knots) must be odd.")
 
     n_freqs = int((nmodes - 1) / 2 * oversample + 1)
     fmax = (nmodes - 1) / Tspan / 2

From 44cbff7e84c1a5234298f3b00e8da7e032c5b364 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 12 Mar 2025 19:46:41 +0100
Subject: [PATCH 24/29] Updated docstring for FFTBasisGP

---
 enterprise/signals/gp_signals.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 698e9acb..cd612da7 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -231,8 +231,7 @@ def FFTBasisGP(
     start_time=None,
     name="red_noise",
 ):
-    """Convenience function to return a BasisGP class with a
-    coarse time basis."""
+    """Function to return a BasisGP class with a coarse time basis."""
 
     if nknots is None:
         nknots = 2 * components + 1

From cd879c9610cab7d8b6b2df8b783bd2da445a2426 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Thu, 13 Mar 2025 12:48:05 +0100
Subject: [PATCH 25/29] Added note to docs about ContitionalGP

---
 docs/fft.rst | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/fft.rst b/docs/fft.rst
index b2c880fa..66380e1d 100644
--- a/docs/fft.rst
+++ b/docs/fft.rst
@@ -117,4 +117,5 @@ would not expect to be identical will be combined.
 The above is not a bug. Combining such bases and the corresponding Phi matrix
 does not matter, because the basis is zero, and that part of the signal has no
 bearing on the data or the model. However, when doing signal reconstruction,
-such as with `la_forge`, make sure to set `combine=False`.
+such as with `la_forge` or `utils.ConditionalGP`, make sure to set
+`combine=False`.

From 742381c73f3f1dbf7e6564f50159602238500de5 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Wed, 19 Mar 2025 12:40:10 +0100
Subject: [PATCH 26/29] Added higher order interpolation methods to the FFT
 correlated GPs

---
 enterprise/signals/gp_bases.py   | 25 ++++++++++---------------
 enterprise/signals/gp_signals.py |  8 ++++++--
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py
index a14ca9f0..44840ebe 100644
--- a/enterprise/signals/gp_bases.py
+++ b/enterprise/signals/gp_bases.py
@@ -5,6 +5,7 @@
 
 import numpy as np
 from enterprise.signals.parameter import function
+import scipy.interpolate as sint
 
 ######################################
 # Fourier-basis signal functions #####
@@ -95,13 +96,14 @@ def createfourierdesignmatrix_red(
 
 
 @function
-def create_fft_time_basis(toas, nknots=30, Tspan=None, start_time=None):
+def create_fft_time_basis(toas, nknots=30, Tspan=None, start_time=None, order=1):
     """
     Construct coarse time-domain design matrix from eq 11 of Chrisostomi et al., 2025
     :param toas: vector of time series in seconds
     :param nknots: number of coarse time samples to use (knots)
     :param Tspan: option to some other Tspan
     :param start_time: option to set some other start epoch of basis
+    :param order: order of the interpolation (1 = linear)
 
     :return B: coarse time-domain design matrix
     :return t_coarse: timestamps of coarse time grid
@@ -120,16 +122,7 @@ def create_fft_time_basis(toas, nknots=30, Tspan=None, start_time=None):
 
     t_fine = toas
     t_coarse = np.linspace(start_time, start_time + Tspan, nknots)
-    dt_coarse = t_coarse[1] - t_coarse[0]
-
-    idx = np.arange(len(t_fine))
-    idy = np.searchsorted(t_coarse, t_fine)
-    idy = np.clip(idy, 1, nknots - 1)
-
-    Bmat = np.zeros((len(t_fine), len(t_coarse)), "d")
-
-    Bmat[idx, idy] = (t_fine - t_coarse[idy - 1]) / dt_coarse
-    Bmat[idx, idy - 1] = (t_coarse[idy] - t_fine) / dt_coarse
+    Bmat = sint.interp1d(t_coarse, np.identity(nknots), kind=order)(t_fine).T
 
     return Bmat, t_coarse
 
@@ -171,7 +164,7 @@ def createfourierdesignmatrix_dm(
 
 
 @function
-def create_fft_time_basis_dm(toas, freqs, nknots=30, Tspan=None, start_time=None, fref=1400):
+def create_fft_time_basis_dm(toas, freqs, nknots=30, Tspan=None, start_time=None, fref=1400, order=1):
     """
     Construct DM-variation linear interpolation design matrix. Current
     normalization expresses DM signal as a deviation [seconds]
@@ -183,13 +176,14 @@ def create_fft_time_basis_dm(toas, freqs, nknots=30, Tspan=None, start_time=None
     :param Tspan: option to some other Tspan
     :param start_time: option to set some other start epoch of basis
     :param fref: reference frequency [MHz]
+    :param order: order of the interpolation (1 = linear)
 
     :return B: coarse time-domain design matrix
     :return t_coarse: timestamps of coarse time grid
     """
 
     # get base course time-domain matrix and times
-    Bmat, t_coarse = create_fft_time_basis(toas, nknots=nknots, Tspan=Tspan, start_time=start_time)
+    Bmat, t_coarse = create_fft_time_basis(toas, nknots=nknots, Tspan=Tspan, start_time=start_time, order=order)
 
     # compute the DM-variation vectors
     Dm = (fref / freqs) ** 2
@@ -363,7 +357,7 @@ def createfourierdesignmatrix_chromatic(
 
 
 @function
-def create_fft_time_basis_chromatic(toas, freqs, nknots=30, Tspan=None, start_time=None, fref=1400, idx=4):
+def create_fft_time_basis_chromatic(toas, freqs, nknots=30, Tspan=None, start_time=None, fref=1400, idx=4, order=1):
     """
     Construct scattering linear interpolation design matrix. Current
     normalization expresses DM signal as a deviation [seconds]
@@ -376,13 +370,14 @@ def create_fft_time_basis_chromatic(toas, freqs, nknots=30, Tspan=None, start_ti
     :param start_time: option to set some other start epoch of basis
     :param fref: reference frequency [MHz]
     :param idx: Index of chromatic effects
+    :param order: order of the interpolation (1 = linear)
 
     :return B: coarse time-domain design matrix
     :return t_coarse: timestamps of coarse time grid
     """
 
     # get base course time-domain matrix and times
-    Bmat, t_coarse = create_fft_time_basis(toas, nknots=nknots, Tspan=Tspan, start_time=start_time)
+    Bmat, t_coarse = create_fft_time_basis(toas, nknots=nknots, Tspan=Tspan, start_time=start_time, order=order)
 
     # compute the DM-variation vectors
     Dm = (fref / freqs) ** idx
diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index cd612da7..583123b5 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -229,6 +229,7 @@ def FFTBasisGP(
     cutbins=1,
     Tspan=None,
     start_time=None,
+    interpolation_order=1,
     name="red_noise",
 ):
     """Function to return a BasisGP class with a coarse time basis."""
@@ -245,7 +246,9 @@ def FFTBasisGP(
         cutbins = int(np.ceil(oversample / cutoff))
 
     if basis is None:
-        basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time)
+        basis = utils.create_fft_time_basis(
+            nknots=nknots, Tspan=Tspan, start_time=start_time, order=interpolation_order
+        )
 
     BaseClass = BasisGP(spectrum, basis, coefficients=coefficients, combine=combine, selection=selection, name=name)
 
@@ -545,6 +548,7 @@ def FFTBasisCommonGP(
     cutoff=None,
     cutbins=1,
     oversample=3,
+    interpolation_order=1,
     name="common_fft",
 ):
     if coefficients and (Tspan is None or start_time is None):
@@ -563,7 +567,7 @@ def FFTBasisCommonGP(
         #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
-    basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time)
+    basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time, order=interpolation_order)
     BaseClass = BasisCommonGP(spectrum, basis, orf, coefficients=coefficients, combine=combine, name=name)
 
     class FFTBasisCommonGP(BaseClass):

From 1eb262bd640b11a2d3aa420009e8dc53de6fc05f Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Fri, 21 Mar 2025 23:09:17 +0100
Subject: [PATCH 27/29] Added caching of the prior/Phi

---
 enterprise/signals/gp_signals.py  | 122 ++++++++++++++++++------------
 enterprise/signals/signal_base.py |   2 +
 2 files changed, 75 insertions(+), 49 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 583123b5..2d2d8dc5 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -82,6 +82,14 @@ def _do_selection(self, psr, priorfn, basisfn, coefficients, selection):
                     self._coefficients[key] = cpar
                     self._params[cpar.name] = cpar
 
+        @property
+        def prior_params(self):
+            """Get any varying prior parameters."""
+            ret = []
+            for prior in self._prior.values():
+                ret.extend([pp.name for pp in prior.params])
+            return ret
+
         @property
         def basis_params(self):
             """Get any varying basis parameters."""
@@ -114,6 +122,12 @@ def _construct_basis(self, params={}):
                 self._slices.update({key: slice(nctot, nn + nctot)})
                 nctot += nn
 
+        @signal_base.cache_call("prior_params")
+        def _construct_phi(self, params):
+            for key, slc in self._slices.items():
+                phislc = self._prior[key](self._labels[key], params=params)
+                self._phi = self._phi.set(phislc, slc)
+
         # this class does different things (and gets different method
         # definitions) if the user wants it to model GP coefficients
         # (e.g., for a hierarchical likelihood) or if they do not
@@ -173,10 +187,8 @@ def get_basis(self, params={}):
 
             def get_phi(self, params):
                 self._construct_basis(params)
+                self._construct_phi(params)
 
-                for key, slc in self._slices.items():
-                    phislc = self._prior[key](self._labels[key], params=params)
-                    self._phi = self._phi.set(phislc, slc)
                 return self._phi
 
             def get_phiinv(self, params):
@@ -257,34 +269,37 @@ class FFTBasisGP(BaseClass):
         signal_name = "red noise"
         signal_id = name
 
-        if coefficients:
-            raise NotImplementedError("Coefficients not supported for FFTBasisGP")
+        @signal_base.cache_call("prior_params")
+        def _construct_phi(self, params):
+            for key, slc in self._slices.items():
+                t_knots = self._labels[key]
 
-        else:
+                freqs = utils.knots_to_freqs(t_knots, oversample=oversample)
 
-            def get_phi(self, params):
-                self._construct_basis(params)
+                # Hack, because Enterprise adds in f=0 and then calculates df,
+                # meaning we cannot simply start freqs from 0. Thus, we use
+                # a modified frequency spacing, such that:
+                # [0, f1, 0, f1, f2, f3] => [df, -df, df, df, df]
+                if cutbins == 0:
+                    freqs_prior = np.concatenate([[freqs[1]], freqs])
+                    psd_prior = self._prior[key](freqs_prior, params=params, components=1)
+                    psd = np.concatenate([[-psd_prior[1]], psd_prior[2:]])
 
-                for key, slc in self._slices.items():
-                    t_knots = self._labels[key]
+                else:
+                    psd_prior = self._prior[key](freqs[1:], params=params, components=1)
+                    psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
 
-                    freqs = utils.knots_to_freqs(t_knots, oversample=oversample)
+                phislc = utils.psd2cov(t_knots, psd)
+                self._phi = self._phi.set(phislc, slc)
 
-                    # Hack, because Enterprise adds in f=0 and then calculates df,
-                    # meaning we cannot simply start freqs from 0. Thus, we use
-                    # a modified frequency spacing, such that:
-                    # [0, f1, 0, f1, f2, f3] => [df, -df, df, df, df]
-                    if cutbins == 0:
-                        freqs_prior = np.concatenate([[freqs[1]], freqs])
-                        psd_prior = self._prior[key](freqs_prior, params=params, components=1)
-                        psd = np.concatenate([[-psd_prior[1]], psd_prior[2:]])
+        if coefficients:
+            raise NotImplementedError("Coefficients not supported for FFTBasisGP")
 
-                    else:
-                        psd_prior = self._prior[key](freqs[1:], params=params, components=1)
-                        psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
+        else:
 
-                    phislc = utils.psd2cov(t_knots, psd)
-                    self._phi = self._phi.set(phislc, slc)
+            def get_phi(self, params):
+                self._construct_basis(params)
+                self._construct_phi(params)
 
                 return self._phi
 
@@ -409,6 +424,11 @@ def __init__(self, psr):
                 self._coefficients[""] = cpar
                 self._params[cpar.name] = cpar
 
+        @property
+        def prior_params(self):
+            """Get any varying prior parameters."""
+            return [pp.name for pp in self._prior.params]
+
         @property
         def basis_params(self):
             """Get any varying basis parameters."""
@@ -421,6 +441,10 @@ def basis_params(self):
         def _construct_basis(self, params={}):
             self._basis, self._labels = self._bases(params=params)
 
+        @signal_base.cache_call("prior_params")
+        def _construct_prior(self, labels, params):
+            return BasisCommonGP._prior(labels, params=params)
+
         if coefficients:
 
             def _get_coefficient_logprior(self, c, **params):
@@ -470,8 +494,7 @@ def get_basis(self, params={}):
 
             def get_phi(self, params):
                 self._construct_basis(params)
-
-                prior = BasisCommonGP._prior(self._labels, params=params)
+                prior = self._construct_prior(self._labels, params)
                 orf = BasisCommonGP._orf(self._psrpos, self._psrpos, params=params)
 
                 return prior * orf
@@ -599,35 +622,36 @@ def _construct_basis(self, params={}):
             freqs = utils.knots_to_freqs(self._t_knots, oversample=oversample)
             self._freqs = freqs
 
-        if coefficients:
-            raise NotImplementedError("Coefficients not supported for FFTBasisCommonGP")
+        @signal_base.cache_call("prior_params")
+        def _construct_prior(self, params):
+            """
+            Compute and cache the time-domain covariance ('phi') for *this* signal's basis.
+            """
+
+            # Hack, because Enterprise adds in f=0 and then calculates df,
+            # meaning we cannot simply start freqs from 0. Thus, we use
+            # a modified frequency spacing, such that:
+            # [0, f1, 0, f1, f2, f3] => [df, -df, df, df, df]
+            if cutbins == 0:
+                freqs_prior = np.concatenate([[self._freqs[1]], self._freqs])
+                psd_prior = FFTBasisCommonGP._prior(freqs_prior, params=params, components=1)
+                psd = np.concatenate([[-psd_prior[1]], psd_prior[2:]])
 
-        else:
-            # @signal_base.cache_call("basis_params", limit=1)
-            def get_phi_matrix(self, params):
-                """
-                Compute and cache the time-domain covariance ('phi') for *this* signal's basis.
-                """
+            else:
+                psd_prior = FFTBasisCommonGP._prior(self._freqs[1:], params=params, components=1)
+                psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
 
-                # Hack, because Enterprise adds in f=0 and then calculates df,
-                # meaning we cannot simply start freqs from 0. Thus, we use
-                # a modified frequency spacing, such that:
-                # [0, f1, 0, f1, f2, f3] => [df, -df, df, df, df]
-                if cutbins == 0:
-                    freqs_prior = np.concatenate([[self._freqs[1]], self._freqs])
-                    psd_prior = FFTBasisCommonGP._prior(freqs_prior, params=params, components=1)
-                    psd = np.concatenate([[-psd_prior[1]], psd_prior[2:]])
+            return utils.psd2cov(self._t_knots, psd)
 
-                else:
-                    psd_prior = FFTBasisCommonGP._prior(self._freqs[1:], params=params, components=1)
-                    psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
+        if coefficients:
+            raise NotImplementedError("Coefficients not supported for FFTBasisCommonGP")
 
-                return utils.psd2cov(self._t_knots, psd)
+        else:
 
             def get_phi(self, params):
                 """Over-load constructing Phi to deal with the FFT"""
                 self._construct_basis(params)
-                phi = self.get_phi_matrix(params)
+                phi = self._construct_prior(params)
 
                 orf = FFTBasisCommonGP._orf(self._psrpos, self._psrpos, params=params)
 
@@ -637,8 +661,8 @@ def get_phi(self, params):
             def get_phicross(cls, signal1, signal2, params):
                 """Use Phi from signal1, ORF from signal1 vs signal2"""
 
-                phi1 = signal1.get_phi_matrix(params)
-                # phi2 = signal2.get_phi_matrix(params)
+                phi1 = signal1._construct_prior(params)
+                # phi2 = signal2._construct_phi(params)
 
                 orf = FFTBasisCommonGP._orf(signal1._psrpos, signal2._psrpos, params=params)
 
diff --git a/enterprise/signals/signal_base.py b/enterprise/signals/signal_base.py
index f6da48ac..2ec768fc 100644
--- a/enterprise/signals/signal_base.py
+++ b/enterprise/signals/signal_base.py
@@ -799,6 +799,7 @@ def _set_cache_parameters(self):
             self.white_params = []
             self.basis_params = []
             self.delay_params = []
+            self.prior_params = []
             for signal in self._signals:
                 if signal.signal_type == "white noise":
                     self.white_params.extend(signal.ndiag_params)
@@ -807,6 +808,7 @@ def _set_cache_parameters(self):
                     # for common GPs, which do not have coefficients yet
                     self.delay_params.extend(getattr(signal, "delay_params", []))
                     self.basis_params.extend(signal.basis_params)
+                    self.prior_params.extend(getattr(signal, "prior_params", []))
                 elif signal.signal_type in ["deterministic"]:
                     self.delay_params.extend(signal.delay_params)
                 else:

From e3c032049f033cfa52c12a8de069f0e55c4384a1 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Fri, 21 Mar 2025 23:35:31 +0100
Subject: [PATCH 28/29] Fixed caching, but untested still

---
 enterprise/signals/gp_signals.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 2d2d8dc5..25409c28 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -123,7 +123,7 @@ def _construct_basis(self, params={}):
                 nctot += nn
 
         @signal_base.cache_call("prior_params")
-        def _construct_phi(self, params):
+        def _construct_prior(self, params):
             for key, slc in self._slices.items():
                 phislc = self._prior[key](self._labels[key], params=params)
                 self._phi = self._phi.set(phislc, slc)
@@ -187,7 +187,7 @@ def get_basis(self, params={}):
 
             def get_phi(self, params):
                 self._construct_basis(params)
-                self._construct_phi(params)
+                self._construct_prior(params)
 
                 return self._phi
 
@@ -270,7 +270,7 @@ class FFTBasisGP(BaseClass):
         signal_id = name
 
         @signal_base.cache_call("prior_params")
-        def _construct_phi(self, params):
+        def _construct_prior(self, params):
             for key, slc in self._slices.items():
                 t_knots = self._labels[key]
 
@@ -299,7 +299,7 @@ def _construct_phi(self, params):
 
             def get_phi(self, params):
                 self._construct_basis(params)
-                self._construct_phi(params)
+                self._construct_prior(params)
 
                 return self._phi
 
@@ -442,8 +442,8 @@ def _construct_basis(self, params={}):
             self._basis, self._labels = self._bases(params=params)
 
         @signal_base.cache_call("prior_params")
-        def _construct_prior(self, labels, params):
-            return BasisCommonGP._prior(labels, params=params)
+        def _construct_prior(self, params):
+            return BasisCommonGP._prior(self._labels, params=params)
 
         if coefficients:
 
@@ -494,7 +494,7 @@ def get_basis(self, params={}):
 
             def get_phi(self, params):
                 self._construct_basis(params)
-                prior = self._construct_prior(self._labels, params)
+                prior = self._construct_prior(params)
                 orf = BasisCommonGP._orf(self._psrpos, self._psrpos, params=params)
 
                 return prior * orf
@@ -662,7 +662,7 @@ def get_phicross(cls, signal1, signal2, params):
                 """Use Phi from signal1, ORF from signal1 vs signal2"""
 
                 phi1 = signal1._construct_prior(params)
-                # phi2 = signal2._construct_phi(params)
+                # phi2 = signal2._construct_prior(params)
 
                 orf = FFTBasisCommonGP._orf(signal1._psrpos, signal2._psrpos, params=params)
 

From e04d9d4fe8c25546647324953e2665be1b2aff71 Mon Sep 17 00:00:00 2001
From: Rutger van Haasteren <vhaasteren@gmail.com>
Date: Sun, 23 Mar 2025 19:50:56 +0100
Subject: [PATCH 29/29] Added fmax_factor to FFT methods

---
 enterprise/signals/gp_signals.py | 12 +++++++---
 enterprise/signals/utils.py      | 41 ++++++++++++++++----------------
 2 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/enterprise/signals/gp_signals.py b/enterprise/signals/gp_signals.py
index 25409c28..c048de49 100644
--- a/enterprise/signals/gp_signals.py
+++ b/enterprise/signals/gp_signals.py
@@ -237,6 +237,7 @@ def FFTBasisGP(
     nknots=None,
     selection=Selection(selections.no_selection),
     oversample=3,
+    fmax_factor=1,
     cutoff=None,
     cutbins=1,
     Tspan=None,
@@ -257,6 +258,8 @@ def FFTBasisGP(
         #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
+    fmax_factor = int(fmax_factor) if fmax_factor >= 1 else 1
+
     if basis is None:
         basis = utils.create_fft_time_basis(
             nknots=nknots, Tspan=Tspan, start_time=start_time, order=interpolation_order
@@ -274,7 +277,7 @@ def _construct_prior(self, params):
             for key, slc in self._slices.items():
                 t_knots = self._labels[key]
 
-                freqs = utils.knots_to_freqs(t_knots, oversample=oversample)
+                freqs = utils.knots_to_freqs(t_knots, oversample=oversample, fmax_factor=fmax_factor)
 
                 # Hack, because Enterprise adds in f=0 and then calculates df,
                 # meaning we cannot simply start freqs from 0. Thus, we use
@@ -289,7 +292,7 @@ def _construct_prior(self, params):
                     psd_prior = self._prior[key](freqs[1:], params=params, components=1)
                     psd = np.concatenate([np.zeros(cutbins), psd_prior[cutbins - 1 :]])
 
-                phislc = utils.psd2cov(t_knots, psd)
+                phislc = utils.psd2cov(t_knots, psd, fmax_factor=fmax_factor)
                 self._phi = self._phi.set(phislc, slc)
 
         if coefficients:
@@ -571,6 +574,7 @@ def FFTBasisCommonGP(
     cutoff=None,
     cutbins=1,
     oversample=3,
+    fmax_factor=1,
     interpolation_order=1,
     name="common_fft",
 ):
@@ -590,6 +594,8 @@ def FFTBasisCommonGP(
         #                low-frequency cut-off of the PSD
         cutbins = int(np.ceil(oversample / cutoff))
 
+    fmax_factor = int(fmax_factor) if fmax_factor >= 1 else 1
+
     basis = utils.create_fft_time_basis(nknots=nknots, Tspan=Tspan, start_time=start_time, order=interpolation_order)
     BaseClass = BasisCommonGP(spectrum, basis, orf, coefficients=coefficients, combine=combine, name=name)
 
@@ -619,7 +625,7 @@ def _construct_basis(self, params={}):
             self._basis, self._labels = self._bases(params=params, Tspan=span, start_time=start)
 
             self._t_knots = self._labels
-            freqs = utils.knots_to_freqs(self._t_knots, oversample=oversample)
+            freqs = utils.knots_to_freqs(self._t_knots, oversample=oversample, fmax_factor=fmax_factor)
             self._freqs = freqs
 
         @signal_base.cache_call("prior_params")
diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index 0d845445..06912f1e 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -842,18 +842,16 @@ def linear_interp_basis(toas, dt=30 * 86400):
     return M[:, idx], x[idx]
 
 
-def psd2cov(
-    t_knots,
-    psd,
-):
+def psd2cov(t_knots, psd, fmax_factor=1):
     """
-    Convert a power spectral density function, defined by (freqs, psd), to a covariance matrix
+    Convert a power spectral density function, defined by (freqs, psd), to a covariance matrix.
 
-    :param t_knots: Timestamps of the coarse time grid
-    :param psd: values of the PSD at frequencies freqs (assumes *delta_f in psd)
-                so psd is assumed to be in units of [s^2]
+    :param t_knots: Timestamps of the coarse time grid.
+    :param psd: PSD values evaluated at frequencies from knots_to_freqs
+                (assumes *delta_f in psd, so units of [s^2]).
+    :param fmax_factor: Integer factor to scale up fmax.
 
-    :return covmat: Covariance matrix at coarse time grid
+    :return covmat: Covariance matrix at the coarse time grid.
     """
 
     def toeplitz(c):
@@ -863,23 +861,25 @@ def toeplitz(c):
         j = np.arange(n).reshape(1, -1)
         return c[np.abs(i - j)]
 
-    def covmat(*args):
-        fullpsd = np.concatenate([psd, psd[-2:0:-1]])
+    # Create the full symmetric PSD (excluding duplicate Nyquist term)
+    fullpsd = np.concatenate([psd, psd[-2:0:-1]])
 
-        Cfreq = np.fft.ifft(fullpsd, norm="backward")
-        Ctau = Cfreq.real * len(fullpsd) / 2
+    # Compute the inverse FFT
+    Cfreq = np.fft.ifft(fullpsd, norm="backward")
+    Ctau = Cfreq.real * len(fullpsd) / 2
 
-        return toeplitz(Ctau[: len(t_knots)])
+    # With fmax_factor > 1, the IFFT time grid is finer by that factor.
+    # Slice out every fmax_factor-th sample to match the coarse grid.
+    return toeplitz(Ctau[::fmax_factor][: len(t_knots)])
 
-    return covmat()
 
-
-def knots_to_freqs(t_knots, oversample=3):
+def knots_to_freqs(t_knots, oversample=3, fmax_factor=1):
     """
     Convert knots of coarse time grid to frequencies
 
     :param t_knots: Timestamps of the coarse time grid
     :param oversample: amount by which to over-sample the frequency grid
+    :param fmax_factor: Integer factor to scale the maximum frequency.
 
     :return freqs: Frequencies, regularly sampled with
                    delta-f = 1/(oversample*T), fmax=1/(2*delta_t_knots)
@@ -890,9 +890,10 @@ def knots_to_freqs(t_knots, oversample=3):
     if nmodes % 2 == 0:
         raise ValueError("len(t_knots) must be odd.")
 
-    n_freqs = int((nmodes - 1) / 2 * oversample + 1)
-    fmax = (nmodes - 1) / Tspan / 2
-    return np.linspace(0, fmax, n_freqs)
+    n_freqs = int((nmodes - 1) / 2 * oversample * fmax_factor + 1)
+    fmax = (nmodes - 1) / (2 * Tspan)
+
+    return np.linspace(0, fmax * fmax_factor, n_freqs)
 
 
 # overlap reduction functions