feat(gw): time-domain whitened-strain PPD plots on Result

.gitignore

@@ -151,4 +151,7 @@ dmypy.json
 # directories with temporary files
 **/tmp_files/*
 
-dingo-venv/
+dingo-venv/
+
+# scratch/output dir for generated images and throwaway artifacts
+local/

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,19 @@ 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.
+
+    Even in-prior posterior draws can hit waveform-model edge cases -- e.g. SEOBNRv5EHM's
+    eccentric interpolation failing on an interior parameter, or simply a waveform backend
+    (pyseobnr / lalsimulation) that differs from the one used when the result was produced.
+    Returning ``None`` lets callers such as :meth:`dingo.gw.result.Result._compute_ppd` drop
+    the offending draw rather than aborting a whole :class:`multiprocessing.Pool` batch.
+    Defined at module level (not a lambda/nested function) so it is picklable for the pool.
+    """
+    try:
+        return gw_signal.signal(theta)
+    except Exception:
+        return None

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, Sequence, Tuple, Union
 
 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.plotting import plot_ppd_td as _plot_ppd_td
 from dingo.gw.prior import build_prior_with_defaults
 from dingo.core.utils.backward_compatibility import check_minimum_version
 
@@ -781,3 +785,260 @@ def pesummary_prior(self):
                 except AttributeError:
                     continue
         return prior
+
+    def _compute_ppd(
+        self,
+        credible_interval: Union[float, Sequence[float]] = 0.9,
+        num_waveforms: int = 1000,
+        num_processes: int = 1,
+        seed: int = RANDOM_STATE,
+    ) -> Tuple[dict, dict, dict, dict]:
+        """Generate whitened waveforms for the time-domain strain PPD.
+
+        Builds the highest-posterior-density credible set at the largest requested
+        ``credible_interval`` level for each available posterior -- ``"dingo"`` (the raw
+        network samples) and, when the result is importance-sampled, ``"dingo-is"`` --
+        generates up to ``num_waveforms`` of each, and projects them onto the detectors via
+        the likelihood. Model waveforms come back already whitened (``h / asd / noise_std``),
+        matching the whitened data. The generated draws are returned ordered by descending
+        density, and per-level counts are returned so the plot can draw nested bands.
+
+        Parameters
+        ----------
+        credible_interval : float or sequence of float
+            Posterior probability of the credible set(s). A sequence (e.g. ``[0.5, 0.9]``)
+            requests nested credible bands; the waveforms are generated for the largest level.
+        num_waveforms : int
+            Maximum number of waveforms drawn from the (largest) credible set.
+        num_processes : int
+            Processes used for waveform generation.
+        seed : int
+            Seed for the uniform subsampling, for deterministic figures.
+
+        Returns
+        -------
+        wf_fd : dict
+            ``{mode: {ifo: complex ndarray (n_kept, n_freq)}}`` whitened model waveforms,
+            where ``mode`` is ``"dingo"`` (always) and ``"dingo-is"`` (present iff the result
+            is importance-sampled). Rows are the successfully generated draws **ordered by
+            descending density**; the rare failed draw is dropped.
+        data_fd : dict
+            ``{ifo: complex ndarray (n_freq,)}`` whitened detector data (shared across modes).
+        map_fd : dict
+            ``{mode: {ifo: complex ndarray (n_freq,)}}`` the single maximum-probability
+            whitened waveform per mode -- max ``log_prob`` for ``"dingo"``, max
+            ``log_likelihood + log_prior`` for ``"dingo-is"`` -- drawn as a line over the band.
+            A mode is absent iff its MAP draw failed generation.
+        level_counts : dict
+            ``{mode: {level: k}}`` -- for each requested credible level, ``k`` is how many of
+            the (density-ordered) generated draws fall in that level's HPD set, so the plot can
+            take the densest ``wf_fd[mode][ifo][:k]`` as the band for that level.
+
+        All are consumed by :func:`dingo.gw.utils.plotting.plot_ppd_td` (together with
+        ``self.domain``).
+
+        Notes
+        -----
+        Per-mode credible set. The ``credible_interval`` (= X) region is the highest-
+        posterior-density set ``R = {theta : p(theta) >= c}`` with ``\\int_R p = X``. It needs
+        *two distinct* quantities per sample -- a probability **mass** and a **ranking key** --
+        not one weight. Given samples ``theta_i ~ g`` (the sampling density) and target density
+        ``p``, self-normalized importance sampling estimates any set's mass as
+        ``P(A) ~= sum_i m_i 1[theta_i in A]`` with ``m_i = (p_i/g_i) / sum_j (p_j/g_j)``
+
+        - ``"dingo"``: ``m_i propto 1/N``; rank by ``log_prob``.
+        - ``"dingo-is"``: ``m_i propto p/q``; rank by ``log_likelihood + log_prior``.
+
+        Ranking by the density and accumulating ``m_i`` until the cumulative mass reaches X
+        keeps exactly the samples above the density threshold ``c``
+
+        Prior filter. Samples are first restricted to the prior support (finite ``log_prior``).
+        The flow proposal is a smooth density over a box and leaks a small fraction of draws
+        outside the prior (e.g. tiny negative spin magnitudes); these are unphysical, carry
+        zero target mass, and account for most waveform-generation failures. This is exactly
+        the in-prior restriction that :meth:`importance_sample` applies.
+
+        Per-sample guard. The prior filter is not sufficient on its own: a waveform model can
+        fail on an *interior* parameter (e.g. SEOBNRv5EHM's eccentric ``CubicSpline`` on a
+        non-monotonic omega array), and this happens even for importance-sampled results
+        whenever the waveform backend here differs from the one used at IS time -- re-plotting
+        a stored result with a newer/older ``pyseobnr``/``lalsimulation`` is enough. So each
+        draw is generated through ``safe_signal``: failures are dropped from the envelope with
+        a warning, and only an all-failed batch raises.
+        """
+        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())
+        rng = np.random.default_rng(seed)
+
+        levels = sorted(
+            {float(credible_interval)}
+            if np.isscalar(credible_interval)
+            else {float(x) for x in credible_interval}
+        )
+
+        # Whitened data, same convention as the (already-whitened) model waveforms.
+        data_fd = {
+            ifo: np.sqrt(4 * domain.delta_f)
+            * self.context["waveform"][ifo]
+            / self.context["asds"][ifo]
+            for ifo in ifos
+        }
+
+        # Restrict to the prior support before generating anything. The flow proposal is a
+        # smooth density over a box and leaks a small fraction of draws outside the prior
+        # (e.g. tiny negative spin magnitudes); these are unphysical and are exactly the
+        # draws that fail waveform generation. Dropping them here -- the same in-prior
+        # restriction importance sampling uses -- removes those failures (guaranteed for an
+        # importance-sampled result, whose in-prior samples all generated during IS) and
+        # keeps unphysical waveforms out of the envelope.
+        non_fixed = [
+            k
+            for k, v in self.prior.items()
+            if not isinstance(v, (Constraint, DeltaFunction))
+        ]
+        in_prior = np.isfinite(
+            np.asarray(self.prior.ln_prob(self.samples[non_fixed], axis=0), dtype=float)
+        )
+
+        # One credible set per available posterior: "dingo" always; "dingo-is" when the
+        # result is importance-sampled (weights present). map_fd holds the single
+        # maximum-probability waveform per mode (drawn as a line over the band).
+        wf_fd = {}
+        map_fd = {}
+        level_counts = {}
+        for mode in ("dingo", "dingo-is"):
+            if mode == "dingo":
+                mass = in_prior.astype(float)
+                rank = self.samples["log_prob"].to_numpy()
+            else:
+                if "weights" not in self.samples:
+                    continue
+                mass = self.samples["weights"].to_numpy().astype(float)
+                rank = (
+                    self.samples["log_likelihood"] + self.samples["log_prior"]
+                ).to_numpy()
+            # Out-of-prior draws get zero mass and bottom rank, so they never enter the
+            # credible set (this also sidesteps NaN log_likelihood on unevaluated samples).
+            mass = np.where(in_prior, mass, 0.0)
+            mass = mass / mass.sum()
+            rank = np.where(in_prior, rank, -np.inf)
+            order = np.argsort(rank)[::-1]
+            cum_mass = np.cumsum(mass[order])
+
+            # The HPD set at level L is {theta : density(theta) >= c_L}. Find each level's
+            # density threshold c_L from the population (rank where the cumulative mass first
+            # reaches L); a smaller level is just a higher threshold on the same ranked draws,
+            # so nested bands are a top-k slice of the density-ordered draws -- no extra
+            # generation. Build (and generate) the set for the largest level.
+            thresholds = {
+                lvl: rank[order[min(int(np.searchsorted(cum_mass, lvl)), len(order) - 1)]]
+                for lvl in levels
+            }
+            n_keep = int(np.searchsorted(cum_mass, levels[-1])) + 1
+            credible_idx = order[:n_keep]
+            if len(credible_idx) > num_waveforms:
+                credible_idx = rng.choice(
+                    credible_idx, size=num_waveforms, replace=False
+                )
+            # Order the draws by descending density so the plot can take the densest top-k as
+            # the inner (tighter) credible level.
+            credible_idx = credible_idx[np.argsort(rank[credible_idx])[::-1]]
+            sel_rank = rank[credible_idx]
+            theta = self.samples.iloc[credible_idx]
+
+            # The prior filter removes out-of-prior draws, but it does not guarantee every
+            # in-prior draw generates: a waveform model can fail on an interior parameter
+            # (e.g. SEOBNRv5EHM's eccentric CubicSpline hitting a non-monotonic omega array).
+            # This bites even for importance-sampled results whenever the waveform backend
+            # here differs from the one used at IS time -- re-plotting a stored result with a
+            # newer/older pyseobnr/lalsimulation is enough to make an IS-accepted draw fail.
+            # So guard per-sample: drop failures, warn with the count, raise only if all fail.
+            signals = apply_func_with_multiprocessing(
+                partial(safe_signal, self.likelihood),
+                theta,
+                num_processes=num_processes,
+            )
+            failed = np.array([s is None for s in signals])
+            n_failed = int(failed.sum())
+            kept_rank = sel_rank[~failed]
+            signals = [s for s in signals if s is not None]
+            if not signals:
+                raise RuntimeError(
+                    f"All '{mode}' PPD waveform generations failed; cannot build the PPD."
+                )
+            if n_failed:
+                warnings.warn(
+                    f"{n_failed} in-prior '{mode}' PPD waveform(s) failed to generate "
+                    "and were dropped from the envelope.",
+                    stacklevel=2,
+                )
+            wf_fd[mode] = {
+                ifo: np.array([s["waveform"][ifo] for s in signals]) for ifo in ifos
+            }
+            # Density-ordered draws already kept; count how many fall in each level's HPD set.
+            level_counts[mode] = {
+                lvl: int(np.sum(kept_rank >= thresholds[lvl])) for lvl in levels
+            }
+
+            # Maximum-probability waveform for this mode: the single highest-ranked
+            # (in-prior) draw -- max log_prob for "dingo", max (log_likelihood + log_prior)
+            # for "dingo-is". Gives a clean representative chirp over the min/max band.
+            map_signal = safe_signal(
+                self.likelihood, self.samples.iloc[order[0]].to_dict()
+            )
+            if map_signal is not None:
+                map_fd[mode] = {ifo: map_signal["waveform"][ifo] for ifo in ifos}
+
+        if wf_fd["dingo"][ifos[0]].shape[1] != len(domain()):
+            raise ValueError(
+                "Whitened waveform length does not match the inverse-FFT domain; "
+                "time-domain PPD requires a uniform frequency domain."
+            )
+        return wf_fd, data_fd, map_fd, level_counts
+
+    def plot_ppd_td(
+        self,
+        filename: str = "ppd_td.png",
+        credible_interval: Union[float, Sequence[float]] = 0.9,
+        num_waveforms: int = 1000,
+        num_processes: int = 1,
+        zoom: Optional[Tuple[float, float]] = None,
+    ) -> Tuple[dict, dict, dict]:
+        """Plot the time-domain whitened-strain posterior-predictive distribution.
+
+        For each mode and detector, inverse-FFTs whitened waveforms to the time domain with
+        the merger at t = 0 and shades the pointwise min/max envelope, with the
+        maximum-probability waveform drawn as a line over it and the whitened data in grey.
+        Panels are stacked vertically, one per (mode, detector): the **Dingo** posterior on
+        top and, when the result is importance-sampled, the **Dingo-IS** posterior below.
+
+        ``credible_interval`` may be a single level (default 0.9) or a sequence such as
+        ``[0.5, 0.9]`` to overlay nested credible bands. ``zoom`` is the (left, right)
+        x-limit in seconds-to-merger; defaults to (-1.0, 0.2). Returns the
+        ``(wf_fd, data_fd, map_fd)`` tuple that was plotted.
+        """
+        wf_fd, data_fd, map_fd, level_counts = self._compute_ppd(
+            credible_interval, num_waveforms, num_processes
+        )
+        domain = (
+            self.domain.base_domain
+            if hasattr(self.domain, "base_domain")
+            else self.domain
+        )
+        _plot_ppd_td(
+            wf_fd,
+            data_fd,
+            domain,
+            map_fd=map_fd,
+            level_counts=level_counts,
+            filename=filename,
+            zoom=zoom,
+        )
+        return wf_fd, data_fd, map_fd

dingo/gw/utils/__init__.py

@@ -0,0 +1 @@
+"""GW-specific utilities."""