Saved notebook contents
Code, text, and outputs are displayed as saved in the repository. This documentation build does not execute or revalidate this study.
This notebook has saved text outputs but no embedded figures. The paper PDFs referenced by its code are not included in this checkout.
The binned likelihood against the exact one, sampled
The paper's pointwise checks compare the hybrid grid with the exact fine-grid likelihood at a point (\(\log\mathcal{L}\) offset \(0.086\)), over a \(3\sigma\) ball (drift \(\le 0.023\)), and at the Laplace level. None of them compares the two sampled posteriors on the same data with everything else held fixed.
This notebook closes that. Same data, same priors, same sampler, same starting procedure --- the only thing that changes is whether the likelihood sums over all \(93\,697\) fine bins or over \(1\,725\) coarse blocks plus a \(384\)-bin window.
| chain | likelihood | TDI | seed | role |
|---|---|---|---|---|
hyb1_a |
hybrid | 1 | 1 | the paper's analysis |
hyb1_b |
hybrid | 1 | 2 | control: the sampler's own scatter |
full1 |
exact | 1 | 1 | the measurement |
hyb2_a |
hybrid | 2 | 1 | the same, in the second generation |
hyb2_b |
hybrid | 2 | 2 | control, TDI-2 |
full2 |
exact | 2 | 1 | the null test, with no binning anywhere |
The control is the point. Two hybrid chains that differ only in their seed give the scale on which any two chains of this length disagree; if hybrid-against-exact sits inside that, the binning is invisible to the sampler and no autocorrelation modelling is needed to say so.
import os, sys, time
os.environ.setdefault("JAX_PLATFORMS", "cpu") # determinism: GPU reductions are not
# bit-reproducible, and a 1e-6 wobble in
# logL flips accept/reject and forks the
# chain. On the CPU this notebook is
# byte-reproducible run to run.
from pathlib import Path
import numpy as np
import jax, jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
import matplotlib.pyplot as plt
REPO = next(p for p in [Path.cwd(), *Path.cwd().parents] if (p / "pyproject.toml").exists())
GB = REPO / "notebooks" / "glitch_GB"
for p in (REPO / "notebooks", GB, REPO / "paper" / "validation"):
sys.path.insert(0, str(p))
import noise as ns
import fd_pipeline as fp
from _style import C, FULL_IN, save, set_source
set_source("notebooks/glitch_only/exact_vs_hybrid.ipynb")
DATASET = GB / "dataset.npz"
print("jax:", jax.default_backend(), jax.devices())
print("dataset:", DATASET.relative_to(REPO))
jax: cpu [CpuDevice(id=0)]
dataset: notebooks/glitch_GB/dataset.npz
The two likelihoods
fp.build_hybrid is the production likelihood of the paper. fp.build_full is the
reference: every bin, no blocks, no stride. Both are built from the one stored data
stream, so the two see identical numbers down to the noise draw.
DS = np.load(DATASET)
grid = fp.make_grid(int(DS["N"]), float(DS["DT"]))
model, n_gb = fp.make_gb_model(grid["T_OBS"], int(DS["N_GB"]))
coords = fp.Coords(free_sky=False)
LABELS = list(coords.names)
DIM = len(LABELS)
th_true = coords.to_sampling(jnp.asarray(DS["gb_true"]), jnp.asarray(DS["glitch_true"]))
BOUNDS = {"log_f0": (np.log(1e-4), np.log(3e-3)), "log_fdot": (np.log(1e-22), np.log(1e-15)),
"log_A_gb": (np.log(1e-25), np.log(1e-20)), "psi": (0.0, float(np.pi)),
"t0": (0.0, float(grid["T_OBS"])),
"log_Ag": (float(np.log(1e-17)), float(np.log(5e-3))),
"log_tau": (float(np.log(0.1)), float(np.log(5e4)))}
log_prior = fp.make_log_prior(coords, BOUNDS)
# TDI-2 is the same physical realisation pushed through h2 = (1 - D^4) h1, signal and
# noise alike, with the PSD transformed to match -- as in run_convergence.py.
TFX1 = -1.0 + jnp.exp(-4j * float(DS["T_ARM"]) * 2.0 * jnp.pi * grid["f_safe"])
data1 = jnp.asarray(DS["data_tdi1"])
LIK = {}
for tdi in (1, 2):
data = data1 if tdi == 1 else (-TFX1[:, None] * data1).at[0].set(0 + 0j)
psd = (ns.psd_tdi1_array if tdi == 1 else ns.psd_tdi2_array)(
grid["f_safe"], t_obs=grid["T_OBS"])
hyb, mh = fp.build_hybrid(grid, data, psd, int(DS["k_min"]), model, n_gb, coords, tdi=tdi)
ful, mf = fp.build_full(grid, data, psd, model, n_gb, coords, tdi=tdi)
LIK[("hybrid", tdi)], LIK[("exact", tdi)] = hyb, ful
if tdi == 1:
N_EVAL_HYB = mh["n_blocks"] + mh["n_win"]
N_EVAL_FUL = mf["n_evals"]
print(f"hybrid: {mh['n_blocks']} blocks + {mh['n_win']} window bins "
f"= {N_EVAL_HYB} evaluations/call")
print(f"exact : {N_EVAL_FUL:,} evaluations/call "
f"(x{N_EVAL_FUL / N_EVAL_HYB:.0f})")
print()
for tdi in (1, 2):
a = float(LIK[("hybrid", tdi)](th_true)); b = float(LIK[("exact", tdi)](th_true))
print(f"TDI-{tdi} at the injection: hybrid {a:14.3f} exact {b:14.3f} "
f"offset {a - b:+.4f}")
/home/giorgio/Desktop/jaxglitches/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
/home/giorgio/Desktop/jaxglitches/.venv/lib/python3.12/site-packages/lisaconstants/compat/astropy.py:252: UserWarning: The following constants differ between lisaconstants and the version of astropy you have installed: VACUUM_PERMEABILITY. The recommended version of astropy is 7.2.0. Use a different one at your own risks.
You may also open an issue at https://gitlab.esa.int/lisa-sgs/commons/lisa-constants to warn that lisaconstants is not compatible with astropy v8.0.1
warnings.warn(
hybrid: 1725 blocks + 384 window bins = 2109 evaluations/call
exact : 93,697 evaluations/call (x44)
TDI-1 at the injection: hybrid -280853.921 exact -280854.007 offset +0.0855
TDI-2 at the injection: hybrid -280853.942 exact -280854.007 offset +0.0647
Cost, measured
The marginal cost per call, after compilation. This is the whole reason the hybrid exists, and it is worth having the number next to the accuracy it buys.
def per_call(f, n=100):
f(th_true).block_until_ready()
best = np.inf
for _ in range(3):
t0 = time.perf_counter()
for _ in range(n):
out = f(th_true)
out.block_until_ready()
best = min(best, (time.perf_counter() - t0) / n)
return best
COST = {k: per_call(v) for k, v in LIK.items() if k[1] == 1}
for (kind, tdi), t in COST.items():
print(f"{kind:7s} {t * 1e6:9.1f} us/call")
print(f"\nratio {COST[('exact', 1)] / COST[('hybrid', 1)]:.1f}x on one call "
f"(the sampler batches 16 walkers, where the ratio is larger)")
hybrid 255.9 us/call
exact 2280.8 us/call
ratio 8.9x on one call (the sampler batches 16 walkers, where the ratio is larger)
Maximum and Laplace width
fp.laplace is the damped Newton iteration used everywhere else in the paper. Running it
on both likelihoods gives the Gaussian-approximation comparison before any sampling, and
the widths it returns are what the walkers are initialised from.
FALLBACK = jnp.asarray([1e-6, 0.3, 0.02, 0.05, 1.0, 0.05, 0.05])
MAP, SIG = {}, {}
for key, lik in LIK.items():
post = jax.jit(lambda t, f=lik: f(t) + log_prior(t))
x, s = fp.laplace(post, th_true, FALLBACK)
MAP[key], SIG[key] = np.asarray(x), np.asarray(s)
print(f"{'parameter':10s}{'sigma exact':>14}{'sigma ratio':>13}{'MAP shift/sigma':>17}")
for i, nm in enumerate(LABELS):
se, sh = SIG[("exact", 1)][i], SIG[("hybrid", 1)][i]
print(f"{nm:10s}{se:14.4g}{sh / se:13.5f}"
f"{(MAP[('hybrid', 1)][i] - MAP[('exact', 1)][i]) / se:17.2e}")
parameter sigma exact sigma ratio MAP shift/sigma
log_f0 1.176e-07 1.00000 1.25e-05
log_fdot 0.2043 1.00000 -1.10e-05
log_A_gb 0.003522 1.00000 7.77e-06
psi 0.003783 1.00000 -1.52e-05
t0 7.135 0.99987 3.01e-03
log_Ag 0.0627 0.99978 8.94e-04
log_tau 0.03085 0.99975 2.61e-04
The chains
Sixteen walkers, \(2\,000\) burn-in and \(20\,000\) retained iterations each, exactly as in
glitch_and_gb.ipynb. The two hybrid TDI-1 chains differ only in their seed, and they
are the control.
NWALK, NBURN, NSAMP = 16, 2_000, 20_000
RUNS = (("hyb1_a", "hybrid", 1, 1), ("hyb1_b", "hybrid", 1, 2), ("full1", "exact", 1, 1),
("hyb2_a", "hybrid", 2, 1), ("hyb2_b", "hybrid", 2, 2), ("full2", "exact", 2, 1))
# Cached per chain, and rewritten after each one. The chains are deterministic on the
# CPU -- same seed, same likelihood, same code -> same samples -- so the cache is a time
# saver and not a second source of truth: deleting it reproduces every number below.
# It also means an interrupted run keeps what it had already finished, which matters
# when two of these six take eleven minutes each.
CACHE = Path("exact_vs_hybrid_chains.npz")
CH = {k: v for k, v in np.load(CACHE).items()} if CACHE.exists() else {}
if CH:
print(f"loaded {len(CH)} chain(s) from {CACHE}: {', '.join(sorted(CH))}\n")
for name, kind, tdi, seed in RUNS:
if name in CH:
print(f"{name:7s} {kind:7s} TDI-{tdi} seed {seed} cached", flush=True)
continue
t0 = time.perf_counter()
CH[name] = np.asarray(fp.run_chain(
LIK[(kind, tdi)], log_prior, jnp.asarray(MAP[(kind, tdi)]),
jnp.asarray(SIG[(kind, tdi)]), DIM, seed,
nwalkers=NWALK, nburn=NBURN, nsamp=NSAMP))
np.savez_compressed(CACHE, **CH)
print(f"{name:7s} {kind:7s} TDI-{tdi} seed {seed} "
f"{CH[name].shape[0]:,} samples in {time.perf_counter() - t0:6.1f} s", flush=True)
loaded 6 chain(s) from exact_vs_hybrid_chains.npz: full1, full2, hyb1_a, hyb1_b, hyb2_a, hyb2_b
hyb1_a hybrid TDI-1 seed 1 cached
hyb1_b hybrid TDI-1 seed 2 cached
full1 exact TDI-1 seed 1 cached
hyb2_a hybrid TDI-2 seed 1 cached
hyb2_b hybrid TDI-2 seed 2 cached
full2 exact TDI-2 seed 1 cached
What the two posteriors do differently
Two statistics per parameter: the ratio of posterior widths, and the shift of the medians in units of the width. Each is reported for the measurement (hybrid against exact) and for the control (hybrid against hybrid, different seed). The control is not an error bar in the formal sense --- it is one draw from the scatter, not its standard deviation --- but it is the right order of magnitude and it is measured rather than modelled.
def compare(a, b):
"""Width ratio and median shift of chain `b` relative to chain `a`."""
sa, sb = a.std(axis=0), b.std(axis=0)
return sb / sa, (np.median(b, axis=0) - np.median(a, axis=0)) / sa
R_MEAS, D_MEAS = compare(CH["hyb1_a"], CH["full1"]) # binning, TDI-1
R_CTRL, D_CTRL = compare(CH["hyb1_a"], CH["hyb1_b"]) # seed only, TDI-1
R_MEA2, D_MEA2 = compare(CH["hyb2_a"], CH["full2"]) # binning, TDI-2
R_CTR2, D_CTR2 = compare(CH["hyb2_a"], CH["hyb2_b"]) # seed only, TDI-2
R_NULL, D_NULL = compare(CH["full1"], CH["full2"]) # TDI-1 vs TDI-2, both exact
# the control band: the scale on which any two chains of this length disagree, taken
# over both generations rather than assumed transferable from one to the other
R_BAND = max(np.abs(np.r_[R_CTRL, R_CTR2] - 1.0))
D_BAND = max(np.abs(np.r_[D_CTRL, D_CTR2]))
hdr = f"{'':10s}{'exact/hyb':>11}{'seed/seed':>11}{'exact/hyb':>11}{'seed/seed':>11}"
print(f"{'':10s}{'width ratio':>22}{'':>0}{'median shift [sigma]':>22}")
print(f"{'parameter':10s}{'TDI-1':>11}{'TDI-1':>11}{'TDI-2':>11}{'TDI-2':>11}")
print(" width ratio")
for i, nm in enumerate(LABELS):
print(f"{nm:10s}{R_MEAS[i]:11.4f}{R_CTRL[i]:11.4f}{R_MEA2[i]:11.4f}{R_CTR2[i]:11.4f}")
print(" median shift [sigma]")
for i, nm in enumerate(LABELS):
print(f"{nm:10s}{D_MEAS[i]:11.4f}{D_CTRL[i]:11.4f}{D_MEA2[i]:11.4f}{D_CTR2[i]:11.4f}")
print(f"\nbinning, TDI-1: widths [{R_MEAS.min():.3f}, {R_MEAS.max():.3f}], "
f"|median shift| <= {np.abs(D_MEAS).max():.3f} sigma")
print(f"control, TDI-1: widths [{R_CTRL.min():.3f}, {R_CTRL.max():.3f}], "
f"|median shift| <= {np.abs(D_CTRL).max():.3f} sigma")
print(f"binning, TDI-2: widths [{R_MEA2.min():.3f}, {R_MEA2.max():.3f}], "
f"|median shift| <= {np.abs(D_MEA2).max():.3f} sigma")
print(f"control, TDI-2: widths [{R_CTR2.min():.3f}, {R_CTR2.max():.3f}], "
f"|median shift| <= {np.abs(D_CTR2).max():.3f} sigma")
print(f"\nTDI-1 vs TDI-2, both exact: widths [{R_NULL.min():.3f}, {R_NULL.max():.3f}], "
f"|median shift| <= {np.abs(D_NULL).max():.3f} sigma")
width ratio median shift [sigma]
parameter TDI-1 TDI-1 TDI-2 TDI-2
width ratio
log_f0 1.0063 1.0017 1.0194 1.0287
log_fdot 1.0141 1.0243 1.0313 1.0598
log_A_gb 0.9859 0.9825 1.0008 0.9997
psi 0.9845 0.9908 0.9975 1.0104
t0 0.9981 1.0020 1.0010 0.9971
log_Ag 0.9949 0.9982 0.9958 1.0187
log_tau 0.9970 1.0070 1.0069 1.0259
median shift [sigma]
log_f0 0.0203 0.0504 -0.0136 -0.0283
log_fdot -0.0108 -0.0388 0.0149 0.0286
log_A_gb -0.0390 -0.0345 0.0290 0.0224
psi -0.0159 -0.0534 -0.0131 0.0068
t0 0.0029 0.0499 -0.0093 0.0238
log_Ag -0.0027 -0.0498 -0.0311 -0.0272
log_tau -0.0117 -0.0537 -0.0256 -0.0116
binning, TDI-1: widths [0.984, 1.014], |median shift| <= 0.039 sigma
control, TDI-1: widths [0.983, 1.024], |median shift| <= 0.054 sigma
binning, TDI-2: widths [0.996, 1.031], |median shift| <= 0.031 sigma
control, TDI-2: widths [0.997, 1.060], |median shift| <= 0.029 sigma
TDI-1 vs TDI-2, both exact: widths [0.995, 1.013], |median shift| <= 0.029 sigma
Figure
Panels (a) and (b) are the two statistics above, measurement against control. Panel (c) is the glitch amplitude--duration block, the most correlated pair in the posterior and therefore the one where a systematic difference would show first.
def contour(ax, ch, ix, iy, color, label, ls="-"):
"""1- and 2-sigma contours of a 2-D marginal, from a histogram of the samples."""
x, y = ch[:, ix], ch[:, iy]
H, xe, ye = np.histogram2d(x, y, bins=60)
Hs = H.T
flat = np.sort(Hs.ravel())[::-1]
csum = np.cumsum(flat) / flat.sum()
lv = [flat[np.searchsorted(csum, q)] for q in (0.393, 0.865)][::-1]
ax.contour(0.5 * (xe[1:] + xe[:-1]), 0.5 * (ye[1:] + ye[:-1]), Hs,
levels=lv, colors=[color], linewidths=1.1, linestyles=[ls, ls])
ax.plot([], [], color=color, lw=1.1, ls=ls, label=label)
SHOW = [r"$\log f_0$", r"$\log\dot f$", r"$\log\mathcal{A}_{\rm GB}$", r"$\psi$",
r"$t_0$", r"$\log A_g$", r"$\log\tau$"]
xx = np.arange(DIM)
fig, ax = plt.subplots(1, 3, figsize=(FULL_IN, 2.35), layout="constrained")
a = ax[0]
a.axhspan(1 - R_BAND, 1 + R_BAND, color=C["grey"], alpha=0.18, lw=0,
label="seed-to-seed control")
a.axhline(1.0, color=C["grey"], lw=0.8, ls=":")
a.plot(xx, R_MEAS, "o", ms=4.5, color=C["blue"], label="TDI-1")
a.plot(xx, R_MEA2, "s", ms=4.0, mfc="none", color=C["orange"], label="TDI-2")
a.set_ylabel(r"$\sigma_{\rm exact}/\sigma_{\rm hybrid}$")
a.set_title("(a) posterior width", fontsize=7, loc="left")
b = ax[1]
b.axhspan(-D_BAND, D_BAND, color=C["grey"], alpha=0.18, lw=0,
label="seed-to-seed control")
b.axhline(0.0, color=C["grey"], lw=0.8, ls=":")
b.plot(xx, D_MEAS, "o", ms=4.5, color=C["blue"], label="TDI-1")
b.plot(xx, D_MEA2, "s", ms=4.0, mfc="none", color=C["orange"], label="TDI-2")
b.set_ylabel(r"$(\mathrm{med}_{\rm exact}-\mathrm{med}_{\rm hybrid})/\sigma$")
b.set_title("(b) posterior location", fontsize=7, loc="left")
for a_ in (ax[0], ax[1]):
a_.set_xticks(xx)
a_.set_xticklabels(SHOW, rotation=55, ha="right", fontsize=6)
a_.legend(fontsize=5.4, handlelength=1.4, borderpad=0.25, labelspacing=0.25)
a_.grid(alpha=0.25, lw=0.4)
c = ax[2]
i_ag, i_tau = LABELS.index("log_Ag"), LABELS.index("log_tau")
contour(c, CH["hyb1_a"], i_ag, i_tau, C["blue"], "hybrid")
contour(c, CH["full1"], i_ag, i_tau, C["red"], "exact", ls="--")
c.plot(float(th_true[i_ag]), float(th_true[i_tau]), "*", color="k", ms=9, zorder=5,
label="injection")
c.set_xlabel(r"$\log A_g$")
c.set_ylabel(r"$\log\tau$")
c.set_title(r"(c) the glitch block, TDI-1", fontsize=7, loc="left")
c.legend(fontsize=5.4, handlelength=1.4, borderpad=0.25, labelspacing=0.25)
c.grid(alpha=0.25, lw=0.4)
save(fig, "fig_exact_vs_hybrid", seed={n: sd for n, _, _, sd in RUNS},
inputs=[DATASET],
note="samples the hybrid and the exact fine-grid likelihood on the same stored "
"data stream, with the same sampler and starting procedure; the grey band is "
"hybrid chains differing only in seed, in both TDI generations. Pinned to the "
"CPU so the chains, and hence the figure, are reproducible")
plt.close(fig)
saved paper/figures/fig_exact_vs_hybrid.pdf
Summary
print("Same data, same priors, same sampler. Only the likelihood grid changes.\n")
print(f" evaluations per call {N_EVAL_HYB:,} (hybrid) vs {N_EVAL_FUL:,} (exact)"
f" -- x{N_EVAL_FUL / N_EVAL_HYB:.0f}")
print(f" logL offset at injection "
f"{float(LIK[('hybrid', 1)](th_true)) - float(LIK[('exact', 1)](th_true)):+.4f} (TDI-1), "
f"{float(LIK[('hybrid', 2)](th_true)) - float(LIK[('exact', 2)](th_true)):+.4f} (TDI-2)")
print(f" Laplace width ratios "
f"[{(SIG[('hybrid', 1)] / SIG[('exact', 1)]).min():.5f}, "
f"{(SIG[('hybrid', 1)] / SIG[('exact', 1)]).max():.5f}]\n")
# Both the worst parameter and the spread over all seven. The maximum of seven noisy
# statistics is itself a very noisy statistic -- comparing one max against another can
# turn over on a single parameter -- so the rms is the fairer of the two and is what the
# verdict below uses. Both are reported, because the max is what a reader will look for.
rms = lambda v: float(np.sqrt(np.mean(v ** 2)))
ROWS = (("binning, TDI-1", R_MEAS, D_MEAS), ("control, TDI-1", R_CTRL, D_CTRL),
("binning, TDI-2", R_MEA2, D_MEA2), ("control, TDI-2", R_CTR2, D_CTR2),
("null, both exact", R_NULL, D_NULL))
print(f"{'':18s}{'max|r-1|':>10}{'rms|r-1|':>10}{'max|d/sig|':>12}{'rms|d/sig|':>12}")
for nm, R, D in ROWS:
print(f"{nm:18s}{np.abs(R - 1).max():10.4f}{rms(R - 1):10.4f}"
f"{np.abs(D).max():12.4f}{rms(D):12.4f}")
print()
for tag, Rm, Dm, Rc, Dc in (("TDI-1", R_MEAS, D_MEAS, R_CTRL, D_CTRL),
("TDI-2", R_MEA2, D_MEA2, R_CTR2, D_CTR2)):
ok = rms(Rm - 1) <= rms(Rc - 1) and rms(Dm) <= rms(Dc)
print(f" {tag}: binning moves the posterior "
+ ("less than" if ok else "MORE than") + " changing the random seed does")
print(f"\n null test with no binning anywhere: "
f"widths [{R_NULL.min():.3f}, {R_NULL.max():.3f}], "
f"medians <= {np.abs(D_NULL).max():.3f} sigma")
Same data, same priors, same sampler. Only the likelihood grid changes.
evaluations per call 2,109 (hybrid) vs 93,697 (exact) -- x44
logL offset at injection +0.0855 (TDI-1), +0.0647 (TDI-2)
Laplace width ratios [0.99975, 1.00000]
max|r-1| rms|r-1| max|d/sig| rms|d/sig|
binning, TDI-1 0.0155 0.0101 0.0390 0.0187
control, TDI-1 0.0243 0.0122 0.0537 0.0477
binning, TDI-2 0.0313 0.0143 0.0311 0.0211
control, TDI-2 0.0598 0.0281 0.0286 0.0227
null, both exact 0.0130 0.0062 0.0286 0.0140
TDI-1: binning moves the posterior less than changing the random seed does
TDI-2: binning moves the posterior less than changing the random seed does
null test with no binning anywhere: widths [0.995, 1.013], medians <= 0.029 sigma