Tracking the herd: Bayesian noise analysis for a single pulsar
¶

In [ ]:
from google.colab import drive
drive.mount('/content/drive')

Prepared by Aurélien Chalumeau and Valentina Di Marco
¶

In the previous tutorial, noise in the wild, we learned to recognise the animals of the noise savanna by injecting each noise process one at a time and watching what it did to the timing residuals.

But identifying an animal from its footprints is not the same as measuring it precisely. In this tutorial, we go further: given a set of timing residuals that contain all the noise at once, can we recover the properties of each individual process? This is the job of Bayesian noise analysis, and it is the foundation of every PTA search for gravitational waves.

The workflow has five stages:

  1. Set up the data: load a simulated pulsar that already has white noise, red noise, and DM variations injected into it.
  2. Build the noise model: use Enterprise to construct the signal model and likelihood.
  3. Sample the posterior: run a Markov Chain Monte Carlo sampler to explore the parameter space.
  4. Diagnose and interpret: check that the chain has converged, then read off the recovered parameter values.
  5. Reconstruct the signal: use the posterior samples to reconstruct the red noise signal in the time domain.

By the end, you will have measured the amplitude and spectral index of the red noise and DM variations in a single pulsar, and you will have seen how well the time-domain reconstruction matches what was injected.


1. Imports¶

In [ ]:
import os
import numpy as np
import scipy.linalg as sl

import libstempo as LT
import libstempo.toasim as LTsim
from enterprise.pulsar import Pulsar
from enterprise.signals.utils import powerlaw, createfourierdesignmatrix_dm
from enterprise.signals import white_signals, gp_signals, parameter
from enterprise.signals.signal_base import PTA
from enterprise_extensions.sampler import JumpProposal

from PTMCMCSampler.PTMCMCSampler import PTSampler as ptmcmc

import matplotlib.pyplot as plt
import corner
/opt/anaconda3/envs/SPNA_env/lib/python3.9/site-packages/enterprise/signals/utils.py:13: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import Requirement, resource_filename
Optional mpi4py package is not installed.  MPI support is not available.
Optional acor package is not installed. Acor is optionally used to calculate the effective chain length for output in the chain file.

2. Prepare the data¶

We need a pulsar with all three noise components already injected. In the noise in the wild tutorial, you built this pulsar from scratch using fakepulsar and injected each noise process yourself. Here, we repeat those injections in a single setup cell so that we start from a fully noisy pulsar and can focus on the inference that follows.

Our pulsar is PSR J1909-3744, a millisecond pulsar timed at three frequencies: 500, 900, and 1400 MHz. We inject:

  • White noise via an EFAC of 1.0
  • Achromatic red noise with amplitude $A_{\rm rn} = 7 \times 10^{-14}$ and spectral index $\gamma_{\rm rn} = 3$,
  • DM variations with amplitude $A_{\rm dm} = 5 \times 10^{-14}$ and spectral index $\gamma_{\rm dm} = 2$.

These are the injected (true) values. The goal of the analysis is to recover them from the data.

In [ ]:
day = 24 * 3600
year = 365.25 * day
import math

def add_efac(psr, efac=1.0, flagid=None, flags=None, seed=None):
    """Adapted from libstempo.
    Add nominal TOA errors multiplied by an EFAC factor."""
    if seed is not None:
        np.random.seed(seed)
    efacvec = np.ones(psr.nobs)
    if flags is None:
        if not np.isscalar(efac):
            raise ValueError("If flags is None, efac must be a scalar")
        efacvec = np.ones(psr.nobs) * efac
    if flags is not None and flagid is not None and not np.isscalar(efac):
        if len(efac) == len(flags):
            for ct, flag in enumerate(flags):
                ind = flag == np.array(psr.flagvals(flagid))
                efacvec[ind] = efac[ct]
    psr.stoas[:] += efacvec * psr.toaerrs * (1e-6 / day) * np.random.randn(psr.nobs)

def add_time_corr_signal(psr, A, gamma, components=10, tspan=None, seed=None, idx=0, factor=1.):
    """Taken from libstempo.
    Add a power-law correlated signal with P(f) = A^2/(12 pi^2) (f year)^-gamma,
    using `components` Fourier modes. For chromatic signals, set idx=2 for DM."""
    if seed is not None:
        np.random.seed(seed)
    t = psr.toas()
    fref = 1400
    v = (fref / psr.freqs) ** idx
    minx, maxx = np.min(t), np.max(t)
    if tspan is None:
        x = (t - minx) / (maxx - minx)
        T = (day / year) * (maxx - minx)
    else:
        x = (t - minx) / tspan
        T = (day / year) * tspan
    size = 2 * components
    F = np.zeros((psr.nobs, size), "d")
    f = np.zeros(size, "d")
    for i in range(components):
        F[:, 2 * i] = np.cos(2 * math.pi * (i + 1) * x)
        F[:, 2 * i + 1] = np.sin(2 * math.pi * (i + 1) * x)
        f[2 * i] = f[2 * i + 1] = (i + 1) / T
    norm = A**2 * year**2 / (12 * math.pi**2 * T)
    prior = norm * f ** (-gamma)
    y = np.sqrt(prior) * np.random.randn(size)
    psr.stoas[:] += (1.0 / day) * np.dot(F, y) * v * factor
In [ ]:
# --- Simulation parameters (the true values we will try to recover) ---
efac   = 1.0
A_rn   = 7e-14;  gamma_rn = 3.0
A_dm   = 5e-14;  gamma_dm = 2.0

# --- Build the fake pulsar ---
parfile = "J1909-3744.par"
cadence = 100.   # days
toas    = np.repeat(np.arange(52000, 59000, cadence), 3)   # MJD, 3 TOAs per epoch
toaerrs = np.random.randn(len(toas)) * 0.2 + 1.            # microseconds
freqs   = np.tile([500, 900, 1400], len(toas) // 3)        # MHz, one freq per band per epoch

ltpsr = LTsim.fakepulsar(parfile, obstimes=toas, toaerr=toaerrs, freq=freqs)

# --- Inject noise processes and keep copies of each contribution ---
res_i = np.copy(ltpsr.residuals())
add_efac(ltpsr, efac=efac)
res_wn = np.copy(ltpsr.residuals()) - res_i

toa_tmp = np.copy(ltpsr.residuals())
add_time_corr_signal(ltpsr, A=A_rn, gamma=gamma_rn, components=30, seed=4105)
res_rn = np.copy(ltpsr.residuals()) - toa_tmp

toa_tmp = np.copy(ltpsr.residuals())
add_time_corr_signal(ltpsr, A=A_dm, gamma=gamma_dm, components=30, idx=2, seed=413405)
res_dm = np.copy(ltpsr.residuals()) - toa_tmp

print("Injected values:")
print(f"  EFAC              = {efac}")
print(f"  RN  log10(A)      = {np.log10(A_rn):.3f},  gamma = {gamma_rn}")
print(f"  DM  log10(A)      = {np.log10(A_dm):.3f},  gamma = {gamma_dm}")
print(f"  Nobs              = {ltpsr.nobs}")
print(f"  Tspan             = {(ltpsr.toas().max()-ltpsr.toas().min()):.0f} days")
[tempo2Util.C:396] Warning: [TIM1] Please place MODE flags in the parameter file 
[preProcess.C:158] Warning: PSR J1909-3744 uses DM2+ but does not define DM_SERIES. Assume Taylor. This has behaviour has changed since June 2020!
See https://bitbucket.org/psrsoft/tempo2/issues/27/tempo2-dm-polynomial-is-not-a-taylor

Injected values:
  EFAC              = 1.0
  RN  log10(A)      = -13.155,  gamma = 3.0
  DM  log10(A)      = -13.301,  gamma = 2.0
  Nobs              = 210
  Tspan             = 6900 days

Let's have a quick look at the residuals. Red noise and DM variations are both present, so the residuals should no longer look like independent draws from a pure $N(0,1)$ distribution.

In [ ]:
f, (a0, a1) = plt.subplots(1, 2, figsize=(22, 7),
                            gridspec_kw={'width_ratios': [3, 1], 'hspace': 0, 'wspace': 0})
bands = [500, 900, 1400]

for band in bands:
    m = ltpsr.freqs == band
    a0.errorbar(
        ltpsr.toas()[m],
        ltpsr.residuals()[m] * 1e6,
        yerr=ltpsr.toaerrs[m],
        fmt='.',
        label=f"{band} MHz"
    )

a0.plot(ltpsr.toas(), res_rn * 1e6, 'r', lw=2, zorder=10, label="Injected red noise")
a0.axhline(0., ls=':', c='k', lw=2, zorder=10)
a0.grid(alpha=.4)
a0.set_xlabel("Epochs [MJD]", fontsize=14)
a0.set_ylabel(r"Timing residuals [$\mu$s]", fontsize=14)
a0.tick_params(labelsize=13)
a0.legend(fontsize=13)

Gauss = np.random.randn(10000)
a1.hist(Gauss, bins=30, color='grey', histtype='step', orientation="horizontal", density=True, label="N(0,1)")
a1.hist(ltpsr.residuals() / (ltpsr.toaerrs * 1e-6), bins=30, histtype='step',
        orientation="horizontal", density=True, label="Weighted residuals")
plt.setp(a1.get_xticklabels(), visible=False)
plt.setp(a1.get_yticklabels(), visible=False)
a1.grid(alpha=.4)
a1.legend(fontsize=11)
plt.suptitle("PSR J1909-3744: simulated residuals with all noise injected", fontsize=14)
plt.show()
Question: The data points do not track the red curve perfectly — they scatter around it rather than following it closely. Why?
Click to reveal the answer
The red curve shows only the injected achromatic red noise. But the data also contains DM variations, which add a radio frequency-dependent offset on top. Because the three observing frequencies (500, 900, 1400 MHz) each experience a different DM delay, the data points scatter around the red curve rather than sitting on it. If you coloured the points by frequency, you would see that the offsets are systematic per frequency band, not random. And, of course there is the scatter due to white noise too.

3. Build the noise model with Enterprise¶

Now comes the tracker's job: given this messy savanna of overlapping signals, how do we measure each animal separately? The answer is to write down a statistical model that describes every noise process mathematically, and then fit all parameters simultaneously using Bayesian inference.

Enterprise is the standard toolkit for this in pulsar timing. It takes a description of each signal component and assembles the full likelihood function automatically.

Our model has four components:

  1. Timing model residuals (TimingModel): marginalises over small errors in the pulsar ephemeris parameters.
  2. EFAC (MeasurementNoise): a single multiplicative scaling of the TOA uncertainties.
  3. Achromatic red noise (FourierBasisGP with a power-law spectrum): the slow, frequency-independent drift.
  4. DM variations (BasisGP with a power-law spectrum and the DM Fourier basis): the frequency-dependent chromatic noise.

The free parameters are: $E$ (EFAC), $\log_{10} A_{\rm rn}$, $\gamma_{\rm rn}$, $\log_{10} A_{\rm dm}$, $\gamma_{\rm dm}$, five parameters in total. We place uniform priors on all of them.

In [ ]:
# Convert the libstempo object to an Enterprise Pulsar object
psr = Pulsar(ltpsr)

For real data, we would create the Enterprise pulsar directly from the .par and .tim files like this:

psr = Pulsar("J1909-3744.par", "J1909-3744.tim")

3.1 Timing model¶

The timing model marginalisation absorbs small offsets in the ephemeris parameters so they do not bias the noise estimates. This is always included.

In [ ]:
s = gp_signals.TimingModel()

3.2 White noise (EFAC)¶

EFAC rescales the formal TOA uncertainties. A value of 1 means the uncertainties are perfectly calibrated. Values above 1 indicate underestimated errors; below 1, overestimated. We place a broad uniform prior $E \in [0.1, 5]$.

In [ ]:
efac_prior = parameter.Uniform(0.1, 5)
s += white_signals.MeasurementNoise(efac=efac_prior)

At this point Entrpise has done nothing else than instantiate a prior class and a signal class.

3.3 Achromatic red noise¶

The red noise is modelled as a Gaussian process with a power-law covariance. We use 30 Fourier modes to represent it. This is sufficient to capture all the power above the $1/T_{\rm span}$ fundamental frequency.

Priors: $\log_{10} A_{\rm rn} \in [-18, -10]$, $\gamma_{\rm rn} \in [0, 7]$.

In [ ]:
log10_A_rn_prior = parameter.Uniform(-18, -10)
gamma_rn_prior   = parameter.Uniform(0, 7)
pl_rn = powerlaw(log10_A=log10_A_rn_prior, gamma=gamma_rn_prior)
s += gp_signals.FourierBasisGP(pl_rn, components=30)

3.4 DM variations¶

DM variations are also a power-law Gaussian process, but they use a chromatic Fourier basis: the time delays scale as $\nu^{-2}$ with observing frequency $\nu$. This is what distinguishes them from achromatic red noise — the same low-frequency power, but frequency-dependent in amplitude.

Priors: $\log_{10} A_{\rm dm} \in [-18, -10]$, $\gamma_{\rm dm} \in [0, 7]$.

In [ ]:
log10_A_dm_prior = parameter.Uniform(-18, -10)
gamma_dm_prior   = parameter.Uniform(0, 7)
pl_dm    = powerlaw(log10_A=log10_A_dm_prior, gamma=gamma_dm_prior)  # <-- NOT pl_rn!
dm_basis = createfourierdesignmatrix_dm(nmodes=30)
s += gp_signals.BasisGP(pl_dm, dm_basis, name="dm_gp")

3.5 Assemble the PTA object¶

We combine the signal model with the pulsar data into a PTA object. This object knows how to compute the log-likelihood and log-prior for any point in parameter space.

In [ ]:
pta = PTA([s(psr)])

print("Free parameters in this model:")
for p in pta.param_names:
    print(" ", p)
Free parameters in this model:
  J1909-3744_dm_gp_gamma
  J1909-3744_dm_gp_log10_A
  J1909-3744_efac
  J1909-3744_red_noise_gamma
  J1909-3744_red_noise_log10_A
Question: Look at the five free parameters listed above. Which one do you expect to be hardest to constrain, and why?
Click to reveal the answer
The DM parameters (dm_gp_log10_A and dm_gp_gamma) are generally harder to constrain than the red noise parameters, because separating DM variations from achromatic red noise relies mainly on the frequency coverage of the dataset. Both processes produce low-frequency correlated noise. The only thing that distinguishes them is that DM delays scale as $\nu^{-2}$. With only three observing frequencies and irregular frequency sampling, this separation is imperfect, and the DM and red noise posteriors can be partially degenerate with each other.

EFAC, by contrast, is often easier tp constrain: it only affects the white noise level, which is easy to read off from the scatter of residuals around the mean on short timescales.

3.6 Sanity check: evaluate the likelihood at the injected values¶

Before running the sampler, it is worth checking that the likelihood function works and that it gives a finite value at the true parameter point. If the log-likelihood is $-\infty$ here, something is wrong with the model setup.

In [ ]:
x_true = {
    "J1909-3744_dm_gp_gamma"    : gamma_dm,
    "J1909-3744_dm_gp_log10_A"  : np.log10(A_dm),
    "J1909-3744_efac"           : efac,
    "J1909-3744_red_noise_gamma": gamma_rn,
    "J1909-3744_red_noise_log10_A": np.log10(A_rn),
}

lnL = pta.get_lnlikelihood(x_true)
lnP = pta.get_lnprior(x_true)
print(f"log-likelihood at injected values : {lnL:.2f}")
print(f"log-prior     at injected values  : {lnP:.2f}")
print()
if np.isfinite(lnL):
    print("✓  Likelihood is finite — model setup looks correct.")
else:
    print("✗  Likelihood is not finite — check the model setup!")
log-likelihood at injected values : 1538.07
log-prior     at injected values  : -9.64

✓  Likelihood is finite — model setup looks correct.
/opt/anaconda3/envs/SPNA_env/lib/python3.9/site-packages/enterprise/signals/utils.py:875: RuntimeWarning: invalid value encountered in divide
  nmat = Mmat / norm

4. Sample the posterior¶

We want the posterior distribution $p(\theta | d) \propto p(d | \theta)\, p(\theta)$: the probability of each set of noise parameters $\theta$ given the data $d$. The standard approach is Markov Chain Monte Carlo (MCMC).

We use PTMCMCSampler, a parallel-tempering MCMC sampler developed specifically for pulsar timing applications. It mixes three types of proposal:

  • SCAM (Single Component Adaptive Metropolis): adapts the step size for each parameter independently.
  • DE (Differential Evolution): uses past chain states to propose large jumps (very effective for correlated parameters).
  • AM (Adaptive Metropolis): global covariance-based proposals.

You can see these algorithms in action on a 2D toy distribution at the MCMC Interactive Gallery from Chi Feng. try switching between DE-MCMC-Z, AdaptiveMH, and RandomWalkMH to see how each one explores the space differently.

We also add draw-from-prior proposals for each correlated signal. These allow the sampler to jump anywhere in the prior volume, which helps it escape local modes.

We will run only 5000 samples (that is not usually enough).

In [ ]:
outdir   = "./chain/"
os.makedirs(outdir, exist_ok=True)
nsamples = 5e4

# Save parameter names for later
with open(os.path.join(outdir, "pars.txt"), "w") as fout:
    for pname in pta.param_names:
        fout.write(pname + "\n")

# Initial point drawn from the prior; small initial covariance
x0   = np.hstack([p.sample() for p in pta.params])
ndim = len(x0)
cov  = np.diag(np.ones(ndim) * 0.01**2)

sampler = ptmcmc(ndim, pta.get_lnlikelihood, pta.get_lnprior, cov, outDir=outdir)
In [ ]:
# Jump proposals
jp = JumpProposal(pta, None, empirical_distr=None)
sampler.addProposalToCycle(jp.draw_from_prior, 5)

sel_sig = {"red_noise": 10, "dm": 10}
for sig_name in sel_sig:
    if any([sig_name in p for p in pta.param_names]):
        sampler.addProposalToCycle(jp.draw_from_par_prior(sig_name), sel_sig[sig_name])
In [ ]:
# Run the sampler  (this will take a few minutes)
np.random.seed(31)
sampler.sample(x0, int(nsamples), SCAMweight=40, DEweight=60, AMweight=20)
Finished 2.00 percent in 0.793704 s Acceptance rate = 0.667
/opt/anaconda3/envs/SPNA_env/lib/python3.9/site-packages/enterprise/signals/parameter.py:62: RuntimeWarning: divide by zero encountered in log
  logpdf = np.log(self.prior(value, **kwargs))
Finished 20.00 percent in 8.640089 s Acceptance rate = 0.230578Adding DE jump with weight 60
Finished 98.00 percent in 35.653018 s Acceptance rate = 0.381633
Run Complete

5. Diagnose the chain¶

Before trusting any posterior results, we must check that the chain has converged: that it has explored the parameter space thoroughly enough that its samples are representative of the true posterior.

The two main diagnostics are:

  • Trace plots: the value of each parameter as a function of iteration. A well-mixed chain looks like white noise around the posterior mean. If it is still drifting, or is stuck in one region, the chain has not converged.
  • Corner plot: the marginalised 1D and 2D posterior distributions for all parameter pairs. The true (injected) values should lie within the credible regions.

We discard the first 50% of the chain as burn-in.

In [ ]:
ch   = np.loadtxt("%s/chain_1.txt" % outdir)
ch   = ch[int(len(ch) * 0.50):]   # discard burn-in
pars = np.loadtxt("%s/pars.txt"   % outdir, dtype=str)

# Injected values in the same order as pars
p_inj = [
    gamma_dm,
    np.log10(A_dm),
    efac,
    gamma_rn,
    np.log10(A_rn),
]

print(f"Post-burn-in chain length: {len(ch)} samples")
Post-burn-in chain length: 2450 samples

5.1 Trace plots¶

Each plot shows the chain for one parameter. The black horizontal line is the injected (true) value. A well-behaved chain should oscillate around the true value with no systematic drift.

In [ ]:
fig, axes = plt.subplots(len(pars), 1, figsize=(14, 3 * len(pars)))
for i, (ax, par) in enumerate(zip(axes, pars)):
    ax.plot(ch[:, i], lw=0.5, color='steelblue', alpha=0.8)
    ax.axhline(p_inj[i], c='k', lw=2, label="Injected")
    ax.set_ylabel(par, fontsize=10)
    ax.grid(alpha=0.3)
    if i == 0:
        ax.legend(fontsize=10)
axes[-1].set_xlabel("Iteration", fontsize=12)
plt.suptitle("Trace plots (post burn-in)", fontsize=13)
plt.tight_layout()
plt.show()
Question: Look at the trace plots above. What would a non-converged chain look like, and what would you do about it?
Click to reveal the answer
A non-converged chain shows up in a few characteristic ways:
  • Drifting: the trace is still moving systematically in one direction rather than oscillating around a stable mean. The sampler has not yet found the high-likelihood region.
  • Stuck: the trace is flat for long stretches, meaning the sampler is repeatedly rejecting proposals and not moving. This usually means the step size is too large, or the sampler is trapped in a local mode.
  • Two distinct levels: the trace jumps between two very different values, suggesting the posterior has multiple modes and the sampler is occasionally tunnelling between them without properly exploring either.
The fixes depend on the symptom. For drifting, run more samples or increase the burn-in fraction you discard. For a stuck chain, check the acceptance rate (PTMCMCSampler prints this) — it should be roughly 20–30%. For multi-modal posteriors, the draw-from-prior proposals help, as does parallel tempering, which is why PTMCMCSampler uses it by default.

5.2 Corner plot¶

The corner plot shows all pairwise 2D marginals and 1D histograms. The black lines mark the injected values. Ideally each 1D histogram peaks near the true value, and the 2D contours are compact and centred near the truth.

Question: Which parameters are most correlated with each other? Does that make physical sense?
Click to reveal the answer
The strongest correlations are typically between the amplitude and spectral index of the same process, for example, red noise amplitude and spectral index are correlated because a steeper spectrum (larger gamma) with a higher amplitude can produce similar residuals to a shallower spectrum with a lower amplitude. The data cannot always uniquely pin down both simultaneously, especially when the observation span is short relative to the signal's characteristic timescale.

You may also see some correlation between the red noise and DM parameters, because both processes contribute low-frequency power to the residuals. The sampler has to disentangle them using only the frequency dependence of the DM signal which is why multi-frequency observations are essential.

EFAC is usually largely uncorrelated with the red noise and DM parameters, because it only affects the white noise level on short timescales, which is a different regime from the correlated signals.
In [ ]:
# Thin by factor 2 to reduce autocorrelation
corner.corner(
    ch[::2, :-4],
    labels=list(pars),
    truths=p_inj,
    truth_color='k',
    hist_kwargs={"density": True},
    show_titles=True,
    title_fmt=".2f",
)
plt.suptitle("Posterior distributions (black lines = injected values)", fontsize=13, y=1.01)
plt.show()

5.3 Posterior summary¶

Let's print the posterior median and the 16th and 84th percentile range for each parameter containing about 68% of the samples, alongside the injected value, to make the comparison explicit. For a Gaussian posterior this corresponds to a one-sigma interval, but percentiles are more robust when the posterior is skewed or non-Gaussian.

In [ ]:
print(f"{'Parameter':<40} {'Injected':>10}  {'Median':>10}  {'16th%':>10}  {'84th%':>10}")
print("-" * 85)
for i, par in enumerate(pars):
    med  = np.median(ch[:, i])
    lo   = np.percentile(ch[:, i], 16)
    hi   = np.percentile(ch[:, i], 84)
    print(f"{par:<40} {p_inj[i]:>10.3f}  {med:>10.3f}  {lo:>10.3f}  {hi:>10.3f}")
Parameter                                  Injected      Median       16th%       84th%
-------------------------------------------------------------------------------------
J1909-3744_dm_gp_gamma                        2.000       1.490       1.210       1.818
J1909-3744_dm_gp_log10_A                    -13.301     -13.243     -13.313     -13.173
J1909-3744_efac                               1.000       1.036       0.968       1.108
J1909-3744_red_noise_gamma                    3.000       3.314       2.673       4.061
J1909-3744_red_noise_log10_A                -13.155     -13.254     -13.523     -13.036

The corner plot shows the recovered parameters in the abstract space $(\log_{10} A, \gamma)$. It is often more informative to translate those parameters into the corresponding power spectrum, because that shows where the variance of the timing residuals lives in frequency.

For a power-law red process, the power spectral density is a power law $P(f) \propto A^2 f^{-\gamma}$. In the plot below we show the equivalent Fourier-bin amplitude, which is the square root of the PSD multiplied by the frequency-bin width.

$\rho(f) = \sqrt{\frac{A^2}{12\pi^2} f_{\rm yr}^{\gamma-3} f^{-\gamma} / T_{\rm span}}$,

where $T_{\rm span}$ is the total observation time in seconds. This is the square root of the power assigned to a Fourier bin, so $\rho$ has units of seconds and can be read as an RMS timing-residual amplitude at that frequency.

We draw 1000 posterior samples, compute $\rho(f)$ for each, and plot the resulting envelope alongside the injected spectrum. We also overlay the expected white-noise level, which sets the approximate noise floor at high frequencies.

The vertical dashed line marks the reference frequency $1\,{\rm yr}^{-1}$.

In [ ]:
def powerlaw_spectrum(f, log10_A, gamma):
    """Power spectral amplitude rho(f) in seconds."""
    fyr = 1.0 / (365.25 * 86400)
    return np.sqrt(
        (10**log10_A)**2 / (12.0 * np.pi**2)
        * fyr**(gamma - 3) * f**(-gamma)
    )
In [ ]:
Tspan = 86400.0 * (ltpsr.toas().max() - ltpsr.toas().min())
freqs_plot = np.linspace(1.0 / Tspan, 30.0 / Tspan, 30)
fyr = 1.0 / (365.25 * 86400)

N_draw = 1000
idxs   = np.random.choice(len(ch), size=N_draw, replace=False)

As_rn = ch[idxs, list(pars).index('J1909-3744_red_noise_log10_A')]
Gs_rn = ch[idxs, list(pars).index('J1909-3744_red_noise_gamma')]

# Compute bounds before any plotting
sigma   = np.median(ltpsr.toaerrs) * 1e-6
cadence = np.median(np.diff(ltpsr.toas()[::3])) * 86400
rho_wn  = np.sqrt(2.0 * sigma**2 * cadence / Tspan)
rho_inj = powerlaw_spectrum(freqs_plot, np.log10(A_rn), gamma_rn) / np.sqrt(Tspan)

fig, ax = plt.subplots(figsize=(9, 5))
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlim([freqs_plot[0] * 0.9, freqs_plot[-1] * 1.1])
ax.set_ylim([rho_wn * 0.5, rho_inj[0] * 10])
ylim_lo = rho_wn * 0.1
ylim_hi = rho_inj[0] * 10

# Posterior envelope — clipped to plot range before drawing
for i in range(N_draw):
    rho_rec = powerlaw_spectrum(freqs_plot, As_rn[i], Gs_rn[i]) / np.sqrt(Tspan)
    rho_rec = np.clip(rho_rec, ylim_lo, ylim_hi)
    lbl = "Recovered (posterior samples)" if i == 0 else None
    ax.plot(freqs_plot, rho_rec, color='coral', alpha=0.02, label=lbl)

# Injected spectrum
ax.plot(freqs_plot, rho_inj, 'k.', zorder=1000, label="Injected power law", ms=6)

# White noise floor
ax.axhline(rho_wn, ls=':', color='k', alpha=0.8, label="Expected white noise floor")

ax.axvline(fyr, ls='--', c='k', lw=1, alpha=0.3)
ax.text(fyr * 0.6, rho_wn * 3, r"$1\,{\rm yr}^{-1}$", fontsize=11)

ax.set_xlabel("Frequency [Hz]", fontsize=14)
ax.set_ylabel(r"$\rho$ [s]", fontsize=14)
ax.tick_params(labelsize=12)
ax.legend(fontsize=11)
ax.set_title("Red noise power spectrum: injected vs recovered", fontsize=13)
plt.tight_layout()
plt.show()

This plot does not estimate the power spectrum directly from the residuals. Instead, it converts the injected and recovered red-noise parameters, (A) and ($\gamma$), into the corresponding model spectrum. The black points show the spectrum implied by the injected parameters, while the coral curves show spectra implied by posterior samples.


7. Reconstruct the signal in the time domain¶

The power spectrum tells us how much power is at each frequency, but it does not show us the actual time-domain waveform. To recover that, we use the Wiener filter .

Given a posterior sample of the noise parameters $\theta$, the conditional mean of the red noise signal $\mathbf{s}$ given the data $\mathbf{d}$ is:

$\langle \mathbf{s} | \mathbf{d}, \theta \rangle = \Phi T^T N^{-1} \left( T \Phi T^T + N \right)^{-1} \mathbf{r}$

where $T$ is the Fourier design matrix, $\Phi$ is the spectral prior covariance, $N$ is the white noise covariance matrix, and $\mathbf{r}$ are the timing residuals. We draw many posterior samples, compute this quantity for each, and plot the resulting envelope.

The result is compared to the injected red noise signal. A good inference should produce an envelope that covers the injected signal for most of its length.

In [ ]:
def get_b(d, TNT, phiinv):
    """Conditional mean and a random draw from the GP posterior.
    Taken from la_forge: https://github.com/nanograv/la_forge"""
    Sigma = TNT + (np.diag(phiinv) if phiinv.ndim == 1 else phiinv)
    try:
        u, s, _ = sl.svd(Sigma)
        mn = np.dot(u, np.dot(u.T, d) / s)
        Li = u * np.sqrt(1.0 / s)
    except np.linalg.LinAlgError:
        Q, R = sl.qr(Sigma)
        Sigi = sl.solve(R, Q.T)
        mn = np.dot(Sigi, d)
        u, s, _ = sl.svd(Sigi)
        Li = u * np.sqrt(1.0 / s)
    return mn + np.dot(Li, np.random.randn(Li.shape[0]))


def get_tdelay_from_chains(pta, psr, ch, pars,
                           ipsr=0, nsamples=100, ch_idxs=None,
                           signames=['all'], separe_signals=True,
                           verbose=True):
    """Draw nsamples posterior realisations of the named signals."""
    if ch_idxs is None:
        ch_idxs = [np.random.choice(len(ch)) for _ in range(nsamples)]

    if list(pars) != pta.param_names:
        print("WARNING: parameter name mismatch between chain and PTA object.")

    sig_idxs = {}
    for sig, idx in pta._signalcollections[ipsr]._idx.items():
        sig_idxs.update({sig.signal_id: idx})

    if separe_signals:
        delays = np.zeros((len(signames), len(ch_idxs), len(psr.toas)))
    else:
        delays = np.zeros((len(ch_idxs), len(psr.toas)))

    pta_signals = pta._signalcollections[ipsr]._signals

    for i, ch_idx in enumerate(ch_idxs):
        if verbose:
            print(f"{i+1}/{len(ch_idxs)}", end="\r")
        post_sample = pta.map_params(ch[ch_idx, :])

        TNrs    = pta.get_TNr(post_sample)
        TNTs    = pta.get_TNT(post_sample)
        phiinvs = pta.get_phiinv(post_sample, logdet=False)
        Ts      = pta.get_basis(post_sample)

        if any([sig.signal_type in ["basis", "common basis"]
                for sig in pta._signalcollections[ipsr]._signals]):
            w = get_b(TNrs[ipsr], TNTs[ipsr], phiinvs[ipsr])

        for j, signame in enumerate(signames):
            delay = 0
            if signame == 'all':
                pta_sig = [sig for sig in pta_signals]
            else:
                pta_sig = [sig for sig in pta_signals if signame == sig.signal_id]

            for sig in pta_sig:
                if sig.signal_type == 'deterministic':
                    delay += sig.get_delay(post_sample)
                if sig.signal_type in ['basis', 'common basis']:
                    idx    = sig_idxs[sig.signal_id]
                    delay += np.dot(Ts[ipsr][:, idx], w[idx])

            if separe_signals:
                delays[j, i, :] = delay
            else:
                delays[i, :] += delay

    return delays
In [ ]:
# Reconstruct the red noise signal from 100 posterior samples
tdelays = get_tdelay_from_chains(
    pta, psr, ch, pars,
    nsamples=100,
    signames=['red_noise'],
    separe_signals=False,
)
tdelays_dm = get_tdelay_from_chains(
    pta, psr, ch, pars,
    nsamples=100,
    signames=['dm_gp'],
    separe_signals=False,
)
100/100
In [ ]:
means = np.mean(tdelays, axis=0)

fig, ax = plt.subplots(figsize=(20, 7))

# Injected signal
ax.plot(ltpsr.toas(), res_rn, c='r', lw=3, zorder=10, label="Injected red noise")

# Posterior envelope
for i in range(len(tdelays)):
    ax.plot(ltpsr.toas(), tdelays[i, :], color='green', alpha=0.1)
ax.plot(ltpsr.toas(), means, '--', color='darkgreen', lw=3,
        label="Recovered (posterior mean)", zorder=9)

# Data
ax.errorbar(ltpsr.toas(), ltpsr.residuals(), yerr=ltpsr.toaerrs * 1e-6,
            fmt='.', lw=0.7, label="Data", color='steelblue', zorder=5)

ax.axhline(0., ls=':', c='k', lw=2)
ax.grid(alpha=0.4)
ax.set_xlabel("Epochs [MJD]", fontsize=15)
ax.set_ylabel("Time delays / Timing residuals [s]", fontsize=14)
ax.tick_params(labelsize=13)
ax.legend(fontsize=13)
ax.set_title("Time-domain reconstruction of red noise", fontsize=14)
plt.tight_layout()
plt.show()
Question: The green envelope does not perfectly trace the red injection everywhere. Why? Where do you expect the reconstruction to be most uncertain, and why?
Click to reveal the answer
The reconstruction is most uncertain at the edges of the dataset. Near the start and end of the observation span, the Fourier modes are less well constrained because there are fewer data points to anchor them. In the middle of the dataset, where data coverage is dense, the reconstruction is tighter. This is a fundamental limitation of the Fourier-basis approach: it cannot extrapolate beyond the data.

8. Prior sensitivity¶

Bayesian inference always depends on the choice of prior. For well-measured parameters (those where the likelihood is much narrower than the prior) the choice of prior matters little. But for weakly constrained parameters, the prior can dominate the posterior.

Let's test this. We will re-run the analysis with a narrow prior on the red noise amplitude, deliberately placing it away from the true value, and compare the recovered posterior to what we got before.

This is a common mistake in practice: if you are not careful about prior boundaries, you can accidentally exclude the true parameter value, or bias the recovery.

In [ ]:
# --- Deliberately narrow prior, shifted away from the injected value ---
log10_A_rn_narrow = parameter.Uniform(-18, -15)   # injected value is ~-13.15, well outside!
gamma_rn_narrow   = parameter.Uniform(0, 7)
pl_rn_narrow      = powerlaw(log10_A=log10_A_rn_narrow, gamma=gamma_rn_narrow)

log10_A_dm_prior2 = parameter.Uniform(-18, -10)
gamma_dm_prior2   = parameter.Uniform(0, 7)
pl_dm2            = powerlaw(log10_A=log10_A_dm_prior2, gamma=gamma_dm_prior2)

s2 = gp_signals.TimingModel()
s2 += white_signals.MeasurementNoise(efac=parameter.Uniform(0.1, 5))
s2 += gp_signals.FourierBasisGP(pl_rn_narrow, components=30)
dm_basis2 = createfourierdesignmatrix_dm(nmodes=30)
s2 += gp_signals.BasisGP(pl_dm2, dm_basis2, name="dm_gp")

pta2 = PTA([s2(psr)])

outdir2 = "./chain_narrow_prior/"
os.makedirs(outdir2, exist_ok=True)
with open(os.path.join(outdir2, "pars.txt"), "w") as fout:
    for pname in pta2.param_names:
        fout.write(pname + "\n")

x0_2   = np.hstack([p.sample() for p in pta2.params])
cov2   = np.diag(np.ones(len(x0_2)) * 0.01**2)
sampler2 = ptmcmc(len(x0_2), pta2.get_lnlikelihood, pta2.get_lnprior, cov2, outDir=outdir2)

jp2 = JumpProposal(pta2, None, empirical_distr=None)
sampler2.addProposalToCycle(jp2.draw_from_prior, 5)
for sig_name in ["red_noise", "dm"]:
    if any([sig_name in p for p in pta2.param_names]):
        sampler2.addProposalToCycle(jp2.draw_from_par_prior(sig_name), 10)

sampler2.sample(x0_2, int(nsamples), SCAMweight=40, DEweight=60, AMweight=20)
/opt/anaconda3/envs/SPNA_env/lib/python3.9/site-packages/enterprise/signals/utils.py:875: RuntimeWarning: invalid value encountered in divide
  nmat = Mmat / norm
/opt/anaconda3/envs/SPNA_env/lib/python3.9/site-packages/enterprise/signals/parameter.py:62: RuntimeWarning: divide by zero encountered in log
  logpdf = np.log(self.prior(value, **kwargs))
Finished 20.00 percent in 8.108380 s Acceptance rate = 0.194256Adding DE jump with weight 60
Finished 98.00 percent in 33.278824 s Acceptance rate = 0.331347
Run Complete
In [ ]:
ch2   = np.loadtxt("%s/chain_1.txt" % outdir2)
ch2   = ch2[int(len(ch2) * 0.5):]
pars2 = np.loadtxt("%s/pars.txt" % outdir2, dtype=str)

# Compare the red noise amplitude posteriors
idx_A_rn  = list(pars).index('J1909-3744_red_noise_log10_A')
idx_A_rn2 = list(pars2).index('J1909-3744_red_noise_log10_A')

fig, ax = plt.subplots(figsize=(9, 5))
ax.hist(ch[:, idx_A_rn],  bins=40, density=True, histtype='step',
        color='steelblue', lw=2, label="Wide prior [-18, -10]")
ax.hist(ch2[:, idx_A_rn2], bins=40, density=True, histtype='step',
        color='coral', lw=2, label="Narrow prior [-18, -15]")
ax.axvline(np.log10(A_rn), c='k', lw=2, ls='--', label=f"Injected = {np.log10(A_rn):.2f}")
ax.set_xlabel(r"$\log_{10} A_{\rm rn}$", fontsize=14)
ax.set_ylabel("Posterior density", fontsize=13)
ax.legend(fontsize=12)
ax.set_title("Effect of prior choice on RN amplitude recovery", fontsize=13)
ax.grid(alpha=0.3)
plt.tight_layout()
plt.show()

print(f"Wide prior   — median log10(A_rn) = {np.median(ch[:, idx_A_rn]):.3f}")
print(f"Narrow prior — median log10(A_rn) = {np.median(ch2[:, idx_A_rn2]):.3f}")
print(f"Injected                          = {np.log10(A_rn):.3f}")
Wide prior   — median log10(A_rn) = -13.284
Narrow prior — median log10(A_rn) = -15.023
Injected                          = -13.155
Question: The narrow prior places a hard boundary at $\log_{10} A = -15$, well below the injected value of $\approx -13.15$. What do you expect the posterior to look like? What does it actually look like, and why?
Click to reveal the answer
Since the prior places a hard upper boundary at $\log_{10} A = -15$, which is well below the injected value of $-13.15$, the likelihood is simply not allowed to explore the region where the true signal lives. The posterior piles up against the prior boundary at $-15$: the sampler is doing its best within the allowed range, but the data cannot pull it toward the true value because the prior forbids it. The result is a posterior that tells you nothing about the signal. It is entirely prior-dominated and prior-truncated.

9. Model comparison: do we need DM variations?¶

So far we have assumed the correct model: EFAC + red noise + DM variations. But in practice, we do not always know which noise components are present. A key question is: does the data actually require a DM term, or can a simpler model without it fit equally well?

We answer this by fitting a reduced model that has no DM component, and comparing the recovered red noise parameters to those from the full model. If the reduced model absorbs the DM power into the red noise term, biasing its amplitude and spectral index upward, that is a sign that the DM component is genuinely required.

This is a simplified version of Bayesian model selection. A full analysis would compute the Bayesian evidence (marginal likelihood) for each model and take their ratio. Here we simply compare posteriors and look for shifts.

In [ ]:
# --- Model without DM variations ---
log10_A_rn_3 = parameter.Uniform(-18, -10)
gamma_rn_3   = parameter.Uniform(0, 7)
pl_rn_3      = powerlaw(log10_A=log10_A_rn_3, gamma=gamma_rn_3)

s3 = gp_signals.TimingModel()
s3 += white_signals.MeasurementNoise(efac=parameter.Uniform(0.1, 5))
s3 += gp_signals.FourierBasisGP(pl_rn_3, components=30)
# <-- no DM component

pta3 = PTA([s3(psr)])

outdir3 = "./chain_no_dm/"
os.makedirs(outdir3, exist_ok=True)
with open(os.path.join(outdir3, "pars.txt"), "w") as fout:
    for pname in pta3.param_names:
        fout.write(pname + "\n")

x0_3     = np.hstack([p.sample() for p in pta3.params])
cov3     = np.diag(np.ones(len(x0_3)) * 0.01**2)
sampler3 = ptmcmc(len(x0_3), pta3.get_lnlikelihood, pta3.get_lnprior, cov3, outDir=outdir3)

jp3 = JumpProposal(pta3, None, empirical_distr=None)
sampler3.addProposalToCycle(jp3.draw_from_prior, 5)
if any(["red_noise" in p for p in pta3.param_names]):
    sampler3.addProposalToCycle(jp3.draw_from_par_prior("red_noise"), 10)

sampler3.sample(x0_3, int(nsamples), SCAMweight=40, DEweight=60, AMweight=20)
Finished 20.00 percent in 4.688360 s Acceptance rate = 0.392822Adding DE jump with weight 60
Finished 98.00 percent in 19.644840 s Acceptance rate = 0.482714
Run Complete
In [ ]:
ch3   = np.loadtxt("%s/chain_1.txt" % outdir3)
ch3   = ch3[int(len(ch3) * 0.5):]
pars3 = np.loadtxt("%s/pars.txt" % outdir3, dtype=str)

idx_A_full = list(pars).index('J1909-3744_red_noise_log10_A')
idx_G_full = list(pars).index('J1909-3744_red_noise_gamma')
idx_A_nodm = list(pars3).index('J1909-3744_red_noise_log10_A')
idx_G_nodm = list(pars3).index('J1909-3744_red_noise_gamma')

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

axes[0].hist(ch[:, idx_A_full],  bins=40, density=True, histtype='step',
             color='steelblue', lw=2, label="Full model (RN + DM)")
axes[0].hist(ch3[:, idx_A_nodm], bins=40, density=True, histtype='step',
             color='coral', lw=2, label="Reduced model (RN only)")
axes[0].axvline(np.log10(A_rn), c='k', lw=2, ls='--', label="Injected RN")
axes[0].set_xlabel(r"$\log_{10} A_{\rm rn}$", fontsize=13)
axes[0].set_ylabel("Posterior density", fontsize=12)
axes[0].legend(fontsize=10)
axes[0].grid(alpha=0.3)

axes[1].hist(ch[:, idx_G_full],  bins=40, density=True, histtype='step',
             color='steelblue', lw=2, label="Full model (RN + DM)")
axes[1].hist(ch3[:, idx_G_nodm], bins=40, density=True, histtype='step',
             color='coral', lw=2, label="Reduced model (RN only)")
axes[1].axvline(gamma_rn, c='k', lw=2, ls='--', label="Injected gamma")
axes[1].set_xlabel(r"$\gamma_{\rm rn}$", fontsize=13)
axes[1].legend(fontsize=10)
axes[1].grid(alpha=0.3)

fig.suptitle("Model comparison: full vs reduced model (no DM)", fontsize=13)
plt.tight_layout()
plt.show()

print("Red noise amplitude (log10 A):")
print(f"  Injected         = {np.log10(A_rn):.3f}")
print(f"  Full model       = {np.median(ch[:,  idx_A_full]):.3f}")
print(f"  Reduced model    = {np.median(ch3[:, idx_A_nodm]):.3f}")
print()
print("Red noise spectral index (gamma):")
print(f"  Injected         = {gamma_rn:.3f}")
print(f"  Full model       = {np.median(ch[:,  idx_G_full]):.3f}")
print(f"  Reduced model    = {np.median(ch3[:, idx_G_nodm]):.3f}")
Red noise amplitude (log10 A):
  Injected         = -13.155
  Full model       = -13.284
  Reduced model    = -12.816

Red noise spectral index (gamma):
  Injected         = 3.000
  Full model       = 3.399
  Reduced model    = 2.544
Question: Does omitting the DM term bias the red noise recovery? In which direction, and why? What would this mean if you were searching for a gravitational wave background with this pulsar?
Click to reveal the answer
Yes. When the DM component is absent from the model, the sampler has to explain the DM-induced residuals using the only other low-frequency term it has: the red noise. This typically drives the red noise amplitude upward (the model absorbs extra power) and may also shift the spectral index, since DM noise and achromatic red noise have different spectral shapes.

For a gravitational wave search, this matters enormously. Any unmodelled noise that leaks into the red noise term will inflate the apparent red noise amplitude, potentially masking or mimicking a gravitational wave signal.

10. Exercises (Bonus activity)¶

These exercises build on the analysis above. They range from quick parameter changes to more open-ended investigations.

Exercise 1: change the injected noise level¶

Re-run the full analysis with a weaker red noise signal ($A_{\rm rn} = 1 \times 10^{-15}$, $\gamma_{\rm rn} = 3$). What happens to the posterior? Is the parameter recovered, or does the posterior just reflect the prior?

Exercise 2: single-frequency data¶

Re-simulate the pulsar using only a single observing frequency (freqs = np.full(len(toas), 1400)). How does this affect the recovery of DM variations? Can you still separate red noise from DM noise? Why or why not?

Exercise 3: EQUAD¶

The current white noise model includes only EFAC. Modify the Enterprise model to also include EQUAD (white_signals.MeasurementNoise(efac=..., equad=...)). Does adding this extra parameter change the recovery of the red noise parameters? Does the EQUAD posterior peak near zero (as expected, since we did not inject any)?

Exercise 4: longer observation span¶

Double the observation span by changing the TOA grid to np.arange(52000, 66000, cadence). How does the longer baseline affect the power spectrum recovery at the lowest frequencies? Does the time-domain reconstruction improve near the edges of the dataset?