Linear interpolated FFT-based time-correlated signals

docs/fft.rst

@@ -0,0 +1,121 @@
+
+.. 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, 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 `nodes=`, where it is understood that
+`nodes=2*n_components+1`. This is because `components` actually means
+frequencies.  In the time-domain, the number of `nodes` 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
+
+    nnodes = 81
+    dm_basis = utils.create_fft_time_basis_dm(nnodes=nnodes)
+    dm_pl = powerlaw(log10_A=dm_log10_A, gamma=dm_gamma)
+    dm_fft = gp_signals.FFTBasisGP(dm_pl, basis=dm_basis, nnodes=nnodes, name='dmgp')
+
+    chrom_basis = utils.create_fft_time_basis_chromatic(nnodes=nnodes, idx=chrom_idx)
+    chrom_pl = powerlaw(log10_A=chrom_log10_A, gamma=chrom_gamma)
+    chrom_fft = gp_signals.FFTBasisGP(chrom_pl, basis=chrom_basis, nnodes=nnodes, 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` or `utils.ConditionalGP`, make sure to set
+`combine=False`.

docs/index.rst

@@ -42,6 +42,7 @@ searches, and timing model analysis.
 
    mdc
    nano9
+   fft
 
 .. toctree::
    :maxdepth: 2

enterprise/signals/gp_bases.py

@@ -5,19 +5,23 @@
 
 import numpy as np
 from enterprise.signals.parameter import function
+import scipy.interpolate as sint
 
 ######################################
 # Fourier-basis signal functions #####
 ######################################
 
 __all__ = [
     "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",
 ]
 
@@ -91,6 +95,38 @@ def createfourierdesignmatrix_red(
     return F, Ffreqs
 
 
+@function
+def create_fft_time_basis(toas, nnodes=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 nnodes: number of coarse time samples to use (nodes)
+    :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
+    """
+    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) - 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, nnodes)
+    Bmat = sint.interp1d(t_coarse, np.identity(nnodes), kind=order)(t_fine).T
+
+    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
@@ -127,6 +163,34 @@ def createfourierdesignmatrix_dm(
     return F * Dm[:, None], Ffreqs
 
 
+@function
+def create_fft_time_basis_dm(toas, freqs, nnodes=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]
+    at fref [MHz]
+
+    :param toas: vector of time series in seconds
+    :param freqs: radio frequencies of observations [MHz]
+    :param nnodes: number of coarse time samples to use (nodes)
+    :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, nnodes=nnodes, Tspan=Tspan, start_time=start_time, order=order)
+
+    # 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
@@ -292,6 +356,35 @@ def createfourierdesignmatrix_chromatic(
     return F * Dm[:, None], Ffreqs
 
 
+@function
+def create_fft_time_basis_chromatic(toas, freqs, nnodes=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]
+    at fref [MHz]
+
+    :param toas: vector of time series in seconds
+    :param freqs: radio frequencies of observations [MHz]
+    :param nnodes: number of coarse time samples to use (nodes)
+    :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
+    :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, nnodes=nnodes, Tspan=Tspan, start_time=start_time, order=order)
+
+    # compute the DM-variation vectors
+    Dm = (fref / freqs) ** idx
+
+    return Bmat * Dm[:, None], t_coarse
+
+
 @function
 def createfourierdesignmatrix_general(
     toas,

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)