diff --git a/dingo/gw/injection.py b/dingo/gw/injection.py index 87f7a066..b6c07111 100644 --- a/dingo/gw/injection.py +++ b/dingo/gw/injection.py @@ -1,3 +1,5 @@ +from typing import Optional + import numpy as np from bilby.gw.detector import InterferometerList from torchvision.transforms import Compose @@ -442,3 +444,15 @@ def random_injection(self): k: float(v) for k, v in theta.items() } # Some parameters are np.float64 return self.injection(theta) + + +def safe_signal(gw_signal: "GWSignal", theta: dict) -> Optional[dict]: + """Call ``gw_signal.signal(theta)``, returning ``None`` if generation fails. + + Lets a caller drop a failed draw rather than abort a whole + :class:`multiprocessing.Pool` batch. Module level, so it is picklable for the pool. + """ + try: + return gw_signal.signal(theta) + except Exception: + return None diff --git a/dingo/gw/result.py b/dingo/gw/result.py index 51616032..e2579c5e 100644 --- a/dingo/gw/result.py +++ b/dingo/gw/result.py @@ -1,6 +1,8 @@ import copy import time -from typing import Optional +import warnings +from functools import partial +from typing import Optional, Tuple import numpy as np import yaml @@ -18,7 +20,9 @@ from dingo.gw.domains import MultibandedFrequencyDomain from dingo.gw.domains import build_domain from dingo.gw.gwutils import get_extrinsic_prior_dict +from dingo.gw.injection import safe_signal from dingo.gw.likelihood import StationaryGaussianGWLikelihood +from dingo.gw.utils import plotting as gw_plotting from dingo.gw.prior import build_prior_with_defaults from dingo.core.utils.backward_compatibility import check_minimum_version @@ -781,3 +785,177 @@ def pesummary_prior(self): except AttributeError: continue return prior + + def _compute_ppd( + self, + num_waveforms: int = 5000, + num_processes: int = 1, + ) -> Tuple[dict, dict, np.ndarray]: + """Generate whitened time-domain waveforms for the pointwise strain PPD. + + Projects posterior draws onto the detectors via the likelihood (already whitened, + ``h / asd / noise_std``) and inverse-FFTs them to the time domain with bilby's + whitened inverse FFT. Each of the two posteriors is represented by an + **equally-weighted** set of draws, as :func:`~dingo.gw.utils.plotting.pointwise_hdi` + requires: ``"dingo"`` is a uniform subsample and, when the result is + importance-sampled, ``"dingo-is"`` is a rejection sample of the weighted draws. + + Parameters + ---------- + num_waveforms : int + Upper bound on the number of posterior draws per mode (a waveform-generation + budget). The ``"dingo-is"`` rejection sample is often smaller, being ESS-limited. + num_processes : int + Processes used for waveform generation. + + Returns + ------- + wf_td : dict + ``{mode: {ifo: (n_kept, n_times) real}}`` whitened time-domain waveforms. + data_td : dict + ``{ifo: (n_times,) real}`` whitened time-domain detector data. + times : numpy.ndarray + ``(n_times,)`` GPS times of the analysed segment less ``t_ref``, so ``t = 0`` + is the trigger and the axis spans ``[time_buffer - duration, time_buffer]``. + """ + if getattr(self, "likelihood", None) is None: + self._build_likelihood() + + domain = ( + self.domain.base_domain + if hasattr(self.domain, "base_domain") + else self.domain + ) + ifos = list(self.context["waveform"].keys()) + + # download_strain_data_in_FD cut the segment as + # [t_ref + time_buffer - duration, t_ref + time_buffer] and then cyclically shifted + # it by time_buffer, so index 0 of the inverse FFT is t_ref itself and everything + # before the event is wrapped around to the end of the array. + duration = 1 / domain.delta_f + sampling_frequency = 2 * domain.f_max + time_buffer = (self.event_metadata or {}).get("time_buffer", duration / 2) + start_time = self.t_ref + time_buffer - duration + + # The likelihood's bilby interferometers are otherwise bare, so bilby's whitened + # inverse FFT cannot be called on them. sampling_frequency = 2 * f_max makes bilby's + # one-sided frequency array coincide with Dingo's domain(). + for ifo in self.likelihood.ifo_list: + ifo.set_strain_data_from_zero_noise( + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + ) + ifo.minimum_frequency = domain.f_min + ifo.maximum_frequency = domain.f_max + interferometers = {ifo.name: ifo for ifo in self.likelihood.ifo_list} + # Undo that cyclic shift, so sample j sits at start_time + j / f_s. + roll = int(round(time_buffer * sampling_frequency)) + + def to_td(fd, ifo): + """Whitened frequency-domain array(s) -> time domain (bilby convention).""" + td = interferometers[ + ifo + ].get_whitened_time_series_from_whitened_frequency_series(np.asarray(fd)) + return np.roll(td, -roll, axis=-1) + + data_fd = self.likelihood.whitened_strains + if len(data_fd[ifos[0]]) != len(domain()): + raise ValueError( + "Whitened data length does not match the inverse-FFT domain; the " + "time-domain PPD requires a uniform frequency domain." + ) + data_td = {ifo: to_td(data_fd[ifo], ifo) for ifo in ifos} + # Absolute GPS times, less the trigger time: t = 0 is the event. + times = interferometers[ifos[0]].time_array - self.t_ref + + # Restrict to the prior support, as importance_sample does: the flow proposal leaks + # a few unphysical draws (e.g. tiny negative spin magnitudes) that fail generation. + in_prior = np.isfinite( + np.asarray( + self.prior.ln_prob(self.samples[self.search_parameter_keys], axis=0), + dtype=float, + ) + ) + + def cap(samples): + """Bound the per-mode waveform-generation cost, without introducing repeats.""" + n = min(num_waveforms, len(samples)) + return samples.sample(n=n, random_state=RANDOM_STATE) + + mode_samples = {"dingo": cap(self.samples[in_prior])} + if "weights" in self.samples: + # max_samples_per_draw=1 => no duplicate waveforms (repeated rows would narrow + # the HDI around whichever draws happened to be duplicated); zero-weight + # (out-of-prior) draws are never accepted. clip_weights is essential, not an + # optimisation: unclipped, the yield + # is sum(w)/max(w), which a single large weight can collapse to a few dozen + # draws regardless of the ESS. See rejection_sample for the clipping scheme. + mode_samples["dingo-is"] = cap( + self.rejection_sample( + max_samples_per_draw=1, clip_weights=True, random_state=RANDOM_STATE + ) + ) + + def mode_td(samples): + signals = apply_func_with_multiprocessing( + partial(safe_signal, self.likelihood), + samples, + num_processes=num_processes, + ) + rows = [s for s in signals if s is not None] + if not rows: + raise RuntimeError( + "All PPD waveform generations failed; cannot build the PPD." + ) + n_failed = len(samples) - len(rows) + if n_failed: + warnings.warn( + f"{n_failed} in-prior PPD waveform(s) failed to generate and were " + "dropped from the distribution.", + stacklevel=2, + ) + return { + ifo: to_td(np.array([s["waveform"][ifo] for s in rows]), ifo) + for ifo in ifos + } + + wf_td = {mode: mode_td(samples) for mode, samples in mode_samples.items()} + return wf_td, data_td, times + + def plot_ppd_td( + self, + filename: str = "ppd_td.png", + num_waveforms: int = 5000, + num_processes: int = 1, + hdi_level: float = 0.9, + num_plotted_draws: int = 0, + **kwargs, + ) -> Tuple[dict, dict, np.ndarray]: + """Plot the time-domain whitened-strain posterior-predictive distribution. + + One stacked panel per (mode, detector), each filling the *pointwise* credible band + of ``p(h(t)|d)`` over the raw whitened data. See + :func:`dingo.gw.utils.plotting.plot_ppd_td`. + + ``hdi_level`` is the credible level of the band (default 90%). + ``num_plotted_draws`` additionally overlays that many individual waveform draws + as faint traces. + The time axis is measured from ``t_ref``, whose GPS value is written into the + title. Remaining keyword arguments (the band and data colours, the number of + overlaid draws) are passed through. Returns the ``(wf_td, data_td, times)`` tuple. + """ + wf_td, data_td, times = self._compute_ppd( + num_waveforms=num_waveforms, num_processes=num_processes + ) + gw_plotting.plot_ppd_td( + wf_td, + data_td, + times, + filename=filename, + hdi_level=hdi_level, + num_plotted_draws=num_plotted_draws, + trigger_time=self.t_ref, + **kwargs, + ) + return wf_td, data_td, times diff --git a/dingo/gw/utils/__init__.py b/dingo/gw/utils/__init__.py new file mode 100644 index 00000000..630ed375 --- /dev/null +++ b/dingo/gw/utils/__init__.py @@ -0,0 +1 @@ +"""GW-specific utilities.""" diff --git a/dingo/gw/utils/plotting.py b/dingo/gw/utils/plotting.py new file mode 100644 index 00000000..1cc497f1 --- /dev/null +++ b/dingo/gw/utils/plotting.py @@ -0,0 +1,194 @@ +"""Time-domain strain posterior-predictive-distribution (PPD) plotting for GW results. + +GW-specific counterpart to :mod:`dingo.core.utils.plotting`. At each time sample it takes +the highest-density interval (HDI) of the whitened strain ``h(t)`` over the +posterior-predictive draws produced by :meth:`dingo.gw.result.Result._compute_ppd`, and +fills between its edges. This is the *pointwise* credible band, as opposed to the central +percentile band shaded by bilby's ``plot_interferometer_waveform_posterior``. +""" + +from typing import Dict, Optional, Tuple + +import numpy as np +from matplotlib import colors as mcolors +from matplotlib import pyplot as plt + + +def pointwise_hdi(td: np.ndarray, level: float) -> np.ndarray: + """Pointwise highest-density interval (HDI) of ``p(h(t) | d)`` at each time sample. + + The narrowest interval containing a ``level`` fraction of the draws, found by sliding a + window of that many draws over the sorted column and keeping the narrowest -- the + unimodal algorithm of :func:`arviz.hdi`, vectorised over the time axis. It assumes + ``p(h(t)|d)`` is unimodal at fixed ``t``, which whitened-strain draws are in practice; + on a multimodal column it bridges the gap between the modes rather than resolving them. + There is no weights argument, so ``td`` must be equally weighted draws. + + Parameters + ---------- + td : numpy.ndarray + ``(n_draws, n_times)`` real whitened time-domain waveforms, equally weighted. + level : float + Credible level in ``(0, 1)``, e.g. ``0.9``. + + Returns + ------- + numpy.ndarray + ``(n_times, 2)`` lower and upper interval edges. + """ + td = np.sort(np.asarray(td, dtype=float), axis=0) + n_draws = td.shape[0] + # Number of sorted draws spanned by the interval; its two edges are that far apart. + span = int(np.floor(level * n_draws)) + widths = td[span:] - td[: n_draws - span] + start = np.argmin(widths, axis=0) + cols = np.arange(td.shape[1]) + return np.stack((td[start, cols], td[start + span, cols]), axis=-1) + + +def plot_ppd_td( + wf_td: Dict[str, Dict[str, np.ndarray]], + data_td: Dict[str, np.ndarray], + times: np.ndarray, + filename: str = "ppd_td.png", + hdi_level: float = 0.9, + num_plotted_draws: int = 0, + trigger_time: Optional[float] = None, + band_colors: Optional[Dict[str, str]] = None, + data_color: str = "#555555", +) -> np.ndarray: + """Plot the time-domain whitened-strain PPD as pointwise credible bands. + + One panel per ``(mode, detector)``, stacked vertically. Each fills the pointwise HDI of + ``p(h(t)|d)`` (:func:`pointwise_hdi`) over the raw whitened data, drawn as a faint grey + trace on the whitened-noise scale (bilby/LVK convention). ``t = 0`` is the trigger. + + Everything is drawn over the full ``times`` array; the axes merely open on the last + second of inspiral through ringdown, with the y-axis auto-scaled to the whitened noise + over that view. Both limits are ordinary matplotlib state, so a caller wanting a + different window sets it on the returned axes rather than through an argument here -- + the figure is closed once saved, but its axes stay live:: + + axes = plot_ppd_td(wf_td, data_td, times) + axes[0].set_xlim(-0.3, 0.1) # sharex, so this moves every panel + axes[0].figure.savefig("ppd_td_mergerzoom.png", dpi=200, bbox_inches="tight") + + Parameters + ---------- + wf_td : dict + ``{mode: {ifo: (n_draws, n_times) real}}`` whitened time-domain draws. + data_td : dict + ``{ifo: (n_times,) real}`` whitened detector data; its keys set the detectors. + times : numpy.ndarray + ``(n_times,)`` time axis in seconds relative to the trigger (``t = 0``). + filename : str + Output path for the saved figure. + hdi_level : float + Credible level of the filled band, in ``(0, 1)``. + num_plotted_draws : int + Overlay this many individual waveform draws, as faint traces underneath the band, + taken as an evenly spaced subsample. Zero (the default) draws none: overlaying + thousands is slow to render and mostly obscures the band. + trigger_time : float or None + GPS time that ``times`` is measured from, written into the title so the axis can be + read back as an absolute time. ``None`` titles the figure generically. + band_colors : dict or None + ``{mode: colour}`` for the band, its draws and the panel label; merged over the + defaults (blue for ``"dingo"``, orange for ``"dingo-is"``), with any further mode + falling back to the matplotlib property cycle. The band's edges are drawn in a + darker shade of it, so they stay legible where overlaid draws saturate the band. + data_color : str + Colour of the whitened detector data trace. + + Returns + ------- + numpy.ndarray of the matplotlib Axes drawn onto (one per stacked (mode, detector) panel). + """ + times = np.asarray(times, dtype=float) + ifos = list(data_td.keys()) + modes = list(wf_td.keys()) + band_colors = {"dingo": "#4C72B0", "dingo-is": "#DD8452", **(band_colors or {})} + colors = { + mode: band_colors.get(mode, f"C{i}") for i, mode in enumerate(modes) + } + + # The window the axes open on. times is monotone, so it is a contiguous slice -- index + # with it rather than a boolean mask, to view the (n_draws, n_times) arrays instead of + # copying them. Only the y-scaling reads it; the traces themselves are drawn in full. + xlim = (-1.0, 0.2) + start, stop = np.searchsorted(times, xlim) + view = slice(start, stop) if stop > start else slice(None) + + panels = [(mode, ifo) for mode in modes for ifo in ifos] + fig, axes = plt.subplots( + len(panels), 1, figsize=(11, 2.6 * len(panels)), sharex=True, squeeze=False + ) + axes = axes[:, 0] + + for ax, (mode, ifo) in zip(axes, panels): + band = np.asarray(wf_td[mode][ifo]) + data = np.asarray(data_td[ifo]) + band_color = colors[mode] + edge_color = _darken(band_color) + + # Faint raw data underneath (grey noise), then the credible band on top. + ax.plot( + times, data, color=data_color, lw=0.5, alpha=0.3, zorder=1, label="data", + ) + if num_plotted_draws > 0: + step = max(1, len(band) // num_plotted_draws) + for i, draw in enumerate(band[::step][:num_plotted_draws]): + ax.plot( + times, draw, color=band_color, lw=0.4, alpha=0.15, zorder=2, + label="draws" if i == 0 else None, + ) + lower, upper = pointwise_hdi(band, hdi_level).T + ax.fill_between( + times, lower, upper, color=band_color, alpha=0.30, lw=0, zorder=3, + label=f"{hdi_level:.0%} HDI", + ) + for edge in (lower, upper): + ax.plot(times, edge, color=edge_color, lw=0.8, alpha=0.9, zorder=4) + + ax.set_xlim(*xlim) + ax.set_ylim(*_strain_range(band[:, view], data[view])) + ax.set_ylabel("whitened strain") + ax.text( + 0.01, 0.95, f"{mode} · {ifo}", transform=ax.transAxes, + va="top", ha="left", fontweight="bold", color=band_color, + ) + # A legend per panel, not just the first: the modes differ in colour, so one + # shared legend would show the wrong swatch for every panel below it. + legend = ax.legend(loc="upper right", fontsize=8, framealpha=0.9) + # The traces are deliberately faint on the axes; opaque keys keep them readable. + for handle in legend.get_lines(): + handle.set_alpha(1.0) + + axes[-1].set_xlabel("time relative to trigger (s)") + fig.suptitle( + "whitened strain PPD" + if trigger_time is None + else f"whitened strain PPD · trigger at GPS {trigger_time:.4f}" + ) + fig.savefig(filename, dpi=200, bbox_inches="tight") + plt.close(fig) + return axes + + +def _darken(color: str, factor: float = 0.65) -> Tuple[float, float, float]: + """Scale a colour's RGB channels towards black, for the band edges.""" + return tuple(factor * channel for channel in mcolors.to_rgb(color)) + + +def _strain_range(band: np.ndarray, data: np.ndarray) -> Tuple[float, float]: + """Robust, symmetric strain range on the whitened-noise scale. + + Set by the data's 99th ``|value|`` percentile (not its max, so a rare noise spike does + not blow it up), widened if the reconstruction band would otherwise clip, padded by 10%. + """ + hi = float(np.percentile(np.abs(np.asarray(data, dtype=float)), 99.0)) + env = np.percentile(np.abs(np.asarray(band, dtype=float)), 97.5, axis=0) + if env.size: + hi = max(hi, float(np.max(env))) + hi = hi * 1.1 if hi > 0 else 1.0 + return -hi, hi diff --git a/tests/gw/test_ppd.py b/tests/gw/test_ppd.py new file mode 100644 index 00000000..34528985 --- /dev/null +++ b/tests/gw/test_ppd.py @@ -0,0 +1,133 @@ +"""Tests for the pointwise strain-PPD helpers in :mod:`dingo.gw.utils.plotting`.""" + +import numpy as np + +from dingo.gw.utils.plotting import plot_ppd_td, pointwise_hdi + + +def test_pointwise_hdi_matches_normal_quantiles() -> None: + """For a symmetric unimodal (normal) column the HDI is centred on the mode.""" + rng = np.random.default_rng(9) + td = rng.normal(size=(20000, 3)) + intervals = pointwise_hdi(td, 0.5) + assert intervals.shape == (3, 2) + # Normal 50% HDI is [-0.674, 0.674] sigma, centred at 0. + assert np.allclose(intervals[:, 0], -0.674, atol=0.06) + assert np.allclose(intervals[:, 1], 0.674, atol=0.06) + + +def test_pointwise_hdi_is_narrowest_interval() -> None: + """On a skewed column the HDI is narrower than the equal-tailed interval.""" + rng = np.random.default_rng(11) + td = rng.exponential(size=(20000, 1)) + lower, upper = pointwise_hdi(td, 0.9)[0] + quantiles = np.percentile(td[:, 0], [5.0, 95.0]) + assert upper - lower < quantiles[1] - quantiles[0] + assert lower < 0.1 # the HDI of an exponential starts at the boundary + + +def test_pointwise_hdi_contains_the_requested_fraction() -> None: + """The interval holds at least ``level`` of the draws at every time sample.""" + rng = np.random.default_rng(5) + td = rng.normal(size=(1000, 25)) + lower, upper = pointwise_hdi(td, 0.9).T + inside = ((td >= lower) & (td <= upper)).mean(axis=0) + assert np.all(inside >= 0.9) + + +def test_pointwise_hdi_nested_levels() -> None: + """The 50% interval is nested inside the 90% interval at every time sample.""" + rng = np.random.default_rng(3) + td = rng.normal(size=(2000, 25)) + i90 = pointwise_hdi(td, 0.9) + i50 = pointwise_hdi(td, 0.5) + assert i90.shape == i50.shape == (25, 2) + assert np.all(i90[:, 0] <= i50[:, 0] + 1e-9) + assert np.all(i50[:, 1] <= i90[:, 1] + 1e-9) + + +def test_plot_ppd_td_renders_all_panels(tmp_path) -> None: + """One panel per (mode, ifo); the figure writes and returns that many axes.""" + rng = np.random.default_rng(4) + n, T = 200, 400 + times = np.linspace(-1.0, 0.2, T) + wf_td = { + "dingo": {"H1": rng.normal(size=(n, T)), "L1": rng.normal(size=(n, T))}, + "dingo-is": {"H1": rng.normal(size=(n, T)), "L1": rng.normal(size=(n, T))}, + } + data_td = {"H1": rng.normal(size=T), "L1": rng.normal(size=T)} + out = tmp_path / "ppd.png" + axes = plot_ppd_td(wf_td, data_td, times, filename=str(out)) + assert out.exists() + assert len(axes) == 4 # 2 modes x 2 detectors + + +def test_plot_ppd_td_labels_the_trigger_time(tmp_path) -> None: + """The trigger GPS time is written into the title, at sub-sample precision.""" + rng = np.random.default_rng(8) + T = 40 + times = np.linspace(-1.0, 0.2, T) + args = ({"dingo": {"H1": rng.normal(size=(20, T))}}, {"H1": rng.normal(size=T)}, times) + + axes = plot_ppd_td( + *args, filename=str(tmp_path / "titled.png"), trigger_time=1264316116.4 + ) + assert axes[0].figure._suptitle.get_text().endswith("GPS 1264316116.4000") + assert axes[-1].get_xlabel() == "time relative to trigger (s)" + + axes = plot_ppd_td(*args, filename=str(tmp_path / "untitled.png")) + assert "GPS" not in axes[0].figure._suptitle.get_text() + + +def test_plot_ppd_td_colors_are_per_mode(tmp_path) -> None: + """Each mode gets its own band colour; dingo is blue and dingo-is orange by default.""" + rng = np.random.default_rng(7) + n, T = 50, 40 + times = np.linspace(-1.0, 0.2, T) + wf_td = { + "dingo": {"H1": rng.normal(size=(n, T))}, + "dingo-is": {"H1": rng.normal(size=(n, T))}, + } + axes = plot_ppd_td( + wf_td, {"H1": rng.normal(size=T)}, times, filename=str(tmp_path / "colors.png") + ) + dingo_fill, dingo_is_fill = (ax.collections[0].get_facecolor()[0] for ax in axes) + assert dingo_fill[2] > dingo_fill[0] # blue channel dominates + assert dingo_is_fill[0] > dingo_is_fill[2] # red channel dominates + + # The band edges are a darker shade of the band, and both are overridable. + axes = plot_ppd_td( + wf_td, + {"H1": rng.normal(size=T)}, + times, + filename=str(tmp_path / "override.png"), + band_colors={"dingo": "#00FF00"}, + data_color="#FF0000", + ) + data_line, *edge_lines = axes[0].lines + assert data_line.get_color() == "#FF0000" + assert axes[0].collections[0].get_facecolor()[0][1] == 1.0 # green band fill + assert all(line.get_color() == (0.0, 0.65, 0.0) for line in edge_lines) + + +def test_plot_ppd_td_draws_overlay(tmp_path) -> None: + """``num_plotted_draws`` overlays that many subsampled draw traces; 0 overlays none.""" + rng = np.random.default_rng(6) + n, T = 200, 100 + times = np.linspace(-1.0, 0.2, T) + wf_td = {"dingo": {"H1": rng.normal(size=(n, T))}} + data_td = {"H1": rng.normal(size=T)} + + plain = plot_ppd_td( + wf_td, data_td, times, filename=str(tmp_path / "plain.png") + ) + with_draws = plot_ppd_td( + wf_td, + data_td, + times, + filename=str(tmp_path / "draws.png"), + num_plotted_draws=20, + ) + # 20 extra draw traces, and each is fainter than the band edges. + assert len(with_draws[0].lines) == len(plain[0].lines) + 20 + assert max(line.get_alpha() for line in with_draws[0].lines[1:21]) < 0.3