From 6529070e5f901dab6b3001b76c122accbaaa434c Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 10:44:40 +0200
Subject: [PATCH 01/65] Add factorized sampler core with plain-NPE parity

Introduce the factorized-sampler spine from the hackathon design
(vault/Hackathon/Factorized_Sampler_Design.md): the posterior as an ordered
product of conditional factors, q(theta|d) = prod_i q_i(theta_i | f_i(theta_<i, d)).

dingo/core/factors.py (domain-agnostic):
- Factor (ABC): physical-in/physical-out sample_and_log_prob / log_prob
- ChainComposer: autoregressive composition, topological-order validation,
  per-factor log-prob summation
- Standardization: per-factor mean/std adapter (both directions + Jacobian)
- ComposedSampler: thin run_sampler facade (batching -> DataFrame)
- SamplerContext protocol

dingo/gw/inference/factors.py (GW plain-NPE path):
- GWSamplerContext.from_model: domain + one-time data prep, cached prepared_data
- FlowFactor: wraps a posterior model, encapsulates standardization
- GWComposedSampler: adds GW _post_process (fixed-param injection + RA correction)
- FixedFactor / SyntheticPhaseFactor stubs for later steps

Plain-NPE parity verified bit-exact against the current GWSampler on a real
GW200129 NPE model, both at the core (_run_sampler) and end-to-end (run_sampler,
batched, with post-processing) levels.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 237 ++++++++++++++++++++++++
 dingo/gw/inference/factors.py | 339 ++++++++++++++++++++++++++++++++++
 2 files changed, 576 insertions(+)
 create mode 100644 dingo/core/factors.py
 create mode 100644 dingo/gw/inference/factors.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
new file mode 100644
index 000000000..4bfb86802
--- /dev/null
+++ b/dingo/core/factors.py
@@ -0,0 +1,237 @@
+"""
+Factorized sampler core.
+
+This module implements the domain-agnostic spine of the factorized-sampler design
+(see vault/Hackathon/Factorized_Sampler_Design.md): the posterior is represented as
+an ordered product of conditional factors
+
+    q(theta_1, ..., theta_n | d) = prod_i q_i(theta_i | f_i(theta_<i, d)),
+
+where each ``Factor`` samples one block of parameters and returns its own log-prob,
+and a composer drives them. The ``ChainComposer`` runs the factors autoregressively
+(exact, importance-sampling-friendly); a Gibbs composer for multi-iteration GNPE is
+added separately.
+
+Key invariant -- *physical in, physical out*: factors expose only physical-space
+parameters. Standardized values exist transiently inside a factor's forward pass; the
+``Standardization`` helper below is the per-factor adapter between a network's internal
+standardized space and physical space. Nothing we store, pass, or serialize is ever
+standardized.
+"""
+
+from __future__ import annotations
+
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional, Protocol, runtime_checkable
+
+import pandas as pd
+import torch
+
+
+class Standardization:
+    """
+    Per-factor parameter standardization: the affine map between a network's internal
+    standardized space and physical space.
+
+    A network is trained on standardized parameters ``z = (theta - mean) / std``. This
+    object holds that network's ``mean``/``std`` (read from its training metadata) and
+    is used in *both* directions -- de-standardizing samples on the way out and
+    standardizing parameters on the way in for ``log_prob`` -- so there is a single
+    source of truth and no "inverted twice" drift.
+
+    Standardization is a property of the *network*, not the event: different factors
+    (e.g. a GNPE init network and the main network) carry different standardizations,
+    and physical space is the shared interchange between them.
+
+    Notes
+    -----
+    The eventual design stores ``mean``/``std`` as registered buffers inside the network
+    module (for NSF, as a fixed affine transform at the flow boundary so its log-Jacobian
+    folds into the density automatically). This pure-Python version is the equivalent
+    seam; moving it into buffers is a follow-up that does not change the interface.
+    """
+
+    def __init__(self, mean: dict[str, float], std: dict[str, float]):
+        self.mean = dict(mean)
+        self.std = dict(std)
+
+    def standardize(
+        self, values: dict[str, torch.Tensor], names: list[str]
+    ) -> torch.Tensor:
+        """Map physical ``values`` (a dict of named tensors) to a standardized tensor
+        with columns in ``names`` order."""
+        cols = [(values[n] - self.mean[n]) / self.std[n] for n in names]
+        return torch.stack(cols, dim=-1)
+
+    def destandardize(
+        self, z: torch.Tensor, names: list[str]
+    ) -> dict[str, torch.Tensor]:
+        """Map a standardized tensor (columns in ``names`` order) back to a dict of
+        named physical tensors."""
+        return {n: z[..., i] * self.std[n] + self.mean[n] for i, n in enumerate(names)}
+
+    def log_det(self, names: list[str]) -> float:
+        """The term to *add* to a network log-prob to express it in physical space:
+        ``log p_theta = log p_z - sum log std``. Same correction in both directions."""
+        return -sum(math.log(self.std[n]) for n in names)
+
+
+@runtime_checkable
+class SamplerContext(Protocol):
+    """
+    Per-event shared state referenced by every factor: the data ``d`` and everything
+    derived from it (the one-time prepared data representation, the likelihood, event
+    metadata). Concrete implementations are domain-specific (see
+    ``dingo.gw.inference.factors.GWSamplerContext``). Serialized as the transport state
+    between pipe stages.
+    """
+
+    event_metadata: Optional[dict]
+
+    def prepared_data(self) -> torch.Tensor:
+        """The one-time data representation (whiten/decimate/repackage/...), computed
+        once and cached. The conditioning input shared by data-conditioned factors."""
+        ...
+
+    def likelihood(self):
+        """The likelihood, for likelihood-based factors (synthetic phase) and IS."""
+        ...
+
+
+@dataclass
+class Conditioning:
+    """
+    What the composer hands a factor: the shared ``context`` (-> prepared data,
+    likelihood, domain) plus the physical values of the earlier-block parameters this
+    factor conditions on.
+    """
+
+    context: SamplerContext
+    given: dict[str, torch.Tensor] = field(default_factory=dict)
+
+
+class Factor(ABC):
+    """
+    A conditional distribution ``q_i(theta_i | f_i(theta_<i, d))`` over one block of
+    parameters. Emits physical-space samples and a physical-space log-prob (its
+    standardization Jacobian already applied), so the chain product ``sum_i log q_i`` is
+    directly the physical posterior log-density.
+
+    Attributes
+    ----------
+    parameters : list[str]
+        The block ``theta_i`` this factor produces.
+    conditioning : list[str]
+        Earlier-block parameter names this factor conditions on. Data dependence is
+        implicit via the shared context.
+    """
+
+    parameters: list[str]
+    conditioning: list[str]
+
+    @abstractmethod
+    def sample_and_log_prob(
+        self, num_samples: int, cond: Conditioning
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        """Draw ``theta_i ~ q_i``; return ``(physical samples, physical-space log q_i)``."""
+
+    @abstractmethod
+    def log_prob(
+        self, theta_i: dict[str, torch.Tensor], cond: Conditioning
+    ) -> torch.Tensor:
+        """Evaluate ``log q_i(theta_i | f_i)`` at given physical ``theta_i`` (for IS /
+        re-plug)."""
+
+
+class ChainComposer:
+    """
+    Autoregressive composer: sample (or evaluate) the factors in declared order, which
+    must be a topological order of the dependency DAG, and sum their log-probs. Exact
+    and importance-sampling-friendly. Covers plain NPE, single-step GNPE, prior
+    conditioning, synthetic phase, and intrinsic/extrinsic splits.
+
+    Multi-iteration (cyclic) GNPE is *not* expressible here and uses a separate Gibbs
+    composer.
+    """
+
+    def __init__(self, factors: list[Factor]):
+        self.factors = factors
+        self._validate()
+
+    def _validate(self):
+        """Check the declared order is a valid topological order: every conditioning
+        name is produced by an earlier factor."""
+        produced: set[str] = set()
+        for factor in self.factors:
+            missing = [c for c in factor.conditioning if c not in produced]
+            if missing:
+                raise ValueError(
+                    f"Factor producing {factor.parameters} conditions on {missing}, "
+                    f"which no earlier factor produces. Check chain order."
+                )
+            produced.update(factor.parameters)
+
+    def sample_and_log_prob(
+        self, num_samples: int, context: SamplerContext
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        samples: dict[str, torch.Tensor] = {}
+        total: torch.Tensor | float = 0.0
+        for factor in self.factors:
+            cond = Conditioning(context, {k: samples[k] for k in factor.conditioning})
+            block, lp = factor.sample_and_log_prob(num_samples, cond)
+            samples.update(block)
+            total = total + lp
+        return samples, total
+
+    def log_prob(
+        self, samples: dict[str, torch.Tensor], context: SamplerContext
+    ) -> torch.Tensor:
+        total: torch.Tensor | float = 0.0
+        for factor in self.factors:
+            cond = Conditioning(context, {k: samples[k] for k in factor.conditioning})
+            theta_i = {k: samples[k] for k in factor.parameters}
+            total = total + factor.log_prob(theta_i, cond)
+        return total
+
+
+class ComposedSampler:
+    """
+    Thin user-facing façade over a ``ChainComposer`` and a ``SamplerContext``: runs the
+    chain with optional batching, applies domain-specific post-processing, and returns
+    samples as a DataFrame. The per-factor compute lives in the composer; this class only
+    handles batching, consolidation, and post-processing -- the role the monolithic
+    ``Sampler`` plays today.
+    """
+
+    def __init__(self, composer: ChainComposer, context: SamplerContext):
+        self.composer = composer
+        self.context = context
+        self.samples: Optional[pd.DataFrame] = None
+
+    def _post_process(self, samples: dict, inverse: bool = False):
+        """Hook for domain-specific post-processing (e.g. GW fixed-parameter injection
+        and reference-time correction). No-op by default."""
+        pass
+
+    def _run_batch(self, num_samples: int) -> dict[str, torch.Tensor]:
+        samples, log_prob = self.composer.sample_and_log_prob(num_samples, self.context)
+        out = dict(samples)
+        out["log_prob"] = log_prob
+        return out
+
+    def run_sampler(
+        self, num_samples: int, batch_size: Optional[int] = None
+    ) -> pd.DataFrame:
+        if batch_size is None:
+            batch_size = num_samples
+        full_batches, remainder = divmod(num_samples, batch_size)
+        batches = [self._run_batch(batch_size) for _ in range(full_batches)]
+        if remainder > 0:
+            batches.append(self._run_batch(remainder))
+        merged = {k: torch.cat([b[k] for b in batches]) for k in batches[0].keys()}
+        merged = {k: v.cpu().numpy() for k, v in merged.items()}
+        self._post_process(merged)
+        self.samples = pd.DataFrame(merged)
+        return self.samples
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
new file mode 100644
index 000000000..11124a216
--- /dev/null
+++ b/dingo/gw/inference/factors.py
@@ -0,0 +1,339 @@
+"""
+Gravitational-wave factors for the factorized sampler.
+
+Concrete ``Factor`` / ``SamplerContext`` implementations for GW inference (see
+vault/Hackathon/Factorized_Sampler_Design.md). This first increment covers the plain-NPE
+path: a single ``FlowFactor`` whose conditioning map ``f_i`` is the one-time data
+preprocessing (decimate / whiten / repackage), run on a ``GWSamplerContext``. The
+fixed, synthetic-phase, and GNPE-proxy factors are stubbed for later steps.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import numpy as np
+import pandas as pd
+import torch
+from astropy.time import Time
+from bilby.core.prior import DeltaFunction
+from torchvision.transforms import Compose
+
+from dingo.core.factors import (
+    ChainComposer,
+    ComposedSampler,
+    Conditioning,
+    Factor,
+    Standardization,
+)
+from dingo.core.posterior_models import BasePosteriorModel
+from dingo.core.transforms import GetItem
+from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
+from dingo.gw.gwutils import get_extrinsic_prior_dict
+from dingo.gw.prior import build_prior_with_defaults
+from dingo.gw.transforms import (
+    DecimateWaveformsAndASDS,
+    RepackageStrainsAndASDS,
+    ToTorch,
+    WhitenAndScaleStrain,
+)
+
+
+def _base_model_metadata(model: BasePosteriorModel) -> dict:
+    """The training metadata describing the data domain / standardization. For an
+    unconditional (density-recovery) model this lives under ``metadata["base"]``."""
+    metadata = model.metadata
+    if metadata["train_settings"]["data"].get("unconditional", False):
+        return metadata["base"]
+    return metadata
+
+
+class GWSamplerContext:
+    """
+    Per-event shared GW state: the data ``d`` and everything derived from it. Referenced
+    by every factor in a chain, and serialized as the transport state between pipe
+    stages.
+
+    This implements the ``dingo.core.factors.SamplerContext`` protocol. It owns the
+    one-time data preprocessing (``prepared_data``) and -- once wired in a later step --
+    the likelihood used by likelihood-based factors and importance sampling.
+
+    Event metadata lives here (not on individual factors): it is a property of the data,
+    and it drives frequency cropping, the t_ref/RA correction, and the likelihood.
+    """
+
+    def __init__(
+        self,
+        domain,
+        detectors: list[str],
+        t_ref: float,
+        data_prep: Compose,
+        raw_context: dict,
+        event_metadata: Optional[dict] = None,
+    ):
+        self.domain = domain
+        self.detectors = detectors
+        self.t_ref = t_ref
+        self._data_prep = data_prep
+        self.raw_context = raw_context
+        self.event_metadata = event_metadata
+        self._prepared: Optional[torch.Tensor] = None
+
+    @classmethod
+    def from_model(
+        cls,
+        model: BasePosteriorModel,
+        raw_context: dict,
+        event_metadata: Optional[dict] = None,
+    ) -> "GWSamplerContext":
+        """Build the context (domain + one-time data-prep chain) from a model's
+        metadata. The data-prep chain reproduces ``GWSampler`` preprocessing for the
+        plain-NPE case (parameter-dependent transforms are handled per-factor)."""
+        meta = _base_model_metadata(model)
+        data_settings = meta["train_settings"]["data"]
+
+        domain = build_domain(meta["dataset_settings"]["domain"])
+        if "domain_update" in data_settings:
+            domain.update(data_settings["domain_update"])
+        detectors = data_settings["detectors"]
+
+        transforms = []
+        # Decimate from the base domain when using a multibanded frequency domain.
+        if isinstance(domain, MultibandedFrequencyDomain):
+            transforms.append(
+                DecimateWaveformsAndASDS(domain, decimation_mode="whitened")
+            )
+        # Whiten and scale (the network expects standardized data).
+        transforms.append(WhitenAndScaleStrain(domain.noise_std))
+        # Repackage strains/ASDs into an array, move to torch, extract the waveform.
+        # TODO: frequency-range cropping (MaskDataForFrequencyRangeUpdate) -- follow-up.
+        transforms += [
+            RepackageStrainsAndASDS(ifos=detectors, first_index=domain.min_idx),
+            ToTorch(device=model.device),
+            GetItem("waveform"),
+        ]
+
+        return cls(
+            domain=domain,
+            detectors=detectors,
+            t_ref=data_settings["ref_time"],
+            data_prep=Compose(transforms),
+            raw_context=raw_context,
+            event_metadata=event_metadata,
+        )
+
+    def prepared_data(self) -> torch.Tensor:
+        """One-time data preprocessing, computed once and cached."""
+        if self._prepared is None:
+            self._prepared = self._data_prep(self.raw_context)
+        return self._prepared
+
+    def likelihood(self):
+        raise NotImplementedError(
+            "Likelihood construction moves onto the context in a later step "
+            "(synthetic-phase / importance-sampling factors)."
+        )
+
+
+class FlowFactor(Factor):
+    """
+    A factor wrapping a posterior model (NPE flow, FMPE, ...). Its conditioning map
+    ``f_i`` applies the cheap parameter-dependent data transform (none for plain NPE;
+    time-shift / heterodyne for GNPE -- follow-up), standardizes, then runs the network.
+
+    Encapsulates the network's own standardization: the public interface is physical-in
+    / physical-out. Standardized values never leave the factor.
+    """
+
+    def __init__(
+        self,
+        model: BasePosteriorModel,
+        parameters: list[str],
+        conditioning: Optional[list[str]] = None,
+        context_parameters: Optional[list[str]] = None,
+    ):
+        self.model = model
+        self.parameters = parameters
+        self.conditioning = conditioning or []
+        # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
+        self.context_parameters = context_parameters or []
+        std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
+        self.standardization = Standardization(std["mean"], std["std"])
+
+    @classmethod
+    def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
+        """Build a plain-NPE factor from a model: ``parameters`` and ``context_parameters``
+        come from the model's training metadata (first-class conditioning)."""
+        data_settings = _base_model_metadata(model)["train_settings"]["data"]
+        context_parameters = data_settings.get("context_parameters") or []
+        return cls(
+            model=model,
+            parameters=data_settings["inference_parameters"],
+            conditioning=list(context_parameters),
+            context_parameters=list(context_parameters),
+        )
+
+    def _network_context(self, cond: Conditioning) -> tuple[torch.Tensor, ...]:
+        """Assemble the model's positional context: the prepared data, plus standardized
+        context parameters when the network conditions on proxies."""
+        data = cond.context.prepared_data()
+        if not self.context_parameters:
+            return (data.unsqueeze(0),)
+        # Standardize the (physical) conditioning values the network was trained on.
+        ctx = self.standardization.standardize(cond.given, self.context_parameters)
+        return data.unsqueeze(0), ctx.unsqueeze(0)
+
+    def sample_and_log_prob(self, num_samples, cond):
+        net_context = self._network_context(cond)
+        self.model.network.eval()
+        with torch.no_grad():
+            z, log_prob = self.model.sample_and_log_prob(
+                *net_context, num_samples=num_samples
+            )
+        # Squeeze the batch dimension added for the (single) context.
+        z = z.squeeze(0)
+        log_prob = log_prob.squeeze(0)
+        theta = self.standardization.destandardize(z, self.parameters)
+        log_prob = log_prob + self.standardization.log_det(self.parameters)
+        return theta, log_prob
+
+    def log_prob(self, theta_i, cond):
+        num_samples = next(iter(theta_i.values())).shape[0]
+        z = self.standardization.standardize(theta_i, self.parameters)
+        data = cond.context.prepared_data()
+        data = data.expand(num_samples, *data.shape)
+        net_context: tuple[torch.Tensor, ...]
+        if not self.context_parameters:
+            net_context = (data,)
+        else:
+            ctx = self.standardization.standardize(cond.given, self.context_parameters)
+            net_context = (data, ctx)
+        self.model.network.eval()
+        with torch.no_grad():
+            log_prob = self.model.log_prob(z, *net_context)
+        return log_prob + self.standardization.log_det(self.parameters)
+
+
+# --- Stubs for later steps -------------------------------------------------------------
+
+
+class FixedFactor(Factor):
+    """``q_i = delta(theta_i - c)``: pins parameters to fixed values (prior-conditioning
+    / known proxies). Subsumes ``FixedInitSampler``. TODO: implement."""
+
+    def __init__(self, values: dict[str, float]):
+        self.values = values
+        self.parameters = list(values)
+        self.conditioning = []
+
+    def sample_and_log_prob(self, num_samples, cond):
+        samples = {
+            p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
+        }
+        return samples, torch.zeros(num_samples)
+
+    def log_prob(self, theta_i, cond):
+        raise NotImplementedError("FixedFactor.log_prob -- later step.")
+
+
+class SyntheticPhaseFactor(Factor):
+    """``q(phase | theta_rest, d)`` from the likelihood on a phase grid; the final,
+    non-NN factor, run in the CPU post-processing stage. TODO: port from
+    ``Result.sample_synthetic_phase``."""
+
+    def __init__(self):
+        self.parameters = ["phase"]
+        self.conditioning = []  # conditions on all preceding params via the likelihood
+
+    def sample_and_log_prob(self, num_samples, cond):
+        raise NotImplementedError("SyntheticPhaseFactor -- later step.")
+
+    def log_prob(self, theta_i, cond):
+        raise NotImplementedError("SyntheticPhaseFactor -- later step.")
+
+
+class GWComposedSampler(ComposedSampler):
+    """
+    GW façade: a ``ComposedSampler`` that adds the gravitational-wave post-processing --
+    injecting fixed (delta-prior) parameters and correcting the sky position for the
+    model's fixed reference time. Ported from ``GWSamplerMixin._post_process``; in the
+    longer run the reference-time correction becomes a deterministic reparametrization
+    factor in the chain (see the design doc, Q#4/Q#7).
+    """
+
+    def __init__(
+        self,
+        composer: ChainComposer,
+        context: GWSamplerContext,
+        metadata: dict,
+        inference_parameters: list[str],
+    ):
+        super().__init__(composer, context)
+        self.metadata = metadata
+        self.inference_parameters = inference_parameters
+
+    @classmethod
+    def from_model(
+        cls,
+        model: BasePosteriorModel,
+        raw_context: dict,
+        event_metadata: Optional[dict] = None,
+    ) -> "GWComposedSampler":
+        """Build a single-factor (plain-NPE) GW sampler from a model and event data."""
+        context = GWSamplerContext.from_model(model, raw_context, event_metadata)
+        factor = FlowFactor.from_model(model)
+        return cls(
+            composer=ChainComposer([factor]),
+            context=context,
+            metadata=_base_model_metadata(model),
+            inference_parameters=factor.parameters,
+        )
+
+    def _correct_reference_time(
+        self, samples: Union[dict, pd.DataFrame], inverse: bool = False
+    ):
+        """Correct the right ascension for the difference between the event time and the
+        model's training reference time (fixed detector positions)."""
+        event_metadata = self.context.event_metadata
+        if event_metadata is None:
+            return
+        t_event = event_metadata.get("time_event")
+        t_ref = self.context.t_ref
+        if t_event is None or t_event == t_ref or "ra" not in samples:
+            return
+        ra = samples["ra"]
+        longitude_event = Time(t_event, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        longitude_reference = Time(t_ref, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        ra_correction = (longitude_event - longitude_reference).rad
+        if not inverse:
+            samples["ra"] = (ra + ra_correction) % (2 * np.pi)
+        else:
+            samples["ra"] = (ra - ra_correction) % (2 * np.pi)
+
+    def _post_process(self, samples: Union[dict, pd.DataFrame], inverse: bool = False):
+        intrinsic_prior = self.metadata["dataset_settings"]["intrinsic_prior"]
+        extrinsic_prior = get_extrinsic_prior_dict(
+            self.metadata["train_settings"]["data"]["extrinsic_prior"]
+        )
+        prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
+
+        if not inverse:
+            # Add fixed (delta-prior) parameters not produced by the chain.
+            num_samples = len(samples[list(samples.keys())[0]])
+            for k, p in prior.items():
+                if isinstance(p, DeltaFunction) and k not in samples:
+                    print(f"Adding fixed parameter {k} = {p.peak} from prior.")
+                    samples[k] = p.peak * np.ones(num_samples)
+        else:
+            drop = [k for k in samples.keys() if k not in self.inference_parameters]
+            if isinstance(samples, pd.DataFrame):
+                samples.drop(columns=drop, inplace=True, errors="ignore")
+            else:
+                for k in drop:
+                    samples.pop(k, None)
+
+        self._correct_reference_time(samples, inverse)

From 2675f68f3abc96026a1354569772828f9a855c2e Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 13:35:55 +0200
Subject: [PATCH 02/65] Move FlowFactor and _base_model_metadata into core

FlowFactor is domain-agnostic -- it reaches data only through the
SamplerContext.prepared_data() protocol -- so it belongs in core alongside the
Factor ABC rather than in the GW module. The GW layer keeps only the GW-specific
pieces: GWSamplerContext (which builds the data-prep conditioning map f_i) and
GWComposedSampler (post-processing). Document GWSamplerContext.raw_context
(= EventDataset.data; retained for the not-yet-wired likelihood and the
serialized transport state).

dingo.gw.inference.factors re-exports FlowFactor, so existing imports keep
working. Plain-NPE parity with GWSampler remains bit-exact.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         |  94 +++++++++++++++++++++++++++++
 dingo/gw/inference/factors.py | 107 ++++------------------------------
 2 files changed, 106 insertions(+), 95 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 4bfb86802..16d33cae6 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -29,6 +29,8 @@
 import pandas as pd
 import torch
 
+from dingo.core.posterior_models import BasePosteriorModel
+
 
 class Standardization:
     """
@@ -145,6 +147,98 @@ def log_prob(
         re-plug)."""
 
 
+def _base_model_metadata(model: BasePosteriorModel) -> dict:
+    """The training metadata describing the parameter standardization / data settings.
+    For an unconditional (density-recovery) model this lives under ``metadata["base"]``.
+    """
+    metadata = model.metadata
+    if metadata["train_settings"]["data"].get("unconditional", False):
+        return metadata["base"]
+    return metadata
+
+
+class FlowFactor(Factor):
+    """
+    A factor wrapping a posterior model (NPE flow, FMPE, ...). Domain-agnostic: it reaches
+    data only through the ``SamplerContext.prepared_data()`` protocol, so the GW-specific
+    conditioning map ``f_i`` (the data preprocessing) lives entirely on the context.
+
+    Encapsulates the network's own standardization: the public interface is physical-in /
+    physical-out -- standardized values never leave the factor.
+
+    The cheap *parameter-dependent* part of ``f_i`` (time-shift / heterodyne for GNPE) is
+    a no-op here; a domain subclass overrides ``_apply_param_transforms`` to add it.
+    """
+
+    def __init__(
+        self,
+        model: BasePosteriorModel,
+        parameters: list[str],
+        conditioning: Optional[list[str]] = None,
+        context_parameters: Optional[list[str]] = None,
+    ):
+        self.model = model
+        self.parameters = parameters
+        self.conditioning = conditioning or []
+        # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
+        self.context_parameters = context_parameters or []
+        std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
+        self.standardization = Standardization(std["mean"], std["std"])
+
+    @classmethod
+    def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
+        """Build a factor from a model: ``parameters`` and ``context_parameters`` come from
+        the model's training metadata (first-class conditioning)."""
+        data_settings = _base_model_metadata(model)["train_settings"]["data"]
+        context_parameters = data_settings.get("context_parameters") or []
+        return cls(
+            model=model,
+            parameters=data_settings["inference_parameters"],
+            conditioning=list(context_parameters),
+            context_parameters=list(context_parameters),
+        )
+
+    def _network_context(self, cond: Conditioning) -> tuple[torch.Tensor, ...]:
+        """Assemble the model's positional context: the prepared data, plus standardized
+        context parameters when the network conditions on proxies."""
+        data = cond.context.prepared_data()
+        if not self.context_parameters:
+            return (data.unsqueeze(0),)
+        # Standardize the (physical) conditioning values the network was trained on.
+        ctx = self.standardization.standardize(cond.given, self.context_parameters)
+        return data.unsqueeze(0), ctx.unsqueeze(0)
+
+    def sample_and_log_prob(self, num_samples, cond):
+        net_context = self._network_context(cond)
+        self.model.network.eval()
+        with torch.no_grad():
+            z, log_prob = self.model.sample_and_log_prob(
+                *net_context, num_samples=num_samples
+            )
+        # Squeeze the batch dimension added for the (single) context.
+        z = z.squeeze(0)
+        log_prob = log_prob.squeeze(0)
+        theta = self.standardization.destandardize(z, self.parameters)
+        log_prob = log_prob + self.standardization.log_det(self.parameters)
+        return theta, log_prob
+
+    def log_prob(self, theta_i, cond):
+        num_samples = next(iter(theta_i.values())).shape[0]
+        z = self.standardization.standardize(theta_i, self.parameters)
+        data = cond.context.prepared_data()
+        data = data.expand(num_samples, *data.shape)
+        net_context: tuple[torch.Tensor, ...]
+        if not self.context_parameters:
+            net_context = (data,)
+        else:
+            ctx = self.standardization.standardize(cond.given, self.context_parameters)
+            net_context = (data, ctx)
+        self.model.network.eval()
+        with torch.no_grad():
+            log_prob = self.model.log_prob(z, *net_context)
+        return log_prob + self.standardization.log_det(self.parameters)
+
+
 class ChainComposer:
     """
     Autoregressive composer: sample (or evaluate) the factors in declared order, which
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 11124a216..79dfce5aa 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -1,11 +1,12 @@
 """
 Gravitational-wave factors for the factorized sampler.
 
-Concrete ``Factor`` / ``SamplerContext`` implementations for GW inference (see
-vault/Hackathon/Factorized_Sampler_Design.md). This first increment covers the plain-NPE
-path: a single ``FlowFactor`` whose conditioning map ``f_i`` is the one-time data
-preprocessing (decimate / whiten / repackage), run on a ``GWSamplerContext``. The
-fixed, synthetic-phase, and GNPE-proxy factors are stubbed for later steps.
+GW-specific pieces of the factorized-sampler design (see
+vault/Hackathon/Factorized_Sampler_Design.md). The generic ``FlowFactor`` is
+domain-agnostic and lives in ``dingo.core.factors``; what is GW-specific is the
+``GWSamplerContext`` (which builds the data-preprocessing conditioning map ``f_i``) and
+the post-processing in ``GWComposedSampler``. This first increment covers the plain-NPE
+path; the fixed, synthetic-phase, and GNPE-proxy factors are stubbed for later steps.
 """
 
 from __future__ import annotations
@@ -22,9 +23,9 @@
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
-    Conditioning,
     Factor,
-    Standardization,
+    FlowFactor,
+    _base_model_metadata,
 )
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem
@@ -39,15 +40,6 @@
 )
 
 
-def _base_model_metadata(model: BasePosteriorModel) -> dict:
-    """The training metadata describing the data domain / standardization. For an
-    unconditional (density-recovery) model this lives under ``metadata["base"]``."""
-    metadata = model.metadata
-    if metadata["train_settings"]["data"].get("unconditional", False):
-        return metadata["base"]
-    return metadata
-
-
 class GWSamplerContext:
     """
     Per-event shared GW state: the data ``d`` and everything derived from it. Referenced
@@ -75,6 +67,10 @@ def __init__(
         self.detectors = detectors
         self.t_ref = t_ref
         self._data_prep = data_prep
+        # The raw event data `d` (strain + ASDs per detector), i.e. EventDataset.data --
+        # the same object set as GWSampler.context today. Consumed lazily by
+        # prepared_data(); also retained for the likelihood (synthetic-phase / IS factors)
+        # and as part of the serialized transport state, neither of which is wired yet.
         self.raw_context = raw_context
         self.event_metadata = event_metadata
         self._prepared: Optional[torch.Tensor] = None
@@ -135,85 +131,6 @@ def likelihood(self):
         )
 
 
-class FlowFactor(Factor):
-    """
-    A factor wrapping a posterior model (NPE flow, FMPE, ...). Its conditioning map
-    ``f_i`` applies the cheap parameter-dependent data transform (none for plain NPE;
-    time-shift / heterodyne for GNPE -- follow-up), standardizes, then runs the network.
-
-    Encapsulates the network's own standardization: the public interface is physical-in
-    / physical-out. Standardized values never leave the factor.
-    """
-
-    def __init__(
-        self,
-        model: BasePosteriorModel,
-        parameters: list[str],
-        conditioning: Optional[list[str]] = None,
-        context_parameters: Optional[list[str]] = None,
-    ):
-        self.model = model
-        self.parameters = parameters
-        self.conditioning = conditioning or []
-        # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
-        self.context_parameters = context_parameters or []
-        std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
-        self.standardization = Standardization(std["mean"], std["std"])
-
-    @classmethod
-    def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
-        """Build a plain-NPE factor from a model: ``parameters`` and ``context_parameters``
-        come from the model's training metadata (first-class conditioning)."""
-        data_settings = _base_model_metadata(model)["train_settings"]["data"]
-        context_parameters = data_settings.get("context_parameters") or []
-        return cls(
-            model=model,
-            parameters=data_settings["inference_parameters"],
-            conditioning=list(context_parameters),
-            context_parameters=list(context_parameters),
-        )
-
-    def _network_context(self, cond: Conditioning) -> tuple[torch.Tensor, ...]:
-        """Assemble the model's positional context: the prepared data, plus standardized
-        context parameters when the network conditions on proxies."""
-        data = cond.context.prepared_data()
-        if not self.context_parameters:
-            return (data.unsqueeze(0),)
-        # Standardize the (physical) conditioning values the network was trained on.
-        ctx = self.standardization.standardize(cond.given, self.context_parameters)
-        return data.unsqueeze(0), ctx.unsqueeze(0)
-
-    def sample_and_log_prob(self, num_samples, cond):
-        net_context = self._network_context(cond)
-        self.model.network.eval()
-        with torch.no_grad():
-            z, log_prob = self.model.sample_and_log_prob(
-                *net_context, num_samples=num_samples
-            )
-        # Squeeze the batch dimension added for the (single) context.
-        z = z.squeeze(0)
-        log_prob = log_prob.squeeze(0)
-        theta = self.standardization.destandardize(z, self.parameters)
-        log_prob = log_prob + self.standardization.log_det(self.parameters)
-        return theta, log_prob
-
-    def log_prob(self, theta_i, cond):
-        num_samples = next(iter(theta_i.values())).shape[0]
-        z = self.standardization.standardize(theta_i, self.parameters)
-        data = cond.context.prepared_data()
-        data = data.expand(num_samples, *data.shape)
-        net_context: tuple[torch.Tensor, ...]
-        if not self.context_parameters:
-            net_context = (data,)
-        else:
-            ctx = self.standardization.standardize(cond.given, self.context_parameters)
-            net_context = (data, ctx)
-        self.model.network.eval()
-        with torch.no_grad():
-            log_prob = self.model.log_prob(z, *net_context)
-        return log_prob + self.standardization.log_det(self.parameters)
-
-
 # --- Stubs for later steps -------------------------------------------------------------
 
 

From 1a6bc918bce08b4601e6ea42c5a194ae96df88da Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 14:02:28 +0200
Subject: [PATCH 03/65] Add GNPE Gibbs composer with bit-exact parity

Add multi-iteration time-GNPE to the factorized sampler:
- core: GibbsComposer (seed from an init factor, iterate a GNPE step to a fixed
  point, return samples without log_prob); ChainComposer.sample; unify
  ComposedSampler._run_batch so the facade drives either composer.
- gw: GNPEFlowFactor (one GNPE iteration -- blur proxies, time-shift strain,
  standardize context, sample, recompute detector times; replicates
  GWSamplerGNPE's transforms) and GWComposedSampler.from_gnpe_models.

Multi-iteration GNPE parity is bit-exact vs GWSamplerGNPE over 4 Gibbs
iterations on the GW200129 init/main pair; plain-NPE parity unchanged.

Single-step GNPE (num_iterations=1, log_prob-preserving / BNS) is the chain
composer path and is not yet implemented.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         |  48 ++++++++++-
 dingo/gw/inference/factors.py | 153 +++++++++++++++++++++++++++++++++-
 2 files changed, 195 insertions(+), 6 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 16d33cae6..ff299a6b4 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -289,6 +289,47 @@ def log_prob(
             total = total + factor.log_prob(theta_i, cond)
         return total
 
+    def sample(
+        self, num_samples: int, context: SamplerContext
+    ) -> dict[str, torch.Tensor]:
+        """Full per-sample dict (parameters + ``log_prob``), for the sampler façade."""
+        samples, log_prob = self.sample_and_log_prob(num_samples, context)
+        return {**samples, "log_prob": log_prob}
+
+
+class GibbsComposer:
+    """
+    Multi-iteration GNPE composition. Seeds the Gibbs chain from an ``init_factor`` (e.g.
+    a network predicting detector coalescence times), then iterates a GNPE step
+    (``gnpe_factor.gibbs_step``) to a fixed point. The dependency is cyclic
+    (theta <-> proxies), so this breaks density access: ``sample`` returns parameters
+    WITHOUT a ``log_prob``. (Single-step GNPE that *does* preserve log_prob is a
+    two-factor ``ChainComposer`` instead.)
+    """
+
+    def __init__(self, init_factor: Factor, gnpe_factor, num_iterations: int):
+        self.init_factor = init_factor
+        self.gnpe_factor = gnpe_factor
+        self.num_iterations = num_iterations
+
+    def sample(
+        self, num_samples: int, context: SamplerContext
+    ) -> dict[str, torch.Tensor]:
+        # Seed the Gibbs chain with the init factor's parameters (e.g. detector times).
+        seed, _ = self.init_factor.sample_and_log_prob(
+            num_samples, Conditioning(context)
+        )
+        extrinsic = dict(seed)
+        parameters: dict[str, torch.Tensor] = {}
+        for _ in range(self.num_iterations):
+            parameters, extrinsic = self.gnpe_factor.gibbs_step(
+                num_samples, context, extrinsic
+            )
+        proxies = {
+            p: extrinsic[p] for p in self.gnpe_factor.proxy_parameters if p in extrinsic
+        }
+        return {**parameters, **proxies}
+
 
 class ComposedSampler:
     """
@@ -310,10 +351,9 @@ def _post_process(self, samples: dict, inverse: bool = False):
         pass
 
     def _run_batch(self, num_samples: int) -> dict[str, torch.Tensor]:
-        samples, log_prob = self.composer.sample_and_log_prob(num_samples, self.context)
-        out = dict(samples)
-        out["log_prob"] = log_prob
-        return out
+        # Uniform across composers: ChainComposer.sample adds log_prob; GibbsComposer
+        # (multi-iteration GNPE) returns parameters + proxies with no log_prob.
+        return self.composer.sample(num_samples, self.context)
 
     def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 79dfce5aa..1c2522689 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -18,6 +18,7 @@
 import torch
 from astropy.time import Time
 from bilby.core.prior import DeltaFunction
+from bilby.gw.detector import InterferometerList
 from torchvision.transforms import Compose
 
 from dingo.core.factors import (
@@ -25,16 +26,24 @@
     ComposedSampler,
     Factor,
     FlowFactor,
+    GibbsComposer,
     _base_model_metadata,
 )
 from dingo.core.posterior_models import BasePosteriorModel
-from dingo.core.transforms import GetItem
+from dingo.core.transforms import GetItem, RenameKey
 from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
 from dingo.gw.gwutils import get_extrinsic_prior_dict
 from dingo.gw.prior import build_prior_with_defaults
 from dingo.gw.transforms import (
+    CopyToExtrinsicParameters,
     DecimateWaveformsAndASDS,
+    GetDetectorTimes,
+    GNPEBase,
+    GNPECoalescenceTimes,
+    PostCorrectGeocentTime,
     RepackageStrainsAndASDS,
+    SelectStandardizeRepackageParameters,
+    TimeShiftStrain,
     ToTorch,
     WhitenAndScaleStrain,
 )
@@ -169,6 +178,122 @@ def log_prob(self, theta_i, cond):
         raise NotImplementedError("SyntheticPhaseFactor -- later step.")
 
 
+class GNPEFlowFactor:
+    """
+    One iteration of time-shift GNPE for the Gibbs composer. Given the current detector-
+    time estimates, it blurs them into proxies, shifts the strain, standardizes the
+    proxies as network context, samples the main network, then recomputes detector times
+    for the next iteration. Wraps the main model and the per-iteration transforms (the
+    same ones ``GWSamplerGNPE`` builds).
+
+    Not a ``Factor``: the ABC's ``sample_and_log_prob`` does not fit a Gibbs step. It
+    exposes ``gibbs_step(num_samples, context, extrinsic) -> (parameters, extrinsic)``.
+    """
+
+    def __init__(
+        self,
+        model: BasePosteriorModel,
+        transform_pre: Compose,
+        transform_post: Compose,
+        gnpe_parameters: list[str],
+        parameters: list[str],
+    ):
+        self.model = model
+        self.transform_pre = transform_pre
+        self.transform_post = transform_post
+        self.gnpe_parameters = gnpe_parameters
+        self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
+        self.parameters = parameters
+
+    @classmethod
+    def from_model(cls, model: BasePosteriorModel) -> "GNPEFlowFactor":
+        """Build the GNPE per-iteration transforms from the main model's metadata,
+        replicating ``GWSamplerGNPE._initialize_transforms`` (time-shift GNPE)."""
+        meta = _base_model_metadata(model)
+        data_settings = meta["train_settings"]["data"]
+        ifo_list = InterferometerList(data_settings["detectors"])
+        domain = build_domain(meta["dataset_settings"]["domain"])
+        if "domain_update" in data_settings:
+            domain.update(data_settings["domain_update"])
+
+        gnpe_time_settings = data_settings.get("gnpe_time_shifts")
+        if not gnpe_time_settings:
+            raise NotImplementedError(
+                "Only time-shift GNPE is supported here so far "
+                "(no gnpe_chirp / gnpe_phase)."
+            )
+
+        transform_pre = [
+            RenameKey("data", "waveform"),
+            GNPECoalescenceTimes(
+                ifo_list,
+                gnpe_time_settings["kernel"],
+                gnpe_time_settings["exact_equiv"],
+                inference=True,
+            ),
+            TimeShiftStrain(ifo_list, domain),
+            SelectStandardizeRepackageParameters(
+                {"context_parameters": data_settings["context_parameters"]},
+                data_settings["standardization"],
+                device=model.device,
+            ),
+            RenameKey("waveform", "data"),
+        ]
+        gnpe_parameters = []
+        for transform in transform_pre:
+            if isinstance(transform, GNPEBase):
+                gnpe_parameters += transform.input_parameter_names
+
+        inference_parameters = data_settings["inference_parameters"]
+        transform_post = [
+            SelectStandardizeRepackageParameters(
+                {"inference_parameters": inference_parameters},
+                data_settings["standardization"],
+                inverse=True,
+                as_type="dict",
+            ),
+            PostCorrectGeocentTime(),
+            CopyToExtrinsicParameters(
+                "ra", "dec", "geocent_time", "chirp_mass", "mass_ratio", "phase"
+            ),
+            GetDetectorTimes(ifo_list, data_settings["ref_time"]),
+        ]
+        return cls(
+            model,
+            Compose(transform_pre),
+            Compose(transform_post),
+            gnpe_parameters,
+            inference_parameters,
+        )
+
+    def gibbs_step(
+        self, num_samples: int, context, extrinsic: dict
+    ) -> tuple[dict, dict]:
+        """One GNPE Gibbs iteration: returns (sampled parameters, updated extrinsic state
+        -- recomputed detector times + proxies for the next iteration)."""
+        x = {
+            "extrinsic_parameters": {k: extrinsic[k] for k in self.gnpe_parameters},
+            "parameters": {},
+        }
+        d = context.prepared_data().clone()
+        x["data"] = d.expand(num_samples, *d.shape)
+        x = self.transform_pre(x)
+        self.model.network.eval()
+        with torch.no_grad():
+            if "context_parameters" in x:
+                y, log_prob = self.model.sample_and_log_prob(
+                    x["data"], x["context_parameters"]
+                )
+            else:
+                y, log_prob = self.model.sample_and_log_prob(x["data"])
+        # sample_and_log_prob(num_samples=1) adds a singleton dim; the batch is the
+        # num_samples Gibbs walkers.
+        x["parameters"] = y.squeeze(1)
+        x["log_prob"] = log_prob.squeeze(1)
+        x = self.transform_post(x)
+        return x["parameters"], x["extrinsic_parameters"]
+
+
 class GWComposedSampler(ComposedSampler):
     """
     GW façade: a ``ComposedSampler`` that adds the gravitational-wave post-processing --
@@ -180,7 +305,7 @@ class GWComposedSampler(ComposedSampler):
 
     def __init__(
         self,
-        composer: ChainComposer,
+        composer: Union[ChainComposer, GibbsComposer],
         context: GWSamplerContext,
         metadata: dict,
         inference_parameters: list[str],
@@ -206,6 +331,30 @@ def from_model(
             inference_parameters=factor.parameters,
         )
 
+    @classmethod
+    def from_gnpe_models(
+        cls,
+        init_model: BasePosteriorModel,
+        main_model: BasePosteriorModel,
+        raw_context: dict,
+        event_metadata: Optional[dict] = None,
+        num_iterations: int = 30,
+    ) -> "GWComposedSampler":
+        """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
+        init model's data preprocessing, an init ``FlowFactor`` to seed, and a
+        ``GNPEFlowFactor`` driven by a ``GibbsComposer``. Returns samples without a
+        log_prob (Gibbs breaks density access)."""
+        context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
+        init_factor = FlowFactor.from_model(init_model)
+        gnpe_factor = GNPEFlowFactor.from_model(main_model)
+        composer = GibbsComposer(init_factor, gnpe_factor, num_iterations)
+        return cls(
+            composer,
+            context,
+            _base_model_metadata(main_model),
+            gnpe_factor.parameters,
+        )
+
     def _correct_reference_time(
         self, samples: Union[dict, pd.DataFrame], inverse: bool = False
     ):

From 365daf101d50336ed5c9f4e35a8c1503cd36deed Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 14:41:14 +0200
Subject: [PATCH 04/65] ComposedSampler: accept GibbsComposer as well as
 ChainComposer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The façade already drives either composer at runtime (it needs only
sample(num_samples, context) -> dict); broaden the type hint and docstring to
match (Union[ChainComposer, GibbsComposer]).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index ff299a6b4..e0df3baf8 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -24,7 +24,7 @@
 import math
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Optional, Protocol, runtime_checkable
+from typing import Optional, Protocol, Union, runtime_checkable
 
 import pandas as pd
 import torch
@@ -333,14 +333,17 @@ def sample(
 
 class ComposedSampler:
     """
-    Thin user-facing façade over a ``ChainComposer`` and a ``SamplerContext``: runs the
-    chain with optional batching, applies domain-specific post-processing, and returns
-    samples as a DataFrame. The per-factor compute lives in the composer; this class only
-    handles batching, consolidation, and post-processing -- the role the monolithic
-    ``Sampler`` plays today.
+    Thin user-facing façade over a composer (``ChainComposer`` or ``GibbsComposer``) and
+    a ``SamplerContext``: runs the composer with optional batching, applies domain-specific
+    post-processing, and returns samples as a DataFrame. The per-factor compute lives in
+    the composer (which need only expose ``sample(num_samples, context) -> dict``); this
+    class only handles batching, consolidation, and post-processing -- the role the
+    monolithic ``Sampler`` plays today.
     """
 
-    def __init__(self, composer: ChainComposer, context: SamplerContext):
+    def __init__(
+        self, composer: Union[ChainComposer, GibbsComposer], context: SamplerContext
+    ):
         self.composer = composer
         self.context = context
         self.samples: Optional[pd.DataFrame] = None

From 4e43778a10c7d752d6ed7dfa84319f0c332738b1 Mon Sep 17 00:00:00 2001
From: Maximilian Dax <maximilian.dax@tuebingen.mpg.de>
Date: Tue, 30 Jun 2026 14:55:24 +0200
Subject: [PATCH 05/65] Add per-factor batching to the factorized sampler
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Move batching from the ComposedSampler façade's single whole-chain loop down
to each Factor: a per-factor `batch_size` plus base-class batching wrappers
(`sample_and_log_prob`/`log_prob`) that chunk `num_samples`, slice the
conditioning to match, call the `_sample_and_log_prob`/`_log_prob` hooks, and
concatenate. The ChainComposer already materializes the full sample set
between factors (cheap parameter vectors), so each factor batches
independently at the size that fits its own memory footprint; the peak is the
single most expensive factor, not the whole chain at one global size.

GNPEFlowFactor.gibbs_step gets the same treatment, chunking the (independent)
Gibbs walkers. The façade keeps `run_sampler(batch_size=...)` as a convenience
that broadcasts a default across factors without their own. `batch_size=None`
short-circuits to the original single-pass path, so default behavior is
unchanged and plain-NPE/GNPE parity is preserved.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 114 +++++++++++++++++++++++++++++-----
 dingo/gw/inference/factors.py |  37 ++++++++---
 2 files changed, 130 insertions(+), 21 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index e0df3baf8..227afda5b 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -114,6 +114,28 @@ class Conditioning:
     given: dict[str, torch.Tensor] = field(default_factory=dict)
 
 
+def _cat_dict(
+    batches: list[dict[str, torch.Tensor]],
+) -> dict[str, torch.Tensor]:
+    """Concatenate a list of ``name -> tensor`` dicts along dim 0 (re-joining batches)."""
+    return {k: torch.cat([b[k] for b in batches]) for k in batches[0]}
+
+
+def _slice_given(cond: Conditioning, start: int, stop: int) -> Conditioning:
+    """A view of ``cond`` whose conditioning values are restricted to rows
+    ``[start:stop]``, so a factor's batch sees the conditioning aligned with the samples
+    it produces. The shared context is unchanged (the prepared data is one event)."""
+    return Conditioning(
+        cond.context, {k: v[start:stop] for k, v in cond.given.items()}
+    )
+
+
+def _batch_bounds(num_samples: int, batch_size: int):
+    """Yield ``(start, stop)`` chunks covering ``range(num_samples)``."""
+    for start in range(0, num_samples, batch_size):
+        yield start, min(start + batch_size, num_samples)
+
+
 class Factor(ABC):
     """
     A conditional distribution ``q_i(theta_i | f_i(theta_<i, d))`` over one block of
@@ -121,6 +143,15 @@ class Factor(ABC):
     standardization Jacobian already applied), so the chain product ``sum_i log q_i`` is
     directly the physical posterior log-density.
 
+    Batching is a per-factor concern: ``sample_and_log_prob`` / ``log_prob`` are
+    batching wrappers that split ``num_samples`` into chunks of ``batch_size`` (slicing
+    the conditioning to match), call the subclass hooks ``_sample_and_log_prob`` /
+    ``_log_prob`` per chunk, and concatenate. Because the chain composer already
+    materializes the full sample set between factors, each factor batches independently
+    at the size that fits *its own* memory footprint -- the peak is the single most
+    expensive factor at its own ``batch_size``, not the whole chain at one size.
+    ``batch_size = None`` runs the whole request in one pass (identical to no batching).
+
     Attributes
     ----------
     parameters : list[str]
@@ -128,23 +159,56 @@ class Factor(ABC):
     conditioning : list[str]
         Earlier-block parameter names this factor conditions on. Data dependence is
         implicit via the shared context.
+    batch_size : int, optional
+        Max number of samples processed in a single forward pass. ``None`` -> no
+        internal batching.
     """
 
     parameters: list[str]
     conditioning: list[str]
+    batch_size: Optional[int] = None
 
-    @abstractmethod
     def sample_and_log_prob(
         self, num_samples: int, cond: Conditioning
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """Draw ``theta_i ~ q_i``; return ``(physical samples, physical-space log q_i)``."""
+        """Draw ``theta_i ~ q_i`` (internally batched by ``batch_size``); return
+        ``(physical samples, physical-space log q_i)``."""
+        if not self.batch_size or self.batch_size >= num_samples:
+            return self._sample_and_log_prob(num_samples, cond)
+        sample_batches, lp_batches = [], []
+        for start, stop in _batch_bounds(num_samples, self.batch_size):
+            block, lp = self._sample_and_log_prob(
+                stop - start, _slice_given(cond, start, stop)
+            )
+            sample_batches.append(block)
+            lp_batches.append(lp)
+        return _cat_dict(sample_batches), torch.cat(lp_batches)
 
-    @abstractmethod
     def log_prob(
         self, theta_i: dict[str, torch.Tensor], cond: Conditioning
     ) -> torch.Tensor:
         """Evaluate ``log q_i(theta_i | f_i)`` at given physical ``theta_i`` (for IS /
-        re-plug)."""
+        re-plug), internally batched by ``batch_size``."""
+        num_samples = next(iter(theta_i.values())).shape[0]
+        if not self.batch_size or self.batch_size >= num_samples:
+            return self._log_prob(theta_i, cond)
+        lp_batches = []
+        for start, stop in _batch_bounds(num_samples, self.batch_size):
+            chunk = {k: v[start:stop] for k, v in theta_i.items()}
+            lp_batches.append(self._log_prob(chunk, _slice_given(cond, start, stop)))
+        return torch.cat(lp_batches)
+
+    @abstractmethod
+    def _sample_and_log_prob(
+        self, num_samples: int, cond: Conditioning
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        """Single-pass draw of ``num_samples`` from ``q_i`` (no internal batching)."""
+
+    @abstractmethod
+    def _log_prob(
+        self, theta_i: dict[str, torch.Tensor], cond: Conditioning
+    ) -> torch.Tensor:
+        """Single-pass evaluation of ``log q_i`` at given physical ``theta_i``."""
 
 
 def _base_model_metadata(model: BasePosteriorModel) -> dict:
@@ -176,17 +240,21 @@ def __init__(
         parameters: list[str],
         conditioning: Optional[list[str]] = None,
         context_parameters: Optional[list[str]] = None,
+        batch_size: Optional[int] = None,
     ):
         self.model = model
         self.parameters = parameters
         self.conditioning = conditioning or []
         # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
         self.context_parameters = context_parameters or []
+        self.batch_size = batch_size
         std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
         self.standardization = Standardization(std["mean"], std["std"])
 
     @classmethod
-    def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
+    def from_model(
+        cls, model: BasePosteriorModel, batch_size: Optional[int] = None
+    ) -> "FlowFactor":
         """Build a factor from a model: ``parameters`` and ``context_parameters`` come from
         the model's training metadata (first-class conditioning)."""
         data_settings = _base_model_metadata(model)["train_settings"]["data"]
@@ -196,6 +264,7 @@ def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
             parameters=data_settings["inference_parameters"],
             conditioning=list(context_parameters),
             context_parameters=list(context_parameters),
+            batch_size=batch_size,
         )
 
     def _network_context(self, cond: Conditioning) -> tuple[torch.Tensor, ...]:
@@ -208,7 +277,7 @@ def _network_context(self, cond: Conditioning) -> tuple[torch.Tensor, ...]:
         ctx = self.standardization.standardize(cond.given, self.context_parameters)
         return data.unsqueeze(0), ctx.unsqueeze(0)
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def _sample_and_log_prob(self, num_samples, cond):
         net_context = self._network_context(cond)
         self.model.network.eval()
         with torch.no_grad():
@@ -222,7 +291,7 @@ def sample_and_log_prob(self, num_samples, cond):
         log_prob = log_prob + self.standardization.log_det(self.parameters)
         return theta, log_prob
 
-    def log_prob(self, theta_i, cond):
+    def _log_prob(self, theta_i, cond):
         num_samples = next(iter(theta_i.values())).shape[0]
         z = self.standardization.standardize(theta_i, self.parameters)
         data = cond.context.prepared_data()
@@ -296,6 +365,14 @@ def sample(
         samples, log_prob = self.sample_and_log_prob(num_samples, context)
         return {**samples, "log_prob": log_prob}
 
+    def set_default_batch_size(self, batch_size: Optional[int]):
+        """Apply ``batch_size`` as the default to every factor that has not set its own
+        (per-factor ``batch_size`` takes precedence). The façade calls this to honour a
+        single ``run_sampler(batch_size=...)`` request across the whole chain."""
+        for factor in self.factors:
+            if factor.batch_size is None:
+                factor.batch_size = batch_size
+
 
 class GibbsComposer:
     """
@@ -330,6 +407,13 @@ def sample(
         }
         return {**parameters, **proxies}
 
+    def set_default_batch_size(self, batch_size: Optional[int]):
+        """Apply ``batch_size`` as the default to the init and GNPE factors that have not
+        set their own (per-factor ``batch_size`` takes precedence)."""
+        for factor in (self.init_factor, self.gnpe_factor):
+            if getattr(factor, "batch_size", None) is None:
+                factor.batch_size = batch_size
+
 
 class ComposedSampler:
     """
@@ -361,13 +445,15 @@ def _run_batch(self, num_samples: int) -> dict[str, torch.Tensor]:
     def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
-        if batch_size is None:
-            batch_size = num_samples
-        full_batches, remainder = divmod(num_samples, batch_size)
-        batches = [self._run_batch(batch_size) for _ in range(full_batches)]
-        if remainder > 0:
-            batches.append(self._run_batch(remainder))
-        merged = {k: torch.cat([b[k] for b in batches]) for k in batches[0].keys()}
+        """Run the chain and return samples as a DataFrame.
+
+        Batching is per-factor (each factor splits its own ``num_samples`` by its
+        ``batch_size``); ``batch_size`` here is a convenience that broadcasts a single
+        default across every factor without its own, reproducing one global batch size.
+        ``None`` -> no batching."""
+        if batch_size is not None:
+            self.composer.set_default_batch_size(batch_size)
+        merged = self._run_batch(num_samples)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
         self._post_process(merged)
         self.samples = pd.DataFrame(merged)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 1c2522689..735d6439e 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -28,6 +28,7 @@
     FlowFactor,
     GibbsComposer,
     _base_model_metadata,
+    _cat_dict,
 )
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
@@ -152,13 +153,13 @@ def __init__(self, values: dict[str, float]):
         self.parameters = list(values)
         self.conditioning = []
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def _sample_and_log_prob(self, num_samples, cond):
         samples = {
             p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
         }
         return samples, torch.zeros(num_samples)
 
-    def log_prob(self, theta_i, cond):
+    def _log_prob(self, theta_i, cond):
         raise NotImplementedError("FixedFactor.log_prob -- later step.")
 
 
@@ -171,10 +172,10 @@ def __init__(self):
         self.parameters = ["phase"]
         self.conditioning = []  # conditions on all preceding params via the likelihood
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def _sample_and_log_prob(self, num_samples, cond):
         raise NotImplementedError("SyntheticPhaseFactor -- later step.")
 
-    def log_prob(self, theta_i, cond):
+    def _log_prob(self, theta_i, cond):
         raise NotImplementedError("SyntheticPhaseFactor -- later step.")
 
 
@@ -197,6 +198,7 @@ def __init__(
         transform_post: Compose,
         gnpe_parameters: list[str],
         parameters: list[str],
+        batch_size: Optional[int] = None,
     ):
         self.model = model
         self.transform_pre = transform_pre
@@ -204,9 +206,12 @@ def __init__(
         self.gnpe_parameters = gnpe_parameters
         self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
         self.parameters = parameters
+        self.batch_size = batch_size
 
     @classmethod
-    def from_model(cls, model: BasePosteriorModel) -> "GNPEFlowFactor":
+    def from_model(
+        cls, model: BasePosteriorModel, batch_size: Optional[int] = None
+    ) -> "GNPEFlowFactor":
         """Build the GNPE per-iteration transforms from the main model's metadata,
         replicating ``GWSamplerGNPE._initialize_transforms`` (time-shift GNPE)."""
         meta = _base_model_metadata(model)
@@ -264,13 +269,31 @@ def from_model(cls, model: BasePosteriorModel) -> "GNPEFlowFactor":
             Compose(transform_post),
             gnpe_parameters,
             inference_parameters,
+            batch_size=batch_size,
         )
 
     def gibbs_step(
         self, num_samples: int, context, extrinsic: dict
     ) -> tuple[dict, dict]:
-        """One GNPE Gibbs iteration: returns (sampled parameters, updated extrinsic state
-        -- recomputed detector times + proxies for the next iteration)."""
+        """One GNPE Gibbs iteration, internally batched by ``batch_size`` over the Gibbs
+        walkers (independent across the batch dimension), to bound the per-pass memory.
+        ``batch_size = None`` runs all walkers in one pass."""
+        if not self.batch_size or self.batch_size >= num_samples:
+            return self._gibbs_step(num_samples, context, extrinsic)
+        param_batches, extrinsic_batches = [], []
+        for start in range(0, num_samples, self.batch_size):
+            stop = min(start + self.batch_size, num_samples)
+            chunk = {k: v[start:stop] for k, v in extrinsic.items()}
+            params, extr = self._gibbs_step(stop - start, context, chunk)
+            param_batches.append(params)
+            extrinsic_batches.append(extr)
+        return _cat_dict(param_batches), _cat_dict(extrinsic_batches)
+
+    def _gibbs_step(
+        self, num_samples: int, context, extrinsic: dict
+    ) -> tuple[dict, dict]:
+        """One GNPE Gibbs iteration (single pass): returns (sampled parameters, updated
+        extrinsic state -- recomputed detector times + proxies for the next iteration)."""
         x = {
             "extrinsic_parameters": {k: extrinsic[k] for k in self.gnpe_parameters},
             "parameters": {},

From f045821063306522b89e66ef415cbbce983603fc Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 14:57:24 +0200
Subject: [PATCH 06/65] Add GWComposedSampler.to_result()

Export the factorized sampler's output to a gw Result so the existing
post-processing pipeline -- synthetic phase, importance sampling, evidence,
plotting -- runs on it unchanged. `context` is the raw event data
(GWSamplerContext.raw_context), which Result needs to rebuild the likelihood.

Validated end-to-end: NPE -> to_result() -> importance_sample reproduces the
current GWSampler pipeline bit-exactly (log_prob, log_likelihood, log_prior,
weights, log_evidence all max|delta|=0); GNPE -> to_result() builds a valid
Result.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 735d6439e..6529d4407 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import copy
 from typing import Optional, Union
 
 import numpy as np
@@ -426,3 +427,24 @@ def _post_process(self, samples: Union[dict, pd.DataFrame], inverse: bool = Fals
                     samples.pop(k, None)
 
         self._correct_reference_time(samples, inverse)
+
+    def to_result(self):
+        """Export to a gw ``Result`` (samples + raw event data + metadata), so the
+        existing post-processing pipeline -- synthetic phase, importance sampling,
+        evidence, plotting -- runs on the factorized sampler's output unchanged.
+
+        ``context`` is the *raw* event-data dict (``GWSamplerContext.raw_context``, what
+        ``Result`` needs to rebuild the likelihood), not the ``SamplerContext`` object.
+        """
+        from dingo.gw.result import Result
+
+        data_dict = {
+            "samples": self.samples,
+            "context": self.context.raw_context,
+            "event_metadata": self.context.event_metadata,
+            "importance_sampling_metadata": None,
+            "log_evidence": None,
+            "log_noise_evidence": None,
+            "settings": copy.deepcopy(self.metadata),
+        }
+        return Result(dictionary=data_dict)

From cd0c502d88ae034f1180b1f292b43684f27a7846 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 15:02:18 +0200
Subject: [PATCH 07/65] Apply black formatting

Two cosmetic fixups in the merged factor files (docstring close + a one-line
call) so the branch is black-clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 4 +---
 dingo/gw/inference/factors.py | 3 ++-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 227afda5b..e3503ea60 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -125,9 +125,7 @@ def _slice_given(cond: Conditioning, start: int, stop: int) -> Conditioning:
     """A view of ``cond`` whose conditioning values are restricted to rows
     ``[start:stop]``, so a factor's batch sees the conditioning aligned with the samples
     it produces. The shared context is unchanged (the prepared data is one event)."""
-    return Conditioning(
-        cond.context, {k: v[start:stop] for k, v in cond.given.items()}
-    )
+    return Conditioning(cond.context, {k: v[start:stop] for k, v in cond.given.items()})
 
 
 def _batch_bounds(num_samples: int, batch_size: int):
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 6529d4407..e6e94c1f2 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -294,7 +294,8 @@ def _gibbs_step(
         self, num_samples: int, context, extrinsic: dict
     ) -> tuple[dict, dict]:
         """One GNPE Gibbs iteration (single pass): returns (sampled parameters, updated
-        extrinsic state -- recomputed detector times + proxies for the next iteration)."""
+        extrinsic state -- recomputed detector times + proxies for the next iteration).
+        """
         x = {
             "extrinsic_parameters": {k: extrinsic[k] for k in self.gnpe_parameters},
             "parameters": {},

From 79108d2f4daf7df73194c9e7966507334afade3b Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 22:59:58 +0200
Subject: [PATCH 08/65] Replace per-factor batching with a vertical chain
 executor

Pull batching out of the factors into a single vertical (depth-first)
executor on ChainComposer, driven by a shared chunk_and_concat helper
(also used by GibbsComposer). Add Stage(factor, fan_out) so a chain can
expand num_samples per conditioning row (M intrinsic x K extrinsic);
num_samples is the base/root count, total rows = num_samples * expansion.

Factors become single-pass: drop batch_size, the batching wrapper methods,
and _slice_given / _batch_bounds / set_default_batch_size / _gibbs_step
(keep _cat_dict); inline ComposedSampler._run_batch into run_sampler. Fix
FlowFactor's conditioned path to draw num_samples per conditioning row
(N-row context + expanded data), matching what GNPEFlowFactor already does.

Rationale: only the vertical executor bounds fan-out memory (you cannot
materialize M*K intermediates for 1000-extrinsic-per-intrinsic), and
chunking the whole chain / Gibbs loop per base-chunk reproduces the old
sampler's batching order, so parity is exact. See the hackathon design doc
section 14.

Bit-exact vs GWSampler (NPE) and GWSamplerGNPE (multi-iteration GNPE),
batched and unbatched, including post-processing. Add tests/core/test_factors.py
(mock-factor unit tests for fan-out expansion + alignment, log-prob
summation, topological validation, and the batching primitive).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 298 +++++++++++++++++++---------------
 dingo/gw/inference/factors.py |  34 +---
 tests/core/test_factors.py    | 148 +++++++++++++++++
 3 files changed, 321 insertions(+), 159 deletions(-)
 create mode 100644 tests/core/test_factors.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index e3503ea60..0e36e4872 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -24,7 +24,7 @@
 import math
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Optional, Protocol, Union, runtime_checkable
+from typing import Callable, Optional, Protocol, Union, runtime_checkable
 
 import pandas as pd
 import torch
@@ -121,17 +121,33 @@ def _cat_dict(
     return {k: torch.cat([b[k] for b in batches]) for k in batches[0]}
 
 
-def _slice_given(cond: Conditioning, start: int, stop: int) -> Conditioning:
-    """A view of ``cond`` whose conditioning values are restricted to rows
-    ``[start:stop]``, so a factor's batch sees the conditioning aligned with the samples
-    it produces. The shared context is unchanged (the prepared data is one event)."""
-    return Conditioning(cond.context, {k: v[start:stop] for k, v in cond.given.items()})
-
-
-def _batch_bounds(num_samples: int, batch_size: int):
-    """Yield ``(start, stop)`` chunks covering ``range(num_samples)``."""
-    for start in range(0, num_samples, batch_size):
-        yield start, min(start + batch_size, num_samples)
+def chunk_and_concat(
+    total: int,
+    batch_size: Optional[int],
+    run_once: Callable[[int], tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]],
+) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
+    """The vertical (depth-first) batching primitive, shared by the chain and Gibbs
+    composers.
+
+    Split ``total`` base samples into chunks of ``batch_size``, run ``run_once(n)`` -- a
+    full pass through the *whole* pipeline (the entire chain, or the entire Gibbs loop) --
+    per chunk, and concatenate. Working memory is bounded to one chunk carried through the
+    pipeline; only the output buffer grows to ``total`` (a real run streams it to CPU per
+    chunk). This is the order the old ``Sampler.run_sampler`` batches in, so parity is
+    exact when ``batch_size`` matches. ``batch_size=None`` runs everything in one pass.
+
+    ``run_once`` returns ``(samples, log_prob)``; ``log_prob`` may be ``None`` (the Gibbs
+    path has no tractable density)."""
+    bs = batch_size or total
+    sample_parts: list[dict[str, torch.Tensor]] = []
+    lp_parts: list[Optional[torch.Tensor]] = []
+    for start in range(0, total, bs):
+        block, lp = run_once(min(bs, total - start))
+        sample_parts.append(block)
+        lp_parts.append(lp)
+    samples = _cat_dict(sample_parts)
+    log_prob = None if lp_parts[0] is None else torch.cat(lp_parts)
+    return samples, log_prob
 
 
 class Factor(ABC):
@@ -141,14 +157,14 @@ class Factor(ABC):
     standardization Jacobian already applied), so the chain product ``sum_i log q_i`` is
     directly the physical posterior log-density.
 
-    Batching is a per-factor concern: ``sample_and_log_prob`` / ``log_prob`` are
-    batching wrappers that split ``num_samples`` into chunks of ``batch_size`` (slicing
-    the conditioning to match), call the subclass hooks ``_sample_and_log_prob`` /
-    ``_log_prob`` per chunk, and concatenate. Because the chain composer already
-    materializes the full sample set between factors, each factor batches independently
-    at the size that fits *its own* memory footprint -- the peak is the single most
-    expensive factor at its own ``batch_size``, not the whole chain at one size.
-    ``batch_size = None`` runs the whole request in one pass (identical to no batching).
+    *Physical in, physical out*, and a **single forward pass per call** -- batching is not
+    a factor concern; the composer drives it (``chunk_and_concat``).
+
+    Contract the composer relies on (this is what makes fan-out work): a call draws
+    ``num_samples`` samples **per conditioning row** and returns ``n_rows * num_samples``
+    rows, flattened row-major, where ``n_rows`` is the number of rows in ``cond.given``
+    (or 1 if unconditioned). This mirrors the network's own
+    ``sample_and_log_prob(*context, num_samples=n) -> (B, n, dim)`` with ``B = n_rows``.
 
     Attributes
     ----------
@@ -157,56 +173,25 @@ class Factor(ABC):
     conditioning : list[str]
         Earlier-block parameter names this factor conditions on. Data dependence is
         implicit via the shared context.
-    batch_size : int, optional
-        Max number of samples processed in a single forward pass. ``None`` -> no
-        internal batching.
     """
 
     parameters: list[str]
     conditioning: list[str]
-    batch_size: Optional[int] = None
 
+    @abstractmethod
     def sample_and_log_prob(
         self, num_samples: int, cond: Conditioning
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """Draw ``theta_i ~ q_i`` (internally batched by ``batch_size``); return
-        ``(physical samples, physical-space log q_i)``."""
-        if not self.batch_size or self.batch_size >= num_samples:
-            return self._sample_and_log_prob(num_samples, cond)
-        sample_batches, lp_batches = [], []
-        for start, stop in _batch_bounds(num_samples, self.batch_size):
-            block, lp = self._sample_and_log_prob(
-                stop - start, _slice_given(cond, start, stop)
-            )
-            sample_batches.append(block)
-            lp_batches.append(lp)
-        return _cat_dict(sample_batches), torch.cat(lp_batches)
+        """Draw ``num_samples`` per conditioning row from ``q_i``; return ``(physical
+        samples, physical-space log q_i)``. The returned dict may carry extra named
+        side-channel columns (e.g. an IS kernel correction) alongside ``parameters``."""
 
+    @abstractmethod
     def log_prob(
         self, theta_i: dict[str, torch.Tensor], cond: Conditioning
     ) -> torch.Tensor:
         """Evaluate ``log q_i(theta_i | f_i)`` at given physical ``theta_i`` (for IS /
-        re-plug), internally batched by ``batch_size``."""
-        num_samples = next(iter(theta_i.values())).shape[0]
-        if not self.batch_size or self.batch_size >= num_samples:
-            return self._log_prob(theta_i, cond)
-        lp_batches = []
-        for start, stop in _batch_bounds(num_samples, self.batch_size):
-            chunk = {k: v[start:stop] for k, v in theta_i.items()}
-            lp_batches.append(self._log_prob(chunk, _slice_given(cond, start, stop)))
-        return torch.cat(lp_batches)
-
-    @abstractmethod
-    def _sample_and_log_prob(
-        self, num_samples: int, cond: Conditioning
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """Single-pass draw of ``num_samples`` from ``q_i`` (no internal batching)."""
-
-    @abstractmethod
-    def _log_prob(
-        self, theta_i: dict[str, torch.Tensor], cond: Conditioning
-    ) -> torch.Tensor:
-        """Single-pass evaluation of ``log q_i`` at given physical ``theta_i``."""
+        re-plug). One conditioning row per ``theta_i`` row (no fan-out)."""
 
 
 def _base_model_metadata(model: BasePosteriorModel) -> dict:
@@ -238,21 +223,17 @@ def __init__(
         parameters: list[str],
         conditioning: Optional[list[str]] = None,
         context_parameters: Optional[list[str]] = None,
-        batch_size: Optional[int] = None,
     ):
         self.model = model
         self.parameters = parameters
         self.conditioning = conditioning or []
         # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
         self.context_parameters = context_parameters or []
-        self.batch_size = batch_size
         std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
         self.standardization = Standardization(std["mean"], std["std"])
 
     @classmethod
-    def from_model(
-        cls, model: BasePosteriorModel, batch_size: Optional[int] = None
-    ) -> "FlowFactor":
+    def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
         """Build a factor from a model: ``parameters`` and ``context_parameters`` come from
         the model's training metadata (first-class conditioning)."""
         data_settings = _base_model_metadata(model)["train_settings"]["data"]
@@ -262,34 +243,41 @@ def from_model(
             parameters=data_settings["inference_parameters"],
             conditioning=list(context_parameters),
             context_parameters=list(context_parameters),
-            batch_size=batch_size,
         )
 
-    def _network_context(self, cond: Conditioning) -> tuple[torch.Tensor, ...]:
-        """Assemble the model's positional context: the prepared data, plus standardized
-        context parameters when the network conditions on proxies."""
+    def sample_and_log_prob(self, num_samples, cond):
+        """Draw ``num_samples`` per conditioning row. Unconditioned (plain NPE): one
+        shared data context, ``num_samples`` draws. Conditioned: ``N`` context rows in
+        (one embedding pass each), ``num_samples`` draws per row -> ``N * num_samples``
+        rows, flattened row-major to align with the composer's ``repeat_interleave``."""
         data = cond.context.prepared_data()
-        if not self.context_parameters:
-            return (data.unsqueeze(0),)
-        # Standardize the (physical) conditioning values the network was trained on.
-        ctx = self.standardization.standardize(cond.given, self.context_parameters)
-        return data.unsqueeze(0), ctx.unsqueeze(0)
-
-    def _sample_and_log_prob(self, num_samples, cond):
-        net_context = self._network_context(cond)
         self.model.network.eval()
-        with torch.no_grad():
-            z, log_prob = self.model.sample_and_log_prob(
-                *net_context, num_samples=num_samples
-            )
-        # Squeeze the batch dimension added for the (single) context.
-        z = z.squeeze(0)
-        log_prob = log_prob.squeeze(0)
+        if not self.context_parameters:
+            with torch.no_grad():
+                z, log_prob = self.model.sample_and_log_prob(
+                    data.unsqueeze(0), num_samples=num_samples
+                )
+            # Squeeze the batch dimension added for the single shared context.
+            z = z.squeeze(0)
+            log_prob = log_prob.squeeze(0)
+        else:
+            # Standardize the (physical) conditioning the network was trained on, giving
+            # B = N context rows; expand the shared data to match. The network draws
+            # num_samples per row -> (N, num_samples, dim).
+            ctx = self.standardization.standardize(cond.given, self.context_parameters)
+            n_rows = ctx.shape[0]
+            data = data.expand(n_rows, *data.shape)
+            with torch.no_grad():
+                z, log_prob = self.model.sample_and_log_prob(
+                    data, ctx, num_samples=num_samples
+                )
+            z = z.reshape(n_rows * num_samples, z.shape[-1])
+            log_prob = log_prob.reshape(n_rows * num_samples)
         theta = self.standardization.destandardize(z, self.parameters)
         log_prob = log_prob + self.standardization.log_det(self.parameters)
         return theta, log_prob
 
-    def _log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, cond):
         num_samples = next(iter(theta_i.values())).shape[0]
         z = self.standardization.standardize(theta_i, self.parameters)
         data = cond.context.prepared_data()
@@ -306,19 +294,37 @@ def _log_prob(self, theta_i, cond):
         return log_prob + self.standardization.log_det(self.parameters)
 
 
+@dataclass
+class Stage:
+    """One factor in the chain plus its **fan-out**: how many samples it draws per
+    incoming conditioning row. The root (first) stage ignores ``fan_out`` -- it draws the
+    base count directly; a 1:1 chain link uses ``fan_out=1``; an intrinsic/extrinsic
+    expansion uses ``fan_out=K`` (``K`` extrinsic draws per intrinsic sample)."""
+
+    factor: Factor
+    fan_out: int = 1
+
+
 class ChainComposer:
     """
-    Autoregressive composer: sample (or evaluate) the factors in declared order, which
-    must be a topological order of the dependency DAG, and sum their log-probs. Exact
-    and importance-sampling-friendly. Covers plain NPE, single-step GNPE, prior
-    conditioning, synthetic phase, and intrinsic/extrinsic splits.
-
-    Multi-iteration (cyclic) GNPE is *not* expressible here and uses a separate Gibbs
-    composer.
+    Autoregressive composer **and vertical (depth-first) executor**: sample the stages in
+    declared order (a topological order of the dependency DAG), expanding by each stage's
+    fan-out, and sum their log-probs. Exact and importance-sampling-friendly. Covers plain
+    NPE, single-step GNPE, prior conditioning, synthetic phase, and intrinsic/extrinsic
+    splits.
+
+    Vertical execution carries a chunk of base samples all the way down the chain, emits
+    it, and discards -- bounding working memory to one chunk (the only way to draw ``M``
+    intrinsic x ``K`` extrinsic without materializing ``M*K`` intermediates) and
+    reproducing the old sampler's batching order (bit-exact when ``batch_size`` matches).
+    Batching lives in ``chunk_and_concat``; the factors are single-pass.
+
+    Multi-iteration (cyclic) GNPE is *not* expressible here and uses ``GibbsComposer``.
+    Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or explicit ``Stage``s.
     """
 
-    def __init__(self, factors: list[Factor]):
-        self.factors = factors
+    def __init__(self, stages: list[Union["Stage", Factor]]):
+        self.stages = [s if isinstance(s, Stage) else Stage(s) for s in stages]
         self._validate()
 
     def _validate(self):
@@ -334,14 +340,51 @@ def _validate(self):
                 )
             produced.update(factor.parameters)
 
+    @property
+    def factors(self) -> list[Factor]:
+        """The factors in order (without fan-out), for callers that only need them."""
+        return [stage.factor for stage in self.stages]
+
+    @property
+    def expansion(self) -> int:
+        """Product of the non-root fan-outs; total rows returned = num_samples *
+        expansion (e.g. M intrinsic x K extrinsic for an extrinsic fan-out of K)."""
+        return math.prod(stage.fan_out for stage in self.stages[1:])
+
     def sample_and_log_prob(
-        self, num_samples: int, context: SamplerContext
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        batch_size: Optional[int] = None,
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        """``num_samples`` is the **base (root) count** -- e.g. ``M`` intrinsic samples.
+        Each fan-out stage multiplies it, so the result has ``num_samples * expansion``
+        rows (``M * K`` for an extrinsic ``fan_out=K``); for a plain chain (no fan-out)
+        that is just ``num_samples``. ``batch_size`` chunks the base count, in the same
+        (root) unit as ``num_samples``."""
+        return chunk_and_concat(
+            num_samples, batch_size, lambda n: self._run_chain_once(n, context)
+        )
+
+    def _run_chain_once(
+        self, base: int, context: SamplerContext
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        """One depth-first pass of the whole chain for ``base`` root samples."""
         samples: dict[str, torch.Tensor] = {}
         total: torch.Tensor | float = 0.0
-        for factor in self.factors:
+        for i, stage in enumerate(self.stages):
+            factor = stage.factor
+            n = base if i == 0 else stage.fan_out
             cond = Conditioning(context, {k: samples[k] for k in factor.conditioning})
-            block, lp = factor.sample_and_log_prob(num_samples, cond)
+            block, lp = factor.sample_and_log_prob(n, cond)
+            # The factor returned (rows_so_far * n) rows. Expand the carried columns to
+            # match -- each upstream row repeated n times (the block is flattened in the
+            # same row-major order). The root has no carried rows, and fan_out=1 stages
+            # are a no-op, so 1:1 chains are untouched.
+            if i > 0 and n > 1:
+                samples = {k: v.repeat_interleave(n, 0) for k, v in samples.items()}
+                if torch.is_tensor(total):
+                    total = total.repeat_interleave(n, 0)
             samples.update(block)
             total = total + lp
         return samples, total
@@ -357,20 +400,15 @@ def log_prob(
         return total
 
     def sample(
-        self, num_samples: int, context: SamplerContext
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        batch_size: Optional[int] = None,
     ) -> dict[str, torch.Tensor]:
         """Full per-sample dict (parameters + ``log_prob``), for the sampler façade."""
-        samples, log_prob = self.sample_and_log_prob(num_samples, context)
+        samples, log_prob = self.sample_and_log_prob(num_samples, context, batch_size)
         return {**samples, "log_prob": log_prob}
 
-    def set_default_batch_size(self, batch_size: Optional[int]):
-        """Apply ``batch_size`` as the default to every factor that has not set its own
-        (per-factor ``batch_size`` takes precedence). The façade calls this to honour a
-        single ``run_sampler(batch_size=...)`` request across the whole chain."""
-        for factor in self.factors:
-            if factor.batch_size is None:
-                factor.batch_size = batch_size
-
 
 class GibbsComposer:
     """
@@ -379,7 +417,10 @@ class GibbsComposer:
     (``gnpe_factor.gibbs_step``) to a fixed point. The dependency is cyclic
     (theta <-> proxies), so this breaks density access: ``sample`` returns parameters
     WITHOUT a ``log_prob``. (Single-step GNPE that *does* preserve log_prob is a
-    two-factor ``ChainComposer`` instead.)
+    two-stage ``ChainComposer`` instead.)
+
+    Batching is vertical, via ``chunk_and_concat``: chunk the Gibbs walkers and run the
+    whole loop per chunk -- the old ``GWSamplerGNPE.run_sampler`` order.
     """
 
     def __init__(self, init_factor: Factor, gnpe_factor, num_iterations: int):
@@ -388,6 +429,17 @@ def __init__(self, init_factor: Factor, gnpe_factor, num_iterations: int):
         self.num_iterations = num_iterations
 
     def sample(
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        batch_size: Optional[int] = None,
+    ) -> dict[str, torch.Tensor]:
+        samples, _ = chunk_and_concat(
+            num_samples, batch_size, lambda n: (self._run_once(n, context), None)
+        )
+        return samples
+
+    def _run_once(
         self, num_samples: int, context: SamplerContext
     ) -> dict[str, torch.Tensor]:
         # Seed the Gibbs chain with the init factor's parameters (e.g. detector times).
@@ -405,22 +457,15 @@ def sample(
         }
         return {**parameters, **proxies}
 
-    def set_default_batch_size(self, batch_size: Optional[int]):
-        """Apply ``batch_size`` as the default to the init and GNPE factors that have not
-        set their own (per-factor ``batch_size`` takes precedence)."""
-        for factor in (self.init_factor, self.gnpe_factor):
-            if getattr(factor, "batch_size", None) is None:
-                factor.batch_size = batch_size
-
 
 class ComposedSampler:
     """
     Thin user-facing façade over a composer (``ChainComposer`` or ``GibbsComposer``) and
-    a ``SamplerContext``: runs the composer with optional batching, applies domain-specific
-    post-processing, and returns samples as a DataFrame. The per-factor compute lives in
-    the composer (which need only expose ``sample(num_samples, context) -> dict``); this
-    class only handles batching, consolidation, and post-processing -- the role the
-    monolithic ``Sampler`` plays today.
+    a ``SamplerContext``: runs the composer (which bounds memory via vertical batching),
+    applies domain-specific post-processing, and returns samples as a DataFrame. The
+    per-factor compute lives in the composer (which need only expose
+    ``sample(num_samples, context, batch_size) -> dict``); this class only handles
+    consolidation and post-processing -- the role the monolithic ``Sampler`` plays today.
     """
 
     def __init__(
@@ -435,23 +480,14 @@ def _post_process(self, samples: dict, inverse: bool = False):
         and reference-time correction). No-op by default."""
         pass
 
-    def _run_batch(self, num_samples: int) -> dict[str, torch.Tensor]:
-        # Uniform across composers: ChainComposer.sample adds log_prob; GibbsComposer
-        # (multi-iteration GNPE) returns parameters + proxies with no log_prob.
-        return self.composer.sample(num_samples, self.context)
-
     def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
-        """Run the chain and return samples as a DataFrame.
-
-        Batching is per-factor (each factor splits its own ``num_samples`` by its
-        ``batch_size``); ``batch_size`` here is a convenience that broadcasts a single
-        default across every factor without its own, reproducing one global batch size.
-        ``None`` -> no batching."""
-        if batch_size is not None:
-            self.composer.set_default_batch_size(batch_size)
-        merged = self._run_batch(num_samples)
+        """Run the composer (vertically batched by ``batch_size``) and return samples as a
+        DataFrame. ``ChainComposer.sample`` adds ``log_prob``; ``GibbsComposer``
+        (multi-iteration GNPE) returns parameters + proxies with no ``log_prob``.
+        ``batch_size=None`` runs in a single pass."""
+        merged = self.composer.sample(num_samples, self.context, batch_size)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
         self._post_process(merged)
         self.samples = pd.DataFrame(merged)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index e6e94c1f2..78c22fc03 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -29,7 +29,6 @@
     FlowFactor,
     GibbsComposer,
     _base_model_metadata,
-    _cat_dict,
 )
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
@@ -154,13 +153,13 @@ def __init__(self, values: dict[str, float]):
         self.parameters = list(values)
         self.conditioning = []
 
-    def _sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, cond):
         samples = {
             p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
         }
         return samples, torch.zeros(num_samples)
 
-    def _log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, cond):
         raise NotImplementedError("FixedFactor.log_prob -- later step.")
 
 
@@ -173,10 +172,10 @@ def __init__(self):
         self.parameters = ["phase"]
         self.conditioning = []  # conditions on all preceding params via the likelihood
 
-    def _sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, cond):
         raise NotImplementedError("SyntheticPhaseFactor -- later step.")
 
-    def _log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, cond):
         raise NotImplementedError("SyntheticPhaseFactor -- later step.")
 
 
@@ -199,7 +198,6 @@ def __init__(
         transform_post: Compose,
         gnpe_parameters: list[str],
         parameters: list[str],
-        batch_size: Optional[int] = None,
     ):
         self.model = model
         self.transform_pre = transform_pre
@@ -207,12 +205,9 @@ def __init__(
         self.gnpe_parameters = gnpe_parameters
         self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
         self.parameters = parameters
-        self.batch_size = batch_size
 
     @classmethod
-    def from_model(
-        cls, model: BasePosteriorModel, batch_size: Optional[int] = None
-    ) -> "GNPEFlowFactor":
+    def from_model(cls, model: BasePosteriorModel) -> "GNPEFlowFactor":
         """Build the GNPE per-iteration transforms from the main model's metadata,
         replicating ``GWSamplerGNPE._initialize_transforms`` (time-shift GNPE)."""
         meta = _base_model_metadata(model)
@@ -270,31 +265,14 @@ def from_model(
             Compose(transform_post),
             gnpe_parameters,
             inference_parameters,
-            batch_size=batch_size,
         )
 
     def gibbs_step(
         self, num_samples: int, context, extrinsic: dict
-    ) -> tuple[dict, dict]:
-        """One GNPE Gibbs iteration, internally batched by ``batch_size`` over the Gibbs
-        walkers (independent across the batch dimension), to bound the per-pass memory.
-        ``batch_size = None`` runs all walkers in one pass."""
-        if not self.batch_size or self.batch_size >= num_samples:
-            return self._gibbs_step(num_samples, context, extrinsic)
-        param_batches, extrinsic_batches = [], []
-        for start in range(0, num_samples, self.batch_size):
-            stop = min(start + self.batch_size, num_samples)
-            chunk = {k: v[start:stop] for k, v in extrinsic.items()}
-            params, extr = self._gibbs_step(stop - start, context, chunk)
-            param_batches.append(params)
-            extrinsic_batches.append(extr)
-        return _cat_dict(param_batches), _cat_dict(extrinsic_batches)
-
-    def _gibbs_step(
-        self, num_samples: int, context, extrinsic: dict
     ) -> tuple[dict, dict]:
         """One GNPE Gibbs iteration (single pass): returns (sampled parameters, updated
         extrinsic state -- recomputed detector times + proxies for the next iteration).
+        Batching over the Gibbs walkers is the composer's job (``chunk_and_concat``).
         """
         x = {
             "extrinsic_parameters": {k: extrinsic[k] for k in self.gnpe_parameters},
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
new file mode 100644
index 000000000..58d7bc3d6
--- /dev/null
+++ b/tests/core/test_factors.py
@@ -0,0 +1,148 @@
+"""
+Unit tests for the factorized-sampler core executor (``dingo.core.factors``).
+
+These exercise the vertical ``ChainComposer`` and ``chunk_and_concat`` with deterministic
+mock factors (no networks), so they are portable and fast: fan-out expansion + conditioning
+alignment, log-prob summation, topological validation, and the batching plumbing. Bit-exact
+parity against the legacy samplers is a model-based check that lives with the GW models, not
+here.
+"""
+
+import pytest
+import torch
+
+from dingo.core.factors import (
+    ChainComposer,
+    Conditioning,
+    Factor,
+    Stage,
+    chunk_and_concat,
+)
+
+
+class _ConstFactor(Factor):
+    """Deterministic factor (no RNG) honoring the executor contract: draws ``n`` samples
+    per conditioning row. Unconditioned, it returns ``arange(n)``; conditioned, it returns
+    ``sum(given) + within_row_index/1000`` so fan-out alignment is checkable by eye. Each
+    factor contributes a constant ``0.5`` to the log-prob."""
+
+    def __init__(self, name, conditioning=()):
+        self.parameters = [name]
+        self.conditioning = list(conditioning)
+        self._name = name
+
+    def sample_and_log_prob(self, n, cond):
+        if self.conditioning:
+            base = sum(cond.given[k] for k in self.conditioning)  # (N,)
+            n_rows = base.shape[0]
+            # Integer-valued within-row offset: exact in float32 (no cancellation when
+            # the test recovers it as b - a), and distinct per fan-out draw.
+            within = torch.arange(n, dtype=base.dtype)
+            vals = (base.unsqueeze(1) + within).reshape(-1)  # (N*n,), row-major
+            lp = torch.full((n_rows * n,), 0.5)
+        else:
+            vals = torch.arange(n, dtype=torch.float32)
+            lp = torch.full((n,), 0.5)
+        return {self._name: vals}, lp
+
+    def log_prob(self, theta_i, cond):
+        return torch.zeros(next(iter(theta_i.values())).shape[0])
+
+
+def test_root_only_draws_num_samples():
+    comp = ChainComposer([_ConstFactor("a")])
+    assert comp.expansion == 1
+    out = comp.sample(5, context=None)
+    assert torch.equal(out["a"], torch.arange(5, dtype=torch.float32))
+    assert torch.allclose(out["log_prob"], torch.full((5,), 0.5))
+
+
+def test_bare_factor_is_wrapped_as_stage():
+    comp = ChainComposer([_ConstFactor("a")])
+    assert isinstance(comp.stages[0], Stage)
+    assert comp.stages[0].fan_out == 1
+    assert comp.factors[0].parameters == ["a"]
+
+
+def test_fan_out_expansion_alignment_and_logprob():
+    m, k = 4, 3
+    comp = ChainComposer(
+        [
+            Stage(_ConstFactor("a")),  # root -> base m
+            Stage(_ConstFactor("b", conditioning=["a"]), fan_out=k),
+        ]
+    )
+    assert comp.expansion == k
+    out = comp.sample(m, context=None)  # num_samples is the base (M intrinsic)
+    assert out["a"].shape == (m * k,) and out["b"].shape == (m * k,)
+
+    # 'a' is each root value repeated k times (repeat_interleave layout).
+    a = out["a"].reshape(m, k)
+    assert torch.all(a == a[:, :1])
+    assert torch.equal(a[:, 0], torch.arange(m, dtype=torch.float32))
+
+    # 'b' conditions on the matching 'a': (b - a) recovers the within-row fan-out index,
+    # which is distinct per draw (0..k-1) and identical across rows.
+    b = out["b"].reshape(m, k)
+    expected_within = torch.arange(k, dtype=torch.float32).repeat(m, 1)
+    assert torch.equal(b - a, expected_within)
+
+    # log_prob is the sum over both factors (0.5 + 0.5).
+    assert torch.allclose(out["log_prob"], torch.full((m * k,), 1.0))
+
+
+def test_chained_fan_out_product():
+    comp = ChainComposer(
+        [
+            Stage(_ConstFactor("a")),
+            Stage(_ConstFactor("b", conditioning=["a"]), fan_out=2),
+            Stage(_ConstFactor("c", conditioning=["b"]), fan_out=5),
+        ]
+    )
+    assert comp.expansion == 10
+    out = comp.sample(3, context=None)  # base 3 -> 3 * 10 = 30 rows
+    assert all(out[p].shape == (30,) for p in ("a", "b", "c"))
+
+
+def test_num_samples_is_the_base_count():
+    # num_samples is the base (root) count, not the total: total rows = num_samples *
+    # expansion, with no divisibility constraint (here 7 is not a multiple of 4).
+    m, k = 7, 4
+    comp = ChainComposer(
+        [
+            Stage(_ConstFactor("a")),
+            Stage(_ConstFactor("b", conditioning=["a"]), fan_out=k),
+        ]
+    )
+    out = comp.sample(m, context=None)
+    assert out["a"].shape == (m * k,) and out["b"].shape == (m * k,)
+
+
+def test_topological_validation():
+    with pytest.raises(ValueError, match="conditions on"):
+        ChainComposer([_ConstFactor("b", conditioning=["a"])])  # 'a' never produced
+
+
+def test_chunk_and_concat_plumbing():
+    # Position-independent run_once -> batched concat is identical to a single pass.
+    def run_once(n):
+        return {"x": torch.full((n,), 7.0)}, torch.ones(n)
+
+    s_full, lp_full = chunk_and_concat(10, None, run_once)
+    s_chunk, lp_chunk = chunk_and_concat(10, 3, run_once)  # 3+3+3+1
+    assert s_full["x"].shape == (10,)
+    assert torch.equal(s_full["x"], s_chunk["x"])
+    assert torch.equal(lp_full, lp_chunk)
+
+
+def test_chunk_and_concat_allows_none_log_prob():
+    # The Gibbs path carries no density.
+    samples, lp = chunk_and_concat(6, 2, lambda n: ({"x": torch.zeros(n)}, None))
+    assert lp is None
+    assert samples["x"].shape == (6,)
+
+
+def test_conditioning_dataclass_defaults():
+    cond = Conditioning(context="ctx")
+    assert cond.context == "ctx"
+    assert cond.given == {}

From 0a2f9f364c59098aab36a7ee2b320f37f3615ad6 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 23:26:23 +0200
Subject: [PATCH 09/65] Rename GibbsComposer to GNPEGibbsComposer; add Composer
 protocol

Gibbs sampling in Dingo is used only for GNPE, and the class's members
(gnpe_factor, proxy_parameters, gibbs_step) are all GNPE-specific, so the
general name was misleading. Rename to GNPEGibbsComposer and keep it in
dingo.gw.inference.factors alongside GNPEFlowFactor.

Core keeps the domain-agnostic pieces -- ChainComposer, chunk_and_concat --
and gains a Composer protocol that ComposedSampler depends on, so core no
longer references the GW-specific composer by name.

GNPE sampling remains bit-exact vs GWSamplerGNPE.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 67 +++++++++--------------------------
 dingo/gw/inference/factors.py | 66 +++++++++++++++++++++++++++++++---
 2 files changed, 79 insertions(+), 54 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 0e36e4872..96b3ac10d 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -319,8 +319,9 @@ class ChainComposer:
     reproducing the old sampler's batching order (bit-exact when ``batch_size`` matches).
     Batching lives in ``chunk_and_concat``; the factors are single-pass.
 
-    Multi-iteration (cyclic) GNPE is *not* expressible here and uses ``GibbsComposer``.
-    Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or explicit ``Stage``s.
+    Multi-iteration (cyclic) GNPE is *not* expressible here and uses the GW
+    ``GNPEGibbsComposer``. Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or
+    explicit ``Stage``s.
     """
 
     def __init__(self, stages: list[Union["Stage", Factor]]):
@@ -410,67 +411,33 @@ def sample(
         return {**samples, "log_prob": log_prob}
 
 
-class GibbsComposer:
-    """
-    Multi-iteration GNPE composition. Seeds the Gibbs chain from an ``init_factor`` (e.g.
-    a network predicting detector coalescence times), then iterates a GNPE step
-    (``gnpe_factor.gibbs_step``) to a fixed point. The dependency is cyclic
-    (theta <-> proxies), so this breaks density access: ``sample`` returns parameters
-    WITHOUT a ``log_prob``. (Single-step GNPE that *does* preserve log_prob is a
-    two-stage ``ChainComposer`` instead.)
-
-    Batching is vertical, via ``chunk_and_concat``: chunk the Gibbs walkers and run the
-    whole loop per chunk -- the old ``GWSamplerGNPE.run_sampler`` order.
-    """
-
-    def __init__(self, init_factor: Factor, gnpe_factor, num_iterations: int):
-        self.init_factor = init_factor
-        self.gnpe_factor = gnpe_factor
-        self.num_iterations = num_iterations
+@runtime_checkable
+class Composer(Protocol):
+    """The minimal interface ``ComposedSampler`` drives: draw a per-sample dict
+    (parameters, plus ``log_prob`` for density-preserving composers). ``ChainComposer``
+    (here, domain-agnostic) satisfies it, as does the GW ``GNPEGibbsComposer`` (multi-iteration
+    GNPE) -- which lives in ``dingo.gw.inference.factors`` because it is built entirely
+    around the GNPE step (its proxies, detector times, time-shift), not in core."""
 
     def sample(
         self,
         num_samples: int,
         context: SamplerContext,
         batch_size: Optional[int] = None,
-    ) -> dict[str, torch.Tensor]:
-        samples, _ = chunk_and_concat(
-            num_samples, batch_size, lambda n: (self._run_once(n, context), None)
-        )
-        return samples
-
-    def _run_once(
-        self, num_samples: int, context: SamplerContext
-    ) -> dict[str, torch.Tensor]:
-        # Seed the Gibbs chain with the init factor's parameters (e.g. detector times).
-        seed, _ = self.init_factor.sample_and_log_prob(
-            num_samples, Conditioning(context)
-        )
-        extrinsic = dict(seed)
-        parameters: dict[str, torch.Tensor] = {}
-        for _ in range(self.num_iterations):
-            parameters, extrinsic = self.gnpe_factor.gibbs_step(
-                num_samples, context, extrinsic
-            )
-        proxies = {
-            p: extrinsic[p] for p in self.gnpe_factor.proxy_parameters if p in extrinsic
-        }
-        return {**parameters, **proxies}
+    ) -> dict[str, torch.Tensor]: ...
 
 
 class ComposedSampler:
     """
-    Thin user-facing façade over a composer (``ChainComposer`` or ``GibbsComposer``) and
-    a ``SamplerContext``: runs the composer (which bounds memory via vertical batching),
-    applies domain-specific post-processing, and returns samples as a DataFrame. The
-    per-factor compute lives in the composer (which need only expose
+    Thin user-facing façade over a ``Composer`` (e.g. ``ChainComposer``, or the GW
+    ``GNPEGibbsComposer``) and a ``SamplerContext``: runs the composer (which bounds memory via
+    vertical batching), applies domain-specific post-processing, and returns samples as a
+    DataFrame. The per-factor compute lives in the composer (which need only expose
     ``sample(num_samples, context, batch_size) -> dict``); this class only handles
     consolidation and post-processing -- the role the monolithic ``Sampler`` plays today.
     """
 
-    def __init__(
-        self, composer: Union[ChainComposer, GibbsComposer], context: SamplerContext
-    ):
+    def __init__(self, composer: Composer, context: SamplerContext):
         self.composer = composer
         self.context = context
         self.samples: Optional[pd.DataFrame] = None
@@ -484,7 +451,7 @@ def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
         """Run the composer (vertically batched by ``batch_size``) and return samples as a
-        DataFrame. ``ChainComposer.sample`` adds ``log_prob``; ``GibbsComposer``
+        DataFrame. ``ChainComposer.sample`` adds ``log_prob``; ``GNPEGibbsComposer``
         (multi-iteration GNPE) returns parameters + proxies with no ``log_prob``.
         ``batch_size=None`` runs in a single pass."""
         merged = self.composer.sample(num_samples, self.context, batch_size)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 78c22fc03..074d9aa6d 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -25,10 +25,12 @@
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
+    Conditioning,
     Factor,
     FlowFactor,
-    GibbsComposer,
+    SamplerContext,
     _base_model_metadata,
+    chunk_and_concat,
 )
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
@@ -297,6 +299,62 @@ def gibbs_step(
         return x["parameters"], x["extrinsic_parameters"]
 
 
+class GNPEGibbsComposer:
+    """
+    The Gibbs sampler for multi-iteration GNPE -- GNPE-specific by name and by guts (it
+    drives the GNPE step: proxies, detector times, time-shift). Gibbs sampling in general
+    is broader than this, but Dingo only ever uses it for GNPE, so the class owns that
+    specificity rather than posing as a general Gibbs facility. The genuinely generic
+    pieces live in core: ``ChainComposer``, the ``chunk_and_concat`` batching primitive,
+    and the ``Composer`` protocol this satisfies.
+
+    Each iteration is a 2-block Gibbs sweep on the joint q(theta, theta_hat | d): blur the
+    detector times into proxies (theta_hat ~ kernel(.|theta)), then sample the network
+    (theta ~ q(.|theta_hat, d)). Both blocks are bundled in ``gnpe_factor.gibbs_step``;
+    the ``init_factor`` seeds the chain and ``num_iterations`` is burn-in. The dependency
+    is cyclic, so this breaks density access: ``sample`` returns parameters WITHOUT a
+    ``log_prob``. (Single-step GNPE that *does* preserve log_prob is a two-stage
+    ``ChainComposer`` instead.)
+
+    Batching is vertical, via ``chunk_and_concat``: chunk the Gibbs walkers and run the
+    whole loop per chunk -- the old ``GWSamplerGNPE.run_sampler`` order.
+    """
+
+    def __init__(self, init_factor: Factor, gnpe_factor, num_iterations: int):
+        self.init_factor = init_factor
+        self.gnpe_factor = gnpe_factor
+        self.num_iterations = num_iterations
+
+    def sample(
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        batch_size: Optional[int] = None,
+    ) -> dict[str, torch.Tensor]:
+        samples, _ = chunk_and_concat(
+            num_samples, batch_size, lambda n: (self._run_once(n, context), None)
+        )
+        return samples
+
+    def _run_once(
+        self, num_samples: int, context: SamplerContext
+    ) -> dict[str, torch.Tensor]:
+        # Seed the Gibbs chain with the init factor's parameters (e.g. detector times).
+        seed, _ = self.init_factor.sample_and_log_prob(
+            num_samples, Conditioning(context)
+        )
+        extrinsic = dict(seed)
+        parameters: dict[str, torch.Tensor] = {}
+        for _ in range(self.num_iterations):
+            parameters, extrinsic = self.gnpe_factor.gibbs_step(
+                num_samples, context, extrinsic
+            )
+        proxies = {
+            p: extrinsic[p] for p in self.gnpe_factor.proxy_parameters if p in extrinsic
+        }
+        return {**parameters, **proxies}
+
+
 class GWComposedSampler(ComposedSampler):
     """
     GW façade: a ``ComposedSampler`` that adds the gravitational-wave post-processing --
@@ -308,7 +366,7 @@ class GWComposedSampler(ComposedSampler):
 
     def __init__(
         self,
-        composer: Union[ChainComposer, GibbsComposer],
+        composer: Union[ChainComposer, GNPEGibbsComposer],
         context: GWSamplerContext,
         metadata: dict,
         inference_parameters: list[str],
@@ -345,12 +403,12 @@ def from_gnpe_models(
     ) -> "GWComposedSampler":
         """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
         init model's data preprocessing, an init ``FlowFactor`` to seed, and a
-        ``GNPEFlowFactor`` driven by a ``GibbsComposer``. Returns samples without a
+        ``GNPEFlowFactor`` driven by a ``GNPEGibbsComposer``. Returns samples without a
         log_prob (Gibbs breaks density access)."""
         context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
         init_factor = FlowFactor.from_model(init_model)
         gnpe_factor = GNPEFlowFactor.from_model(main_model)
-        composer = GibbsComposer(init_factor, gnpe_factor, num_iterations)
+        composer = GNPEGibbsComposer(init_factor, gnpe_factor, num_iterations)
         return cls(
             composer,
             context,

From 51b969cff3b1329005350a8583ece45ae2c73bcf Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 30 Jun 2026 23:32:40 +0200
Subject: [PATCH 10/65] Tighten core factor docstrings

Rewrite the dingo.core.factors docstrings as reference documentation: state
what each class and function is and does, with parameters and returns. Drop
design-justification, codebase-history narration, comparisons to the old
samplers, and conversational asides -- that rationale lives in the design doc.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py | 212 ++++++++++++++++--------------------------
 1 file changed, 78 insertions(+), 134 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 96b3ac10d..40b84fa15 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -1,22 +1,16 @@
 """
-Factorized sampler core.
+Factorized sampler core: the domain-agnostic spine of the factorized-sampler design
+(see vault/Hackathon/Factorized_Sampler_Design.md).
 
-This module implements the domain-agnostic spine of the factorized-sampler design
-(see vault/Hackathon/Factorized_Sampler_Design.md): the posterior is represented as
-an ordered product of conditional factors
+The posterior is an ordered product of conditional factors,
 
     q(theta_1, ..., theta_n | d) = prod_i q_i(theta_i | f_i(theta_<i, d)),
 
-where each ``Factor`` samples one block of parameters and returns its own log-prob,
-and a composer drives them. The ``ChainComposer`` runs the factors autoregressively
-(exact, importance-sampling-friendly); a Gibbs composer for multi-iteration GNPE is
-added separately.
+where each ``Factor`` samples one parameter block and returns its own log-prob, and a
+composer evaluates them. ``ChainComposer`` runs the factors autoregressively.
 
-Key invariant -- *physical in, physical out*: factors expose only physical-space
-parameters. Standardized values exist transiently inside a factor's forward pass; the
-``Standardization`` helper below is the per-factor adapter between a network's internal
-standardized space and physical space. Nothing we store, pass, or serialize is ever
-standardized.
+Factors work in physical parameter space; a network's standardized space exists only
+inside its forward pass, mediated by ``Standardization``.
 """
 
 from __future__ import annotations
@@ -34,25 +28,13 @@
 
 class Standardization:
     """
-    Per-factor parameter standardization: the affine map between a network's internal
-    standardized space and physical space.
-
-    A network is trained on standardized parameters ``z = (theta - mean) / std``. This
-    object holds that network's ``mean``/``std`` (read from its training metadata) and
-    is used in *both* directions -- de-standardizing samples on the way out and
-    standardizing parameters on the way in for ``log_prob`` -- so there is a single
-    source of truth and no "inverted twice" drift.
-
-    Standardization is a property of the *network*, not the event: different factors
-    (e.g. a GNPE init network and the main network) carry different standardizations,
-    and physical space is the shared interchange between them.
-
-    Notes
-    -----
-    The eventual design stores ``mean``/``std`` as registered buffers inside the network
-    module (for NSF, as a fixed affine transform at the flow boundary so its log-Jacobian
-    folds into the density automatically). This pure-Python version is the equivalent
-    seam; moving it into buffers is a follow-up that does not change the interface.
+    Affine map between a network's standardized space (``z = (theta - mean) / std``) and
+    physical parameter space.
+
+    Holds one network's ``mean`` and ``std`` and applies them in both directions:
+    de-standardizing network outputs to physical samples, and standardizing physical
+    parameters for ``log_prob``. Different factors (e.g. a GNPE init and main network)
+    carry different instances.
     """
 
     def __init__(self, mean: dict[str, float], std: dict[str, float]):
@@ -83,18 +65,17 @@ def log_det(self, names: list[str]) -> float:
 @runtime_checkable
 class SamplerContext(Protocol):
     """
-    Per-event shared state referenced by every factor: the data ``d`` and everything
-    derived from it (the one-time prepared data representation, the likelihood, event
-    metadata). Concrete implementations are domain-specific (see
-    ``dingo.gw.inference.factors.GWSamplerContext``). Serialized as the transport state
-    between pipe stages.
+    Per-event shared state referenced by every factor: the data ``d`` and quantities
+    derived from it (the prepared-data representation, the likelihood, event metadata).
+    Concrete implementations are domain-specific; see
+    ``dingo.gw.inference.factors.GWSamplerContext``.
     """
 
     event_metadata: Optional[dict]
 
     def prepared_data(self) -> torch.Tensor:
-        """The one-time data representation (whiten/decimate/repackage/...), computed
-        once and cached. The conditioning input shared by data-conditioned factors."""
+        """The data representation the factors condition on
+        (whiten/decimate/repackage/...), computed once and cached."""
         ...
 
     def likelihood(self):
@@ -105,9 +86,8 @@ def likelihood(self):
 @dataclass
 class Conditioning:
     """
-    What the composer hands a factor: the shared ``context`` (-> prepared data,
-    likelihood, domain) plus the physical values of the earlier-block parameters this
-    factor conditions on.
+    The conditioning passed to a factor: the shared ``context``, and ``given``, the
+    physical values of the earlier-block parameters the factor conditions on.
     """
 
     context: SamplerContext
@@ -126,18 +106,12 @@ def chunk_and_concat(
     batch_size: Optional[int],
     run_once: Callable[[int], tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]],
 ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
-    """The vertical (depth-first) batching primitive, shared by the chain and Gibbs
-    composers.
-
-    Split ``total`` base samples into chunks of ``batch_size``, run ``run_once(n)`` -- a
-    full pass through the *whole* pipeline (the entire chain, or the entire Gibbs loop) --
-    per chunk, and concatenate. Working memory is bounded to one chunk carried through the
-    pipeline; only the output buffer grows to ``total`` (a real run streams it to CPU per
-    chunk). This is the order the old ``Sampler.run_sampler`` batches in, so parity is
-    exact when ``batch_size`` matches. ``batch_size=None`` runs everything in one pass.
-
-    ``run_once`` returns ``(samples, log_prob)``; ``log_prob`` may be ``None`` (the Gibbs
-    path has no tractable density)."""
+    """Run ``run_once`` over batches of ``total`` and concatenate the results.
+
+    Splits ``total`` into chunks of ``batch_size``, calls ``run_once(n) -> (samples,
+    log_prob)`` per chunk, and concatenates the sample dicts and log-probs. Caps peak
+    memory at one chunk. ``batch_size=None`` runs ``total`` in a single call. ``log_prob``
+    may be ``None`` (for a composer without a tractable density, i.e. Gibbs)."""
     bs = batch_size or total
     sample_parts: list[dict[str, torch.Tensor]] = []
     lp_parts: list[Optional[torch.Tensor]] = []
@@ -152,27 +126,20 @@ def chunk_and_concat(
 
 class Factor(ABC):
     """
-    A conditional distribution ``q_i(theta_i | f_i(theta_<i, d))`` over one block of
-    parameters. Emits physical-space samples and a physical-space log-prob (its
-    standardization Jacobian already applied), so the chain product ``sum_i log q_i`` is
-    directly the physical posterior log-density.
-
-    *Physical in, physical out*, and a **single forward pass per call** -- batching is not
-    a factor concern; the composer drives it (``chunk_and_concat``).
+    A conditional distribution ``q_i(theta_i | f_i(theta_<i, d))`` over one parameter
+    block, emitting physical-space samples and a physical-space log-prob.
 
-    Contract the composer relies on (this is what makes fan-out work): a call draws
-    ``num_samples`` samples **per conditioning row** and returns ``n_rows * num_samples``
-    rows, flattened row-major, where ``n_rows`` is the number of rows in ``cond.given``
-    (or 1 if unconditioned). This mirrors the network's own
-    ``sample_and_log_prob(*context, num_samples=n) -> (B, n, dim)`` with ``B = n_rows``.
+    A call draws ``num_samples`` samples per conditioning row and returns ``n_rows *
+    num_samples`` rows in row-major order, where ``n_rows`` is the number of rows in
+    ``cond.given`` (1 if unconditioned). This mirrors the network's
+    ``sample_and_log_prob(*context, num_samples=n) -> (n_rows, n, dim)``.
 
     Attributes
     ----------
     parameters : list[str]
-        The block ``theta_i`` this factor produces.
+        The parameter block this factor produces.
     conditioning : list[str]
-        Earlier-block parameter names this factor conditions on. Data dependence is
-        implicit via the shared context.
+        Earlier-block parameters it conditions on (data dependence is via the context).
     """
 
     parameters: list[str]
@@ -182,16 +149,16 @@ class Factor(ABC):
     def sample_and_log_prob(
         self, num_samples: int, cond: Conditioning
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """Draw ``num_samples`` per conditioning row from ``q_i``; return ``(physical
-        samples, physical-space log q_i)``. The returned dict may carry extra named
-        side-channel columns (e.g. an IS kernel correction) alongside ``parameters``."""
+        """Draw ``num_samples`` samples per conditioning row; return ``(samples,
+        log_prob)`` in physical space. The samples dict may include named columns beyond
+        ``parameters``."""
 
     @abstractmethod
     def log_prob(
         self, theta_i: dict[str, torch.Tensor], cond: Conditioning
     ) -> torch.Tensor:
-        """Evaluate ``log q_i(theta_i | f_i)`` at given physical ``theta_i`` (for IS /
-        re-plug). One conditioning row per ``theta_i`` row (no fan-out)."""
+        """Evaluate ``log q_i`` at given physical ``theta_i`` (one conditioning row per
+        row)."""
 
 
 def _base_model_metadata(model: BasePosteriorModel) -> dict:
@@ -206,15 +173,12 @@ def _base_model_metadata(model: BasePosteriorModel) -> dict:
 
 class FlowFactor(Factor):
     """
-    A factor wrapping a posterior model (NPE flow, FMPE, ...). Domain-agnostic: it reaches
-    data only through the ``SamplerContext.prepared_data()`` protocol, so the GW-specific
-    conditioning map ``f_i`` (the data preprocessing) lives entirely on the context.
+    Factor wrapping a posterior model (NPE flow, FMPE, ...). Reaches data only through
+    ``SamplerContext.prepared_data()`` and encapsulates the network's standardization, so
+    its interface is in physical parameter space.
 
-    Encapsulates the network's own standardization: the public interface is physical-in /
-    physical-out -- standardized values never leave the factor.
-
-    The cheap *parameter-dependent* part of ``f_i`` (time-shift / heterodyne for GNPE) is
-    a no-op here; a domain subclass overrides ``_apply_param_transforms`` to add it.
+    An unconditioned model draws from the shared data context; a model with
+    ``context_parameters`` conditions on the values in ``cond.given``.
     """
 
     def __init__(
@@ -234,8 +198,8 @@ def __init__(
 
     @classmethod
     def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
-        """Build a factor from a model: ``parameters`` and ``context_parameters`` come from
-        the model's training metadata (first-class conditioning)."""
+        """Build a factor from a model, reading ``parameters`` and ``context_parameters``
+        from its training metadata."""
         data_settings = _base_model_metadata(model)["train_settings"]["data"]
         context_parameters = data_settings.get("context_parameters") or []
         return cls(
@@ -246,10 +210,9 @@ def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
         )
 
     def sample_and_log_prob(self, num_samples, cond):
-        """Draw ``num_samples`` per conditioning row. Unconditioned (plain NPE): one
-        shared data context, ``num_samples`` draws. Conditioned: ``N`` context rows in
-        (one embedding pass each), ``num_samples`` draws per row -> ``N * num_samples``
-        rows, flattened row-major to align with the composer's ``repeat_interleave``."""
+        """Draw ``num_samples`` samples per conditioning row. Unconditioned: ``num_samples``
+        draws from the shared data context. Conditioned: ``num_samples`` per context row,
+        returning ``n_rows * num_samples`` rows in row-major order."""
         data = cond.context.prepared_data()
         self.model.network.eval()
         if not self.context_parameters:
@@ -296,10 +259,8 @@ def log_prob(self, theta_i, cond):
 
 @dataclass
 class Stage:
-    """One factor in the chain plus its **fan-out**: how many samples it draws per
-    incoming conditioning row. The root (first) stage ignores ``fan_out`` -- it draws the
-    base count directly; a 1:1 chain link uses ``fan_out=1``; an intrinsic/extrinsic
-    expansion uses ``fan_out=K`` (``K`` extrinsic draws per intrinsic sample)."""
+    """A chain factor and its fan-out: the number of samples drawn per incoming
+    conditioning row. The root stage draws the base count and ignores ``fan_out``."""
 
     factor: Factor
     fan_out: int = 1
@@ -307,21 +268,14 @@ class Stage:
 
 class ChainComposer:
     """
-    Autoregressive composer **and vertical (depth-first) executor**: sample the stages in
-    declared order (a topological order of the dependency DAG), expanding by each stage's
-    fan-out, and sum their log-probs. Exact and importance-sampling-friendly. Covers plain
-    NPE, single-step GNPE, prior conditioning, synthetic phase, and intrinsic/extrinsic
-    splits.
-
-    Vertical execution carries a chunk of base samples all the way down the chain, emits
-    it, and discards -- bounding working memory to one chunk (the only way to draw ``M``
-    intrinsic x ``K`` extrinsic without materializing ``M*K`` intermediates) and
-    reproducing the old sampler's batching order (bit-exact when ``batch_size`` matches).
-    Batching lives in ``chunk_and_concat``; the factors are single-pass.
-
-    Multi-iteration (cyclic) GNPE is *not* expressible here and uses the GW
-    ``GNPEGibbsComposer``. Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or
-    explicit ``Stage``s.
+    Autoregressive composer over an ordered list of ``Stage``s.
+
+    Draws the stages in declared order -- a topological order of the conditioning DAG --
+    expanding each stage by its fan-out and summing the log-probs. Covers plain NPE,
+    single-step GNPE, prior conditioning, synthetic phase, and intrinsic/extrinsic splits.
+    Multi-iteration (cyclic) GNPE uses the GW ``GNPEGibbsComposer`` instead.
+
+    Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or explicit ``Stage``s.
     """
 
     def __init__(self, stages: list[Union["Stage", Factor]]):
@@ -343,13 +297,13 @@ def _validate(self):
 
     @property
     def factors(self) -> list[Factor]:
-        """The factors in order (without fan-out), for callers that only need them."""
+        """The stage factors, in order."""
         return [stage.factor for stage in self.stages]
 
     @property
     def expansion(self) -> int:
-        """Product of the non-root fan-outs; total rows returned = num_samples *
-        expansion (e.g. M intrinsic x K extrinsic for an extrinsic fan-out of K)."""
+        """Product of the non-root fan-outs; ``sample`` returns ``num_samples *
+        expansion`` rows."""
         return math.prod(stage.fan_out for stage in self.stages[1:])
 
     def sample_and_log_prob(
@@ -358,11 +312,9 @@ def sample_and_log_prob(
         context: SamplerContext,
         batch_size: Optional[int] = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """``num_samples`` is the **base (root) count** -- e.g. ``M`` intrinsic samples.
-        Each fan-out stage multiplies it, so the result has ``num_samples * expansion``
-        rows (``M * K`` for an extrinsic ``fan_out=K``); for a plain chain (no fan-out)
-        that is just ``num_samples``. ``batch_size`` chunks the base count, in the same
-        (root) unit as ``num_samples``."""
+        """Draw samples. ``num_samples`` is the base (root) count; the result has
+        ``num_samples * expansion`` rows. ``batch_size`` chunks the base count (``None``
+        draws in one pass)."""
         return chunk_and_concat(
             num_samples, batch_size, lambda n: self._run_chain_once(n, context)
         )
@@ -370,7 +322,7 @@ def sample_and_log_prob(
     def _run_chain_once(
         self, base: int, context: SamplerContext
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """One depth-first pass of the whole chain for ``base`` root samples."""
+        """One pass of the whole chain for ``base`` root samples."""
         samples: dict[str, torch.Tensor] = {}
         total: torch.Tensor | float = 0.0
         for i, stage in enumerate(self.stages):
@@ -406,18 +358,16 @@ def sample(
         context: SamplerContext,
         batch_size: Optional[int] = None,
     ) -> dict[str, torch.Tensor]:
-        """Full per-sample dict (parameters + ``log_prob``), for the sampler façade."""
+        """Per-sample dict of parameters plus ``log_prob``."""
         samples, log_prob = self.sample_and_log_prob(num_samples, context, batch_size)
         return {**samples, "log_prob": log_prob}
 
 
 @runtime_checkable
 class Composer(Protocol):
-    """The minimal interface ``ComposedSampler`` drives: draw a per-sample dict
-    (parameters, plus ``log_prob`` for density-preserving composers). ``ChainComposer``
-    (here, domain-agnostic) satisfies it, as does the GW ``GNPEGibbsComposer`` (multi-iteration
-    GNPE) -- which lives in ``dingo.gw.inference.factors`` because it is built entirely
-    around the GNPE step (its proxies, detector times, time-shift), not in core."""
+    """Interface required by ``ComposedSampler``: ``sample`` returns a per-sample dict of
+    parameters (plus ``log_prob`` for density-preserving composers). Satisfied by
+    ``ChainComposer`` and the GW ``GNPEGibbsComposer``."""
 
     def sample(
         self,
@@ -429,12 +379,8 @@ def sample(
 
 class ComposedSampler:
     """
-    Thin user-facing façade over a ``Composer`` (e.g. ``ChainComposer``, or the GW
-    ``GNPEGibbsComposer``) and a ``SamplerContext``: runs the composer (which bounds memory via
-    vertical batching), applies domain-specific post-processing, and returns samples as a
-    DataFrame. The per-factor compute lives in the composer (which need only expose
-    ``sample(num_samples, context, batch_size) -> dict``); this class only handles
-    consolidation and post-processing -- the role the monolithic ``Sampler`` plays today.
+    Façade over a ``Composer`` and a ``SamplerContext``. Runs the composer, applies
+    domain-specific post-processing, and returns the samples as a DataFrame.
     """
 
     def __init__(self, composer: Composer, context: SamplerContext):
@@ -443,17 +389,15 @@ def __init__(self, composer: Composer, context: SamplerContext):
         self.samples: Optional[pd.DataFrame] = None
 
     def _post_process(self, samples: dict, inverse: bool = False):
-        """Hook for domain-specific post-processing (e.g. GW fixed-parameter injection
-        and reference-time correction). No-op by default."""
+        """Hook for domain-specific post-processing; no-op by default."""
         pass
 
     def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
-        """Run the composer (vertically batched by ``batch_size``) and return samples as a
-        DataFrame. ``ChainComposer.sample`` adds ``log_prob``; ``GNPEGibbsComposer``
-        (multi-iteration GNPE) returns parameters + proxies with no ``log_prob``.
-        ``batch_size=None`` runs in a single pass."""
+        """Draw ``num_samples`` samples (chunked by ``batch_size``), post-process, and
+        return them as a DataFrame. ``ChainComposer`` includes ``log_prob``;
+        ``GNPEGibbsComposer`` does not."""
         merged = self.composer.sample(num_samples, self.context, batch_size)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
         self._post_process(merged)

From 9a7152e3b4796608757e49776baae89536c82201 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 00:51:07 +0200
Subject: [PATCH 11/65] Split GNPE main factor into kernel + flow; generalize
 Gibbs composer

Replace the bundled GNPEFlowFactor (a non-Factor Gibbs step) and the
single-step GNPEChainFactor with two conditional Factors:

- GNPEKernelFactor: the perturbation kernel p(theta_hat | theta). Blurs
  detector times into proxies; its log_prob is the delta_log_prob_target
  importance-sampling correction.
- GNPEFlowFactor: the main network q(theta | theta_hat, d), now a real
  Factor, shared by both composers.

GNPEGibbsComposer becomes the generic GibbsComposer in core.factors
(cycles a factor list, no GNPE specifics). Multi-iteration GNPE is
GibbsComposer([kernel, flow]); single-step GNPE is
ChainComposer([proxy_source, flow]) with the kernel correction applied
out of the proposal sum.

ComposedSampler keeps the Composer protocol, now used consistently by
GWComposedSampler too (dropped the dead @runtime_checkable).

Bit-exact vs the old GWSampler / GWSamplerGNPE for NPE, multi-iteration
GNPE, and single-step GNPE across both verification harnesses.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         |  58 +++++-
 dingo/gw/inference/factors.py | 334 +++++++++++++++++++---------------
 2 files changed, 240 insertions(+), 152 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 40b84fa15..c545fc160 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -273,7 +273,7 @@ class ChainComposer:
     Draws the stages in declared order -- a topological order of the conditioning DAG --
     expanding each stage by its fan-out and summing the log-probs. Covers plain NPE,
     single-step GNPE, prior conditioning, synthetic phase, and intrinsic/extrinsic splits.
-    Multi-iteration (cyclic) GNPE uses the GW ``GNPEGibbsComposer`` instead.
+    Multi-iteration (cyclic) GNPE uses ``GibbsComposer`` instead.
 
     Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or explicit ``Stage``s.
     """
@@ -363,11 +363,61 @@ def sample(
         return {**samples, "log_prob": log_prob}
 
 
-@runtime_checkable
+class GibbsComposer:
+    """
+    Blocked Gibbs composer over a cyclic conditioning dependency.
+
+    Seeds the chain with an init factor, then sweeps the factor list in order for
+    ``num_iterations`` iterations; each factor conditions on the current state and
+    overwrites its own block. The cyclic dependency has no tractable marginal, so
+    ``sample`` returns parameters without a ``log_prob`` (recoverable by fitting an
+    unconditional density to the samples and taking one ``ChainComposer`` step). Dingo uses
+    this only for multi-iteration GNPE (the GNPE factors in ``dingo.gw.inference.factors``),
+    but the loop is generic.
+
+    Batching is vertical, via ``chunk_and_concat``: chunk the walkers and run the whole
+    loop per chunk.
+    """
+
+    def __init__(self, init_factor: Factor, factors: list[Factor], num_iterations: int):
+        self.init_factor = init_factor
+        self.factors = list(factors)
+        self.num_iterations = num_iterations
+
+    def sample(
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        batch_size: Optional[int] = None,
+    ) -> dict[str, torch.Tensor]:
+        samples, _ = chunk_and_concat(
+            num_samples, batch_size, lambda n: (self._run_once(n, context), None)
+        )
+        return samples
+
+    def _run_once(
+        self, num_samples: int, context: SamplerContext
+    ) -> dict[str, torch.Tensor]:
+        # Seed the chain (e.g. an init network's detector times); the walkers are the rows.
+        seed, _ = self.init_factor.sample_and_log_prob(
+            num_samples, Conditioning(context)
+        )
+        state = dict(seed)
+        for _ in range(self.num_iterations):
+            for factor in self.factors:
+                cond = Conditioning(context, {k: state[k] for k in factor.conditioning})
+                # One sample per walker (Gibbs is 1:1); walkers are the conditioning rows.
+                block, _ = factor.sample_and_log_prob(1, cond)
+                state.update(block)
+        # Each factor's parameter block (e.g. proxies + inference parameters), dropping
+        # side-channel columns such as the recomputed detector times.
+        return {p: state[p] for factor in self.factors for p in factor.parameters}
+
+
 class Composer(Protocol):
     """Interface required by ``ComposedSampler``: ``sample`` returns a per-sample dict of
     parameters (plus ``log_prob`` for density-preserving composers). Satisfied by
-    ``ChainComposer`` and the GW ``GNPEGibbsComposer``."""
+    ``ChainComposer`` and ``GibbsComposer``."""
 
     def sample(
         self,
@@ -397,7 +447,7 @@ def run_sampler(
     ) -> pd.DataFrame:
         """Draw ``num_samples`` samples (chunked by ``batch_size``), post-process, and
         return them as a DataFrame. ``ChainComposer`` includes ``log_prob``;
-        ``GNPEGibbsComposer`` does not."""
+        ``GibbsComposer`` does not."""
         merged = self.composer.sample(num_samples, self.context, batch_size)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
         self._post_process(merged)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 074d9aa6d..d085a6972 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -5,8 +5,10 @@
 vault/Hackathon/Factorized_Sampler_Design.md). The generic ``FlowFactor`` is
 domain-agnostic and lives in ``dingo.core.factors``; what is GW-specific is the
 ``GWSamplerContext`` (which builds the data-preprocessing conditioning map ``f_i``) and
-the post-processing in ``GWComposedSampler``. This first increment covers the plain-NPE
-path; the fixed, synthetic-phase, and GNPE-proxy factors are stubbed for later steps.
+the post-processing in ``GWComposedSampler``. Covers plain NPE (``FlowFactor`` via a
+``ChainComposer``) and multi-iteration time-shift GNPE (``GNPEKernelFactor`` +
+``GNPEFlowFactor`` cycled by a ``GibbsComposer``); the synthetic-phase factor and
+single-step GNPE importance sampling remain for later steps.
 """
 
 from __future__ import annotations
@@ -18,19 +20,18 @@
 import pandas as pd
 import torch
 from astropy.time import Time
-from bilby.core.prior import DeltaFunction
+from bilby.core.prior import DeltaFunction, PriorDict
 from bilby.gw.detector import InterferometerList
 from torchvision.transforms import Compose
 
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
-    Conditioning,
+    Composer,
     Factor,
     FlowFactor,
-    SamplerContext,
+    GibbsComposer,
     _base_model_metadata,
-    chunk_and_concat,
 )
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
@@ -181,16 +182,151 @@ def log_prob(self, theta_i, cond):
         raise NotImplementedError("SyntheticPhaseFactor -- later step.")
 
 
-class GNPEFlowFactor:
+def _build_gnpe_transforms(model: BasePosteriorModel):
+    """Build the time-shift GNPE per-step transforms from a model's metadata (the same
+    chains as ``GWSamplerGNPE._initialize_transforms``).
+
+    Returns
+    -------
+    transform_pre, transform_post : Compose
+    gnpe_parameters : list[str]
+        The GNPE input parameters (detector times).
+    inference_parameters : list[str]
+    kernel : PriorDict
+        The proxy perturbation kernel.
+    gnpe_transform : GNPECoalescenceTimes
+        The blur transform itself, shared so the kernel factor can call ``sample_proxies``.
     """
-    One iteration of time-shift GNPE for the Gibbs composer. Given the current detector-
-    time estimates, it blurs them into proxies, shifts the strain, standardizes the
-    proxies as network context, samples the main network, then recomputes detector times
-    for the next iteration. Wraps the main model and the per-iteration transforms (the
-    same ones ``GWSamplerGNPE`` builds).
-
-    Not a ``Factor``: the ABC's ``sample_and_log_prob`` does not fit a Gibbs step. It
-    exposes ``gibbs_step(num_samples, context, extrinsic) -> (parameters, extrinsic)``.
+    meta = _base_model_metadata(model)
+    data_settings = meta["train_settings"]["data"]
+    ifo_list = InterferometerList(data_settings["detectors"])
+    domain = build_domain(meta["dataset_settings"]["domain"])
+    if "domain_update" in data_settings:
+        domain.update(data_settings["domain_update"])
+
+    gnpe_time_settings = data_settings.get("gnpe_time_shifts")
+    if not gnpe_time_settings:
+        raise NotImplementedError(
+            "Only time-shift GNPE is supported here so far (no gnpe_chirp / gnpe_phase)."
+        )
+
+    gnpe_transform = GNPECoalescenceTimes(
+        ifo_list,
+        gnpe_time_settings["kernel"],
+        gnpe_time_settings["exact_equiv"],
+        inference=True,
+    )
+    transform_pre = [
+        RenameKey("data", "waveform"),
+        gnpe_transform,
+        TimeShiftStrain(ifo_list, domain),
+        SelectStandardizeRepackageParameters(
+            {"context_parameters": data_settings["context_parameters"]},
+            data_settings["standardization"],
+            device=model.device,
+        ),
+        RenameKey("waveform", "data"),
+    ]
+    gnpe_parameters: list[str] = []
+    kernel = PriorDict()
+    for transform in transform_pre:
+        if isinstance(transform, GNPEBase):
+            gnpe_parameters += transform.input_parameter_names
+            for k, v in transform.kernel.items():
+                kernel[k] = v
+
+    inference_parameters = data_settings["inference_parameters"]
+    transform_post = [
+        SelectStandardizeRepackageParameters(
+            {"inference_parameters": inference_parameters},
+            data_settings["standardization"],
+            inverse=True,
+            as_type="dict",
+        ),
+        PostCorrectGeocentTime(),
+        CopyToExtrinsicParameters(
+            "ra", "dec", "geocent_time", "chirp_mass", "mass_ratio", "phase"
+        ),
+        GetDetectorTimes(ifo_list, data_settings["ref_time"]),
+    ]
+    return (
+        Compose(transform_pre),
+        Compose(transform_post),
+        gnpe_parameters,
+        inference_parameters,
+        kernel,
+        gnpe_transform,
+    )
+
+
+class GNPEKernelFactor(Factor):
+    """
+    The GNPE perturbation kernel p(theta_hat | theta) as a factor (non-NN).
+
+    GNPE (arXiv:2111.13139) conditions the main network on "proxy" parameters theta_hat
+    drawn from a fixed kernel p(theta_hat | theta). Here theta are the detector coalescence
+    times and the kernel adds a bounded perturbation to each. Conditioning on the proxies
+    frees the network input to be simplified by them -- ``GNPEFlowFactor`` shifts each
+    detector's strain to its proxy time -- at the price of inferring theta and theta_hat
+    jointly.
+
+    The parameter block is the proxies; the conditioning is the detector times.
+    ``sample_and_log_prob`` blurs the times into proxies (the proxy update of a Gibbs
+    sweep). ``log_prob`` returns the kernel density log p(theta_hat | theta); for
+    single-step GNPE this is the ``delta_log_prob_target`` importance-sampling correction,
+    evaluated at the proxies and the detector times recomputed from theta. One proxy per
+    detector-time row.
+    """
+
+    def __init__(self, gnpe_transform: GNPEBase, gnpe_parameters: list[str]):
+        self.gnpe = gnpe_transform
+        self.gnpe_parameters = gnpe_parameters
+        self.parameters = [p + "_proxy" for p in gnpe_parameters]
+        self.conditioning = list(gnpe_parameters)
+        self.kernel = gnpe_transform.kernel
+
+    @classmethod
+    def from_model(cls, model: BasePosteriorModel) -> "GNPEKernelFactor":
+        """Build from the main model's metadata (the kernel / blur transform)."""
+        _, _, gnpe_parameters, _, _, gnpe_transform = _build_gnpe_transforms(model)
+        return cls(gnpe_transform, gnpe_parameters)
+
+    def sample_and_log_prob(self, num_samples, cond):
+        """Blur the conditioning detector times into proxies; ``num_samples`` must be 1
+        (GNPE is 1:1). Returns the proxies and their kernel log-prob."""
+        if num_samples != 1:
+            raise ValueError("GNPE proxy is 1:1; use fan_out=1.")
+        times = {k: cond.given[k] for k in self.gnpe_parameters}
+        proxies = self.gnpe.sample_proxies(times)
+        return proxies, self.log_prob(proxies, cond)
+
+    def log_prob(self, theta_i, cond):
+        """``log p(theta_hat | theta)`` from the kernel, at the proxies (``theta_i``) and
+        the detector times (``cond.given``)."""
+        diffs = {}
+        for k in self.kernel.keys():
+            diff = cond.given[k] - theta_i[f"{k}_proxy"]
+            if torch.is_tensor(diff):
+                diff = diff.detach().cpu().numpy()
+            diffs[k] = np.asarray(diff)
+        return torch.as_tensor(self.kernel.ln_prob(diffs, axis=0), dtype=torch.float32)
+
+
+class GNPEFlowFactor(Factor):
+    """
+    The GNPE main network q(theta | theta_hat, d) as a factor.
+
+    Conditions on the detector-time proxies from ``GNPEKernelFactor``: it shifts each
+    detector's strain by the corresponding proxy time (standardizing the network input),
+    samples the network, and recomputes the detector times from the sampled sky position
+    and geocent time. The proxies are supplied, so no blurring happens here. Wraps the main
+    model and the per-iteration transforms built for ``GWSamplerGNPE``.
+
+    The single network factor in either GNPE mode: cycled by ``GibbsComposer`` for
+    multi-iteration GNPE, or a ``ChainComposer`` factor for single-step GNPE. The recomputed
+    detector times are returned as extra columns: the next Gibbs iteration blurs them into
+    fresh proxies, and single-step GNPE evaluates the kernel correction at them. One sample
+    per proxy row.
     """
 
     def __init__(
@@ -207,81 +343,26 @@ def __init__(
         self.gnpe_parameters = gnpe_parameters
         self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
         self.parameters = parameters
+        self.conditioning = list(self.proxy_parameters)
 
     @classmethod
     def from_model(cls, model: BasePosteriorModel) -> "GNPEFlowFactor":
-        """Build the GNPE per-iteration transforms from the main model's metadata,
-        replicating ``GWSamplerGNPE._initialize_transforms`` (time-shift GNPE)."""
-        meta = _base_model_metadata(model)
-        data_settings = meta["train_settings"]["data"]
-        ifo_list = InterferometerList(data_settings["detectors"])
-        domain = build_domain(meta["dataset_settings"]["domain"])
-        if "domain_update" in data_settings:
-            domain.update(data_settings["domain_update"])
-
-        gnpe_time_settings = data_settings.get("gnpe_time_shifts")
-        if not gnpe_time_settings:
-            raise NotImplementedError(
-                "Only time-shift GNPE is supported here so far "
-                "(no gnpe_chirp / gnpe_phase)."
-            )
-
-        transform_pre = [
-            RenameKey("data", "waveform"),
-            GNPECoalescenceTimes(
-                ifo_list,
-                gnpe_time_settings["kernel"],
-                gnpe_time_settings["exact_equiv"],
-                inference=True,
-            ),
-            TimeShiftStrain(ifo_list, domain),
-            SelectStandardizeRepackageParameters(
-                {"context_parameters": data_settings["context_parameters"]},
-                data_settings["standardization"],
-                device=model.device,
-            ),
-            RenameKey("waveform", "data"),
-        ]
-        gnpe_parameters = []
-        for transform in transform_pre:
-            if isinstance(transform, GNPEBase):
-                gnpe_parameters += transform.input_parameter_names
-
-        inference_parameters = data_settings["inference_parameters"]
-        transform_post = [
-            SelectStandardizeRepackageParameters(
-                {"inference_parameters": inference_parameters},
-                data_settings["standardization"],
-                inverse=True,
-                as_type="dict",
-            ),
-            PostCorrectGeocentTime(),
-            CopyToExtrinsicParameters(
-                "ra", "dec", "geocent_time", "chirp_mass", "mass_ratio", "phase"
-            ),
-            GetDetectorTimes(ifo_list, data_settings["ref_time"]),
-        ]
-        return cls(
-            model,
-            Compose(transform_pre),
-            Compose(transform_post),
-            gnpe_parameters,
-            inference_parameters,
+        """Build the GNPE per-iteration transforms from the main model's metadata."""
+        pre, post, gnpe_parameters, inference_parameters, _, _ = _build_gnpe_transforms(
+            model
         )
+        return cls(model, pre, post, gnpe_parameters, inference_parameters)
 
-    def gibbs_step(
-        self, num_samples: int, context, extrinsic: dict
-    ) -> tuple[dict, dict]:
-        """One GNPE Gibbs iteration (single pass): returns (sampled parameters, updated
-        extrinsic state -- recomputed detector times + proxies for the next iteration).
-        Batching over the Gibbs walkers is the composer's job (``chunk_and_concat``).
-        """
-        x = {
-            "extrinsic_parameters": {k: extrinsic[k] for k in self.gnpe_parameters},
-            "parameters": {},
-        }
-        d = context.prepared_data().clone()
-        x["data"] = d.expand(num_samples, *d.shape)
+    def sample_and_log_prob(self, num_samples, cond):
+        """Sample one parameter set per proxy row; ``num_samples`` must be 1 (GNPE is 1:1).
+        Returns theta plus the recomputed detector times, and the network log-prob."""
+        if num_samples != 1:
+            raise ValueError("GNPE is 1:1; draw one sample per proxy (fan_out=1).")
+        proxies = {p: cond.given[p] for p in self.proxy_parameters}
+        n_rows = next(iter(proxies.values())).shape[0]
+        x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
+        d = cond.context.prepared_data().clone()
+        x["data"] = d.expand(n_rows, *d.shape)
         x = self.transform_pre(x)
         self.model.network.eval()
         with torch.no_grad():
@@ -291,68 +372,22 @@ def gibbs_step(
                 )
             else:
                 y, log_prob = self.model.sample_and_log_prob(x["data"])
-        # sample_and_log_prob(num_samples=1) adds a singleton dim; the batch is the
-        # num_samples Gibbs walkers.
+        # sample_and_log_prob(num_samples=1) adds a singleton dim; the batch is the proxy
+        # rows.
         x["parameters"] = y.squeeze(1)
         x["log_prob"] = log_prob.squeeze(1)
         x = self.transform_post(x)
-        return x["parameters"], x["extrinsic_parameters"]
-
-
-class GNPEGibbsComposer:
-    """
-    The Gibbs sampler for multi-iteration GNPE -- GNPE-specific by name and by guts (it
-    drives the GNPE step: proxies, detector times, time-shift). Gibbs sampling in general
-    is broader than this, but Dingo only ever uses it for GNPE, so the class owns that
-    specificity rather than posing as a general Gibbs facility. The genuinely generic
-    pieces live in core: ``ChainComposer``, the ``chunk_and_concat`` batching primitive,
-    and the ``Composer`` protocol this satisfies.
-
-    Each iteration is a 2-block Gibbs sweep on the joint q(theta, theta_hat | d): blur the
-    detector times into proxies (theta_hat ~ kernel(.|theta)), then sample the network
-    (theta ~ q(.|theta_hat, d)). Both blocks are bundled in ``gnpe_factor.gibbs_step``;
-    the ``init_factor`` seeds the chain and ``num_iterations`` is burn-in. The dependency
-    is cyclic, so this breaks density access: ``sample`` returns parameters WITHOUT a
-    ``log_prob``. (Single-step GNPE that *does* preserve log_prob is a two-stage
-    ``ChainComposer`` instead.)
-
-    Batching is vertical, via ``chunk_and_concat``: chunk the Gibbs walkers and run the
-    whole loop per chunk -- the old ``GWSamplerGNPE.run_sampler`` order.
-    """
-
-    def __init__(self, init_factor: Factor, gnpe_factor, num_iterations: int):
-        self.init_factor = init_factor
-        self.gnpe_factor = gnpe_factor
-        self.num_iterations = num_iterations
+        params = dict(x["parameters"])
+        # Surface the recomputed detector times: the next Gibbs iteration's input, and the
+        # evaluation point for GNPEKernelFactor's importance-sampling correction.
+        for k in self.gnpe_parameters:
+            params[k] = x["extrinsic_parameters"][k]
+        return params, x["log_prob"]
 
-    def sample(
-        self,
-        num_samples: int,
-        context: SamplerContext,
-        batch_size: Optional[int] = None,
-    ) -> dict[str, torch.Tensor]:
-        samples, _ = chunk_and_concat(
-            num_samples, batch_size, lambda n: (self._run_once(n, context), None)
-        )
-        return samples
-
-    def _run_once(
-        self, num_samples: int, context: SamplerContext
-    ) -> dict[str, torch.Tensor]:
-        # Seed the Gibbs chain with the init factor's parameters (e.g. detector times).
-        seed, _ = self.init_factor.sample_and_log_prob(
-            num_samples, Conditioning(context)
+    def log_prob(self, theta_i, cond):
+        raise NotImplementedError(
+            "GNPEFlowFactor.log_prob -- later step (single-step GNPE importance sampling)."
         )
-        extrinsic = dict(seed)
-        parameters: dict[str, torch.Tensor] = {}
-        for _ in range(self.num_iterations):
-            parameters, extrinsic = self.gnpe_factor.gibbs_step(
-                num_samples, context, extrinsic
-            )
-        proxies = {
-            p: extrinsic[p] for p in self.gnpe_factor.proxy_parameters if p in extrinsic
-        }
-        return {**parameters, **proxies}
 
 
 class GWComposedSampler(ComposedSampler):
@@ -366,7 +401,7 @@ class GWComposedSampler(ComposedSampler):
 
     def __init__(
         self,
-        composer: Union[ChainComposer, GNPEGibbsComposer],
+        composer: Composer,
         context: GWSamplerContext,
         metadata: dict,
         inference_parameters: list[str],
@@ -403,17 +438,20 @@ def from_gnpe_models(
     ) -> "GWComposedSampler":
         """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
         init model's data preprocessing, an init ``FlowFactor`` to seed, and a
-        ``GNPEFlowFactor`` driven by a ``GNPEGibbsComposer``. Returns samples without a
-        log_prob (Gibbs breaks density access)."""
+        ``GibbsComposer`` cycling the GNPE kernel and main-network factors. Returns samples
+        without a log_prob (Gibbs breaks density access)."""
         context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
         init_factor = FlowFactor.from_model(init_model)
-        gnpe_factor = GNPEFlowFactor.from_model(main_model)
-        composer = GNPEGibbsComposer(init_factor, gnpe_factor, num_iterations)
+        kernel_factor = GNPEKernelFactor.from_model(main_model)
+        flow_factor = GNPEFlowFactor.from_model(main_model)
+        composer = GibbsComposer(
+            init_factor, [kernel_factor, flow_factor], num_iterations
+        )
         return cls(
             composer,
             context,
             _base_model_metadata(main_model),
-            gnpe_factor.parameters,
+            flow_factor.parameters,
         )
 
     def _correct_reference_time(

From 700b7d064dd0e920530d522ca5cd4dac71de7599 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 09:43:31 +0200
Subject: [PATCH 12/65] Add single-step GNPE sampler on the factorized spine

GWComposedSampler.from_singlestep_gnpe builds a ChainComposer of
[proxy_source, GNPEFlowFactor] for density-preserving single-step GNPE.
proxy_source supplies the detector-time proxies: a FixedFactor for prior
conditioning (BNS), or an unconditional NDE for density recovery (which
supersedes prepare_log_prob's single-step half).

The kernel correction delta_log_prob_target = log p(theta_hat | theta) is
evaluated in post-processing (via a held GNPEKernelFactor) at the proxies
and the detector times recomputed from theta, then the intermediate
detector times are dropped. It is kept out of the proposal log_prob; the
column-driven importance-sampling layer applies it to the joint target.

Bit-exact vs GWSamplerGNPE(num_iterations=1) end to end through
run_sampler + to_result (fixed-proxy case).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py | 49 +++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index d085a6972..6d60bb2ca 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -28,6 +28,7 @@
     ChainComposer,
     ComposedSampler,
     Composer,
+    Conditioning,
     Factor,
     FlowFactor,
     GibbsComposer,
@@ -405,10 +406,15 @@ def __init__(
         context: GWSamplerContext,
         metadata: dict,
         inference_parameters: list[str],
+        kernel_factor: Optional[GNPEKernelFactor] = None,
     ):
         super().__init__(composer, context)
         self.metadata = metadata
         self.inference_parameters = inference_parameters
+        # Set for single-step GNPE. Supplies the delta_log_prob_target kernel correction
+        # (evaluated in post-processing, out of the proposal sum) that importance sampling
+        # applies to the joint target q(theta, theta_hat | d).
+        self.kernel_factor = kernel_factor
 
     @classmethod
     def from_model(
@@ -454,6 +460,33 @@ def from_gnpe_models(
             flow_factor.parameters,
         )
 
+    @classmethod
+    def from_singlestep_gnpe(
+        cls,
+        main_model: BasePosteriorModel,
+        proxy_source: Factor,
+        raw_context: dict,
+        event_metadata: Optional[dict] = None,
+    ) -> "GWComposedSampler":
+        """Build a single-step (density-preserving) time-GNPE sampler: a ``ChainComposer``
+        of ``[proxy_source, GNPEFlowFactor]``. ``proxy_source`` supplies the detector-time
+        proxies -- a ``FixedFactor`` for prior conditioning (BNS), or an unconditional NDE
+        for density recovery. Unlike multi-iteration GNPE the chain is autoregressive, so
+        log_prob is preserved; the ``GNPEKernelFactor`` supplies the
+        ``delta_log_prob_target`` correction that importance sampling adds to the target.
+        """
+        context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
+        flow_factor = GNPEFlowFactor.from_model(main_model)
+        kernel_factor = GNPEKernelFactor.from_model(main_model)
+        composer = ChainComposer([proxy_source, flow_factor])
+        return cls(
+            composer,
+            context,
+            _base_model_metadata(main_model),
+            flow_factor.parameters,
+            kernel_factor=kernel_factor,
+        )
+
     def _correct_reference_time(
         self, samples: Union[dict, pd.DataFrame], inverse: bool = False
     ):
@@ -479,7 +512,23 @@ def _correct_reference_time(
         else:
             samples["ra"] = (ra - ra_correction) % (2 * np.pi)
 
+    def _add_kernel_correction(self, samples: dict):
+        """Single-step GNPE: add ``delta_log_prob_target = log p(theta_hat | theta)``,
+        evaluated at the proxies and the detector times recomputed from theta, then drop
+        those intermediate detector times. This is the joint-target correction importance
+        sampling applies; it is deliberately kept out of the proposal ``log_prob``."""
+        kf = self.kernel_factor
+        proxies = {p: samples[p] for p in kf.parameters}
+        times = {k: samples[k] for k in kf.gnpe_parameters}
+        correction = kf.log_prob(proxies, Conditioning(self.context, times))
+        samples["delta_log_prob_target"] = np.asarray(correction)
+        for k in kf.gnpe_parameters:
+            samples.pop(k, None)
+
     def _post_process(self, samples: Union[dict, pd.DataFrame], inverse: bool = False):
+        if self.kernel_factor is not None and not inverse:
+            self._add_kernel_correction(samples)
+
         intrinsic_prior = self.metadata["dataset_settings"]["intrinsic_prior"]
         extrinsic_prior = get_extrinsic_prior_dict(
             self.metadata["train_settings"]["data"]["extrinsic_prior"]

From f0220b52b607de2a004d543734aab2f84f1a7270 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 11:36:37 +0200
Subject: [PATCH 13/65] Generalize ChainComposer to a Step list; drop the
 Composer protocol

- Add a `Step` protocol (parameters / conditioning / sample_and_log_prob ->
  (dict, Optional log_prob)); Stage.factor -> Stage.step, .factors -> .steps.
- ChainComposer's fold is None-aware: a single density-free step nulls the
  chain's log_prob, and sample() then omits it.
- Replace GibbsComposer with GibbsBlock, a density-free Step; multi-iteration
  GNPE is now ChainComposer([GibbsBlock(...)]). The Gibbs loop is unchanged,
  so parity stays bit-exact.
- Drop the now one-implementer Composer protocol.

Parity: NPE, multi-iteration GNPE, single-step, and packaged to_result all
bit-exact (max|delta|=0); tests/core/test_factors.py 12/12.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 171 +++++++++++++++++++---------------
 dingo/gw/inference/factors.py |  22 ++---
 tests/core/test_factors.py    |  59 +++++++++++-
 3 files changed, 162 insertions(+), 90 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index c545fc160..6b6f94e14 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -257,12 +257,30 @@ def log_prob(self, theta_i, cond):
         return log_prob + self.standardization.log_det(self.parameters)
 
 
+class Step(Protocol):
+    """
+    One entry a ``ChainComposer`` folds over: it emits a parameter block and an optional
+    contribution to the proposal ``log_prob``.
+
+    Density-contributing steps (``Factor``) return a tensor; density-free sampling blocks
+    (``GibbsBlock``) return ``None``. ``parameters`` names the block produced, ``conditioning``
+    the earlier-block parameters read from the chain (data is implicit via the context).
+    """
+
+    parameters: list[str]
+    conditioning: list[str]
+
+    def sample_and_log_prob(
+        self, num_samples: int, cond: "Conditioning"
+    ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]: ...
+
+
 @dataclass
 class Stage:
-    """A chain factor and its fan-out: the number of samples drawn per incoming
+    """A chain ``Step`` and its fan-out: the number of samples drawn per incoming
     conditioning row. The root stage draws the base count and ignores ``fan_out``."""
 
-    factor: Factor
+    step: Step
     fan_out: int = 1
 
 
@@ -270,35 +288,38 @@ class ChainComposer:
     """
     Autoregressive composer over an ordered list of ``Stage``s.
 
-    Draws the stages in declared order -- a topological order of the conditioning DAG --
-    expanding each stage by its fan-out and summing the log-probs. Covers plain NPE,
-    single-step GNPE, prior conditioning, synthetic phase, and intrinsic/extrinsic splits.
-    Multi-iteration (cyclic) GNPE uses ``GibbsComposer`` instead.
+    Folds the steps in declared order -- a topological order of the conditioning DAG --
+    expanding each by its fan-out and summing the proposal log-probs. A step is a
+    ``Factor`` (contributes ``log q_i``) or a density-free sampling block (``GibbsBlock``,
+    contributes ``None``); if any step is density-free the chain has no tractable density
+    and ``sample`` omits ``log_prob``. Covers plain NPE, single-step GNPE, prior
+    conditioning, synthetic phase, intrinsic/extrinsic splits, and -- via ``GibbsBlock`` --
+    multi-iteration GNPE.
 
-    Accepts bare factors (wrapped as ``Stage(factor, fan_out=1)``) or explicit ``Stage``s.
+    Accepts bare steps (wrapped as ``Stage(step, fan_out=1)``) or explicit ``Stage``s.
     """
 
-    def __init__(self, stages: list[Union["Stage", Factor]]):
+    def __init__(self, stages: list[Union["Stage", Step]]):
         self.stages = [s if isinstance(s, Stage) else Stage(s) for s in stages]
         self._validate()
 
     def _validate(self):
         """Check the declared order is a valid topological order: every conditioning
-        name is produced by an earlier factor."""
+        name is produced by an earlier step."""
         produced: set[str] = set()
-        for factor in self.factors:
-            missing = [c for c in factor.conditioning if c not in produced]
+        for step in self.steps:
+            missing = [c for c in step.conditioning if c not in produced]
             if missing:
                 raise ValueError(
-                    f"Factor producing {factor.parameters} conditions on {missing}, "
-                    f"which no earlier factor produces. Check chain order."
+                    f"A step producing {step.parameters} conditions on {missing}, "
+                    f"which no earlier step produces. Check chain order."
                 )
-            produced.update(factor.parameters)
+            produced.update(step.parameters)
 
     @property
-    def factors(self) -> list[Factor]:
-        """The stage factors, in order."""
-        return [stage.factor for stage in self.stages]
+    def steps(self) -> list[Step]:
+        """The stage steps, in order."""
+        return [stage.step for stage in self.stages]
 
     @property
     def expansion(self) -> int:
@@ -311,45 +332,53 @@ def sample_and_log_prob(
         num_samples: int,
         context: SamplerContext,
         batch_size: Optional[int] = None,
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+    ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
         """Draw samples. ``num_samples`` is the base (root) count; the result has
         ``num_samples * expansion`` rows. ``batch_size`` chunks the base count (``None``
-        draws in one pass)."""
+        draws in one pass). The log-prob is ``None`` if any step is density-free."""
         return chunk_and_concat(
             num_samples, batch_size, lambda n: self._run_chain_once(n, context)
         )
 
     def _run_chain_once(
         self, base: int, context: SamplerContext
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """One pass of the whole chain for ``base`` root samples."""
+    ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
+        """One pass of the whole chain for ``base`` root samples. Returns the samples and
+        the summed proposal log-prob, or ``None`` if any step is density-free."""
         samples: dict[str, torch.Tensor] = {}
         total: torch.Tensor | float = 0.0
+        has_density = True
         for i, stage in enumerate(self.stages):
-            factor = stage.factor
+            step = stage.step
             n = base if i == 0 else stage.fan_out
-            cond = Conditioning(context, {k: samples[k] for k in factor.conditioning})
-            block, lp = factor.sample_and_log_prob(n, cond)
-            # The factor returned (rows_so_far * n) rows. Expand the carried columns to
-            # match -- each upstream row repeated n times (the block is flattened in the
-            # same row-major order). The root has no carried rows, and fan_out=1 stages
-            # are a no-op, so 1:1 chains are untouched.
+            cond = Conditioning(context, {k: samples[k] for k in step.conditioning})
+            block, lp = step.sample_and_log_prob(n, cond)
+            # The step returned (rows_so_far * n) rows. Expand the carried columns to match
+            # -- each upstream row repeated n times (the block is flattened in the same
+            # row-major order). The root has no carried rows, and fan_out=1 stages are a
+            # no-op, so 1:1 chains are untouched.
             if i > 0 and n > 1:
                 samples = {k: v.repeat_interleave(n, 0) for k, v in samples.items()}
                 if torch.is_tensor(total):
                     total = total.repeat_interleave(n, 0)
             samples.update(block)
-            total = total + lp
-        return samples, total
+            # A single density-free step (Gibbs) makes the whole chain density-free.
+            if lp is None:
+                has_density = False
+            elif has_density:
+                total = total + lp
+        return samples, (total if has_density else None)
 
     def log_prob(
         self, samples: dict[str, torch.Tensor], context: SamplerContext
     ) -> torch.Tensor:
+        """Sum each step's ``log_prob`` at the given samples. Valid only for an all-density
+        chain (every step a ``Factor``)."""
         total: torch.Tensor | float = 0.0
-        for factor in self.factors:
-            cond = Conditioning(context, {k: samples[k] for k in factor.conditioning})
-            theta_i = {k: samples[k] for k in factor.parameters}
-            total = total + factor.log_prob(theta_i, cond)
+        for step in self.steps:
+            cond = Conditioning(context, {k: samples[k] for k in step.conditioning})
+            theta_i = {k: samples[k] for k in step.parameters}
+            total = total + step.log_prob(theta_i, cond)
         return total
 
     def sample(
@@ -358,42 +387,45 @@ def sample(
         context: SamplerContext,
         batch_size: Optional[int] = None,
     ) -> dict[str, torch.Tensor]:
-        """Per-sample dict of parameters plus ``log_prob``."""
+        """Per-sample dict of parameters, plus ``log_prob`` for an all-density chain."""
         samples, log_prob = self.sample_and_log_prob(num_samples, context, batch_size)
+        if log_prob is None:
+            return dict(samples)
         return {**samples, "log_prob": log_prob}
 
 
-class GibbsComposer:
+class GibbsBlock:
     """
-    Blocked Gibbs composer over a cyclic conditioning dependency.
+    A density-free sampling-block ``Step``: runs blocked Gibbs internally and yields no
+    proposal log-prob.
 
     Seeds the chain with an init factor, then sweeps the factor list in order for
     ``num_iterations`` iterations; each factor conditions on the current state and
-    overwrites its own block. The cyclic dependency has no tractable marginal, so
-    ``sample`` returns parameters without a ``log_prob`` (recoverable by fitting an
-    unconditional density to the samples and taking one ``ChainComposer`` step). Dingo uses
-    this only for multi-iteration GNPE (the GNPE factors in ``dingo.gw.inference.factors``),
-    but the loop is generic.
-
-    Batching is vertical, via ``chunk_and_concat``: chunk the walkers and run the whole
-    loop per chunk.
+    overwrites its own block. As a chain ``Step`` it produces the swept parameter blocks
+    and returns ``None`` for the log-prob -- the cyclic dependency has no tractable marginal
+    (recoverable by fitting an unconditional density to the samples and taking one
+    ``ChainComposer`` step). Dingo uses this only for multi-iteration GNPE (the GNPE factors
+    in ``dingo.gw.inference.factors``), but the loop is generic.
+
+    Batching is handled by the enclosing ``ChainComposer``: it chunks the walkers and runs
+    the whole loop per chunk (``chunk_and_concat``).
     """
 
     def __init__(self, init_factor: Factor, factors: list[Factor], num_iterations: int):
         self.init_factor = init_factor
         self.factors = list(factors)
         self.num_iterations = num_iterations
+        # The blocks this step produces (proxies + inference parameters), dropping
+        # side-channel columns such as the recomputed detector times.
+        self.parameters = [p for factor in self.factors for p in factor.parameters]
+        self.conditioning: list[str] = []
 
-    def sample(
-        self,
-        num_samples: int,
-        context: SamplerContext,
-        batch_size: Optional[int] = None,
-    ) -> dict[str, torch.Tensor]:
-        samples, _ = chunk_and_concat(
-            num_samples, batch_size, lambda n: (self._run_once(n, context), None)
-        )
-        return samples
+    def sample_and_log_prob(
+        self, num_samples: int, cond: Conditioning
+    ) -> tuple[dict[str, torch.Tensor], None]:
+        """Run the Gibbs loop for ``num_samples`` walkers; return ``(samples, None)``.
+        ``num_samples`` is the walker (root) count -- Gibbs does not fan out."""
+        return self._run_once(num_samples, cond.context), None
 
     def _run_once(
         self, num_samples: int, context: SamplerContext
@@ -405,35 +437,20 @@ def _run_once(
         state = dict(seed)
         for _ in range(self.num_iterations):
             for factor in self.factors:
-                cond = Conditioning(context, {k: state[k] for k in factor.conditioning})
+                c = Conditioning(context, {k: state[k] for k in factor.conditioning})
                 # One sample per walker (Gibbs is 1:1); walkers are the conditioning rows.
-                block, _ = factor.sample_and_log_prob(1, cond)
+                block, _ = factor.sample_and_log_prob(1, c)
                 state.update(block)
-        # Each factor's parameter block (e.g. proxies + inference parameters), dropping
-        # side-channel columns such as the recomputed detector times.
-        return {p: state[p] for factor in self.factors for p in factor.parameters}
-
-
-class Composer(Protocol):
-    """Interface required by ``ComposedSampler``: ``sample`` returns a per-sample dict of
-    parameters (plus ``log_prob`` for density-preserving composers). Satisfied by
-    ``ChainComposer`` and ``GibbsComposer``."""
-
-    def sample(
-        self,
-        num_samples: int,
-        context: SamplerContext,
-        batch_size: Optional[int] = None,
-    ) -> dict[str, torch.Tensor]: ...
+        return {p: state[p] for p in self.parameters}
 
 
 class ComposedSampler:
     """
-    Façade over a ``Composer`` and a ``SamplerContext``. Runs the composer, applies
+    Façade over a ``ChainComposer`` and a ``SamplerContext``. Runs the composer, applies
     domain-specific post-processing, and returns the samples as a DataFrame.
     """
 
-    def __init__(self, composer: Composer, context: SamplerContext):
+    def __init__(self, composer: ChainComposer, context: SamplerContext):
         self.composer = composer
         self.context = context
         self.samples: Optional[pd.DataFrame] = None
@@ -446,8 +463,8 @@ def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
         """Draw ``num_samples`` samples (chunked by ``batch_size``), post-process, and
-        return them as a DataFrame. ``ChainComposer`` includes ``log_prob``;
-        ``GibbsComposer`` does not."""
+        return them as a DataFrame. An all-density chain includes ``log_prob``; a chain
+        with a ``GibbsBlock`` step does not."""
         merged = self.composer.sample(num_samples, self.context, batch_size)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
         self._post_process(merged)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 6d60bb2ca..36f84f674 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -7,7 +7,7 @@
 ``GWSamplerContext`` (which builds the data-preprocessing conditioning map ``f_i``) and
 the post-processing in ``GWComposedSampler``. Covers plain NPE (``FlowFactor`` via a
 ``ChainComposer``) and multi-iteration time-shift GNPE (``GNPEKernelFactor`` +
-``GNPEFlowFactor`` cycled by a ``GibbsComposer``); the synthetic-phase factor and
+``GNPEFlowFactor`` cycled by a ``GibbsBlock``); the synthetic-phase factor and
 single-step GNPE importance sampling remain for later steps.
 """
 
@@ -27,11 +27,10 @@
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
-    Composer,
     Conditioning,
     Factor,
     FlowFactor,
-    GibbsComposer,
+    GibbsBlock,
     _base_model_metadata,
 )
 from dingo.core.posterior_models import BasePosteriorModel
@@ -323,7 +322,7 @@ class GNPEFlowFactor(Factor):
     and geocent time. The proxies are supplied, so no blurring happens here. Wraps the main
     model and the per-iteration transforms built for ``GWSamplerGNPE``.
 
-    The single network factor in either GNPE mode: cycled by ``GibbsComposer`` for
+    The single network factor in either GNPE mode: cycled by a ``GibbsBlock`` for
     multi-iteration GNPE, or a ``ChainComposer`` factor for single-step GNPE. The recomputed
     detector times are returned as extra columns: the next Gibbs iteration blurs them into
     fresh proxies, and single-step GNPE evaluates the kernel correction at them. One sample
@@ -402,7 +401,7 @@ class GWComposedSampler(ComposedSampler):
 
     def __init__(
         self,
-        composer: Composer,
+        composer: ChainComposer,
         context: GWSamplerContext,
         metadata: dict,
         inference_parameters: list[str],
@@ -443,18 +442,17 @@ def from_gnpe_models(
         num_iterations: int = 30,
     ) -> "GWComposedSampler":
         """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
-        init model's data preprocessing, an init ``FlowFactor`` to seed, and a
-        ``GibbsComposer`` cycling the GNPE kernel and main-network factors. Returns samples
-        without a log_prob (Gibbs breaks density access)."""
+        init model's data preprocessing, an init ``FlowFactor`` to seed, and a single
+        ``GibbsBlock`` step -- cycling the GNPE kernel and main-network factors -- in a
+        ``ChainComposer``. Returns samples without a log_prob (Gibbs breaks density
+        access)."""
         context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
         init_factor = FlowFactor.from_model(init_model)
         kernel_factor = GNPEKernelFactor.from_model(main_model)
         flow_factor = GNPEFlowFactor.from_model(main_model)
-        composer = GibbsComposer(
-            init_factor, [kernel_factor, flow_factor], num_iterations
-        )
+        gibbs = GibbsBlock(init_factor, [kernel_factor, flow_factor], num_iterations)
         return cls(
-            composer,
+            ChainComposer([gibbs]),
             context,
             _base_model_metadata(main_model),
             flow_factor.parameters,
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 58d7bc3d6..8042bbbf2 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -15,6 +15,7 @@
     ChainComposer,
     Conditioning,
     Factor,
+    GibbsBlock,
     Stage,
     chunk_and_concat,
 )
@@ -61,7 +62,7 @@ def test_bare_factor_is_wrapped_as_stage():
     comp = ChainComposer([_ConstFactor("a")])
     assert isinstance(comp.stages[0], Stage)
     assert comp.stages[0].fan_out == 1
-    assert comp.factors[0].parameters == ["a"]
+    assert comp.steps[0].parameters == ["a"]
 
 
 def test_fan_out_expansion_alignment_and_logprob():
@@ -146,3 +147,59 @@ def test_conditioning_dataclass_defaults():
     cond = Conditioning(context="ctx")
     assert cond.context == "ctx"
     assert cond.given == {}
+
+
+class _NoDensityStep:
+    """A density-free step (like ``GibbsBlock``): emits a block but returns ``None`` for the
+    log-prob. Honors the per-row fan-out contract so it composes like any step."""
+
+    def __init__(self, name, conditioning=()):
+        self.parameters = [name]
+        self.conditioning = list(conditioning)
+        self._name = name
+
+    def sample_and_log_prob(self, n, cond):
+        if self.conditioning:
+            base = sum(cond.given[k] for k in self.conditioning)
+            within = torch.arange(n, dtype=base.dtype)
+            vals = (base.unsqueeze(1) + within).reshape(-1)
+        else:
+            vals = torch.arange(n, dtype=torch.float32)
+        return {self._name: vals}, None
+
+
+def test_density_free_step_omits_log_prob():
+    # A density-free root step (Gibbs) yields no proposal density.
+    comp = ChainComposer([_NoDensityStep("a")])
+    _, lp = comp.sample_and_log_prob(5, context=None)
+    assert lp is None
+    out = comp.sample(5, context=None)
+    assert "log_prob" not in out
+    assert torch.equal(out["a"], torch.arange(5, dtype=torch.float32))
+
+
+def test_one_density_free_step_nulls_the_whole_chain():
+    # A single None step makes the chain density-free even alongside a density factor.
+    comp = ChainComposer(
+        [_ConstFactor("a"), Stage(_NoDensityStep("b", conditioning=["a"]))]
+    )
+    _, lp = comp.sample_and_log_prob(4, context=None)
+    assert lp is None
+    assert "log_prob" not in comp.sample(4, context=None)
+
+
+def test_gibbs_block_runs_as_a_density_free_step():
+    # GibbsBlock seeds, sweeps its factors num_iterations times, and yields no density;
+    # seed-only params (not reproduced by a factor) are dropped, matching the old composer.
+    block = GibbsBlock(
+        init_factor=_ConstFactor("proxy"),
+        factors=[_ConstFactor("theta", conditioning=["proxy"])],
+        num_iterations=3,
+    )
+    assert block.parameters == ["theta"] and block.conditioning == []
+    comp = ChainComposer([block])
+    out = comp.sample(6, context=None)
+    assert "log_prob" not in out
+    assert "proxy" not in out  # seed-only, dropped
+    # theta | proxy with Gibbs' one-per-walker draw: theta == proxy == arange(6).
+    assert torch.equal(out["theta"], torch.arange(6, dtype=torch.float32))

From be8a74ca45383ba37c12a9addeda85a0afad1050 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 14:33:24 +0200
Subject: [PATCH 14/65] Move RA frame and fixed params into chain steps; drop
 _post_process re-plug

Reparametrizations (Increment 1b):
- Add Reparametrization, a core bijector Step (forward/inverse/log_det,
  contributing -log|det J|, 1:1). RAReparam ports GWSamplerMixin's t_ref/RA
  sky rotation into the chain, via a ra@t_ref -> ra factor-boundary alias, and
  deletes _correct_reference_time. ra is stored float32 (a bounded angle; the
  old float64 was an incidental numpy promotion). The correction is computed in
  float64 (absolute GPS times); harnesses compare ra at float32.
- FlowFactor / GNPEFlowFactor gain an `aliases` map exposing trained names under
  canonical names at the factor boundary.
- The composer makes a Reparametrization consume its input (in-place bijection),
  so the network-frame intermediate never reaches post-processing.

Fixed parameters (Increment 2):
- Delta-prior parameters move from _post_process into a FixedFactor chain step
  (_fixed_prior_steps, resolved once at build time).
- The composer supports an unconditioned non-root step (a filler): it draws one
  value per current row rather than fanning out. FixedFactor also remains the
  chain root for prior-conditioning / known proxies.

Drop the dead _post_process inverse branch: its job was to strip the log_prob
column so old Sampler.log_prob's in-band `+=` would not double-count. The
factorized log_prob sums per-factor densities explicitly, so both the `+=` and
the strip are obsolete. _post_process is now just the single-step-GNPE kernel
correction; the `inverse` parameter is gone.

Parity: NPE, multi-iteration GNPE, single-step, and packaged all bit-exact
(ra at float32). tests/core/test_factors.py (reparam consume + -log_det +
round-trip; unconditioned filler incl. as root) and tests/gw/test_ra_reparam.py
(forward-inverse round trip + no-op) green; 21 tests total.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 116 +++++++++++++---
 dingo/gw/inference/factors.py | 244 +++++++++++++++++++++++-----------
 tests/core/test_factors.py    | 103 ++++++++++++++
 tests/gw/test_ra_reparam.py   |  44 ++++++
 4 files changed, 410 insertions(+), 97 deletions(-)
 create mode 100644 tests/gw/test_ra_reparam.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 6b6f94e14..77e788db7 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -187,9 +187,15 @@ def __init__(
         parameters: list[str],
         conditioning: Optional[list[str]] = None,
         context_parameters: Optional[list[str]] = None,
+        aliases: Optional[dict[str, str]] = None,
     ):
         self.model = model
-        self.parameters = parameters
+        # The network's trained parameter names -- standardization is keyed by these. The
+        # factor exposes them under canonical aliases (e.g. ra -> ra@t_ref), so a downstream
+        # reparametrization can convert frames by name without retraining (design Q#7).
+        self._net_parameters = parameters
+        self.aliases = aliases or {}
+        self.parameters = [self.aliases.get(p, p) for p in parameters]
         self.conditioning = conditioning or []
         # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
         self.context_parameters = context_parameters or []
@@ -197,9 +203,12 @@ def __init__(
         self.standardization = Standardization(std["mean"], std["std"])
 
     @classmethod
-    def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
+    def from_model(
+        cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
+    ) -> "FlowFactor":
         """Build a factor from a model, reading ``parameters`` and ``context_parameters``
-        from its training metadata."""
+        from its training metadata. ``aliases`` maps trained names to exposed canonical
+        names (e.g. ``{"ra": "ra@t_ref"}``) at the factor boundary."""
         data_settings = _base_model_metadata(model)["train_settings"]["data"]
         context_parameters = data_settings.get("context_parameters") or []
         return cls(
@@ -207,6 +216,7 @@ def from_model(cls, model: BasePosteriorModel) -> "FlowFactor":
             parameters=data_settings["inference_parameters"],
             conditioning=list(context_parameters),
             context_parameters=list(context_parameters),
+            aliases=aliases,
         )
 
     def sample_and_log_prob(self, num_samples, cond):
@@ -236,13 +246,19 @@ def sample_and_log_prob(self, num_samples, cond):
                 )
             z = z.reshape(n_rows * num_samples, z.shape[-1])
             log_prob = log_prob.reshape(n_rows * num_samples)
-        theta = self.standardization.destandardize(z, self.parameters)
-        log_prob = log_prob + self.standardization.log_det(self.parameters)
+        theta = self.standardization.destandardize(z, self._net_parameters)
+        log_prob = log_prob + self.standardization.log_det(self._net_parameters)
+        # Expose trained names under their canonical aliases at the factor boundary.
+        theta = {self.aliases.get(k, k): v for k, v in theta.items()}
         return theta, log_prob
 
     def log_prob(self, theta_i, cond):
-        num_samples = next(iter(theta_i.values())).shape[0]
-        z = self.standardization.standardize(theta_i, self.parameters)
+        # theta_i uses exposed (aliased) names; map back to the network's trained names.
+        theta_net = {
+            net: theta_i[self.aliases.get(net, net)] for net in self._net_parameters
+        }
+        num_samples = next(iter(theta_net.values())).shape[0]
+        z = self.standardization.standardize(theta_net, self._net_parameters)
         data = cond.context.prepared_data()
         data = data.expand(num_samples, *data.shape)
         net_context: tuple[torch.Tensor, ...]
@@ -254,7 +270,60 @@ def log_prob(self, theta_i, cond):
         self.model.network.eval()
         with torch.no_grad():
             log_prob = self.model.log_prob(z, *net_context)
-        return log_prob + self.standardization.log_det(self.parameters)
+        return log_prob + self.standardization.log_det(self._net_parameters)
+
+
+class Reparametrization(ABC):
+    """
+    A deterministic bijection ``Step``: it transforms existing parameters (no sampling)
+    and contributes ``-log|det J|`` to the proposal density.
+
+    Unlike a ``Factor`` it is 1:1 and invertible -- ``forward`` maps the conditioning block
+    to the ``parameters`` block, ``inverse`` maps back (for re-plug / importance sampling),
+    and its density contribution is a Jacobian, not a sampled log-prob. Used to relate a
+    network's coordinates to physical ones (e.g. right ascension from the training reference
+    frame to the event frame). Subclasses implement ``forward`` / ``inverse`` and, where the
+    map is not measure-preserving, ``log_det``.
+    """
+
+    parameters: list[str]
+    conditioning: list[str]
+
+    @abstractmethod
+    def forward(
+        self, given: dict[str, torch.Tensor], context: "SamplerContext"
+    ) -> dict[str, torch.Tensor]:
+        """Map the conditioning block to the ``parameters`` block."""
+
+    @abstractmethod
+    def inverse(
+        self, params: dict[str, torch.Tensor], context: "SamplerContext"
+    ) -> dict[str, torch.Tensor]:
+        """Map the ``parameters`` block back to the conditioning block."""
+
+    def log_det(
+        self, given: dict[str, torch.Tensor], context: "SamplerContext"
+    ) -> torch.Tensor:
+        """``log|det J|`` of ``forward``, per row. Default 0 (measure-preserving)."""
+        n = next(iter(given.values())).shape[0]
+        return torch.zeros(n)
+
+    def sample_and_log_prob(
+        self, num_samples: int, cond: "Conditioning"
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        """Apply ``forward`` to the conditioning; contribute ``-log|det J|``. ``num_samples``
+        must be 1 (a reparametrization is 1:1)."""
+        if num_samples != 1:
+            raise ValueError("A reparametrization is 1:1; use fan_out=1.")
+        out = self.forward(cond.given, cond.context)
+        return out, -self.log_det(cond.given, cond.context)
+
+    def log_prob(
+        self, theta_i: dict[str, torch.Tensor], cond: "Conditioning"
+    ) -> torch.Tensor:
+        raise NotImplementedError(
+            "Reparametrization.log_prob (re-plug through the inverse map) -- later step."
+        )
 
 
 class Step(Protocol):
@@ -350,18 +419,31 @@ def _run_chain_once(
         has_density = True
         for i, stage in enumerate(self.stages):
             step = stage.step
-            n = base if i == 0 else stage.fan_out
+            if i == 0:
+                n = base  # root: draw the base count
+            elif step.conditioning:
+                n = stage.fan_out  # fan_out samples per conditioning row
+            else:
+                # Unconditioned non-root step (e.g. a fixed/delta factor): draw one value
+                # per current row -- it fills the batch rather than fanning out.
+                n = len(next(iter(samples.values())))
             cond = Conditioning(context, {k: samples[k] for k in step.conditioning})
             block, lp = step.sample_and_log_prob(n, cond)
-            # The step returned (rows_so_far * n) rows. Expand the carried columns to match
-            # -- each upstream row repeated n times (the block is flattened in the same
-            # row-major order). The root has no carried rows, and fan_out=1 stages are a
-            # no-op, so 1:1 chains are untouched.
-            if i > 0 and n > 1:
-                samples = {k: v.repeat_interleave(n, 0) for k, v in samples.items()}
+            # A conditioned fan-out (fan_out > 1) expands the batch: repeat each carried row
+            # to align with the step's fan_out sub-rows (the block is flattened row-major).
+            # 1:1 stages and unconditioned fillers leave the batch untouched.
+            if i > 0 and step.conditioning and stage.fan_out > 1:
+                fan = stage.fan_out
+                samples = {k: v.repeat_interleave(fan, 0) for k, v in samples.items()}
                 if torch.is_tensor(total):
-                    total = total.repeat_interleave(n, 0)
+                    total = total.repeat_interleave(fan, 0)
             samples.update(block)
+            # A reparametrization is an in-place bijection: it replaces the columns it
+            # consumed with its outputs, so drop the consumed (non-reproduced) inputs.
+            if isinstance(step, Reparametrization):
+                for k in step.conditioning:
+                    if k not in step.parameters:
+                        samples.pop(k, None)
             # A single density-free step (Gibbs) makes the whole chain density-free.
             if lp is None:
                 has_density = False
@@ -455,7 +537,7 @@ def __init__(self, composer: ChainComposer, context: SamplerContext):
         self.context = context
         self.samples: Optional[pd.DataFrame] = None
 
-    def _post_process(self, samples: dict, inverse: bool = False):
+    def _post_process(self, samples: dict):
         """Hook for domain-specific post-processing; no-op by default."""
         pass
 
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 36f84f674..575cdc1f9 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -14,10 +14,9 @@
 from __future__ import annotations
 
 import copy
-from typing import Optional, Union
+from typing import Optional
 
 import numpy as np
-import pandas as pd
 import torch
 from astropy.time import Time
 from bilby.core.prior import DeltaFunction, PriorDict
@@ -31,6 +30,7 @@
     Factor,
     FlowFactor,
     GibbsBlock,
+    Reparametrization,
     _base_model_metadata,
 )
 from dingo.core.posterior_models import BasePosteriorModel
@@ -144,12 +144,15 @@ def likelihood(self):
         )
 
 
-# --- Stubs for later steps -------------------------------------------------------------
-
-
 class FixedFactor(Factor):
-    """``q_i = delta(theta_i - c)``: pins parameters to fixed values (prior-conditioning
-    / known proxies). Subsumes ``FixedInitSampler``. TODO: implement."""
+    """``q_i = delta(theta_i - c)``: pins parameters to fixed values, contributing 0 to the
+    proposal log-prob.
+
+    Serves two roles in a chain: the **root**, for prior-conditioning / known proxies
+    (single-step GNPE, where later factors condition on the pinned values -- subsumes
+    ``FixedInitSampler``), and a non-root **filler** for delta-prior parameters the network
+    does not infer (one constant per current row). ``log_prob`` (re-plug) is a later step.
+    """
 
     def __init__(self, values: dict[str, float]):
         self.values = values
@@ -166,6 +169,9 @@ def log_prob(self, theta_i, cond):
         raise NotImplementedError("FixedFactor.log_prob -- later step.")
 
 
+# --- Stubs for later steps -------------------------------------------------------------
+
+
 class SyntheticPhaseFactor(Factor):
     """``q(phase | theta_rest, d)`` from the likelihood on a phase grid; the final,
     non-NN factor, run in the CPU post-processing stage. TODO: port from
@@ -336,22 +342,31 @@ def __init__(
         transform_post: Compose,
         gnpe_parameters: list[str],
         parameters: list[str],
+        aliases: Optional[dict[str, str]] = None,
     ):
         self.model = model
         self.transform_pre = transform_pre
         self.transform_post = transform_post
         self.gnpe_parameters = gnpe_parameters
         self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
-        self.parameters = parameters
+        self.aliases = aliases or {}
+        self._net_parameters = parameters
+        self.parameters = [self.aliases.get(p, p) for p in parameters]
         self.conditioning = list(self.proxy_parameters)
 
     @classmethod
-    def from_model(cls, model: BasePosteriorModel) -> "GNPEFlowFactor":
-        """Build the GNPE per-iteration transforms from the main model's metadata."""
+    def from_model(
+        cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
+    ) -> "GNPEFlowFactor":
+        """Build the GNPE per-iteration transforms from the main model's metadata.
+        ``aliases`` maps trained names to canonical names (e.g. ``{"ra": "ra@t_ref"}``).
+        """
         pre, post, gnpe_parameters, inference_parameters, _, _ = _build_gnpe_transforms(
             model
         )
-        return cls(model, pre, post, gnpe_parameters, inference_parameters)
+        return cls(
+            model, pre, post, gnpe_parameters, inference_parameters, aliases=aliases
+        )
 
     def sample_and_log_prob(self, num_samples, cond):
         """Sample one parameter set per proxy row; ``num_samples`` must be 1 (GNPE is 1:1).
@@ -378,6 +393,8 @@ def sample_and_log_prob(self, num_samples, cond):
         x["log_prob"] = log_prob.squeeze(1)
         x = self.transform_post(x)
         params = dict(x["parameters"])
+        # Expose trained names under their canonical aliases (e.g. ra -> ra@t_ref).
+        params = {self.aliases.get(k, k): v for k, v in params.items()}
         # Surface the recomputed detector times: the next Gibbs iteration's input, and the
         # evaluation point for GNPEKernelFactor's importance-sampling correction.
         for k in self.gnpe_parameters:
@@ -390,6 +407,86 @@ def log_prob(self, theta_i, cond):
         )
 
 
+class RAReparam(Reparametrization):
+    """
+    Rotate right ascension from the network's training reference frame (``ra@t_ref``) to the
+    event frame (``ra``).
+
+    The network is trained at a fixed reference time; an event at a different GPS time needs
+    the sky rotated by the sidereal-time difference. This is a measure-preserving shift
+    modulo 2*pi (``log_det = 0``), so it contributes nothing to the density. Ported from
+    ``GWSamplerMixin._correct_reference_time``: ``forward`` produces the event-frame ``ra``
+    for downstream, ``inverse`` recovers ``ra@t_ref`` for re-plug / importance sampling
+    (design Q#7). The sidereal correction is read from the shared context (``t_ref`` and the
+    event time).
+    """
+
+    def __init__(self):
+        self.conditioning = ["ra@t_ref"]
+        self.parameters = ["ra"]
+
+    @staticmethod
+    def _correction(context) -> float:
+        """Sidereal-time difference (event minus reference) in radians; 0 when the event
+        time is unset or equal to the reference time."""
+        event_metadata = context.event_metadata
+        t_event = None if event_metadata is None else event_metadata.get("time_event")
+        t_ref = context.t_ref
+        if t_event is None or t_event == t_ref:
+            return 0.0
+        longitude_event = Time(t_event, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        longitude_reference = Time(t_ref, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        return (longitude_event - longitude_reference).rad
+
+    def forward(self, given, context):
+        correction = self._correction(context)
+        if correction == 0.0:
+            return {"ra": given["ra@t_ref"]}
+        # ra is a bounded angle -> float32 is plenty. The correction is a difference of
+        # absolute GPS times, so compute it in float64, but store the wrapped angle float32.
+        ra = (given["ra@t_ref"].double() + correction) % (2 * np.pi)
+        return {"ra": ra.float()}
+
+    def inverse(self, params, context):
+        correction = self._correction(context)
+        if correction == 0.0:
+            return {"ra@t_ref": params["ra"]}
+        ra_tref = (params["ra"].double() - correction) % (2 * np.pi)
+        return {"ra@t_ref": ra_tref.float()}
+
+
+def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
+    """The RA frame alias (``ra`` -> ``ra@t_ref``), applied only when the model infers
+    ``ra``; paired with an ``RAReparam`` step that maps it back to the event frame."""
+    return {"ra": "ra@t_ref"} if "ra" in inference_parameters else {}
+
+
+def _ra_reparam_steps(inference_parameters: list[str]) -> list:
+    """The ``RAReparam`` step, appended to a chain only when the model infers ``ra``."""
+    return [RAReparam()] if "ra" in inference_parameters else []
+
+
+def _fixed_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
+    """Delta-prior parameters the chain does not produce, as a single ``FixedFactor`` step
+    (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0); the
+    old sampler injected them in ``_post_process``."""
+    intrinsic_prior = metadata["dataset_settings"]["intrinsic_prior"]
+    extrinsic_prior = get_extrinsic_prior_dict(
+        metadata["train_settings"]["data"]["extrinsic_prior"]
+    )
+    prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
+    fixed = {
+        k: p.peak
+        for k, p in prior.items()
+        if isinstance(p, DeltaFunction) and k not in inference_parameters
+    }
+    return [FixedFactor(fixed)] if fixed else []
+
+
 class GWComposedSampler(ComposedSampler):
     """
     GW façade: a ``ComposedSampler`` that adds the gravitational-wave post-processing --
@@ -422,14 +519,24 @@ def from_model(
         raw_context: dict,
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
-        """Build a single-factor (plain-NPE) GW sampler from a model and event data."""
+        """Build a plain-NPE GW sampler from a model and event data: the flow exposes
+        ``ra`` as ``ra@t_ref``, followed by an ``RAReparam`` to the event frame."""
         context = GWSamplerContext.from_model(model, raw_context, event_metadata)
-        factor = FlowFactor.from_model(model)
+        metadata = _base_model_metadata(model)
+        inference_parameters = metadata["train_settings"]["data"][
+            "inference_parameters"
+        ]
+        factor = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
+        steps = (
+            [factor]
+            + _ra_reparam_steps(inference_parameters)
+            + _fixed_prior_steps(metadata, inference_parameters)
+        )
         return cls(
-            composer=ChainComposer([factor]),
+            composer=ChainComposer(steps),
             context=context,
-            metadata=_base_model_metadata(model),
-            inference_parameters=factor.parameters,
+            metadata=metadata,
+            inference_parameters=inference_parameters,
         )
 
     @classmethod
@@ -444,18 +551,29 @@ def from_gnpe_models(
         """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
         init model's data preprocessing, an init ``FlowFactor`` to seed, and a single
         ``GibbsBlock`` step -- cycling the GNPE kernel and main-network factors -- in a
-        ``ChainComposer``. Returns samples without a log_prob (Gibbs breaks density
-        access)."""
+        ``ChainComposer``, then an ``RAReparam`` to the event frame. Returns samples without
+        a log_prob (Gibbs breaks density access)."""
         context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
+        metadata = _base_model_metadata(main_model)
+        inference_parameters = metadata["train_settings"]["data"][
+            "inference_parameters"
+        ]
         init_factor = FlowFactor.from_model(init_model)
         kernel_factor = GNPEKernelFactor.from_model(main_model)
-        flow_factor = GNPEFlowFactor.from_model(main_model)
+        flow_factor = GNPEFlowFactor.from_model(
+            main_model, aliases=_ra_aliases(inference_parameters)
+        )
         gibbs = GibbsBlock(init_factor, [kernel_factor, flow_factor], num_iterations)
+        steps = (
+            [gibbs]
+            + _ra_reparam_steps(inference_parameters)
+            + _fixed_prior_steps(metadata, inference_parameters)
+        )
         return cls(
-            ChainComposer([gibbs]),
+            ChainComposer(steps),
             context,
-            _base_model_metadata(main_model),
-            flow_factor.parameters,
+            metadata,
+            inference_parameters,
         )
 
     @classmethod
@@ -467,49 +585,35 @@ def from_singlestep_gnpe(
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
         """Build a single-step (density-preserving) time-GNPE sampler: a ``ChainComposer``
-        of ``[proxy_source, GNPEFlowFactor]``. ``proxy_source`` supplies the detector-time
-        proxies -- a ``FixedFactor`` for prior conditioning (BNS), or an unconditional NDE
-        for density recovery. Unlike multi-iteration GNPE the chain is autoregressive, so
-        log_prob is preserved; the ``GNPEKernelFactor`` supplies the
+        of ``[proxy_source, GNPEFlowFactor, RAReparam]``. ``proxy_source`` supplies the
+        detector-time proxies -- a ``FixedFactor`` for prior conditioning (BNS), or an
+        unconditional NDE for density recovery. Unlike multi-iteration GNPE the chain is
+        autoregressive, so log_prob is preserved; the ``GNPEKernelFactor`` supplies the
         ``delta_log_prob_target`` correction that importance sampling adds to the target.
         """
         context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
-        flow_factor = GNPEFlowFactor.from_model(main_model)
+        metadata = _base_model_metadata(main_model)
+        inference_parameters = metadata["train_settings"]["data"][
+            "inference_parameters"
+        ]
+        flow_factor = GNPEFlowFactor.from_model(
+            main_model, aliases=_ra_aliases(inference_parameters)
+        )
         kernel_factor = GNPEKernelFactor.from_model(main_model)
-        composer = ChainComposer([proxy_source, flow_factor])
+        steps = (
+            [proxy_source, flow_factor]
+            + _ra_reparam_steps(inference_parameters)
+            + _fixed_prior_steps(metadata, inference_parameters)
+        )
+        composer = ChainComposer(steps)
         return cls(
             composer,
             context,
-            _base_model_metadata(main_model),
-            flow_factor.parameters,
+            metadata,
+            inference_parameters,
             kernel_factor=kernel_factor,
         )
 
-    def _correct_reference_time(
-        self, samples: Union[dict, pd.DataFrame], inverse: bool = False
-    ):
-        """Correct the right ascension for the difference between the event time and the
-        model's training reference time (fixed detector positions)."""
-        event_metadata = self.context.event_metadata
-        if event_metadata is None:
-            return
-        t_event = event_metadata.get("time_event")
-        t_ref = self.context.t_ref
-        if t_event is None or t_event == t_ref or "ra" not in samples:
-            return
-        ra = samples["ra"]
-        longitude_event = Time(t_event, format="gps", scale="utc").sidereal_time(
-            "apparent", "greenwich"
-        )
-        longitude_reference = Time(t_ref, format="gps", scale="utc").sidereal_time(
-            "apparent", "greenwich"
-        )
-        ra_correction = (longitude_event - longitude_reference).rad
-        if not inverse:
-            samples["ra"] = (ra + ra_correction) % (2 * np.pi)
-        else:
-            samples["ra"] = (ra - ra_correction) % (2 * np.pi)
-
     def _add_kernel_correction(self, samples: dict):
         """Single-step GNPE: add ``delta_log_prob_target = log p(theta_hat | theta)``,
         evaluated at the proxies and the detector times recomputed from theta, then drop
@@ -523,33 +627,13 @@ def _add_kernel_correction(self, samples: dict):
         for k in kf.gnpe_parameters:
             samples.pop(k, None)
 
-    def _post_process(self, samples: Union[dict, pd.DataFrame], inverse: bool = False):
-        if self.kernel_factor is not None and not inverse:
+    def _post_process(self, samples: dict):
+        # Fixed (delta-prior) parameters are a FixedFactor in the chain and the RA frame
+        # rotation is a chain reparametrization; only the single-step-GNPE kernel correction
+        # remains here (moves to a chain step next).
+        if self.kernel_factor is not None:
             self._add_kernel_correction(samples)
 
-        intrinsic_prior = self.metadata["dataset_settings"]["intrinsic_prior"]
-        extrinsic_prior = get_extrinsic_prior_dict(
-            self.metadata["train_settings"]["data"]["extrinsic_prior"]
-        )
-        prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
-
-        if not inverse:
-            # Add fixed (delta-prior) parameters not produced by the chain.
-            num_samples = len(samples[list(samples.keys())[0]])
-            for k, p in prior.items():
-                if isinstance(p, DeltaFunction) and k not in samples:
-                    print(f"Adding fixed parameter {k} = {p.peak} from prior.")
-                    samples[k] = p.peak * np.ones(num_samples)
-        else:
-            drop = [k for k in samples.keys() if k not in self.inference_parameters]
-            if isinstance(samples, pd.DataFrame):
-                samples.drop(columns=drop, inplace=True, errors="ignore")
-            else:
-                for k in drop:
-                    samples.pop(k, None)
-
-        self._correct_reference_time(samples, inverse)
-
     def to_result(self):
         """Export to a gw ``Result`` (samples + raw event data + metadata), so the
         existing post-processing pipeline -- synthetic phase, importance sampling,
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 8042bbbf2..7e01cfee0 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -16,6 +16,7 @@
     Conditioning,
     Factor,
     GibbsBlock,
+    Reparametrization,
     Stage,
     chunk_and_concat,
 )
@@ -203,3 +204,105 @@ def test_gibbs_block_runs_as_a_density_free_step():
     assert "proxy" not in out  # seed-only, dropped
     # theta | proxy with Gibbs' one-per-walker draw: theta == proxy == arange(6).
     assert torch.equal(out["theta"], torch.arange(6, dtype=torch.float32))
+
+
+class _MockReparam(Reparametrization):
+    """A mock bijection ``u -> v = u + shift`` with a constant nonzero ``log|det J|`` so the
+    density contribution (``-log_det``) is checkable."""
+
+    def __init__(self, shift=10.0, log_det_val=0.5):
+        self.conditioning = ["u"]
+        self.parameters = ["v"]
+        self._shift = shift
+        self._ld = log_det_val
+
+    def forward(self, given, context):
+        return {"v": given["u"] + self._shift}
+
+    def inverse(self, params, context):
+        return {"u": params["v"] - self._shift}
+
+    def log_det(self, given, context):
+        n = next(iter(given.values())).shape[0]
+        return torch.full((n,), self._ld)
+
+
+def test_reparam_step_consumes_input_and_contributes_neg_logdet():
+    # A reparam is an in-place bijection: it replaces its input column and contributes
+    # -log_det to the chain density.
+    comp = ChainComposer([_ConstFactor("u"), _MockReparam(shift=10.0, log_det_val=0.5)])
+    out = comp.sample(4, context=None)
+    assert "u" not in out  # consumed
+    assert torch.equal(out["v"], torch.arange(4, dtype=torch.float32) + 10.0)
+    # chain log_prob = factor(+0.5) + reparam(-0.5) = 0.
+    assert torch.allclose(out["log_prob"], torch.zeros(4))
+
+
+def test_reparam_rejects_fan_out():
+    rp = _MockReparam()
+    with pytest.raises(ValueError, match="1:1"):
+        rp.sample_and_log_prob(
+            2, Conditioning(context=None, given={"u": torch.zeros(3)})
+        )
+
+
+def test_reparam_forward_inverse_round_trip():
+    rp = _MockReparam(shift=10.0)
+    u = torch.arange(5, dtype=torch.float32)
+    v = rp.forward({"u": u}, context=None)["v"]
+    u2 = rp.inverse({"v": v}, context=None)["u"]
+    assert torch.equal(u, u2)
+
+
+class _ConstFillFactor(Factor):
+    """An unconditioned factor emitting a constant (like a ``FixedFactor``). As the root it
+    draws the base count; as a non-root step it fills one value per current row."""
+
+    def __init__(self, name, value=3.0):
+        self.parameters = [name]
+        self.conditioning = []
+        self._name, self._v = name, value
+
+    def sample_and_log_prob(self, n, cond):
+        return {self._name: torch.full((n,), self._v)}, torch.zeros(n)
+
+    def log_prob(self, theta_i, cond):
+        return torch.zeros(next(iter(theta_i.values())).shape[0])
+
+
+def test_unconditioned_filler_fills_the_current_batch():
+    # A non-root unconditioned factor (fixed/delta filler) emits one value per current row,
+    # not a single row, and contributes 0 to the density.
+    comp = ChainComposer([_ConstFactor("a"), _ConstFillFactor("c", 3.0)])
+    out = comp.sample(5, context=None)
+    assert out["a"].shape == (5,) and out["c"].shape == (5,)
+    assert torch.equal(out["c"], torch.full((5,), 3.0))
+    assert torch.allclose(
+        out["log_prob"], torch.full((5,), 0.5)
+    )  # factor(0.5) + fill(0)
+
+
+def test_unconditioned_filler_fills_after_fan_out():
+    # The filler matches the expanded batch, not the base count.
+    comp = ChainComposer(
+        [
+            _ConstFactor("a"),
+            Stage(_ConstFactor("b", conditioning=["a"]), fan_out=3),
+            _ConstFillFactor("c", 7.0),
+        ]
+    )
+    out = comp.sample(4, context=None)  # 4 * 3 = 12 rows
+    assert out["c"].shape == (12,)
+    assert torch.equal(out["c"], torch.full((12,), 7.0))
+
+
+def test_unconditioned_factor_as_root_draws_base_count():
+    # As the chain root (prior-conditioning / FixedFactor proxy), an unconditioned factor
+    # draws the base count; a downstream factor conditions on the pinned values.
+    comp = ChainComposer(
+        [_ConstFillFactor("c", 2.0), _ConstFactor("d", conditioning=["c"])]
+    )
+    out = comp.sample(6, context=None)
+    assert out["c"].shape == (6,) and out["d"].shape == (6,)
+    assert torch.equal(out["c"], torch.full((6,), 2.0))
+    assert torch.equal(out["d"], torch.full((6,), 2.0))  # d | c: sum(c) + 0 == c
diff --git a/tests/gw/test_ra_reparam.py b/tests/gw/test_ra_reparam.py
new file mode 100644
index 000000000..83b4c8b70
--- /dev/null
+++ b/tests/gw/test_ra_reparam.py
@@ -0,0 +1,44 @@
+"""RAReparam: the right-ascension frame reparametrization (network training frame
+``ra@t_ref`` <-> event frame ``ra``). Covers the forward/inverse round trip -- the inverse
+direction the model-based parity harnesses do not exercise -- and the no-op case (event
+time equal to, or absent from, the reference time)."""
+
+import types
+
+import numpy as np
+import torch
+
+from dingo.gw.inference.factors import RAReparam
+
+
+def _context(t_ref, t_event):
+    """A minimal stand-in for the SamplerContext fields RAReparam reads."""
+    event_metadata = None if t_event is None else {"time_event": t_event}
+    return types.SimpleNamespace(t_ref=t_ref, event_metadata=event_metadata)
+
+
+def test_ra_reparam_round_trip():
+    rp = RAReparam()
+    ctx = _context(t_ref=1126259462.4, t_event=1264316116.4)  # differ -> real rotation
+    ra_tref = torch.rand(1000, dtype=torch.float32) * (2 * np.pi)
+
+    ra = rp.forward({"ra@t_ref": ra_tref}, ctx)["ra"]
+    back = rp.inverse({"ra": ra}, ctx)["ra@t_ref"]
+
+    assert ra.dtype == torch.float32  # a bounded angle; float32 is sufficient
+    assert not torch.allclose(ra, ra_tref)  # the sidereal rotation actually moved it
+    assert torch.allclose(back, ra_tref, atol=1e-6)  # forward then inverse recovers it
+
+
+def test_ra_reparam_is_noop_without_a_distinct_event_time():
+    rp = RAReparam()
+    ra_tref = torch.rand(100, dtype=torch.float32) * (2 * np.pi)
+    for ctx in (_context(1126259462.4, 1126259462.4), _context(1126259462.4, None)):
+        out = rp.forward({"ra@t_ref": ra_tref}, ctx)["ra"]
+        assert torch.equal(out, ra_tref)  # untouched: no cast, no modulo
+
+
+def test_ra_reparam_declared_names():
+    rp = RAReparam()
+    assert rp.conditioning == ["ra@t_ref"]
+    assert rp.parameters == ["ra"]

From b0bc6b5f7484ad635699dffb41e9942b47c85ad3 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 14:56:41 +0200
Subject: [PATCH 15/65] Move the GNPE kernel correction into a chain step;
 remove _post_process

- Add TargetCorrection, a core kind-3 Step: it emits an importance-sampling
  target-side column and contributes 0 to the proposal density (1:1).
  GNPEKernelCorrection emits delta_log_prob_target = log p(theta_hat | theta) and
  consumes the intermediate detector times.
- Generalize the composer's consume to a `consumes` attribute (used by both
  reparametrizations and target corrections), and add `produces` so a step's
  side-channel columns (e.g. GNPEFlowFactor's recomputed detector times) satisfy
  the topological check.
- from_singlestep_gnpe adds GNPEKernelCorrection to the chain; drop the
  kernel_factor __init__ argument and _add_kernel_correction.
- GWComposedSampler no longer overrides _post_process (all GW-specific processing
  is now chain steps); the base ComposedSampler hook stays as a no-op.
- Tighten docstrings and NotImplementedError messages.

Parity: NPE, multi-iteration GNPE, single-step, and packaged all bit-exact;
tests/core/test_factors.py + tests/gw/test_ra_reparam.py green (24 tests).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         |  59 +++++++++---
 dingo/gw/inference/factors.py | 171 +++++++++++++++-------------------
 tests/core/test_factors.py    |  67 +++++++++++++
 3 files changed, 190 insertions(+), 107 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 77e788db7..5443352bc 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -318,12 +318,50 @@ def sample_and_log_prob(
         out = self.forward(cond.given, cond.context)
         return out, -self.log_det(cond.given, cond.context)
 
+    @property
+    def consumes(self) -> list[str]:
+        """Conditioning the bijection replaces with its outputs (dropped after the step)."""
+        return [c for c in self.conditioning if c not in self.parameters]
+
     def log_prob(
         self, theta_i: dict[str, torch.Tensor], cond: "Conditioning"
     ) -> torch.Tensor:
-        raise NotImplementedError(
-            "Reparametrization.log_prob (re-plug through the inverse map) -- later step."
-        )
+        raise NotImplementedError("Reparametrization.log_prob")
+
+
+class TargetCorrection(ABC):
+    """
+    A ``Step`` that emits an importance-sampling target correction as a side-channel column
+    and contributes nothing to the proposal density.
+
+    Its value belongs to the IS target, not the proposal: it reads earlier blocks, emits a
+    named column, optionally consumes intermediates (``consumes``), and adds 0 to the
+    proposal log-prob. 1:1.
+    """
+
+    parameters: list[str]
+    conditioning: list[str]
+    consumes: list[str]
+
+    @abstractmethod
+    def correction(
+        self, given: dict[str, torch.Tensor], context: "SamplerContext"
+    ) -> dict[str, torch.Tensor]:
+        """The side-channel column(s), one value per conditioning row."""
+
+    def sample_and_log_prob(
+        self, num_samples: int, cond: "Conditioning"
+    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+        if num_samples != 1:
+            raise ValueError("A target correction is 1:1; use fan_out=1.")
+        out = self.correction(cond.given, cond.context)
+        n = next(iter(out.values())).shape[0]
+        return out, torch.zeros(n)
+
+    def log_prob(
+        self, theta_i: dict[str, torch.Tensor], cond: "Conditioning"
+    ) -> torch.Tensor:
+        raise NotImplementedError("TargetCorrection.log_prob")
 
 
 class Step(Protocol):
@@ -374,7 +412,8 @@ def __init__(self, stages: list[Union["Stage", Step]]):
 
     def _validate(self):
         """Check the declared order is a valid topological order: every conditioning
-        name is produced by an earlier step."""
+        name is produced by an earlier step. A step's emitted columns default to its
+        ``parameters``, but a step may emit side-channel columns too (``produces``)."""
         produced: set[str] = set()
         for step in self.steps:
             missing = [c for c in step.conditioning if c not in produced]
@@ -383,7 +422,7 @@ def _validate(self):
                     f"A step producing {step.parameters} conditions on {missing}, "
                     f"which no earlier step produces. Check chain order."
                 )
-            produced.update(step.parameters)
+            produced.update(getattr(step, "produces", step.parameters))
 
     @property
     def steps(self) -> list[Step]:
@@ -438,12 +477,10 @@ def _run_chain_once(
                 if torch.is_tensor(total):
                     total = total.repeat_interleave(fan, 0)
             samples.update(block)
-            # A reparametrization is an in-place bijection: it replaces the columns it
-            # consumed with its outputs, so drop the consumed (non-reproduced) inputs.
-            if isinstance(step, Reparametrization):
-                for k in step.conditioning:
-                    if k not in step.parameters:
-                        samples.pop(k, None)
+            # Drop any intermediates the step consumed: a reparametrization replaces its
+            # inputs (in-place bijection); a target correction drops the columns it read.
+            for k in getattr(step, "consumes", ()):
+                samples.pop(k, None)
             # A single density-free step (Gibbs) makes the whole chain density-free.
             if lp is None:
                 has_density = False
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 575cdc1f9..ef32b5450 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -1,14 +1,12 @@
 """
-Gravitational-wave factors for the factorized sampler.
-
-GW-specific pieces of the factorized-sampler design (see
-vault/Hackathon/Factorized_Sampler_Design.md). The generic ``FlowFactor`` is
-domain-agnostic and lives in ``dingo.core.factors``; what is GW-specific is the
-``GWSamplerContext`` (which builds the data-preprocessing conditioning map ``f_i``) and
-the post-processing in ``GWComposedSampler``. Covers plain NPE (``FlowFactor`` via a
-``ChainComposer``) and multi-iteration time-shift GNPE (``GNPEKernelFactor`` +
-``GNPEFlowFactor`` cycled by a ``GibbsBlock``); the synthetic-phase factor and
-single-step GNPE importance sampling remain for later steps.
+Gravitational-wave factors, steps, and samplers for the factorized sampler.
+
+The domain-agnostic pieces (``Factor``, ``FlowFactor``, ``Reparametrization``,
+``TargetCorrection``, and the composers) live in ``dingo.core.factors``. This module adds
+the GW-specific ones: ``GWSamplerContext`` (the data-preprocessing conditioning map), the
+GNPE factors and steps (``GNPEKernelFactor``, ``GNPEFlowFactor``, ``GNPEKernelCorrection``),
+the ``RAReparam`` sky-frame reparametrization, and ``GWComposedSampler``, which assembles
+the chain for plain NPE, multi-iteration GNPE, or single-step GNPE.
 """
 
 from __future__ import annotations
@@ -31,6 +29,7 @@
     FlowFactor,
     GibbsBlock,
     Reparametrization,
+    TargetCorrection,
     _base_model_metadata,
 )
 from dingo.core.posterior_models import BasePosteriorModel
@@ -138,20 +137,16 @@ def prepared_data(self) -> torch.Tensor:
         return self._prepared
 
     def likelihood(self):
-        raise NotImplementedError(
-            "Likelihood construction moves onto the context in a later step "
-            "(synthetic-phase / importance-sampling factors)."
-        )
+        raise NotImplementedError("GWSamplerContext.likelihood")
 
 
 class FixedFactor(Factor):
     """``q_i = delta(theta_i - c)``: pins parameters to fixed values, contributing 0 to the
     proposal log-prob.
 
-    Serves two roles in a chain: the **root**, for prior-conditioning / known proxies
-    (single-step GNPE, where later factors condition on the pinned values -- subsumes
-    ``FixedInitSampler``), and a non-root **filler** for delta-prior parameters the network
-    does not infer (one constant per current row). ``log_prob`` (re-plug) is a later step.
+    Used as the chain root for prior-conditioning or known proxies (single-step GNPE, where
+    later factors condition on the pinned values), and as a non-root filler for delta-prior
+    parameters the network does not infer (one constant per current row).
     """
 
     def __init__(self, values: dict[str, float]):
@@ -166,26 +161,25 @@ def sample_and_log_prob(self, num_samples, cond):
         return samples, torch.zeros(num_samples)
 
     def log_prob(self, theta_i, cond):
-        raise NotImplementedError("FixedFactor.log_prob -- later step.")
+        raise NotImplementedError("FixedFactor.log_prob")
 
 
 # --- Stubs for later steps -------------------------------------------------------------
 
 
 class SyntheticPhaseFactor(Factor):
-    """``q(phase | theta_rest, d)`` from the likelihood on a phase grid; the final,
-    non-NN factor, run in the CPU post-processing stage. TODO: port from
-    ``Result.sample_synthetic_phase``."""
+    """``q(phase | theta_rest, d)`` from the likelihood on a phase grid (a non-network
+    factor, the final step of the chain). Not yet implemented."""
 
     def __init__(self):
         self.parameters = ["phase"]
         self.conditioning = []  # conditions on all preceding params via the likelihood
 
     def sample_and_log_prob(self, num_samples, cond):
-        raise NotImplementedError("SyntheticPhaseFactor -- later step.")
+        raise NotImplementedError("SyntheticPhaseFactor")
 
     def log_prob(self, theta_i, cond):
-        raise NotImplementedError("SyntheticPhaseFactor -- later step.")
+        raise NotImplementedError("SyntheticPhaseFactor")
 
 
 def _build_gnpe_transforms(model: BasePosteriorModel):
@@ -267,21 +261,14 @@ def _build_gnpe_transforms(model: BasePosteriorModel):
 
 class GNPEKernelFactor(Factor):
     """
-    The GNPE perturbation kernel p(theta_hat | theta) as a factor (non-NN).
-
-    GNPE (arXiv:2111.13139) conditions the main network on "proxy" parameters theta_hat
-    drawn from a fixed kernel p(theta_hat | theta). Here theta are the detector coalescence
-    times and the kernel adds a bounded perturbation to each. Conditioning on the proxies
-    frees the network input to be simplified by them -- ``GNPEFlowFactor`` shifts each
-    detector's strain to its proxy time -- at the price of inferring theta and theta_hat
-    jointly.
+    The GNPE perturbation kernel ``p(theta_hat | theta)`` as a non-network factor.
 
-    The parameter block is the proxies; the conditioning is the detector times.
+    ``theta`` are the detector coalescence times; the kernel adds a bounded perturbation to
+    each, giving the proxies ``theta_hat`` the main network conditions on. The parameter
+    block is the proxies, the conditioning is the detector times.
     ``sample_and_log_prob`` blurs the times into proxies (the proxy update of a Gibbs
-    sweep). ``log_prob`` returns the kernel density log p(theta_hat | theta); for
-    single-step GNPE this is the ``delta_log_prob_target`` importance-sampling correction,
-    evaluated at the proxies and the detector times recomputed from theta. One proxy per
-    detector-time row.
+    sweep); ``log_prob`` returns the kernel density ``log p(theta_hat | theta)`` at the
+    proxies and the detector times. One proxy per detector-time row.
     """
 
     def __init__(self, gnpe_transform: GNPEBase, gnpe_parameters: list[str]):
@@ -320,19 +307,18 @@ def log_prob(self, theta_i, cond):
 
 class GNPEFlowFactor(Factor):
     """
-    The GNPE main network q(theta | theta_hat, d) as a factor.
+    The GNPE main network ``q(theta | theta_hat, d)`` as a factor.
 
     Conditions on the detector-time proxies from ``GNPEKernelFactor``: it shifts each
     detector's strain by the corresponding proxy time (standardizing the network input),
     samples the network, and recomputes the detector times from the sampled sky position
-    and geocent time. The proxies are supplied, so no blurring happens here. Wraps the main
-    model and the per-iteration transforms built for ``GWSamplerGNPE``.
+    and geocent time. The proxies are supplied, so no blurring happens here.
 
     The single network factor in either GNPE mode: cycled by a ``GibbsBlock`` for
     multi-iteration GNPE, or a ``ChainComposer`` factor for single-step GNPE. The recomputed
-    detector times are returned as extra columns: the next Gibbs iteration blurs them into
-    fresh proxies, and single-step GNPE evaluates the kernel correction at them. One sample
-    per proxy row.
+    detector times are emitted as extra columns (``produces``): the next Gibbs iteration
+    blurs them into fresh proxies, and single-step GNPE evaluates the kernel correction at
+    them. One sample per proxy row.
     """
 
     def __init__(
@@ -354,6 +340,11 @@ def __init__(
         self.parameters = [self.aliases.get(p, p) for p in parameters]
         self.conditioning = list(self.proxy_parameters)
 
+    @property
+    def produces(self) -> list[str]:
+        """Emitted columns: the inference block plus the recomputed detector times."""
+        return self.parameters + self.gnpe_parameters
+
     @classmethod
     def from_model(
         cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
@@ -402,9 +393,7 @@ def sample_and_log_prob(self, num_samples, cond):
         return params, x["log_prob"]
 
     def log_prob(self, theta_i, cond):
-        raise NotImplementedError(
-            "GNPEFlowFactor.log_prob -- later step (single-step GNPE importance sampling)."
-        )
+        raise NotImplementedError("GNPEFlowFactor.log_prob")
 
 
 class RAReparam(Reparametrization):
@@ -414,11 +403,9 @@ class RAReparam(Reparametrization):
 
     The network is trained at a fixed reference time; an event at a different GPS time needs
     the sky rotated by the sidereal-time difference. This is a measure-preserving shift
-    modulo 2*pi (``log_det = 0``), so it contributes nothing to the density. Ported from
-    ``GWSamplerMixin._correct_reference_time``: ``forward`` produces the event-frame ``ra``
-    for downstream, ``inverse`` recovers ``ra@t_ref`` for re-plug / importance sampling
-    (design Q#7). The sidereal correction is read from the shared context (``t_ref`` and the
-    event time).
+    modulo 2*pi (``log_det = 0``), so it contributes nothing to the density. ``forward``
+    produces the event-frame ``ra``, ``inverse`` recovers ``ra@t_ref``. The sidereal
+    correction is read from the shared context (``t_ref`` and the event time).
     """
 
     def __init__(self):
@@ -459,6 +446,31 @@ def inverse(self, params, context):
         return {"ra@t_ref": ra_tref.float()}
 
 
+class GNPEKernelCorrection(TargetCorrection):
+    """
+    The single-step GNPE kernel correction as a target-side chain step.
+
+    Emits ``delta_log_prob_target = log p(theta_hat | theta)`` -- the GNPE kernel evaluated
+    at the proxies and the detector times recomputed from theta -- for importance sampling
+    on the joint proposal ``q(theta, theta_hat | d)``. Contributes 0 to the proposal
+    density and consumes the intermediate detector times.
+    """
+
+    def __init__(self, kernel_factor: GNPEKernelFactor):
+        self.kernel_factor = kernel_factor
+        self.parameters = ["delta_log_prob_target"]
+        self.conditioning = list(kernel_factor.parameters) + list(
+            kernel_factor.gnpe_parameters
+        )
+        self.consumes = list(kernel_factor.gnpe_parameters)
+
+    def correction(self, given, context):
+        proxies = {p: given[p] for p in self.kernel_factor.parameters}
+        times = {k: given[k] for k in self.kernel_factor.gnpe_parameters}
+        correction = self.kernel_factor.log_prob(proxies, Conditioning(context, times))
+        return {"delta_log_prob_target": correction}
+
+
 def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
     """The RA frame alias (``ra`` -> ``ra@t_ref``), applied only when the model infers
     ``ra``; paired with an ``RAReparam`` step that maps it back to the event frame."""
@@ -472,8 +484,8 @@ def _ra_reparam_steps(inference_parameters: list[str]) -> list:
 
 def _fixed_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
     """Delta-prior parameters the chain does not produce, as a single ``FixedFactor`` step
-    (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0); the
-    old sampler injected them in ``_post_process``."""
+    (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0).
+    """
     intrinsic_prior = metadata["dataset_settings"]["intrinsic_prior"]
     extrinsic_prior = get_extrinsic_prior_dict(
         metadata["train_settings"]["data"]["extrinsic_prior"]
@@ -489,11 +501,10 @@ def _fixed_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
 
 class GWComposedSampler(ComposedSampler):
     """
-    GW façade: a ``ComposedSampler`` that adds the gravitational-wave post-processing --
-    injecting fixed (delta-prior) parameters and correcting the sky position for the
-    model's fixed reference time. Ported from ``GWSamplerMixin._post_process``; in the
-    longer run the reference-time correction becomes a deterministic reparametrization
-    factor in the chain (see the design doc, Q#4/Q#7).
+    GW façade over a ``ChainComposer``. Builds the chain for plain NPE, multi-iteration
+    GNPE, or single-step GNPE from model metadata, and exports the samples to a gw
+    ``Result``. All GW-specific processing (RA frame, fixed parameters, kernel correction)
+    is expressed as chain steps, so there is no post-processing step.
     """
 
     def __init__(
@@ -502,15 +513,10 @@ def __init__(
         context: GWSamplerContext,
         metadata: dict,
         inference_parameters: list[str],
-        kernel_factor: Optional[GNPEKernelFactor] = None,
     ):
         super().__init__(composer, context)
         self.metadata = metadata
         self.inference_parameters = inference_parameters
-        # Set for single-step GNPE. Supplies the delta_log_prob_target kernel correction
-        # (evaluated in post-processing, out of the proposal sum) that importance sampling
-        # applies to the joint target q(theta, theta_hat | d).
-        self.kernel_factor = kernel_factor
 
     @classmethod
     def from_model(
@@ -585,10 +591,10 @@ def from_singlestep_gnpe(
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
         """Build a single-step (density-preserving) time-GNPE sampler: a ``ChainComposer``
-        of ``[proxy_source, GNPEFlowFactor, RAReparam]``. ``proxy_source`` supplies the
-        detector-time proxies -- a ``FixedFactor`` for prior conditioning (BNS), or an
-        unconditional NDE for density recovery. Unlike multi-iteration GNPE the chain is
-        autoregressive, so log_prob is preserved; the ``GNPEKernelFactor`` supplies the
+        of ``[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]``.
+        ``proxy_source`` supplies the detector-time proxies -- a ``FixedFactor`` for prior
+        conditioning (BNS), or an unconditional NDE for density recovery. The chain is
+        autoregressive, so log_prob is preserved, and ``GNPEKernelCorrection`` emits the
         ``delta_log_prob_target`` correction that importance sampling adds to the target.
         """
         context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
@@ -601,38 +607,11 @@ def from_singlestep_gnpe(
         )
         kernel_factor = GNPEKernelFactor.from_model(main_model)
         steps = (
-            [proxy_source, flow_factor]
+            [proxy_source, flow_factor, GNPEKernelCorrection(kernel_factor)]
             + _ra_reparam_steps(inference_parameters)
             + _fixed_prior_steps(metadata, inference_parameters)
         )
-        composer = ChainComposer(steps)
-        return cls(
-            composer,
-            context,
-            metadata,
-            inference_parameters,
-            kernel_factor=kernel_factor,
-        )
-
-    def _add_kernel_correction(self, samples: dict):
-        """Single-step GNPE: add ``delta_log_prob_target = log p(theta_hat | theta)``,
-        evaluated at the proxies and the detector times recomputed from theta, then drop
-        those intermediate detector times. This is the joint-target correction importance
-        sampling applies; it is deliberately kept out of the proposal ``log_prob``."""
-        kf = self.kernel_factor
-        proxies = {p: samples[p] for p in kf.parameters}
-        times = {k: samples[k] for k in kf.gnpe_parameters}
-        correction = kf.log_prob(proxies, Conditioning(self.context, times))
-        samples["delta_log_prob_target"] = np.asarray(correction)
-        for k in kf.gnpe_parameters:
-            samples.pop(k, None)
-
-    def _post_process(self, samples: dict):
-        # Fixed (delta-prior) parameters are a FixedFactor in the chain and the RA frame
-        # rotation is a chain reparametrization; only the single-step-GNPE kernel correction
-        # remains here (moves to a chain step next).
-        if self.kernel_factor is not None:
-            self._add_kernel_correction(samples)
+        return cls(ChainComposer(steps), context, metadata, inference_parameters)
 
     def to_result(self):
         """Export to a gw ``Result`` (samples + raw event data + metadata), so the
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 7e01cfee0..942910f3d 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -18,6 +18,7 @@
     GibbsBlock,
     Reparametrization,
     Stage,
+    TargetCorrection,
     chunk_and_concat,
 )
 
@@ -306,3 +307,69 @@ def test_unconditioned_factor_as_root_draws_base_count():
     assert out["c"].shape == (6,) and out["d"].shape == (6,)
     assert torch.equal(out["c"], torch.full((6,), 2.0))
     assert torch.equal(out["d"], torch.full((6,), 2.0))  # d | c: sum(c) + 0 == c
+
+
+class _MockTargetCorrection(TargetCorrection):
+    """A mock kind-3 correction: emits ``corr = 2 * x``, contributes 0 to the proposal,
+    and consumes its input."""
+
+    def __init__(self, reads="x", emits="corr"):
+        self.conditioning = [reads]
+        self.parameters = [emits]
+        self.consumes = [reads]
+        self._reads, self._emits = reads, emits
+
+    def correction(self, given, context):
+        return {self._emits: given[self._reads] * 2}
+
+
+def test_target_correction_emits_side_channel_and_contributes_zero():
+    comp = ChainComposer([_ConstFactor("x"), _MockTargetCorrection(reads="x")])
+    out = comp.sample(5, context=None)
+    assert "x" not in out  # consumed
+    assert torch.equal(out["corr"], 2 * torch.arange(5, dtype=torch.float32))
+    # A correction adds 0 to the proposal density: log_prob = factor(0.5) + corr(0).
+    assert torch.allclose(out["log_prob"], torch.full((5,), 0.5))
+
+
+def test_target_correction_rejects_fan_out():
+    tc = _MockTargetCorrection()
+    with pytest.raises(ValueError, match="1:1"):
+        tc.sample_and_log_prob(
+            2, Conditioning(context=None, given={"x": torch.zeros(3)})
+        )
+
+
+class _SideChannelFactor(Factor):
+    """A factor emitting an extra side-channel column beyond ``parameters``, declared via
+    ``produces`` so the topological check sees it (like GNPEFlowFactor's detector times).
+    """
+
+    def __init__(self, name, side):
+        self.parameters = [name]
+        self.conditioning = []
+        self._name, self._side = name, side
+
+    @property
+    def produces(self):
+        return self.parameters + [self._side]
+
+    def sample_and_log_prob(self, n, cond):
+        block = {
+            self._name: torch.arange(n, dtype=torch.float32),
+            self._side: torch.zeros(n),
+        }
+        return block, torch.zeros(n)
+
+    def log_prob(self, theta_i, cond):
+        return torch.zeros(next(iter(theta_i.values())).shape[0])
+
+
+def test_side_channel_produces_satisfies_topological_check():
+    # A later step may condition on a side-channel column declared via `produces`; without
+    # `produces` the topological check would reject the chain.
+    comp = ChainComposer(
+        [_SideChannelFactor("a", side="s"), _ConstFactor("b", conditioning=["s"])]
+    )
+    out = comp.sample(4, context=None)
+    assert out["b"].shape == (4,) and "s" in out

From 73912279220bf7bbce691d42d5510593d818b580 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 15:00:34 +0200
Subject: [PATCH 16/65] Rename FixedFactor to DeltaFactor

The factor is a point mass q = delta(theta - c). DeltaFactor names the math --
matching bilby's DeltaFunction, the prior these pinned parameters come from --
and frees "fixed" for a future unconditional / data-independent flow factor.
Also renames the _fixed_prior_steps helper to _delta_prior_steps.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py | 22 +++++++++++-----------
 tests/core/test_factors.py    |  4 ++--
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index ef32b5450..0cfcaa72a 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -140,9 +140,9 @@ def likelihood(self):
         raise NotImplementedError("GWSamplerContext.likelihood")
 
 
-class FixedFactor(Factor):
-    """``q_i = delta(theta_i - c)``: pins parameters to fixed values, contributing 0 to the
-    proposal log-prob.
+class DeltaFactor(Factor):
+    """``q_i = delta(theta_i - c)``: a point mass pinning parameters to fixed values,
+    contributing 0 to the proposal log-prob.
 
     Used as the chain root for prior-conditioning or known proxies (single-step GNPE, where
     later factors condition on the pinned values), and as a non-root filler for delta-prior
@@ -161,7 +161,7 @@ def sample_and_log_prob(self, num_samples, cond):
         return samples, torch.zeros(num_samples)
 
     def log_prob(self, theta_i, cond):
-        raise NotImplementedError("FixedFactor.log_prob")
+        raise NotImplementedError("DeltaFactor.log_prob")
 
 
 # --- Stubs for later steps -------------------------------------------------------------
@@ -482,8 +482,8 @@ def _ra_reparam_steps(inference_parameters: list[str]) -> list:
     return [RAReparam()] if "ra" in inference_parameters else []
 
 
-def _fixed_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
-    """Delta-prior parameters the chain does not produce, as a single ``FixedFactor`` step
+def _delta_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
+    """Delta-prior parameters the chain does not produce, as a single ``DeltaFactor`` step
     (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0).
     """
     intrinsic_prior = metadata["dataset_settings"]["intrinsic_prior"]
@@ -496,7 +496,7 @@ def _fixed_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
         for k, p in prior.items()
         if isinstance(p, DeltaFunction) and k not in inference_parameters
     }
-    return [FixedFactor(fixed)] if fixed else []
+    return [DeltaFactor(fixed)] if fixed else []
 
 
 class GWComposedSampler(ComposedSampler):
@@ -536,7 +536,7 @@ def from_model(
         steps = (
             [factor]
             + _ra_reparam_steps(inference_parameters)
-            + _fixed_prior_steps(metadata, inference_parameters)
+            + _delta_prior_steps(metadata, inference_parameters)
         )
         return cls(
             composer=ChainComposer(steps),
@@ -573,7 +573,7 @@ def from_gnpe_models(
         steps = (
             [gibbs]
             + _ra_reparam_steps(inference_parameters)
-            + _fixed_prior_steps(metadata, inference_parameters)
+            + _delta_prior_steps(metadata, inference_parameters)
         )
         return cls(
             ChainComposer(steps),
@@ -592,7 +592,7 @@ def from_singlestep_gnpe(
     ) -> "GWComposedSampler":
         """Build a single-step (density-preserving) time-GNPE sampler: a ``ChainComposer``
         of ``[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]``.
-        ``proxy_source`` supplies the detector-time proxies -- a ``FixedFactor`` for prior
+        ``proxy_source`` supplies the detector-time proxies -- a ``DeltaFactor`` for prior
         conditioning (BNS), or an unconditional NDE for density recovery. The chain is
         autoregressive, so log_prob is preserved, and ``GNPEKernelCorrection`` emits the
         ``delta_log_prob_target`` correction that importance sampling adds to the target.
@@ -609,7 +609,7 @@ def from_singlestep_gnpe(
         steps = (
             [proxy_source, flow_factor, GNPEKernelCorrection(kernel_factor)]
             + _ra_reparam_steps(inference_parameters)
-            + _fixed_prior_steps(metadata, inference_parameters)
+            + _delta_prior_steps(metadata, inference_parameters)
         )
         return cls(ChainComposer(steps), context, metadata, inference_parameters)
 
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 942910f3d..74f22ac11 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -256,7 +256,7 @@ def test_reparam_forward_inverse_round_trip():
 
 
 class _ConstFillFactor(Factor):
-    """An unconditioned factor emitting a constant (like a ``FixedFactor``). As the root it
+    """An unconditioned factor emitting a constant (like a ``DeltaFactor``). As the root it
     draws the base count; as a non-root step it fills one value per current row."""
 
     def __init__(self, name, value=3.0):
@@ -298,7 +298,7 @@ def test_unconditioned_filler_fills_after_fan_out():
 
 
 def test_unconditioned_factor_as_root_draws_base_count():
-    # As the chain root (prior-conditioning / FixedFactor proxy), an unconditioned factor
+    # As the chain root (prior-conditioning / DeltaFactor proxy), an unconditioned factor
     # draws the base count; a downstream factor conditions on the pinned values.
     comp = ChainComposer(
         [_ConstFillFactor("c", 2.0), _ConstFactor("d", conditioning=["c"])]

From b36c5c6f345b4313504c0cca56531371a188df9c Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 15:07:09 +0200
Subject: [PATCH 17/65] Remove the vestigial _post_process hook;
 ComposedSampler is a generic runner

With all GW-specific processing expressed as chain steps (RA reparametrization,
delta factors, kernel correction), the ComposedSampler._post_process hook is a
dead no-op. Remove it and its run_sampler call: ComposedSampler is now a
domain-agnostic runner, and GWComposedSampler is a GW builder (the from_*
constructors) and exporter (to_result) on top of it.

Parity: NPE, multi-iteration GNPE, single-step, and packaged all bit-exact;
tests/core/test_factors.py + tests/gw/test_ra_reparam.py green (24 tests).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 16 ++++++----------
 dingo/gw/inference/factors.py |  9 +++++----
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 5443352bc..3926d170a 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -565,8 +565,9 @@ def _run_once(
 
 class ComposedSampler:
     """
-    Façade over a ``ChainComposer`` and a ``SamplerContext``. Runs the composer, applies
-    domain-specific post-processing, and returns the samples as a DataFrame.
+    Generic runner over a ``ChainComposer`` and a ``SamplerContext``: draws samples and
+    returns them as a DataFrame. Domain-specific processing lives in the chain's steps, so
+    the runner is domain-agnostic.
     """
 
     def __init__(self, composer: ChainComposer, context: SamplerContext):
@@ -574,18 +575,13 @@ def __init__(self, composer: ChainComposer, context: SamplerContext):
         self.context = context
         self.samples: Optional[pd.DataFrame] = None
 
-    def _post_process(self, samples: dict):
-        """Hook for domain-specific post-processing; no-op by default."""
-        pass
-
     def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
-        """Draw ``num_samples`` samples (chunked by ``batch_size``), post-process, and
-        return them as a DataFrame. An all-density chain includes ``log_prob``; a chain
-        with a ``GibbsBlock`` step does not."""
+        """Draw ``num_samples`` samples (chunked by ``batch_size``) and return them as a
+        DataFrame. An all-density chain includes ``log_prob``; a chain with a ``GibbsBlock``
+        step does not."""
         merged = self.composer.sample(num_samples, self.context, batch_size)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
-        self._post_process(merged)
         self.samples = pd.DataFrame(merged)
         return self.samples
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 0cfcaa72a..608da45f6 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -501,10 +501,11 @@ def _delta_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
 
 class GWComposedSampler(ComposedSampler):
     """
-    GW façade over a ``ChainComposer``. Builds the chain for plain NPE, multi-iteration
-    GNPE, or single-step GNPE from model metadata, and exports the samples to a gw
-    ``Result``. All GW-specific processing (RA frame, fixed parameters, kernel correction)
-    is expressed as chain steps, so there is no post-processing step.
+    GW builder and exporter over the generic ``ComposedSampler`` runner. The ``from_*``
+    constructors assemble the chain for plain NPE, multi-iteration GNPE, or single-step
+    GNPE from model metadata; ``to_result`` exports the samples to a gw ``Result``. All
+    GW-specific processing (RA frame, fixed parameters, kernel correction) is expressed as
+    chain steps, so there is no post-processing.
     """
 
     def __init__(

From 17f49a6096d05625dd30f8dd23bb002c96457290 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Wed, 1 Jul 2026 15:35:23 +0200
Subject: [PATCH 18/65] Drop the Conditioning wrapper for explicit (context,
 given) args

Conditioning was a 2-field dataclass (context, given) with no behaviour. Steps now
take sample_and_log_prob(n, context, given) / log_prob(theta_i, context, given)
directly, and the composer passes the two. Any parameter-dependent transform lives
inside the factor (as GNPE shows), so the conditioning never needed to carry more
than these two fields.

Parity: NPE, multi-iteration GNPE, single-step, and packaged all bit-exact;
tests/core/test_factors.py + tests/gw/test_ra_reparam.py green (23 tests).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 95 +++++++++++++++++++----------------
 dingo/gw/inference/factors.py | 31 ++++++------
 tests/core/test_factors.py    | 33 ++++--------
 3 files changed, 79 insertions(+), 80 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 3926d170a..038ea33f7 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -17,7 +17,7 @@
 
 import math
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Callable, Optional, Protocol, Union, runtime_checkable
 
 import pandas as pd
@@ -83,17 +83,6 @@ def likelihood(self):
         ...
 
 
-@dataclass
-class Conditioning:
-    """
-    The conditioning passed to a factor: the shared ``context``, and ``given``, the
-    physical values of the earlier-block parameters the factor conditions on.
-    """
-
-    context: SamplerContext
-    given: dict[str, torch.Tensor] = field(default_factory=dict)
-
-
 def _cat_dict(
     batches: list[dict[str, torch.Tensor]],
 ) -> dict[str, torch.Tensor]:
@@ -131,7 +120,7 @@ class Factor(ABC):
 
     A call draws ``num_samples`` samples per conditioning row and returns ``n_rows *
     num_samples`` rows in row-major order, where ``n_rows`` is the number of rows in
-    ``cond.given`` (1 if unconditioned). This mirrors the network's
+    ``given`` (1 if unconditioned). This mirrors the network's
     ``sample_and_log_prob(*context, num_samples=n) -> (n_rows, n, dim)``.
 
     Attributes
@@ -147,7 +136,10 @@ class Factor(ABC):
 
     @abstractmethod
     def sample_and_log_prob(
-        self, num_samples: int, cond: Conditioning
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
         """Draw ``num_samples`` samples per conditioning row; return ``(samples,
         log_prob)`` in physical space. The samples dict may include named columns beyond
@@ -155,7 +147,10 @@ def sample_and_log_prob(
 
     @abstractmethod
     def log_prob(
-        self, theta_i: dict[str, torch.Tensor], cond: Conditioning
+        self,
+        theta_i: dict[str, torch.Tensor],
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
         """Evaluate ``log q_i`` at given physical ``theta_i`` (one conditioning row per
         row)."""
@@ -178,7 +173,7 @@ class FlowFactor(Factor):
     its interface is in physical parameter space.
 
     An unconditioned model draws from the shared data context; a model with
-    ``context_parameters`` conditions on the values in ``cond.given``.
+    ``context_parameters`` conditions on the values in ``given``.
     """
 
     def __init__(
@@ -219,11 +214,11 @@ def from_model(
             aliases=aliases,
         )
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, context, given=None):
         """Draw ``num_samples`` samples per conditioning row. Unconditioned: ``num_samples``
         draws from the shared data context. Conditioned: ``num_samples`` per context row,
         returning ``n_rows * num_samples`` rows in row-major order."""
-        data = cond.context.prepared_data()
+        data = context.prepared_data()
         self.model.network.eval()
         if not self.context_parameters:
             with torch.no_grad():
@@ -237,7 +232,7 @@ def sample_and_log_prob(self, num_samples, cond):
             # Standardize the (physical) conditioning the network was trained on, giving
             # B = N context rows; expand the shared data to match. The network draws
             # num_samples per row -> (N, num_samples, dim).
-            ctx = self.standardization.standardize(cond.given, self.context_parameters)
+            ctx = self.standardization.standardize(given, self.context_parameters)
             n_rows = ctx.shape[0]
             data = data.expand(n_rows, *data.shape)
             with torch.no_grad():
@@ -252,20 +247,20 @@ def sample_and_log_prob(self, num_samples, cond):
         theta = {self.aliases.get(k, k): v for k, v in theta.items()}
         return theta, log_prob
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         # theta_i uses exposed (aliased) names; map back to the network's trained names.
         theta_net = {
             net: theta_i[self.aliases.get(net, net)] for net in self._net_parameters
         }
         num_samples = next(iter(theta_net.values())).shape[0]
         z = self.standardization.standardize(theta_net, self._net_parameters)
-        data = cond.context.prepared_data()
+        data = context.prepared_data()
         data = data.expand(num_samples, *data.shape)
         net_context: tuple[torch.Tensor, ...]
         if not self.context_parameters:
             net_context = (data,)
         else:
-            ctx = self.standardization.standardize(cond.given, self.context_parameters)
+            ctx = self.standardization.standardize(given, self.context_parameters)
             net_context = (data, ctx)
         self.model.network.eval()
         with torch.no_grad():
@@ -309,22 +304,28 @@ def log_det(
         return torch.zeros(n)
 
     def sample_and_log_prob(
-        self, num_samples: int, cond: "Conditioning"
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
         """Apply ``forward`` to the conditioning; contribute ``-log|det J|``. ``num_samples``
         must be 1 (a reparametrization is 1:1)."""
         if num_samples != 1:
             raise ValueError("A reparametrization is 1:1; use fan_out=1.")
-        out = self.forward(cond.given, cond.context)
-        return out, -self.log_det(cond.given, cond.context)
+        out = self.forward(given, context)
+        return out, -self.log_det(given, context)
 
     @property
     def consumes(self) -> list[str]:
-        """Conditioning the bijection replaces with its outputs (dropped after the step)."""
+        """Inputs the bijection replaces with its outputs (dropped after the step)."""
         return [c for c in self.conditioning if c not in self.parameters]
 
     def log_prob(
-        self, theta_i: dict[str, torch.Tensor], cond: "Conditioning"
+        self,
+        theta_i: dict[str, torch.Tensor],
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
         raise NotImplementedError("Reparametrization.log_prob")
 
@@ -350,16 +351,22 @@ def correction(
         """The side-channel column(s), one value per conditioning row."""
 
     def sample_and_log_prob(
-        self, num_samples: int, cond: "Conditioning"
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
         if num_samples != 1:
             raise ValueError("A target correction is 1:1; use fan_out=1.")
-        out = self.correction(cond.given, cond.context)
+        out = self.correction(given, context)
         n = next(iter(out.values())).shape[0]
         return out, torch.zeros(n)
 
     def log_prob(
-        self, theta_i: dict[str, torch.Tensor], cond: "Conditioning"
+        self,
+        theta_i: dict[str, torch.Tensor],
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
         raise NotImplementedError("TargetCorrection.log_prob")
 
@@ -378,7 +385,10 @@ class Step(Protocol):
     conditioning: list[str]
 
     def sample_and_log_prob(
-        self, num_samples: int, cond: "Conditioning"
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]: ...
 
 
@@ -466,8 +476,8 @@ def _run_chain_once(
                 # Unconditioned non-root step (e.g. a fixed/delta factor): draw one value
                 # per current row -- it fills the batch rather than fanning out.
                 n = len(next(iter(samples.values())))
-            cond = Conditioning(context, {k: samples[k] for k in step.conditioning})
-            block, lp = step.sample_and_log_prob(n, cond)
+            given = {k: samples[k] for k in step.conditioning}
+            block, lp = step.sample_and_log_prob(n, context, given)
             # A conditioned fan-out (fan_out > 1) expands the batch: repeat each carried row
             # to align with the step's fan_out sub-rows (the block is flattened row-major).
             # 1:1 stages and unconditioned fillers leave the batch untouched.
@@ -495,9 +505,9 @@ def log_prob(
         chain (every step a ``Factor``)."""
         total: torch.Tensor | float = 0.0
         for step in self.steps:
-            cond = Conditioning(context, {k: samples[k] for k in step.conditioning})
+            given = {k: samples[k] for k in step.conditioning}
             theta_i = {k: samples[k] for k in step.parameters}
-            total = total + step.log_prob(theta_i, cond)
+            total = total + step.log_prob(theta_i, context, given)
         return total
 
     def sample(
@@ -540,25 +550,26 @@ def __init__(self, init_factor: Factor, factors: list[Factor], num_iterations: i
         self.conditioning: list[str] = []
 
     def sample_and_log_prob(
-        self, num_samples: int, cond: Conditioning
+        self,
+        num_samples: int,
+        context: SamplerContext,
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], None]:
         """Run the Gibbs loop for ``num_samples`` walkers; return ``(samples, None)``.
         ``num_samples`` is the walker (root) count -- Gibbs does not fan out."""
-        return self._run_once(num_samples, cond.context), None
+        return self._run_once(num_samples, context), None
 
     def _run_once(
         self, num_samples: int, context: SamplerContext
     ) -> dict[str, torch.Tensor]:
         # Seed the chain (e.g. an init network's detector times); the walkers are the rows.
-        seed, _ = self.init_factor.sample_and_log_prob(
-            num_samples, Conditioning(context)
-        )
+        seed, _ = self.init_factor.sample_and_log_prob(num_samples, context)
         state = dict(seed)
         for _ in range(self.num_iterations):
             for factor in self.factors:
-                c = Conditioning(context, {k: state[k] for k in factor.conditioning})
+                given = {k: state[k] for k in factor.conditioning}
                 # One sample per walker (Gibbs is 1:1); walkers are the conditioning rows.
-                block, _ = factor.sample_and_log_prob(1, c)
+                block, _ = factor.sample_and_log_prob(1, context, given)
                 state.update(block)
         return {p: state[p] for p in self.parameters}
 
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 608da45f6..624eaec05 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -24,7 +24,6 @@
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
-    Conditioning,
     Factor,
     FlowFactor,
     GibbsBlock,
@@ -154,13 +153,13 @@ def __init__(self, values: dict[str, float]):
         self.parameters = list(values)
         self.conditioning = []
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, context, given=None):
         samples = {
             p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
         }
         return samples, torch.zeros(num_samples)
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         raise NotImplementedError("DeltaFactor.log_prob")
 
 
@@ -175,10 +174,10 @@ def __init__(self):
         self.parameters = ["phase"]
         self.conditioning = []  # conditions on all preceding params via the likelihood
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, context, given=None):
         raise NotImplementedError("SyntheticPhaseFactor")
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         raise NotImplementedError("SyntheticPhaseFactor")
 
 
@@ -284,21 +283,21 @@ def from_model(cls, model: BasePosteriorModel) -> "GNPEKernelFactor":
         _, _, gnpe_parameters, _, _, gnpe_transform = _build_gnpe_transforms(model)
         return cls(gnpe_transform, gnpe_parameters)
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, context, given=None):
         """Blur the conditioning detector times into proxies; ``num_samples`` must be 1
         (GNPE is 1:1). Returns the proxies and their kernel log-prob."""
         if num_samples != 1:
             raise ValueError("GNPE proxy is 1:1; use fan_out=1.")
-        times = {k: cond.given[k] for k in self.gnpe_parameters}
+        times = {k: given[k] for k in self.gnpe_parameters}
         proxies = self.gnpe.sample_proxies(times)
-        return proxies, self.log_prob(proxies, cond)
+        return proxies, self.log_prob(proxies, context, given)
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         """``log p(theta_hat | theta)`` from the kernel, at the proxies (``theta_i``) and
-        the detector times (``cond.given``)."""
+        the detector times (``given``)."""
         diffs = {}
         for k in self.kernel.keys():
-            diff = cond.given[k] - theta_i[f"{k}_proxy"]
+            diff = given[k] - theta_i[f"{k}_proxy"]
             if torch.is_tensor(diff):
                 diff = diff.detach().cpu().numpy()
             diffs[k] = np.asarray(diff)
@@ -359,15 +358,15 @@ def from_model(
             model, pre, post, gnpe_parameters, inference_parameters, aliases=aliases
         )
 
-    def sample_and_log_prob(self, num_samples, cond):
+    def sample_and_log_prob(self, num_samples, context, given=None):
         """Sample one parameter set per proxy row; ``num_samples`` must be 1 (GNPE is 1:1).
         Returns theta plus the recomputed detector times, and the network log-prob."""
         if num_samples != 1:
             raise ValueError("GNPE is 1:1; draw one sample per proxy (fan_out=1).")
-        proxies = {p: cond.given[p] for p in self.proxy_parameters}
+        proxies = {p: given[p] for p in self.proxy_parameters}
         n_rows = next(iter(proxies.values())).shape[0]
         x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = cond.context.prepared_data().clone()
+        d = context.prepared_data().clone()
         x["data"] = d.expand(n_rows, *d.shape)
         x = self.transform_pre(x)
         self.model.network.eval()
@@ -392,7 +391,7 @@ def sample_and_log_prob(self, num_samples, cond):
             params[k] = x["extrinsic_parameters"][k]
         return params, x["log_prob"]
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         raise NotImplementedError("GNPEFlowFactor.log_prob")
 
 
@@ -467,7 +466,7 @@ def __init__(self, kernel_factor: GNPEKernelFactor):
     def correction(self, given, context):
         proxies = {p: given[p] for p in self.kernel_factor.parameters}
         times = {k: given[k] for k in self.kernel_factor.gnpe_parameters}
-        correction = self.kernel_factor.log_prob(proxies, Conditioning(context, times))
+        correction = self.kernel_factor.log_prob(proxies, context, times)
         return {"delta_log_prob_target": correction}
 
 
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 74f22ac11..8cfe748c3 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -13,7 +13,6 @@
 
 from dingo.core.factors import (
     ChainComposer,
-    Conditioning,
     Factor,
     GibbsBlock,
     Reparametrization,
@@ -34,9 +33,9 @@ def __init__(self, name, conditioning=()):
         self.conditioning = list(conditioning)
         self._name = name
 
-    def sample_and_log_prob(self, n, cond):
+    def sample_and_log_prob(self, n, context, given=None):
         if self.conditioning:
-            base = sum(cond.given[k] for k in self.conditioning)  # (N,)
+            base = sum(given[k] for k in self.conditioning)  # (N,)
             n_rows = base.shape[0]
             # Integer-valued within-row offset: exact in float32 (no cancellation when
             # the test recovers it as b - a), and distinct per fan-out draw.
@@ -48,7 +47,7 @@ def sample_and_log_prob(self, n, cond):
             lp = torch.full((n,), 0.5)
         return {self._name: vals}, lp
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         return torch.zeros(next(iter(theta_i.values())).shape[0])
 
 
@@ -145,12 +144,6 @@ def test_chunk_and_concat_allows_none_log_prob():
     assert samples["x"].shape == (6,)
 
 
-def test_conditioning_dataclass_defaults():
-    cond = Conditioning(context="ctx")
-    assert cond.context == "ctx"
-    assert cond.given == {}
-
-
 class _NoDensityStep:
     """A density-free step (like ``GibbsBlock``): emits a block but returns ``None`` for the
     log-prob. Honors the per-row fan-out contract so it composes like any step."""
@@ -160,9 +153,9 @@ def __init__(self, name, conditioning=()):
         self.conditioning = list(conditioning)
         self._name = name
 
-    def sample_and_log_prob(self, n, cond):
+    def sample_and_log_prob(self, n, context, given=None):
         if self.conditioning:
-            base = sum(cond.given[k] for k in self.conditioning)
+            base = sum(given[k] for k in self.conditioning)
             within = torch.arange(n, dtype=base.dtype)
             vals = (base.unsqueeze(1) + within).reshape(-1)
         else:
@@ -242,9 +235,7 @@ def test_reparam_step_consumes_input_and_contributes_neg_logdet():
 def test_reparam_rejects_fan_out():
     rp = _MockReparam()
     with pytest.raises(ValueError, match="1:1"):
-        rp.sample_and_log_prob(
-            2, Conditioning(context=None, given={"u": torch.zeros(3)})
-        )
+        rp.sample_and_log_prob(2, None, {"u": torch.zeros(3)})
 
 
 def test_reparam_forward_inverse_round_trip():
@@ -264,10 +255,10 @@ def __init__(self, name, value=3.0):
         self.conditioning = []
         self._name, self._v = name, value
 
-    def sample_and_log_prob(self, n, cond):
+    def sample_and_log_prob(self, n, context, given=None):
         return {self._name: torch.full((n,), self._v)}, torch.zeros(n)
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         return torch.zeros(next(iter(theta_i.values())).shape[0])
 
 
@@ -335,9 +326,7 @@ def test_target_correction_emits_side_channel_and_contributes_zero():
 def test_target_correction_rejects_fan_out():
     tc = _MockTargetCorrection()
     with pytest.raises(ValueError, match="1:1"):
-        tc.sample_and_log_prob(
-            2, Conditioning(context=None, given={"x": torch.zeros(3)})
-        )
+        tc.sample_and_log_prob(2, None, {"x": torch.zeros(3)})
 
 
 class _SideChannelFactor(Factor):
@@ -354,14 +343,14 @@ def __init__(self, name, side):
     def produces(self):
         return self.parameters + [self._side]
 
-    def sample_and_log_prob(self, n, cond):
+    def sample_and_log_prob(self, n, context, given=None):
         block = {
             self._name: torch.arange(n, dtype=torch.float32),
             self._side: torch.zeros(n),
         }
         return block, torch.zeros(n)
 
-    def log_prob(self, theta_i, cond):
+    def log_prob(self, theta_i, context, given=None):
         return torch.zeros(next(iter(theta_i.values())).shape[0])
 
 

From 4420dcb110ecb24521b7b129f103704bec64f1b1 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Thu, 2 Jul 2026 13:58:54 +0200
Subject: [PATCH 19/65] Add likelihood-in-context and the synthetic-phase
 factor

Give GWSamplerContext a likelihood() that builds the exact GW likelihood on the
event's raw data (the network's decimated/whitened view stays in prepared_data()),
and add SyntheticPhaseFactor, a terminal chain factor that reconstructs the
coalescence phase for a phase-marginalized network from that likelihood on a phase
grid, returning phase and its proposal log-prob.

- GWSamplerContext gains base_metadata and likelihood(...), porting
  Result._build_likelihood. The likelihood reference time is the event time (the
  training-frame RA correction is already applied to the samples).
- SyntheticPhaseFactor supports the (2, 2) approximation and the exact
  mode-decomposed grid, with a uniform floor for mass coverage.

Verified bit-exact against Result._build_likelihood and
Result.sample_synthetic_phase via model-based harnesses; 7 mock-based CI unit
tests added in tests/gw/test_synthetic_phase.py.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py    | 255 +++++++++++++++++++++++++++++--
 tests/gw/test_synthetic_phase.py | 129 ++++++++++++++++
 2 files changed, 372 insertions(+), 12 deletions(-)
 create mode 100644 tests/gw/test_synthetic_phase.py

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 624eaec05..4c930ea53 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -12,15 +12,22 @@
 from __future__ import annotations
 
 import copy
+import logging
+import time
 from typing import Optional
 
 import numpy as np
+import pandas as pd
 import torch
 from astropy.time import Time
-from bilby.core.prior import DeltaFunction, PriorDict
+from bilby.core.prior import DeltaFunction, PriorDict, Uniform
 from bilby.gw.detector import InterferometerList
 from torchvision.transforms import Compose
 
+from dingo.core.density import (
+    interpolated_log_prob_multi,
+    interpolated_sample_and_log_prob_multi,
+)
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
@@ -31,10 +38,12 @@
     TargetCorrection,
     _base_model_metadata,
 )
+from dingo.core.multiprocessing import apply_func_with_multiprocessing
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
 from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
 from dingo.gw.gwutils import get_extrinsic_prior_dict
+from dingo.gw.likelihood import StationaryGaussianGWLikelihood
 from dingo.gw.prior import build_prior_with_defaults
 from dingo.gw.transforms import (
     CopyToExtrinsicParameters,
@@ -50,6 +59,8 @@
     WhitenAndScaleStrain,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class GWSamplerContext:
     """
@@ -58,8 +69,9 @@ class GWSamplerContext:
     stages.
 
     This implements the ``dingo.core.factors.SamplerContext`` protocol. It owns the
-    one-time data preprocessing (``prepared_data``) and -- once wired in a later step --
-    the likelihood used by likelihood-based factors and importance sampling.
+    one-time data preprocessing (``prepared_data``) and builds the exact likelihood
+    (``likelihood``) used by likelihood-based factors (synthetic phase) and importance
+    sampling.
 
     Event metadata lives here (not on individual factors): it is a property of the data,
     and it drives frequency cropping, the t_ref/RA correction, and the likelihood.
@@ -73,11 +85,14 @@ def __init__(
         data_prep: Compose,
         raw_context: dict,
         event_metadata: Optional[dict] = None,
+        base_metadata: Optional[dict] = None,
     ):
         self.domain = domain
         self.detectors = detectors
         self.t_ref = t_ref
         self._data_prep = data_prep
+        # The base (parameter/waveform/domain) metadata, used to build the likelihood.
+        self.base_metadata = base_metadata
         # The raw event data `d` (strain + ASDs per detector), i.e. EventDataset.data --
         # the same object set as GWSampler.context today. Consumed lazily by
         # prepared_data(); also retained for the likelihood (synthetic-phase / IS factors)
@@ -127,6 +142,7 @@ def from_model(
             data_prep=Compose(transforms),
             raw_context=raw_context,
             event_metadata=event_metadata,
+            base_metadata=meta,
         )
 
     def prepared_data(self) -> torch.Tensor:
@@ -135,8 +151,125 @@ def prepared_data(self) -> torch.Tensor:
             self._prepared = self._data_prep(self.raw_context)
         return self._prepared
 
-    def likelihood(self):
-        raise NotImplementedError("GWSamplerContext.likelihood")
+    def likelihood(
+        self,
+        time_marginalization_kwargs: Optional[dict] = None,
+        phase_marginalization_kwargs: Optional[dict] = None,
+        calibration_marginalization_kwargs: Optional[dict] = None,
+        phase_grid: Optional[np.ndarray] = None,
+        use_base_domain: bool = False,
+    ) -> StationaryGaussianGWLikelihood:
+        """
+        Build the exact GW likelihood on this event's data, in physical parameter space.
+
+        The network's standardized, decimated view of the data is `prepared_data()`; the
+        likelihood instead takes the raw event data (`raw_context`) and builds its own
+        representation -- decimating to the multibanded domain unless `use_base_domain`,
+        and masking the ASDs to the event's frequency range. Its reference time is the
+        event time (the training-frame right-ascension correction is already applied to the
+        samples), or the training reference time when no event time is set.
+
+        Importance-sampling domain rebuilds (an updated `T` / `delta_f` or frequency range)
+        are not wired here yet: the likelihood uses this context's domain.
+
+        Parameters
+        ----------
+        time_marginalization_kwargs : dict, optional
+            Analytically marginalize over `geocent_time`; `t_lower` / `t_upper` are filled
+            from the network's (uniform) time prior. Requires a time-marginalized network.
+        phase_marginalization_kwargs : dict, optional
+            Analytically marginalize over `phase`. Requires a uniform [0, 2 pi) phase prior.
+        calibration_marginalization_kwargs : dict, optional
+            Marginalize over detector calibration uncertainty.
+        phase_grid : np.ndarray, optional
+            Phase grid for the exact phase-marginalized / synthetic-phase likelihood.
+        use_base_domain : bool, default False
+            For a multibanded domain, evaluate on the base (undecimated) frequency domain
+            rather than the decimated one.
+
+        Returns
+        -------
+        StationaryGaussianGWLikelihood
+        """
+        # Marginalizing over a parameter needs the (uniform) prior the network
+        # marginalized over; use it to parameterize the requested marginalization.
+        if time_marginalization_kwargs is not None:
+            time_prior = self._marginalized_prior("geocent_time")
+            if time_prior is None:
+                raise NotImplementedError(
+                    "Time marginalization is not compatible with a non-marginalized "
+                    "network."
+                )
+            if type(time_prior) != Uniform:
+                raise NotImplementedError(
+                    "Only uniform time prior is supported for time marginalization."
+                )
+            time_marginalization_kwargs = {
+                **time_marginalization_kwargs,
+                "t_lower": time_prior.minimum,
+                "t_upper": time_prior.maximum,
+            }
+        if phase_marginalization_kwargs is not None:
+            phase_prior = self._marginalized_prior("phase")
+            if not (
+                isinstance(phase_prior, Uniform)
+                and (phase_prior._minimum, phase_prior._maximum) == (0, 2 * np.pi)
+            ):
+                raise ValueError(
+                    f"Phase prior should be uniform [0, 2pi) for phase marginalization, "
+                    f"but is {phase_prior}."
+                )
+
+        dataset_settings = self.base_metadata["dataset_settings"]
+        # WaveformGenerator domain -- generally the dataset domain (it may be wider than
+        # the data domain for generation). delta_f / T importance-sampling updates: TODO.
+        wfg_domain = build_domain(dataset_settings["domain"])
+
+        # Likelihood reference time: the event time (the training-frame RA correction has
+        # already been applied to the samples), falling back to the training reference.
+        if self.event_metadata is not None and "time_event" in self.event_metadata:
+            t_ref = self.event_metadata["time_event"]
+        else:
+            t_ref = self.t_ref
+
+        return StationaryGaussianGWLikelihood(
+            wfg_kwargs=dataset_settings["waveform_generator"],
+            wfg_domain=wfg_domain,
+            data_domain=self.domain,
+            event_data=self.raw_context,
+            t_ref=t_ref,
+            time_marginalization_kwargs=time_marginalization_kwargs,
+            phase_marginalization_kwargs=phase_marginalization_kwargs,
+            calibration_marginalization_kwargs=calibration_marginalization_kwargs,
+            phase_grid=phase_grid,
+            use_base_domain=use_base_domain,
+            frequency_update=dict(
+                minimum_frequency=self._frequency(
+                    "minimum_frequency", self.domain.f_min
+                ),
+                maximum_frequency=self._frequency(
+                    "maximum_frequency", self.domain.f_max
+                ),
+            ),
+        )
+
+    def _frequency(self, key: str, default: float):
+        """The event's frequency-range override for `key` (min/max), else `default`."""
+        if self.event_metadata is None:
+            return default
+        return self.event_metadata.get(key, default)
+
+    def _marginalized_prior(self, name: str):
+        """The prior over `name` if the network marginalized it (i.e. `name` is not an
+        inference parameter), else `None` -- used to parameterize likelihood
+        marginalization over time / phase."""
+        data_settings = self.base_metadata["train_settings"]["data"]
+        if name in data_settings["inference_parameters"]:
+            return None
+        intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
+        extrinsic_prior = get_extrinsic_prior_dict(data_settings["extrinsic_prior"])
+        prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
+        return prior.get(name)
 
 
 class DeltaFactor(Factor):
@@ -163,22 +296,120 @@ def log_prob(self, theta_i, context, given=None):
         raise NotImplementedError("DeltaFactor.log_prob")
 
 
-# --- Stubs for later steps -------------------------------------------------------------
+def _to_numpy(v) -> np.ndarray:
+    """Detach a torch tensor (or coerce anything) to a numpy array."""
+    if torch.is_tensor(v):
+        return v.detach().cpu().numpy()
+    return np.asarray(v)
 
 
 class SyntheticPhaseFactor(Factor):
-    """``q(phase | theta_rest, d)`` from the likelihood on a phase grid (a non-network
-    factor, the final step of the chain). Not yet implemented."""
+    """
+    Reconstruct the coalescence phase for a phase-marginalized network:
+    `q(phase | theta_rest, d)` from the likelihood on a phase grid. The terminal factor of
+    the chain.
+
+    For each incoming `theta_rest` it builds the phase-full likelihood
+    (`context.likelihood()`) and evaluates `log L` on a grid over `[0, 2 pi)`, exploiting
+    that the waveform modes computed once at `phase = 0` each transform as `exp(-i m phase)`
+    -- so the whole grid follows from a single waveform evaluation. The grid is
+    exponentiated into a conditional phase distribution, a uniform floor (weight
+    `uniform_weight`) is added to keep it mass-covering, and one phase is drawn per sample
+    from the interpolated distribution. The returned proposal log-prob
+    `log q(phase | theta_rest, d)` joins the chain's proposal density; importance sampling
+    then targets the phase-full posterior (with the phase prior re-added).
+
+    Two grid modes: `approximation_22_mode=True` assumes a (2, 2)-dominated signal (the
+    whole waveform transforms as `exp(2i phase)`), giving `log L` from the complex overlap
+    `Re[(d | h(phase=0)) exp(2i phase)]`; `False` (the production default) sums the modes
+    exactly and requires the waveform generator's `spin_conversion_phase = 0`.
+    """
 
-    def __init__(self):
+    def __init__(
+        self,
+        conditioning: list[str],
+        n_grid: int = 5001,
+        approximation_22_mode: bool = False,
+        uniform_weight: float = 0.01,
+        num_processes: int = 1,
+    ):
+        """
+        Parameters
+        ----------
+        conditioning : list[str]
+            The physical parameters the likelihood needs to generate the waveform
+            (everything the chain has produced except `phase`).
+        n_grid : int, default 5001
+            Number of phase grid points on `[0, 2 pi)`.
+        approximation_22_mode : bool, default False
+            Use the (2, 2)-mode approximation instead of the exact mode sum.
+        uniform_weight : float, default 0.01
+            Weight of the uniform floor added to the phase distribution for mass coverage.
+        num_processes : int, default 1
+            Parallel processes for the per-sample likelihood evaluation and phase sampling.
+        """
         self.parameters = ["phase"]
-        self.conditioning = []  # conditions on all preceding params via the likelihood
+        self.conditioning = list(conditioning)
+        self.n_grid = n_grid
+        self.approximation_22_mode = approximation_22_mode
+        self.uniform_weight = uniform_weight
+        self.num_processes = num_processes
 
     def sample_and_log_prob(self, num_samples, context, given=None):
-        raise NotImplementedError("SyntheticPhaseFactor")
+        """Draw one phase per `theta_rest` row (`num_samples` must be 1); return the phases
+        and their proposal log-prob `log q(phase | theta_rest, d)`."""
+        if num_samples != 1:
+            raise ValueError(
+                "Synthetic phase is 1:1; draw one phase per sample (fan_out=1)."
+            )
+        n = len(next(iter(given.values())))
+        logger.info(f"Estimating synthetic phase for {n} samples.")
+        t0 = time.time()
+        phases, phase_posterior = self._phase_profile(given, context)
+        new_phase, log_prob = interpolated_sample_and_log_prob_multi(
+            phases, phase_posterior, self.num_processes
+        )
+        logger.info(f"Done. This took {time.time() - t0:.2f} s.")
+        return {"phase": torch.as_tensor(new_phase)}, torch.as_tensor(log_prob)
 
     def log_prob(self, theta_i, context, given=None):
-        raise NotImplementedError("SyntheticPhaseFactor")
+        """Evaluate `log q(phase | theta_rest, d)` at the given phases (re-plug / IS)."""
+        phases, phase_posterior = self._phase_profile(given, context)
+        log_prob = interpolated_log_prob_multi(
+            phases, phase_posterior, _to_numpy(theta_i["phase"]), self.num_processes
+        )
+        return torch.as_tensor(log_prob)
+
+    def _phase_profile(self, given, context):
+        """The phase grid and the mass-covered (un-normalized) phase distribution, one row
+        per sample: evaluate `log L` on the grid, exponentiate (shifted by the per-row
+        max), and add the uniform floor."""
+        theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
+        likelihood = context.likelihood()
+        phases = np.linspace(0, 2 * np.pi, self.n_grid)
+        if self.approximation_22_mode:
+            # Assume the waveform is (2, 2)-dominated (transforms as exp(2i phase)), so the
+            # phase-dependent log-posterior is Re[(d | h(phase=0)) exp(2i phase)].
+            theta = theta.copy()
+            theta["phase"] = 0.0
+            d_inner_h = likelihood.d_inner_h_complex_multi(theta, self.num_processes)
+            phase_log_posterior = np.outer(d_inner_h, np.exp(2j * phases)).real
+        else:
+            # Exact: each mode m contributes exp(-i m phase); needs spin_conversion_phase=0.
+            likelihood.phase_grid = phases
+            phase_log_posterior = apply_func_with_multiprocessing(
+                likelihood.log_likelihood_phase_grid,
+                theta,
+                num_processes=self.num_processes,
+            )
+        phase_posterior = np.exp(
+            phase_log_posterior - np.amax(phase_log_posterior, axis=1, keepdims=True)
+        )
+        # Uniform floor: keep q(phase) > 0 everywhere so importance sampling stays finite.
+        phase_posterior += (
+            phase_posterior.mean(axis=-1, keepdims=True) * self.uniform_weight
+        )
+        return phases, phase_posterior
 
 
 def _build_gnpe_transforms(model: BasePosteriorModel):
diff --git a/tests/gw/test_synthetic_phase.py b/tests/gw/test_synthetic_phase.py
new file mode 100644
index 000000000..cc45224aa
--- /dev/null
+++ b/tests/gw/test_synthetic_phase.py
@@ -0,0 +1,129 @@
+"""CI unit tests for SyntheticPhaseFactor (and a GWSamplerContext helper), using a mock
+context / likelihood so no waveform models or LAL calls are needed. End-to-end parity
+against Result.sample_synthetic_phase is covered by the model-based harness."""
+
+import numpy as np
+import pytest
+import torch
+from bilby.core.utils import random as bilby_random
+
+from dingo.gw.inference.factors import GWSamplerContext, SyntheticPhaseFactor
+
+
+class _MockLikelihood:
+    """Exposes the two methods the factor uses, deterministic in `chirp_mass`."""
+
+    def __init__(self):
+        self.phase_grid = None
+
+    def d_inner_h_complex_multi(self, theta, num_processes=1):
+        # (2, 2)-approx path: one complex overlap (d | h) per row.
+        return np.array([complex(cm, 0.5) for cm in theta["chirp_mass"].to_numpy()])
+
+    def log_likelihood_phase_grid(self, theta):
+        # exact path: a (n_grid,) log-likelihood per row (theta is a row dict).
+        return np.cos(self.phase_grid) * float(theta["chirp_mass"])
+
+
+class _MockContext:
+    def likelihood(self, **kwargs):
+        return _MockLikelihood()
+
+
+def _given(n=5):
+    return {"chirp_mass": torch.linspace(20.0, 40.0, n, dtype=torch.float64)}
+
+
+def _seed(s=0):
+    np.random.seed(s)
+    bilby_random.seed(s)
+
+
+def test_parameters_and_conditioning():
+    factor = SyntheticPhaseFactor(conditioning=["chirp_mass", "theta_jn"])
+    assert factor.parameters == ["phase"]
+    assert factor.conditioning == ["chirp_mass", "theta_jn"]
+
+
+def test_synthetic_phase_is_one_to_one():
+    factor = SyntheticPhaseFactor(conditioning=["chirp_mass"])
+    with pytest.raises(ValueError):
+        factor.sample_and_log_prob(2, _MockContext(), _given())
+
+
+def test_profile_approx_matches_formula():
+    n, n_grid, weight = 5, 257, 0.01
+    factor = SyntheticPhaseFactor(
+        conditioning=["chirp_mass"],
+        n_grid=n_grid,
+        approximation_22_mode=True,
+        uniform_weight=weight,
+    )
+    given = _given(n)
+    phases, profile = factor._phase_profile(given, _MockContext())
+
+    kappa = np.array([complex(cm, 0.5) for cm in given["chirp_mass"].numpy()])
+    log_posterior = np.outer(kappa, np.exp(2j * phases)).real
+    expected = np.exp(log_posterior - log_posterior.max(axis=1, keepdims=True))
+    expected += expected.mean(axis=1, keepdims=True) * weight
+
+    assert phases.shape == (n_grid,)
+    assert profile.shape == (n, n_grid)
+    assert np.allclose(profile, expected)
+    assert (profile > 0).all()  # uniform floor keeps it mass-covering
+
+
+def test_profile_exact_mode_runs():
+    n, n_grid = 4, 129
+    factor = SyntheticPhaseFactor(
+        conditioning=["chirp_mass"], n_grid=n_grid, approximation_22_mode=False
+    )
+    phases, profile = factor._phase_profile(_given(n), _MockContext())
+    assert phases.shape == (n_grid,)
+    assert profile.shape == (n, n_grid)
+    assert (profile > 0).all()
+
+
+def test_sample_shapes_and_range():
+    n = 6
+    factor = SyntheticPhaseFactor(
+        conditioning=["chirp_mass"], approximation_22_mode=True
+    )
+    _seed(0)
+    block, log_prob = factor.sample_and_log_prob(1, _MockContext(), _given(n))
+    phase = block["phase"].numpy()
+    assert set(block) == {"phase"}
+    assert phase.shape == (n,) and log_prob.shape == (n,)
+    assert (phase >= 0).all() and (phase <= 2 * np.pi).all()
+    assert np.isfinite(log_prob.numpy()).all()
+
+
+def test_log_prob_replug_matches_sample():
+    # factor.log_prob at the drawn phase equals the sampled log q (same deterministic
+    # profile, no re-draw).
+    n = 6
+    factor = SyntheticPhaseFactor(
+        conditioning=["chirp_mass"], approximation_22_mode=True
+    )
+    context, given = _MockContext(), _given(n)
+    _seed(1)
+    block, log_prob = factor.sample_and_log_prob(1, context, given)
+    log_prob_replug = factor.log_prob({"phase": block["phase"]}, context, given)
+    assert np.allclose(log_prob.numpy(), log_prob_replug.numpy())
+
+
+def test_context_frequency_override():
+    def context(event_metadata):
+        return GWSamplerContext(
+            domain=None,
+            detectors=[],
+            t_ref=0.0,
+            data_prep=None,
+            raw_context={},
+            event_metadata=event_metadata,
+        )
+
+    with_override = context({"minimum_frequency": 25.0})
+    assert with_override._frequency("minimum_frequency", 20.0) == 25.0
+    assert with_override._frequency("maximum_frequency", 1024.0) == 1024.0
+    assert context(None)._frequency("minimum_frequency", 20.0) == 20.0

From b5fc80810370c20e77e3530b29acf3e281b14c8d Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Thu, 2 Jul 2026 14:21:36 +0200
Subject: [PATCH 20/65] Convert factor docstrings to single backticks and numpy
 style

For the quartodoc (markdown) docs migration: replace RST double-backticks with
markdown single backticks across dingo/core/factors.py and
dingo/gw/inference/factors.py, and add numpy-style Parameters/Returns blocks to
the public constructors and builders. Docstrings only -- no code changes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/factors.py         | 222 ++++++++++++++++++------------
 dingo/gw/inference/factors.py | 249 +++++++++++++++++++++++++---------
 2 files changed, 325 insertions(+), 146 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 038ea33f7..68a806714 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -6,11 +6,11 @@
 
     q(theta_1, ..., theta_n | d) = prod_i q_i(theta_i | f_i(theta_<i, d)),
 
-where each ``Factor`` samples one parameter block and returns its own log-prob, and a
-composer evaluates them. ``ChainComposer`` runs the factors autoregressively.
+where each `Factor` samples one parameter block and returns its own log-prob, and a
+composer evaluates them. `ChainComposer` runs the factors autoregressively.
 
 Factors work in physical parameter space; a network's standardized space exists only
-inside its forward pass, mediated by ``Standardization``.
+inside its forward pass, mediated by `Standardization`.
 """
 
 from __future__ import annotations
@@ -28,12 +28,12 @@
 
 class Standardization:
     """
-    Affine map between a network's standardized space (``z = (theta - mean) / std``) and
+    Affine map between a network's standardized space (`z = (theta - mean) / std`) and
     physical parameter space.
 
-    Holds one network's ``mean`` and ``std`` and applies them in both directions:
+    Holds one network's `mean` and `std` and applies them in both directions:
     de-standardizing network outputs to physical samples, and standardizing physical
-    parameters for ``log_prob``. Different factors (e.g. a GNPE init and main network)
+    parameters for `log_prob`. Different factors (e.g. a GNPE init and main network)
     carry different instances.
     """
 
@@ -44,31 +44,31 @@ def __init__(self, mean: dict[str, float], std: dict[str, float]):
     def standardize(
         self, values: dict[str, torch.Tensor], names: list[str]
     ) -> torch.Tensor:
-        """Map physical ``values`` (a dict of named tensors) to a standardized tensor
-        with columns in ``names`` order."""
+        """Map physical `values` (a dict of named tensors) to a standardized tensor
+        with columns in `names` order."""
         cols = [(values[n] - self.mean[n]) / self.std[n] for n in names]
         return torch.stack(cols, dim=-1)
 
     def destandardize(
         self, z: torch.Tensor, names: list[str]
     ) -> dict[str, torch.Tensor]:
-        """Map a standardized tensor (columns in ``names`` order) back to a dict of
+        """Map a standardized tensor (columns in `names` order) back to a dict of
         named physical tensors."""
         return {n: z[..., i] * self.std[n] + self.mean[n] for i, n in enumerate(names)}
 
     def log_det(self, names: list[str]) -> float:
         """The term to *add* to a network log-prob to express it in physical space:
-        ``log p_theta = log p_z - sum log std``. Same correction in both directions."""
+        `log p_theta = log p_z - sum log std`. Same correction in both directions."""
         return -sum(math.log(self.std[n]) for n in names)
 
 
 @runtime_checkable
 class SamplerContext(Protocol):
     """
-    Per-event shared state referenced by every factor: the data ``d`` and quantities
+    Per-event shared state referenced by every factor: the data `d` and quantities
     derived from it (the prepared-data representation, the likelihood, event metadata).
     Concrete implementations are domain-specific; see
-    ``dingo.gw.inference.factors.GWSamplerContext``.
+    `dingo.gw.inference.factors.GWSamplerContext`.
     """
 
     event_metadata: Optional[dict]
@@ -86,7 +86,7 @@ def likelihood(self):
 def _cat_dict(
     batches: list[dict[str, torch.Tensor]],
 ) -> dict[str, torch.Tensor]:
-    """Concatenate a list of ``name -> tensor`` dicts along dim 0 (re-joining batches)."""
+    """Concatenate a list of `name -> tensor` dicts along dim 0 (re-joining batches)."""
     return {k: torch.cat([b[k] for b in batches]) for k in batches[0]}
 
 
@@ -95,12 +95,27 @@ def chunk_and_concat(
     batch_size: Optional[int],
     run_once: Callable[[int], tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]],
 ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
-    """Run ``run_once`` over batches of ``total`` and concatenate the results.
+    """Run `run_once` over batches of `total` and concatenate the results.
 
-    Splits ``total`` into chunks of ``batch_size``, calls ``run_once(n) -> (samples,
-    log_prob)`` per chunk, and concatenates the sample dicts and log-probs. Caps peak
-    memory at one chunk. ``batch_size=None`` runs ``total`` in a single call. ``log_prob``
-    may be ``None`` (for a composer without a tractable density, i.e. Gibbs)."""
+    Caps peak memory at one chunk. `log_prob` may be `None` (for a composer without a
+    tractable density, i.e. Gibbs).
+
+    Parameters
+    ----------
+    total : int
+        Total number of samples to produce.
+    batch_size : int, optional
+        Chunk size; `None` runs `total` in a single call.
+    run_once : callable
+        `run_once(n) -> (samples, log_prob)` for one chunk of `n` samples.
+
+    Returns
+    -------
+    samples : dict[str, torch.Tensor]
+        The concatenated per-name sample tensors.
+    log_prob : torch.Tensor or None
+        The concatenated log-probs, or `None` when `run_once` returns `None`.
+    """
     bs = batch_size or total
     sample_parts: list[dict[str, torch.Tensor]] = []
     lp_parts: list[Optional[torch.Tensor]] = []
@@ -115,13 +130,13 @@ def chunk_and_concat(
 
 class Factor(ABC):
     """
-    A conditional distribution ``q_i(theta_i | f_i(theta_<i, d))`` over one parameter
+    A conditional distribution `q_i(theta_i | f_i(theta_<i, d))` over one parameter
     block, emitting physical-space samples and a physical-space log-prob.
 
-    A call draws ``num_samples`` samples per conditioning row and returns ``n_rows *
-    num_samples`` rows in row-major order, where ``n_rows`` is the number of rows in
-    ``given`` (1 if unconditioned). This mirrors the network's
-    ``sample_and_log_prob(*context, num_samples=n) -> (n_rows, n, dim)``.
+    A call draws `num_samples` samples per conditioning row and returns `n_rows *
+    num_samples` rows in row-major order, where `n_rows` is the number of rows in
+    `given` (1 if unconditioned). This mirrors the network's
+    `sample_and_log_prob(*context, num_samples=n) -> (n_rows, n, dim)`.
 
     Attributes
     ----------
@@ -141,9 +156,9 @@ def sample_and_log_prob(
         context: SamplerContext,
         given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """Draw ``num_samples`` samples per conditioning row; return ``(samples,
-        log_prob)`` in physical space. The samples dict may include named columns beyond
-        ``parameters``."""
+        """Draw `num_samples` samples per conditioning row; return `(samples,
+        log_prob)` in physical space. The samples dict may include named columns beyond
+        `parameters`."""
 
     @abstractmethod
     def log_prob(
@@ -152,13 +167,13 @@ def log_prob(
         context: SamplerContext,
         given: Optional[dict[str, torch.Tensor]] = None,
     ) -> torch.Tensor:
-        """Evaluate ``log q_i`` at given physical ``theta_i`` (one conditioning row per
+        """Evaluate `log q_i` at given physical `theta_i` (one conditioning row per
         row)."""
 
 
 def _base_model_metadata(model: BasePosteriorModel) -> dict:
     """The training metadata describing the parameter standardization / data settings.
-    For an unconditional (density-recovery) model this lives under ``metadata["base"]``.
+    For an unconditional (density-recovery) model this lives under `metadata["base"]`.
     """
     metadata = model.metadata
     if metadata["train_settings"]["data"].get("unconditional", False):
@@ -169,11 +184,11 @@ def _base_model_metadata(model: BasePosteriorModel) -> dict:
 class FlowFactor(Factor):
     """
     Factor wrapping a posterior model (NPE flow, FMPE, ...). Reaches data only through
-    ``SamplerContext.prepared_data()`` and encapsulates the network's standardization, so
+    `SamplerContext.prepared_data()` and encapsulates the network's standardization, so
     its interface is in physical parameter space.
 
     An unconditioned model draws from the shared data context; a model with
-    ``context_parameters`` conditions on the values in ``given``.
+    `context_parameters` conditions on the values in `given`.
     """
 
     def __init__(
@@ -184,6 +199,22 @@ def __init__(
         context_parameters: Optional[list[str]] = None,
         aliases: Optional[dict[str, str]] = None,
     ):
+        """
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The posterior model (NPE flow, FMPE, ...) wrapped by this factor.
+        parameters : list[str]
+            The network's trained parameter names (standardization is keyed by these).
+        conditioning : list[str], optional
+            Earlier-block parameters this factor conditions on.
+        context_parameters : list[str], optional
+            Network conditioning inputs (GNPE proxies); empty for plain NPE.
+        aliases : dict[str, str], optional
+            Trained-name to exposed-name map at the factor boundary (e.g.
+            `{"ra": "ra@t_ref"}`), so a downstream reparametrization can convert frames
+            by name without retraining.
+        """
         self.model = model
         # The network's trained parameter names -- standardization is keyed by these. The
         # factor exposes them under canonical aliases (e.g. ra -> ra@t_ref), so a downstream
@@ -201,9 +232,21 @@ def __init__(
     def from_model(
         cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
     ) -> "FlowFactor":
-        """Build a factor from a model, reading ``parameters`` and ``context_parameters``
-        from its training metadata. ``aliases`` maps trained names to exposed canonical
-        names (e.g. ``{"ra": "ra@t_ref"}``) at the factor boundary."""
+        """Build a factor from a model, reading `parameters` and `context_parameters`
+        from its training metadata.
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The posterior model to wrap.
+        aliases : dict[str, str], optional
+            Trained-name to exposed canonical-name map (e.g. `{"ra": "ra@t_ref"}`) at
+            the factor boundary.
+
+        Returns
+        -------
+        FlowFactor
+        """
         data_settings = _base_model_metadata(model)["train_settings"]["data"]
         context_parameters = data_settings.get("context_parameters") or []
         return cls(
@@ -215,9 +258,9 @@ def from_model(
         )
 
     def sample_and_log_prob(self, num_samples, context, given=None):
-        """Draw ``num_samples`` samples per conditioning row. Unconditioned: ``num_samples``
-        draws from the shared data context. Conditioned: ``num_samples`` per context row,
-        returning ``n_rows * num_samples`` rows in row-major order."""
+        """Draw `num_samples` samples per conditioning row. Unconditioned: `num_samples`
+        draws from the shared data context. Conditioned: `num_samples` per context row,
+        returning `n_rows * num_samples` rows in row-major order."""
         data = context.prepared_data()
         self.model.network.eval()
         if not self.context_parameters:
@@ -270,15 +313,15 @@ def log_prob(self, theta_i, context, given=None):
 
 class Reparametrization(ABC):
     """
-    A deterministic bijection ``Step``: it transforms existing parameters (no sampling)
-    and contributes ``-log|det J|`` to the proposal density.
+    A deterministic bijection `Step`: it transforms existing parameters (no sampling)
+    and contributes `-log|det J|` to the proposal density.
 
-    Unlike a ``Factor`` it is 1:1 and invertible -- ``forward`` maps the conditioning block
-    to the ``parameters`` block, ``inverse`` maps back (for re-plug / importance sampling),
+    Unlike a `Factor` it is 1:1 and invertible -- `forward` maps the conditioning block
+    to the `parameters` block, `inverse` maps back (for re-plug / importance sampling),
     and its density contribution is a Jacobian, not a sampled log-prob. Used to relate a
     network's coordinates to physical ones (e.g. right ascension from the training reference
-    frame to the event frame). Subclasses implement ``forward`` / ``inverse`` and, where the
-    map is not measure-preserving, ``log_det``.
+    frame to the event frame). Subclasses implement `forward` / `inverse` and, where the
+    map is not measure-preserving, `log_det`.
     """
 
     parameters: list[str]
@@ -288,18 +331,18 @@ class Reparametrization(ABC):
     def forward(
         self, given: dict[str, torch.Tensor], context: "SamplerContext"
     ) -> dict[str, torch.Tensor]:
-        """Map the conditioning block to the ``parameters`` block."""
+        """Map the conditioning block to the `parameters` block."""
 
     @abstractmethod
     def inverse(
         self, params: dict[str, torch.Tensor], context: "SamplerContext"
     ) -> dict[str, torch.Tensor]:
-        """Map the ``parameters`` block back to the conditioning block."""
+        """Map the `parameters` block back to the conditioning block."""
 
     def log_det(
         self, given: dict[str, torch.Tensor], context: "SamplerContext"
     ) -> torch.Tensor:
-        """``log|det J|`` of ``forward``, per row. Default 0 (measure-preserving)."""
+        """`log|det J|` of `forward`, per row. Default 0 (measure-preserving)."""
         n = next(iter(given.values())).shape[0]
         return torch.zeros(n)
 
@@ -309,7 +352,7 @@ def sample_and_log_prob(
         context: SamplerContext,
         given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """Apply ``forward`` to the conditioning; contribute ``-log|det J|``. ``num_samples``
+        """Apply `forward` to the conditioning; contribute `-log|det J|`. `num_samples`
         must be 1 (a reparametrization is 1:1)."""
         if num_samples != 1:
             raise ValueError("A reparametrization is 1:1; use fan_out=1.")
@@ -332,11 +375,11 @@ def log_prob(
 
 class TargetCorrection(ABC):
     """
-    A ``Step`` that emits an importance-sampling target correction as a side-channel column
+    A `Step` that emits an importance-sampling target correction as a side-channel column
     and contributes nothing to the proposal density.
 
     Its value belongs to the IS target, not the proposal: it reads earlier blocks, emits a
-    named column, optionally consumes intermediates (``consumes``), and adds 0 to the
+    named column, optionally consumes intermediates (`consumes`), and adds 0 to the
     proposal log-prob. 1:1.
     """
 
@@ -373,11 +416,11 @@ def log_prob(
 
 class Step(Protocol):
     """
-    One entry a ``ChainComposer`` folds over: it emits a parameter block and an optional
-    contribution to the proposal ``log_prob``.
+    One entry a `ChainComposer` folds over: it emits a parameter block and an optional
+    contribution to the proposal `log_prob`.
 
-    Density-contributing steps (``Factor``) return a tensor; density-free sampling blocks
-    (``GibbsBlock``) return ``None``. ``parameters`` names the block produced, ``conditioning``
+    Density-contributing steps (`Factor`) return a tensor; density-free sampling blocks
+    (`GibbsBlock`) return `None`. `parameters` names the block produced, `conditioning`
     the earlier-block parameters read from the chain (data is implicit via the context).
     """
 
@@ -394,8 +437,8 @@ def sample_and_log_prob(
 
 @dataclass
 class Stage:
-    """A chain ``Step`` and its fan-out: the number of samples drawn per incoming
-    conditioning row. The root stage draws the base count and ignores ``fan_out``."""
+    """A chain `Step` and its fan-out: the number of samples drawn per incoming
+    conditioning row. The root stage draws the base count and ignores `fan_out`."""
 
     step: Step
     fan_out: int = 1
@@ -403,17 +446,17 @@ class Stage:
 
 class ChainComposer:
     """
-    Autoregressive composer over an ordered list of ``Stage``s.
+    Autoregressive composer over an ordered list of `Stage`s.
 
     Folds the steps in declared order -- a topological order of the conditioning DAG --
     expanding each by its fan-out and summing the proposal log-probs. A step is a
-    ``Factor`` (contributes ``log q_i``) or a density-free sampling block (``GibbsBlock``,
-    contributes ``None``); if any step is density-free the chain has no tractable density
-    and ``sample`` omits ``log_prob``. Covers plain NPE, single-step GNPE, prior
-    conditioning, synthetic phase, intrinsic/extrinsic splits, and -- via ``GibbsBlock`` --
+    `Factor` (contributes `log q_i`) or a density-free sampling block (`GibbsBlock`,
+    contributes `None`); if any step is density-free the chain has no tractable density
+    and `sample` omits `log_prob`. Covers plain NPE, single-step GNPE, prior
+    conditioning, synthetic phase, intrinsic/extrinsic splits, and -- via `GibbsBlock` --
     multi-iteration GNPE.
 
-    Accepts bare steps (wrapped as ``Stage(step, fan_out=1)``) or explicit ``Stage``s.
+    Accepts bare steps (wrapped as `Stage(step, fan_out=1)`) or explicit `Stage`s.
     """
 
     def __init__(self, stages: list[Union["Stage", Step]]):
@@ -423,7 +466,7 @@ def __init__(self, stages: list[Union["Stage", Step]]):
     def _validate(self):
         """Check the declared order is a valid topological order: every conditioning
         name is produced by an earlier step. A step's emitted columns default to its
-        ``parameters``, but a step may emit side-channel columns too (``produces``)."""
+        `parameters`, but a step may emit side-channel columns too (`produces`)."""
         produced: set[str] = set()
         for step in self.steps:
             missing = [c for c in step.conditioning if c not in produced]
@@ -441,8 +484,8 @@ def steps(self) -> list[Step]:
 
     @property
     def expansion(self) -> int:
-        """Product of the non-root fan-outs; ``sample`` returns ``num_samples *
-        expansion`` rows."""
+        """Product of the non-root fan-outs; `sample` returns `num_samples *
+        expansion` rows."""
         return math.prod(stage.fan_out for stage in self.stages[1:])
 
     def sample_and_log_prob(
@@ -451,9 +494,9 @@ def sample_and_log_prob(
         context: SamplerContext,
         batch_size: Optional[int] = None,
     ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
-        """Draw samples. ``num_samples`` is the base (root) count; the result has
-        ``num_samples * expansion`` rows. ``batch_size`` chunks the base count (``None``
-        draws in one pass). The log-prob is ``None`` if any step is density-free."""
+        """Draw samples. `num_samples` is the base (root) count; the result has
+        `num_samples * expansion` rows. `batch_size` chunks the base count (`None`
+        draws in one pass). The log-prob is `None` if any step is density-free."""
         return chunk_and_concat(
             num_samples, batch_size, lambda n: self._run_chain_once(n, context)
         )
@@ -461,8 +504,8 @@ def sample_and_log_prob(
     def _run_chain_once(
         self, base: int, context: SamplerContext
     ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
-        """One pass of the whole chain for ``base`` root samples. Returns the samples and
-        the summed proposal log-prob, or ``None`` if any step is density-free."""
+        """One pass of the whole chain for `base` root samples. Returns the samples and
+        the summed proposal log-prob, or `None` if any step is density-free."""
         samples: dict[str, torch.Tensor] = {}
         total: torch.Tensor | float = 0.0
         has_density = True
@@ -501,8 +544,8 @@ def _run_chain_once(
     def log_prob(
         self, samples: dict[str, torch.Tensor], context: SamplerContext
     ) -> torch.Tensor:
-        """Sum each step's ``log_prob`` at the given samples. Valid only for an all-density
-        chain (every step a ``Factor``)."""
+        """Sum each step's `log_prob` at the given samples. Valid only for an all-density
+        chain (every step a `Factor`)."""
         total: torch.Tensor | float = 0.0
         for step in self.steps:
             given = {k: samples[k] for k in step.conditioning}
@@ -516,7 +559,7 @@ def sample(
         context: SamplerContext,
         batch_size: Optional[int] = None,
     ) -> dict[str, torch.Tensor]:
-        """Per-sample dict of parameters, plus ``log_prob`` for an all-density chain."""
+        """Per-sample dict of parameters, plus `log_prob` for an all-density chain."""
         samples, log_prob = self.sample_and_log_prob(num_samples, context, batch_size)
         if log_prob is None:
             return dict(samples)
@@ -525,22 +568,33 @@ def sample(
 
 class GibbsBlock:
     """
-    A density-free sampling-block ``Step``: runs blocked Gibbs internally and yields no
+    A density-free sampling-block `Step`: runs blocked Gibbs internally and yields no
     proposal log-prob.
 
     Seeds the chain with an init factor, then sweeps the factor list in order for
-    ``num_iterations`` iterations; each factor conditions on the current state and
-    overwrites its own block. As a chain ``Step`` it produces the swept parameter blocks
-    and returns ``None`` for the log-prob -- the cyclic dependency has no tractable marginal
+    `num_iterations` iterations; each factor conditions on the current state and
+    overwrites its own block. As a chain `Step` it produces the swept parameter blocks
+    and returns `None` for the log-prob -- the cyclic dependency has no tractable marginal
     (recoverable by fitting an unconditional density to the samples and taking one
-    ``ChainComposer`` step). Dingo uses this only for multi-iteration GNPE (the GNPE factors
-    in ``dingo.gw.inference.factors``), but the loop is generic.
+    `ChainComposer` step). Dingo uses this only for multi-iteration GNPE (the GNPE factors
+    in `dingo.gw.inference.factors`), but the loop is generic.
 
-    Batching is handled by the enclosing ``ChainComposer``: it chunks the walkers and runs
-    the whole loop per chunk (``chunk_and_concat``).
+    Batching is handled by the enclosing `ChainComposer`: it chunks the walkers and runs
+    the whole loop per chunk (`chunk_and_concat`).
     """
 
     def __init__(self, init_factor: Factor, factors: list[Factor], num_iterations: int):
+        """
+        Parameters
+        ----------
+        init_factor : Factor
+            Seeds the chain (e.g. an init network's detector times).
+        factors : list[Factor]
+            The factors swept in order each iteration; each conditions on the current
+            state and overwrites its own block.
+        num_iterations : int
+            Number of Gibbs sweeps.
+        """
         self.init_factor = init_factor
         self.factors = list(factors)
         self.num_iterations = num_iterations
@@ -555,8 +609,8 @@ def sample_and_log_prob(
         context: SamplerContext,
         given: Optional[dict[str, torch.Tensor]] = None,
     ) -> tuple[dict[str, torch.Tensor], None]:
-        """Run the Gibbs loop for ``num_samples`` walkers; return ``(samples, None)``.
-        ``num_samples`` is the walker (root) count -- Gibbs does not fan out."""
+        """Run the Gibbs loop for `num_samples` walkers; return `(samples, None)`.
+        `num_samples` is the walker (root) count -- Gibbs does not fan out."""
         return self._run_once(num_samples, context), None
 
     def _run_once(
@@ -576,7 +630,7 @@ def _run_once(
 
 class ComposedSampler:
     """
-    Generic runner over a ``ChainComposer`` and a ``SamplerContext``: draws samples and
+    Generic runner over a `ChainComposer` and a `SamplerContext`: draws samples and
     returns them as a DataFrame. Domain-specific processing lives in the chain's steps, so
     the runner is domain-agnostic.
     """
@@ -589,8 +643,8 @@ def __init__(self, composer: ChainComposer, context: SamplerContext):
     def run_sampler(
         self, num_samples: int, batch_size: Optional[int] = None
     ) -> pd.DataFrame:
-        """Draw ``num_samples`` samples (chunked by ``batch_size``) and return them as a
-        DataFrame. An all-density chain includes ``log_prob``; a chain with a ``GibbsBlock``
+        """Draw `num_samples` samples (chunked by `batch_size`) and return them as a
+        DataFrame. An all-density chain includes `log_prob`; a chain with a `GibbsBlock`
         step does not."""
         merged = self.composer.sample(num_samples, self.context, batch_size)
         merged = {k: v.cpu().numpy() for k, v in merged.items()}
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 4c930ea53..2dd76195f 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -1,11 +1,11 @@
 """
 Gravitational-wave factors, steps, and samplers for the factorized sampler.
 
-The domain-agnostic pieces (``Factor``, ``FlowFactor``, ``Reparametrization``,
-``TargetCorrection``, and the composers) live in ``dingo.core.factors``. This module adds
-the GW-specific ones: ``GWSamplerContext`` (the data-preprocessing conditioning map), the
-GNPE factors and steps (``GNPEKernelFactor``, ``GNPEFlowFactor``, ``GNPEKernelCorrection``),
-the ``RAReparam`` sky-frame reparametrization, and ``GWComposedSampler``, which assembles
+The domain-agnostic pieces (`Factor`, `FlowFactor`, `Reparametrization`,
+`TargetCorrection`, and the composers) live in `dingo.core.factors`. This module adds
+the GW-specific ones: `GWSamplerContext` (the data-preprocessing conditioning map), the
+GNPE factors and steps (`GNPEKernelFactor`, `GNPEFlowFactor`, `GNPEKernelCorrection`),
+the `RAReparam` sky-frame reparametrization, and `GWComposedSampler`, which assembles
 the chain for plain NPE, multi-iteration GNPE, or single-step GNPE.
 """
 
@@ -64,13 +64,13 @@
 
 class GWSamplerContext:
     """
-    Per-event shared GW state: the data ``d`` and everything derived from it. Referenced
+    Per-event shared GW state: the data `d` and everything derived from it. Referenced
     by every factor in a chain, and serialized as the transport state between pipe
     stages.
 
-    This implements the ``dingo.core.factors.SamplerContext`` protocol. It owns the
-    one-time data preprocessing (``prepared_data``) and builds the exact likelihood
-    (``likelihood``) used by likelihood-based factors (synthetic phase) and importance
+    This implements the `dingo.core.factors.SamplerContext` protocol. It owns the
+    one-time data preprocessing (`prepared_data`) and builds the exact likelihood
+    (`likelihood`) used by likelihood-based factors (synthetic phase) and importance
     sampling.
 
     Event metadata lives here (not on individual factors): it is a property of the data,
@@ -87,16 +87,33 @@ def __init__(
         event_metadata: Optional[dict] = None,
         base_metadata: Optional[dict] = None,
     ):
+        """
+        Parameters
+        ----------
+        domain : Domain
+            The data frequency domain.
+        detectors : list[str]
+            Detector names.
+        t_ref : float
+            Training reference GPS time.
+        data_prep : Compose
+            The one-time data-preprocessing transform chain (whiten / decimate /
+            repackage).
+        raw_context : dict
+            The raw event data `d` (strain + ASDs per detector), i.e. `EventDataset.data`.
+            Consumed lazily by `prepared_data()` and reused for the likelihood.
+        event_metadata : dict, optional
+            Per-event metadata; drives frequency cropping, the RA correction, and the
+            likelihood reference time.
+        base_metadata : dict, optional
+            The base (parameter / waveform / domain) metadata, used to build the
+            likelihood.
+        """
         self.domain = domain
         self.detectors = detectors
         self.t_ref = t_ref
         self._data_prep = data_prep
-        # The base (parameter/waveform/domain) metadata, used to build the likelihood.
         self.base_metadata = base_metadata
-        # The raw event data `d` (strain + ASDs per detector), i.e. EventDataset.data --
-        # the same object set as GWSampler.context today. Consumed lazily by
-        # prepared_data(); also retained for the likelihood (synthetic-phase / IS factors)
-        # and as part of the serialized transport state, neither of which is wired yet.
         self.raw_context = raw_context
         self.event_metadata = event_metadata
         self._prepared: Optional[torch.Tensor] = None
@@ -108,9 +125,23 @@ def from_model(
         raw_context: dict,
         event_metadata: Optional[dict] = None,
     ) -> "GWSamplerContext":
-        """Build the context (domain + one-time data-prep chain) from a model's
-        metadata. The data-prep chain reproduces ``GWSampler`` preprocessing for the
-        plain-NPE case (parameter-dependent transforms are handled per-factor)."""
+        """Build the context (domain + one-time data-prep chain) from a model's metadata.
+        The data-prep chain reproduces `GWSampler` preprocessing for the plain-NPE case
+        (parameter-dependent transforms are handled per-factor).
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The model whose metadata defines the domain and preprocessing.
+        raw_context : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+
+        Returns
+        -------
+        GWSamplerContext
+        """
         meta = _base_model_metadata(model)
         data_settings = meta["train_settings"]["data"]
 
@@ -273,7 +304,7 @@ def _marginalized_prior(self, name: str):
 
 
 class DeltaFactor(Factor):
-    """``q_i = delta(theta_i - c)``: a point mass pinning parameters to fixed values,
+    """`q_i = delta(theta_i - c)`: a point mass pinning parameters to fixed values,
     contributing 0 to the proposal log-prob.
 
     Used as the chain root for prior-conditioning or known proxies (single-step GNPE, where
@@ -414,7 +445,7 @@ def _phase_profile(self, given, context):
 
 def _build_gnpe_transforms(model: BasePosteriorModel):
     """Build the time-shift GNPE per-step transforms from a model's metadata (the same
-    chains as ``GWSamplerGNPE._initialize_transforms``).
+    chains as `GWSamplerGNPE._initialize_transforms`).
 
     Returns
     -------
@@ -425,7 +456,7 @@ def _build_gnpe_transforms(model: BasePosteriorModel):
     kernel : PriorDict
         The proxy perturbation kernel.
     gnpe_transform : GNPECoalescenceTimes
-        The blur transform itself, shared so the kernel factor can call ``sample_proxies``.
+        The blur transform itself, shared so the kernel factor can call `sample_proxies`.
     """
     meta = _base_model_metadata(model)
     data_settings = meta["train_settings"]["data"]
@@ -491,17 +522,25 @@ def _build_gnpe_transforms(model: BasePosteriorModel):
 
 class GNPEKernelFactor(Factor):
     """
-    The GNPE perturbation kernel ``p(theta_hat | theta)`` as a non-network factor.
+    The GNPE perturbation kernel `p(theta_hat | theta)` as a non-network factor.
 
-    ``theta`` are the detector coalescence times; the kernel adds a bounded perturbation to
-    each, giving the proxies ``theta_hat`` the main network conditions on. The parameter
+    `theta` are the detector coalescence times; the kernel adds a bounded perturbation to
+    each, giving the proxies `theta_hat` the main network conditions on. The parameter
     block is the proxies, the conditioning is the detector times.
-    ``sample_and_log_prob`` blurs the times into proxies (the proxy update of a Gibbs
-    sweep); ``log_prob`` returns the kernel density ``log p(theta_hat | theta)`` at the
+    `sample_and_log_prob` blurs the times into proxies (the proxy update of a Gibbs
+    sweep); `log_prob` returns the kernel density `log p(theta_hat | theta)` at the
     proxies and the detector times. One proxy per detector-time row.
     """
 
     def __init__(self, gnpe_transform: GNPEBase, gnpe_parameters: list[str]):
+        """
+        Parameters
+        ----------
+        gnpe_transform : GNPEBase
+            The blur transform supplying the kernel and `sample_proxies`.
+        gnpe_parameters : list[str]
+            The detector-time parameters perturbed into proxies.
+        """
         self.gnpe = gnpe_transform
         self.gnpe_parameters = gnpe_parameters
         self.parameters = [p + "_proxy" for p in gnpe_parameters]
@@ -515,7 +554,7 @@ def from_model(cls, model: BasePosteriorModel) -> "GNPEKernelFactor":
         return cls(gnpe_transform, gnpe_parameters)
 
     def sample_and_log_prob(self, num_samples, context, given=None):
-        """Blur the conditioning detector times into proxies; ``num_samples`` must be 1
+        """Blur the conditioning detector times into proxies; `num_samples` must be 1
         (GNPE is 1:1). Returns the proxies and their kernel log-prob."""
         if num_samples != 1:
             raise ValueError("GNPE proxy is 1:1; use fan_out=1.")
@@ -524,8 +563,8 @@ def sample_and_log_prob(self, num_samples, context, given=None):
         return proxies, self.log_prob(proxies, context, given)
 
     def log_prob(self, theta_i, context, given=None):
-        """``log p(theta_hat | theta)`` from the kernel, at the proxies (``theta_i``) and
-        the detector times (``given``)."""
+        """`log p(theta_hat | theta)` from the kernel, at the proxies (`theta_i`) and
+        the detector times (`given`)."""
         diffs = {}
         for k in self.kernel.keys():
             diff = given[k] - theta_i[f"{k}_proxy"]
@@ -537,16 +576,16 @@ def log_prob(self, theta_i, context, given=None):
 
 class GNPEFlowFactor(Factor):
     """
-    The GNPE main network ``q(theta | theta_hat, d)`` as a factor.
+    The GNPE main network `q(theta | theta_hat, d)` as a factor.
 
-    Conditions on the detector-time proxies from ``GNPEKernelFactor``: it shifts each
+    Conditions on the detector-time proxies from `GNPEKernelFactor`: it shifts each
     detector's strain by the corresponding proxy time (standardizing the network input),
     samples the network, and recomputes the detector times from the sampled sky position
     and geocent time. The proxies are supplied, so no blurring happens here.
 
-    The single network factor in either GNPE mode: cycled by a ``GibbsBlock`` for
-    multi-iteration GNPE, or a ``ChainComposer`` factor for single-step GNPE. The recomputed
-    detector times are emitted as extra columns (``produces``): the next Gibbs iteration
+    The single network factor in either GNPE mode: cycled by a `GibbsBlock` for
+    multi-iteration GNPE, or a `ChainComposer` factor for single-step GNPE. The recomputed
+    detector times are emitted as extra columns (`produces`): the next Gibbs iteration
     blurs them into fresh proxies, and single-step GNPE evaluates the kernel correction at
     them. One sample per proxy row.
     """
@@ -560,6 +599,23 @@ def __init__(
         parameters: list[str],
         aliases: Optional[dict[str, str]] = None,
     ):
+        """
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The GNPE main network.
+        transform_pre : Compose
+            Per-iteration pre-network transforms (proxy bookkeeping, time shift,
+            standardization).
+        transform_post : Compose
+            Post-network transforms (de-standardize, recompute detector times).
+        gnpe_parameters : list[str]
+            The detector-time parameters.
+        parameters : list[str]
+            The network's trained inference parameters.
+        aliases : dict[str, str], optional
+            Trained-name to exposed-name map (e.g. `{"ra": "ra@t_ref"}`).
+        """
         self.model = model
         self.transform_pre = transform_pre
         self.transform_post = transform_post
@@ -580,7 +636,17 @@ def from_model(
         cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
     ) -> "GNPEFlowFactor":
         """Build the GNPE per-iteration transforms from the main model's metadata.
-        ``aliases`` maps trained names to canonical names (e.g. ``{"ra": "ra@t_ref"}``).
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The GNPE main model.
+        aliases : dict[str, str], optional
+            Trained-name to canonical-name map (e.g. `{"ra": "ra@t_ref"}`).
+
+        Returns
+        -------
+        GNPEFlowFactor
         """
         pre, post, gnpe_parameters, inference_parameters, _, _ = _build_gnpe_transforms(
             model
@@ -590,7 +656,7 @@ def from_model(
         )
 
     def sample_and_log_prob(self, num_samples, context, given=None):
-        """Sample one parameter set per proxy row; ``num_samples`` must be 1 (GNPE is 1:1).
+        """Sample one parameter set per proxy row; `num_samples` must be 1 (GNPE is 1:1).
         Returns theta plus the recomputed detector times, and the network log-prob."""
         if num_samples != 1:
             raise ValueError("GNPE is 1:1; draw one sample per proxy (fan_out=1).")
@@ -628,14 +694,14 @@ def log_prob(self, theta_i, context, given=None):
 
 class RAReparam(Reparametrization):
     """
-    Rotate right ascension from the network's training reference frame (``ra@t_ref``) to the
-    event frame (``ra``).
+    Rotate right ascension from the network's training reference frame (`ra@t_ref`) to the
+    event frame (`ra`).
 
     The network is trained at a fixed reference time; an event at a different GPS time needs
     the sky rotated by the sidereal-time difference. This is a measure-preserving shift
-    modulo 2*pi (``log_det = 0``), so it contributes nothing to the density. ``forward``
-    produces the event-frame ``ra``, ``inverse`` recovers ``ra@t_ref``. The sidereal
-    correction is read from the shared context (``t_ref`` and the event time).
+    modulo 2*pi (`log_det = 0`), so it contributes nothing to the density. `forward`
+    produces the event-frame `ra`, `inverse` recovers `ra@t_ref`. The sidereal
+    correction is read from the shared context (`t_ref` and the event time).
     """
 
     def __init__(self):
@@ -680,9 +746,9 @@ class GNPEKernelCorrection(TargetCorrection):
     """
     The single-step GNPE kernel correction as a target-side chain step.
 
-    Emits ``delta_log_prob_target = log p(theta_hat | theta)`` -- the GNPE kernel evaluated
+    Emits `delta_log_prob_target = log p(theta_hat | theta)` -- the GNPE kernel evaluated
     at the proxies and the detector times recomputed from theta -- for importance sampling
-    on the joint proposal ``q(theta, theta_hat | d)``. Contributes 0 to the proposal
+    on the joint proposal `q(theta, theta_hat | d)`. Contributes 0 to the proposal
     density and consumes the intermediate detector times.
     """
 
@@ -702,18 +768,18 @@ def correction(self, given, context):
 
 
 def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
-    """The RA frame alias (``ra`` -> ``ra@t_ref``), applied only when the model infers
-    ``ra``; paired with an ``RAReparam`` step that maps it back to the event frame."""
+    """The RA frame alias (`ra` -> `ra@t_ref`), applied only when the model infers
+    `ra`; paired with an `RAReparam` step that maps it back to the event frame."""
     return {"ra": "ra@t_ref"} if "ra" in inference_parameters else {}
 
 
 def _ra_reparam_steps(inference_parameters: list[str]) -> list:
-    """The ``RAReparam`` step, appended to a chain only when the model infers ``ra``."""
+    """The `RAReparam` step, appended to a chain only when the model infers `ra`."""
     return [RAReparam()] if "ra" in inference_parameters else []
 
 
 def _delta_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
-    """Delta-prior parameters the chain does not produce, as a single ``DeltaFactor`` step
+    """Delta-prior parameters the chain does not produce, as a single `DeltaFactor` step
     (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0).
     """
     intrinsic_prior = metadata["dataset_settings"]["intrinsic_prior"]
@@ -731,9 +797,9 @@ def _delta_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
 
 class GWComposedSampler(ComposedSampler):
     """
-    GW builder and exporter over the generic ``ComposedSampler`` runner. The ``from_*``
+    GW builder and exporter over the generic `ComposedSampler` runner. The `from_*`
     constructors assemble the chain for plain NPE, multi-iteration GNPE, or single-step
-    GNPE from model metadata; ``to_result`` exports the samples to a gw ``Result``. All
+    GNPE from model metadata; `to_result` exports the samples to a gw `Result`. All
     GW-specific processing (RA frame, fixed parameters, kernel correction) is expressed as
     chain steps, so there is no post-processing.
     """
@@ -745,6 +811,18 @@ def __init__(
         metadata: dict,
         inference_parameters: list[str],
     ):
+        """
+        Parameters
+        ----------
+        composer : ChainComposer
+            The assembled chain of steps.
+        context : GWSamplerContext
+            Per-event shared state.
+        metadata : dict
+            Model metadata, carried through to the exported `Result`.
+        inference_parameters : list[str]
+            The inferred parameter names.
+        """
         super().__init__(composer, context)
         self.metadata = metadata
         self.inference_parameters = inference_parameters
@@ -757,7 +835,21 @@ def from_model(
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
         """Build a plain-NPE GW sampler from a model and event data: the flow exposes
-        ``ra`` as ``ra@t_ref``, followed by an ``RAReparam`` to the event frame."""
+        `ra` as `ra@t_ref`, followed by an `RAReparam` to the event frame.
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The NPE model.
+        raw_context : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+
+        Returns
+        -------
+        GWComposedSampler
+        """
         context = GWSamplerContext.from_model(model, raw_context, event_metadata)
         metadata = _base_model_metadata(model)
         inference_parameters = metadata["train_settings"]["data"][
@@ -786,10 +878,29 @@ def from_gnpe_models(
         num_iterations: int = 30,
     ) -> "GWComposedSampler":
         """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
-        init model's data preprocessing, an init ``FlowFactor`` to seed, and a single
-        ``GibbsBlock`` step -- cycling the GNPE kernel and main-network factors -- in a
-        ``ChainComposer``, then an ``RAReparam`` to the event frame. Returns samples without
-        a log_prob (Gibbs breaks density access)."""
+        init model's data preprocessing, an init `FlowFactor` to seed, and a single
+        `GibbsBlock` step -- cycling the GNPE kernel and main-network factors -- in a
+        `ChainComposer`, then an `RAReparam` to the event frame. Returns samples without
+        a log_prob (Gibbs breaks density access).
+
+        Parameters
+        ----------
+        init_model : BasePosteriorModel
+            The init network (detector times); seeds the Gibbs loop and defines the data
+            preprocessing.
+        main_model : BasePosteriorModel
+            The GNPE main network.
+        raw_context : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+        num_iterations : int, default 30
+            Number of Gibbs sweeps.
+
+        Returns
+        -------
+        GWComposedSampler
+        """
         context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
         metadata = _base_model_metadata(main_model)
         inference_parameters = metadata["train_settings"]["data"][
@@ -821,12 +932,26 @@ def from_singlestep_gnpe(
         raw_context: dict,
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
-        """Build a single-step (density-preserving) time-GNPE sampler: a ``ChainComposer``
-        of ``[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]``.
-        ``proxy_source`` supplies the detector-time proxies -- a ``DeltaFactor`` for prior
-        conditioning (BNS), or an unconditional NDE for density recovery. The chain is
-        autoregressive, so log_prob is preserved, and ``GNPEKernelCorrection`` emits the
-        ``delta_log_prob_target`` correction that importance sampling adds to the target.
+        """Build a single-step (density-preserving) time-GNPE sampler: a `ChainComposer`
+        of `[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]`. The chain is
+        autoregressive, so log_prob is preserved, and `GNPEKernelCorrection` emits the
+        `delta_log_prob_target` correction that importance sampling adds to the target.
+
+        Parameters
+        ----------
+        main_model : BasePosteriorModel
+            The GNPE main network.
+        proxy_source : Factor
+            Supplies the detector-time proxies -- a `DeltaFactor` for prior conditioning
+            (BNS), or an unconditional NDE for density recovery.
+        raw_context : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+
+        Returns
+        -------
+        GWComposedSampler
         """
         context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
         metadata = _base_model_metadata(main_model)
@@ -845,12 +970,12 @@ def from_singlestep_gnpe(
         return cls(ChainComposer(steps), context, metadata, inference_parameters)
 
     def to_result(self):
-        """Export to a gw ``Result`` (samples + raw event data + metadata), so the
+        """Export to a gw `Result` (samples + raw event data + metadata), so the
         existing post-processing pipeline -- synthetic phase, importance sampling,
         evidence, plotting -- runs on the factorized sampler's output unchanged.
 
-        ``context`` is the *raw* event-data dict (``GWSamplerContext.raw_context``, what
-        ``Result`` needs to rebuild the likelihood), not the ``SamplerContext`` object.
+        `context` is the *raw* event-data dict (`GWSamplerContext.raw_context`, what
+        `Result` needs to rebuild the likelihood), not the `SamplerContext` object.
         """
         from dingo.gw.result import Result
 

From 211c390c50c121281f9b70f09e08f63fd969398e Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Fri, 3 Jul 2026 17:15:08 +0200
Subject: [PATCH 21/65] Add GWSamplerContext.prior as the static prior owner

Introduce a cached prior property on GWSamplerContext that builds the static intrinsic+extrinsic prior from model metadata, and route _marginalized_prior through it. First step of the context / Result-diet consolidation: the static prior gets a single home on the per-event context.

Defer the _delta_prior_steps and Result._build_prior call sites to the Result-diet step; in the GNPE builder the context is built from the init model while those read the main model metadata, so unifying them needs that resolved first.

Bit-exact: 30/30 factor CI, plus verify_likelihood_context and verify_integrated (NPE+GNPE) at max|delta|=0.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 2dd76195f..ecf0314e6 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -117,6 +117,7 @@ def __init__(
         self.raw_context = raw_context
         self.event_metadata = event_metadata
         self._prepared: Optional[torch.Tensor] = None
+        self._prior: Optional[PriorDict] = None
 
     @classmethod
     def from_model(
@@ -182,6 +183,24 @@ def prepared_data(self) -> torch.Tensor:
             self._prepared = self._data_prep(self.raw_context)
         return self._prepared
 
+    @property
+    def prior(self) -> PriorDict:
+        """The static prior over all parameters, built once from the model metadata
+        (intrinsic + extrinsic priors with Dingo defaults).
+
+        This is the event-independent prior fixed at training time. Importance-sampling
+        prior-bound updates and the time / phase split-off for marginalized networks are
+        applied downstream (they depend on the evolving analysis state), not here.
+        """
+        if self._prior is None:
+            data_settings = self.base_metadata["train_settings"]["data"]
+            intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
+            extrinsic_prior = get_extrinsic_prior_dict(data_settings["extrinsic_prior"])
+            self._prior = build_prior_with_defaults(
+                {**intrinsic_prior, **extrinsic_prior}
+            )
+        return self._prior
+
     def likelihood(
         self,
         time_marginalization_kwargs: Optional[dict] = None,
@@ -294,13 +313,9 @@ def _marginalized_prior(self, name: str):
         """The prior over `name` if the network marginalized it (i.e. `name` is not an
         inference parameter), else `None` -- used to parameterize likelihood
         marginalization over time / phase."""
-        data_settings = self.base_metadata["train_settings"]["data"]
-        if name in data_settings["inference_parameters"]:
+        if name in self.base_metadata["train_settings"]["data"]["inference_parameters"]:
             return None
-        intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
-        extrinsic_prior = get_extrinsic_prior_dict(data_settings["extrinsic_prior"])
-        prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
-        return prior.get(name)
+        return self.prior.get(name)
 
 
 class DeltaFactor(Factor):

From 028108c4962b4ff724d27fa11a4630c7d7f55956 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Fri, 3 Jul 2026 18:51:10 +0200
Subject: [PATCH 22/65] Route the prior through GWSamplerContext as the single
 owner

Complete the prior consolidation started in 2b (context.prior), part of the Result diet (2c):

- Build the multi-iteration GNPE context from the main model (it owns the analysis), guarded by a check that the init and main models agree on the shared data-prep view (domain, detectors, reference time).

- _delta_prior_steps and Result._build_prior now source the static prior from context.prior. to_result passes the live GWSamplerContext as a new sampler_context arg; Result deepcopies it (the marginalization split-off mutates the prior) and still self-builds when loaded from file (no context).

Bit-exact / pass: verify_integrated (NPE+GNPE), verify_singlestep_packaged, verify_prior_delegation, 30/30 factor CI.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 dingo/core/result.py          |  9 ++--
 dingo/gw/inference/factors.py | 81 +++++++++++++++++++++++++++++------
 dingo/gw/result.py            | 40 +++++++++++------
 3 files changed, 101 insertions(+), 29 deletions(-)

diff --git a/dingo/core/result.py b/dingo/core/result.py
index 67017271a..8f9696052 100644
--- a/dingo/core/result.py
+++ b/dingo/core/result.py
@@ -82,11 +82,16 @@ class Result(DingoDataset):
 
     dataset_type = "core_result"
 
-    def __init__(self, file_name=None, dictionary=None):
+    def __init__(self, file_name=None, dictionary=None, sampler_context=None):
         self.event_metadata = None
         self.context = None
         self.samples = None
         self.log_noise_evidence = None
+        # Live sampler context (a GWSamplerContext), set when a sampler builds the
+        # Result rather than loading it from file. When present, _build_prior (and,
+        # later, _build_likelihood) delegate to it; it is not serialized, so a Result
+        # loaded from file has sampler_context=None and self-builds as before.
+        self.sampler_context = sampler_context
         super().__init__(
             file_name=file_name,
             dictionary=dictionary,
@@ -1090,5 +1095,3 @@ def freeze(d):
     elif isinstance(d, list):
         return tuple(freeze(value) for value in d)
     return d
-
-
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index ecf0314e6..18757cf5a 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -793,15 +793,18 @@ def _ra_reparam_steps(inference_parameters: list[str]) -> list:
     return [RAReparam()] if "ra" in inference_parameters else []
 
 
-def _delta_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
+def _delta_prior_steps(prior, inference_parameters: list[str]) -> list:
     """Delta-prior parameters the chain does not produce, as a single `DeltaFactor` step
     (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0).
+
+    Parameters
+    ----------
+    prior : PriorDict
+        The static prior (`GWSamplerContext.prior`); its delta-function entries that are
+        not inference parameters become the pinned constants.
+    inference_parameters : list of str
+        The inferred parameter names.
     """
-    intrinsic_prior = metadata["dataset_settings"]["intrinsic_prior"]
-    extrinsic_prior = get_extrinsic_prior_dict(
-        metadata["train_settings"]["data"]["extrinsic_prior"]
-    )
-    prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
     fixed = {
         k: p.peak
         for k, p in prior.items()
@@ -810,6 +813,52 @@ def _delta_prior_steps(metadata: dict, inference_parameters: list[str]) -> list:
     return [DeltaFactor(fixed)] if fixed else []
 
 
+def _assert_consistent_gnpe_data_prep(init_model, main_model):
+    """Assert the init and main GNPE models agree on the data-preprocessing view.
+
+    Multi-iteration GNPE shares one `GWSamplerContext` (built from the main model)
+    between the init and main factors, so both read the same `prepared_data()` and
+    reference time. That is only valid when the two models agree on everything that
+    determines those: the domain, the detectors, and the reference time. Raises
+    `ValueError` on any mismatch.
+
+    Parameters
+    ----------
+    init_model, main_model : BasePosteriorModel
+        The GNPE init and main networks.
+    """
+    init = _base_model_metadata(init_model)
+    main = _base_model_metadata(main_model)
+    fields = {
+        "domain": (
+            init["dataset_settings"]["domain"],
+            main["dataset_settings"]["domain"],
+        ),
+        "domain_update": (
+            init["train_settings"]["data"].get("domain_update"),
+            main["train_settings"]["data"].get("domain_update"),
+        ),
+        "detectors": (
+            init["train_settings"]["data"]["detectors"],
+            main["train_settings"]["data"]["detectors"],
+        ),
+        "ref_time": (
+            init["train_settings"]["data"]["ref_time"],
+            main["train_settings"]["data"]["ref_time"],
+        ),
+    }
+    mismatched = {k: (i, m) for k, (i, m) in fields.items() if i != m}
+    if mismatched:
+        details = "; ".join(
+            f"{k}: init={i!r} vs main={m!r}" for k, (i, m) in mismatched.items()
+        )
+        raise ValueError(
+            f"GNPE init and main models disagree on the data-preprocessing view "
+            f"({details}). They share one context, so they must agree on the domain, "
+            f"detectors, and reference time."
+        )
+
+
 class GWComposedSampler(ComposedSampler):
     """
     GW builder and exporter over the generic `ComposedSampler` runner. The `from_*`
@@ -874,7 +923,7 @@ def from_model(
         steps = (
             [factor]
             + _ra_reparam_steps(inference_parameters)
-            + _delta_prior_steps(metadata, inference_parameters)
+            + _delta_prior_steps(context.prior, inference_parameters)
         )
         return cls(
             composer=ChainComposer(steps),
@@ -916,7 +965,11 @@ def from_gnpe_models(
         -------
         GWComposedSampler
         """
-        context = GWSamplerContext.from_model(init_model, raw_context, event_metadata)
+        _assert_consistent_gnpe_data_prep(init_model, main_model)
+        # Build the context from the main model: it owns the analysis (likelihood,
+        # prior, inference parameters). The init model shares the data domain and
+        # preprocessing (asserted above), so prepared_data() is identical either way.
+        context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
         metadata = _base_model_metadata(main_model)
         inference_parameters = metadata["train_settings"]["data"][
             "inference_parameters"
@@ -930,7 +983,7 @@ def from_gnpe_models(
         steps = (
             [gibbs]
             + _ra_reparam_steps(inference_parameters)
-            + _delta_prior_steps(metadata, inference_parameters)
+            + _delta_prior_steps(context.prior, inference_parameters)
         )
         return cls(
             ChainComposer(steps),
@@ -980,7 +1033,7 @@ def from_singlestep_gnpe(
         steps = (
             [proxy_source, flow_factor, GNPEKernelCorrection(kernel_factor)]
             + _ra_reparam_steps(inference_parameters)
-            + _delta_prior_steps(metadata, inference_parameters)
+            + _delta_prior_steps(context.prior, inference_parameters)
         )
         return cls(ChainComposer(steps), context, metadata, inference_parameters)
 
@@ -989,8 +1042,10 @@ def to_result(self):
         existing post-processing pipeline -- synthetic phase, importance sampling,
         evidence, plotting -- runs on the factorized sampler's output unchanged.
 
-        `context` is the *raw* event-data dict (`GWSamplerContext.raw_context`, what
-        `Result` needs to rebuild the likelihood), not the `SamplerContext` object.
+        The raw event-data dict (`GWSamplerContext.raw_context`) is stored as the
+        `Result` context (serialized), and the live `GWSamplerContext` is passed as
+        `sampler_context` so `Result` can pull the prior (and, later, the likelihood)
+        from it rather than rebuilding them from metadata.
         """
         from dingo.gw.result import Result
 
@@ -1003,4 +1058,4 @@ def to_result(self):
             "log_noise_evidence": None,
             "settings": copy.deepcopy(self.metadata),
         }
-        return Result(dictionary=data_dict)
+        return Result(dictionary=data_dict, sampler_context=self.context)
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 51616032d..6e3b3bf0f 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -214,11 +214,18 @@ def _rebuild_domain(self, verbose=False):
 
     def _build_prior(self):
         """Build the prior based on model metadata. Called by __init__()."""
-        intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
-        extrinsic_prior = get_extrinsic_prior_dict(
-            self.base_metadata["train_settings"]["data"]["extrinsic_prior"]
-        )
-        self.prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
+        if self.sampler_context is not None:
+            # Delegate the static prior to the sampler context (the single owner).
+            # Deepcopy because the marginalization split-off below mutates it.
+            self.prior = copy.deepcopy(self.sampler_context.prior)
+        else:
+            intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
+            extrinsic_prior = get_extrinsic_prior_dict(
+                self.base_metadata["train_settings"]["data"]["extrinsic_prior"]
+            )
+            self.prior = build_prior_with_defaults(
+                {**intrinsic_prior, **extrinsic_prior}
+            )
 
         prior_update = self.importance_sampling_metadata.get("prior_update")
         if prior_update is not None:
@@ -415,13 +422,18 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict):
         self.calibration_sampling_kwargs = calibration_sampling_kwargs
 
         # Handle correction_type defaults
-        correction_type = self.calibration_sampling_kwargs.get("correction_type", "data")
+        correction_type = self.calibration_sampling_kwargs.get(
+            "correction_type", "data"
+        )
         if correction_type is None:
             correction_type_dict = {
-                ifo: CALIBRATION_CORRECTION_TYPE_LOOKUP[ifo] for ifo in self.interferometers
+                ifo: CALIBRATION_CORRECTION_TYPE_LOOKUP[ifo]
+                for ifo in self.interferometers
             }
         elif correction_type == "data" or correction_type == "template":
-            correction_type_dict = {ifo: correction_type for ifo in self.interferometers}
+            correction_type_dict = {
+                ifo: correction_type for ifo in self.interferometers
+            }
         elif isinstance(correction_type, dict):
             correction_type_dict = correction_type
         else:
@@ -440,7 +452,7 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict):
             )
 
         # Removing the delta function priors on the frequency nodes, amplitude and phase.
-        # Usually the frequency nodes are set to delta functions, but we also remove the 
+        # Usually the frequency nodes are set to delta functions, but we also remove the
         # the amplitude and phase delta functions if present.
         # This avoids large log probs and log priors, since the density of a delta function
         # at the sampled point is infinite. The delta functions do not affect the sampling,
@@ -457,9 +469,9 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict):
         delta_log_prob = np.zeros(num_samples)
 
         # Here we will sample the calibration parameters from the prior.
-        # We treat the *prior as the proposal* distribution and 
-        # therefore add the log_prob of the sampled calibration parameters 
-        # to the existing log_prob. We also will update the prior 
+        # We treat the *prior as the proposal* distribution and
+        # therefore add the log_prob of the sampled calibration parameters
+        # to the existing log_prob. We also will update the prior
         # to include the calibration priors using the importance_sampling_metadata
         prior_update = self.importance_sampling_metadata.get("prior_update", {})
         for ifo, prior in calibration_priors.items():
@@ -690,7 +702,9 @@ def get_samples_bilby_phase(self, num_processes=1):
             num_processes=num_processes,
         )
 
-    def get_pesummary_samples(self, num_processes=1, resampling_method="clip+rejection"):
+    def get_pesummary_samples(
+        self, num_processes=1, resampling_method="clip+rejection"
+    ):
         """Samples in a form suitable for PESummary.
 
         These samples are adjusted to undo certain conventions used internally by

From 292f557c1af7f42a07e66e74811c450511db85ca Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 5 Jul 2026 18:34:39 +0100
Subject: [PATCH 23/65] Make ChainComposer.log_prob a kind-aware reverse fold;
 add missing step log_probs

ChainComposer.log_prob previously called step.log_prob uniformly, which raised
for every real chain (reparams and corrections had raising stubs, and consumed
columns like ra@t_ref were absent from the samples). Now the steps fold in
exact reverse chain order, restoring the column state each step saw during
sampling: a Reparametrization rebuilds its inputs via inverse and contributes
-log|det J|, a Factor contributes its log_prob, a TargetCorrection contributes
nothing, and a density-free chain (GibbsBlock) raises.

- _validate now rejects a step overwriting an existing column, except a
  Reparametrization replacing its own inputs (invertible, so log_prob can
  restore the state); this is the invariant the reverse fold relies on.
- GNPEFlowFactor.log_prob: run the proxies-present data prep, undo the
  post-network geocent_time proxy shift via PostCorrectGeocentTime(inverse=True),
  standardize, and score the network; carries its own Standardization now.
- DeltaFactor moves to core (nothing GW about a point mass) and gains
  log_prob = 0 per row on the evaluated block's device; the gw module re-exports
  it for the builders.
- Drop the dead raising log_prob stubs on Reparametrization/TargetCorrection;
  document that a correction must only consume side-channel intermediates.
- RAReparam: document the principal-branch semantics of the modulo for samples
  drawn outside [0, 2pi).
- Tests: analytic-density re-plug through reparam + correction chains (incl.
  an in-place same-name reparam), overwrite validation, density-free raise.
  Model-based re-plug harness (NPE + single-step GNPE) agrees to float roundoff
  on in-range rows (max ~2e-4, median ~4e-6); sampling parity stays bit-exact.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py         | 121 ++++++++++++++++++++++++++--------
 dingo/gw/inference/factors.py |  67 ++++++++++++-------
 tests/core/test_factors.py    | 119 ++++++++++++++++++++++++++++++++-
 3 files changed, 253 insertions(+), 54 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 68a806714..fd8a36b04 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -311,6 +311,36 @@ def log_prob(self, theta_i, context, given=None):
         return log_prob + self.standardization.log_det(self._net_parameters)
 
 
+class DeltaFactor(Factor):
+    """`q_i = delta(theta_i - c)`: a point mass pinning parameters to fixed values,
+    contributing 0 to the proposal log-prob.
+
+    Used as the chain root for prior-conditioning or known proxies (where later factors
+    condition on the pinned values), and as a non-root filler for delta-prior parameters
+    a model does not infer (one constant per current row).
+    """
+
+    def __init__(self, values: dict[str, float]):
+        self.values = values
+        self.parameters = list(values)
+        self.conditioning = []
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        samples = {
+            p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
+        }
+        return samples, torch.zeros(num_samples)
+
+    def log_prob(self, theta_i, context, given=None):
+        """0 per row, matching sample time: the point mass is evaluated on its own
+        support (the chain only re-plugs its own samples), so the pinned block is
+        conditioned on, not integrated over. Off-support densities are not
+        represented."""
+        # One zero per row, on the same device/dtype as the evaluated block.
+        reference_column = next(iter(theta_i.values()))
+        return torch.zeros_like(reference_column)
+
+
 class Reparametrization(ABC):
     """
     A deterministic bijection `Step`: it transforms existing parameters (no sampling)
@@ -361,17 +391,10 @@ def sample_and_log_prob(
 
     @property
     def consumes(self) -> list[str]:
-        """Inputs the bijection replaces with its outputs (dropped after the step)."""
+        """Inputs the bijection replaces with its outputs (dropped after the step).
+        `ChainComposer.log_prob` rebuilds them via `inverse`."""
         return [c for c in self.conditioning if c not in self.parameters]
 
-    def log_prob(
-        self,
-        theta_i: dict[str, torch.Tensor],
-        context: SamplerContext,
-        given: Optional[dict[str, torch.Tensor]] = None,
-    ) -> torch.Tensor:
-        raise NotImplementedError("Reparametrization.log_prob")
-
 
 class TargetCorrection(ABC):
     """
@@ -381,6 +404,10 @@ class TargetCorrection(ABC):
     Its value belongs to the IS target, not the proposal: it reads earlier blocks, emits a
     named column, optionally consumes intermediates (`consumes`), and adds 0 to the
     proposal log-prob. 1:1.
+
+    `consumes` must name only side-channel intermediates (e.g. recomputed detector
+    times), never a sampled parameter: unlike a `Reparametrization`, a correction has no
+    inverse, so anything it consumes cannot be rebuilt by `ChainComposer.log_prob`.
     """
 
     parameters: list[str]
@@ -405,14 +432,6 @@ def sample_and_log_prob(
         n = next(iter(out.values())).shape[0]
         return out, torch.zeros(n)
 
-    def log_prob(
-        self,
-        theta_i: dict[str, torch.Tensor],
-        context: SamplerContext,
-        given: Optional[dict[str, torch.Tensor]] = None,
-    ) -> torch.Tensor:
-        raise NotImplementedError("TargetCorrection.log_prob")
-
 
 class Step(Protocol):
     """
@@ -465,8 +484,11 @@ def __init__(self, stages: list[Union["Stage", Step]]):
 
     def _validate(self):
         """Check the declared order is a valid topological order: every conditioning
-        name is produced by an earlier step. A step's emitted columns default to its
-        `parameters`, but a step may emit side-channel columns too (`produces`)."""
+        name is produced by an earlier step, and no step overwrites an existing
+        column -- except a `Reparametrization` replacing its own inputs, which is
+        invertible, so `log_prob` can restore the overwritten state. A step's
+        emitted columns default to its `parameters`, but a step may emit
+        side-channel columns too (`produces`)."""
         produced: set[str] = set()
         for step in self.steps:
             missing = [c for c in step.conditioning if c not in produced]
@@ -475,7 +497,18 @@ def _validate(self):
                     f"A step producing {step.parameters} conditions on {missing}, "
                     f"which no earlier step produces. Check chain order."
                 )
-            produced.update(getattr(step, "produces", step.parameters))
+            emitted = set(getattr(step, "produces", step.parameters))
+            replaceable = (
+                set(step.conditioning) if isinstance(step, Reparametrization) else set()
+            )
+            clobbered = (emitted & produced) - replaceable
+            if clobbered:
+                raise ValueError(
+                    f"A step producing {step.parameters} would overwrite existing "
+                    f"column(s) {sorted(clobbered)}. Only a Reparametrization may "
+                    f"replace columns (its inverse can rebuild them for log_prob)."
+                )
+            produced.update(emitted)
 
     @property
     def steps(self) -> list[Step]:
@@ -544,13 +577,49 @@ def _run_chain_once(
     def log_prob(
         self, samples: dict[str, torch.Tensor], context: SamplerContext
     ) -> torch.Tensor:
-        """Sum each step's `log_prob` at the given samples. Valid only for an all-density
-        chain (every step a `Factor`)."""
+        """Evaluate the chain's proposal log-density at given physical samples
+        (re-plug / importance sampling).
+
+        The steps are folded in exact reverse chain order, so the columns are
+        restored to the state each step saw during sampling: a `Reparametrization`
+        rebuilds its inputs via `inverse` (e.g. `ra@t_ref` from the event-frame
+        `ra`) and contributes `-log|det J|`, a `Factor` contributes its `log_prob`,
+        and a `TargetCorrection` contributes nothing (target-side only). Raises for
+        a density-free chain (one containing a `GibbsBlock`).
+
+        Parameters
+        ----------
+        samples : dict[str, torch.Tensor]
+            The chain's emitted columns (one value per row). Consumed intermediates
+            are rebuilt via the reparametrization inverses and need not be present.
+        context : SamplerContext
+            Per-event shared state.
+
+        Returns
+        -------
+        torch.Tensor
+            The proposal log-density per row.
+        """
+        if any(isinstance(step, GibbsBlock) for step in self.steps):
+            raise ValueError(
+                "The chain contains a density-free step (GibbsBlock), so its "
+                "log_prob is unavailable; recover a density first (fit an "
+                "unconditional model and take a chain step)."
+            )
+        values = dict(samples)
         total: torch.Tensor | float = 0.0
-        for step in self.steps:
-            given = {k: samples[k] for k in step.conditioning}
-            theta_i = {k: samples[k] for k in step.parameters}
-            total = total + step.log_prob(theta_i, context, given)
+        for step in reversed(self.steps):
+            if isinstance(step, TargetCorrection):
+                continue
+            if isinstance(step, Reparametrization):
+                params = {k: values[k] for k in step.parameters}
+                values.update(step.inverse(params, context))
+                given = {k: values[k] for k in step.conditioning}
+                total = total - step.log_det(given, context)
+            else:
+                given = {k: values[k] for k in step.conditioning}
+                theta_i = {k: values[k] for k in step.parameters}
+                total = total + step.log_prob(theta_i, context, given)
         return total
 
     def sample(
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 18757cf5a..97f6bfdfd 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -31,10 +31,12 @@
 from dingo.core.factors import (
     ChainComposer,
     ComposedSampler,
+    DeltaFactor,
     Factor,
     FlowFactor,
     GibbsBlock,
     Reparametrization,
+    Standardization,
     TargetCorrection,
     _base_model_metadata,
 )
@@ -318,30 +320,6 @@ def _marginalized_prior(self, name: str):
         return self.prior.get(name)
 
 
-class DeltaFactor(Factor):
-    """`q_i = delta(theta_i - c)`: a point mass pinning parameters to fixed values,
-    contributing 0 to the proposal log-prob.
-
-    Used as the chain root for prior-conditioning or known proxies (single-step GNPE, where
-    later factors condition on the pinned values), and as a non-root filler for delta-prior
-    parameters the network does not infer (one constant per current row).
-    """
-
-    def __init__(self, values: dict[str, float]):
-        self.values = values
-        self.parameters = list(values)
-        self.conditioning = []
-
-    def sample_and_log_prob(self, num_samples, context, given=None):
-        samples = {
-            p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
-        }
-        return samples, torch.zeros(num_samples)
-
-    def log_prob(self, theta_i, context, given=None):
-        raise NotImplementedError("DeltaFactor.log_prob")
-
-
 def _to_numpy(v) -> np.ndarray:
     """Detach a torch tensor (or coerce anything) to a numpy array."""
     if torch.is_tensor(v):
@@ -640,6 +618,10 @@ def __init__(
         self._net_parameters = parameters
         self.parameters = [self.aliases.get(p, p) for p in parameters]
         self.conditioning = list(self.proxy_parameters)
+        # For log_prob: the sampling path de-standardizes (and corrects the log-prob)
+        # inside transform_post, but evaluating at a point needs the forward map too.
+        std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
+        self.standardization = Standardization(std["mean"], std["std"])
 
     @property
     def produces(self) -> list[str]:
@@ -704,7 +686,37 @@ def sample_and_log_prob(self, num_samples, context, given=None):
         return params, x["log_prob"]
 
     def log_prob(self, theta_i, context, given=None):
-        raise NotImplementedError("GNPEFlowFactor.log_prob")
+        """Evaluate the network density `log q(theta | theta_hat, d)` in physical space
+        at given `theta_i` (exposed / aliased names), one row per proxy row in `given`.
+        Runs the same proxies-present data preparation as sampling (time-shift by the
+        proxies, standardize the conditioning), then scores the standardized parameters
+        under the network."""
+        proxies = {p: given[p] for p in self.proxy_parameters}
+        n_rows = next(iter(proxies.values())).shape[0]
+        x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
+        d = context.prepared_data().clone()
+        x["data"] = d.expand(n_rows, *d.shape)
+        x = self.transform_pre(x)
+        theta_net = {
+            net: theta_i[self.aliases.get(net, net)] for net in self._net_parameters
+        }
+        # Mirror transform_post: sampling shifts geocent_time by the preferred proxy
+        # after the network (PostCorrectGeocentTime, using the extrinsic geocent_time
+        # set up by the GNPE transform), so score the network in its own output frame
+        # by applying the inverse correction first.
+        y = {
+            "parameters": dict(theta_net),
+            "extrinsic_parameters": dict(x["extrinsic_parameters"]),
+        }
+        theta_net = PostCorrectGeocentTime(inverse=True)(y)["parameters"]
+        z = self.standardization.standardize(theta_net, self._net_parameters)
+        self.model.network.eval()
+        with torch.no_grad():
+            if "context_parameters" in x:
+                log_prob = self.model.log_prob(z, x["data"], x["context_parameters"])
+            else:
+                log_prob = self.model.log_prob(z, x["data"])
+        return log_prob + self.standardization.log_det(self._net_parameters)
 
 
 class RAReparam(Reparametrization):
@@ -717,6 +729,11 @@ class RAReparam(Reparametrization):
     modulo 2*pi (`log_det = 0`), so it contributes nothing to the density. `forward`
     produces the event-frame `ra`, `inverse` recovers `ra@t_ref`. The sidereal
     correction is read from the shared context (`t_ref` and the event time).
+
+    The modulo makes the map a bijection on the circle, while the flow's density lives
+    on the real line: a sample drawn outside `[0, 2 pi)` is wrapped, so `inverse`
+    recovers its principal-branch representative and a re-evaluated `log_prob` refers
+    to that branch. Only tail samples outside the bounded `ra` prior are affected.
     """
 
     def __init__(self):
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 8cfe748c3..1a9d65101 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -8,6 +8,8 @@
 here.
 """
 
+import math
+
 import pytest
 import torch
 
@@ -302,12 +304,12 @@ def test_unconditioned_factor_as_root_draws_base_count():
 
 class _MockTargetCorrection(TargetCorrection):
     """A mock kind-3 correction: emits ``corr = 2 * x``, contributes 0 to the proposal,
-    and consumes its input."""
+    and (by default) consumes its input."""
 
-    def __init__(self, reads="x", emits="corr"):
+    def __init__(self, reads="x", emits="corr", consume=True):
         self.conditioning = [reads]
         self.parameters = [emits]
-        self.consumes = [reads]
+        self.consumes = [reads] if consume else []
         self._reads, self._emits = reads, emits
 
     def correction(self, given, context):
@@ -362,3 +364,114 @@ def test_side_channel_produces_satisfies_topological_check():
     )
     out = comp.sample(4, context=None)
     assert out["b"].shape == (4,) and "s" in out
+
+
+class _GaussFactor(Factor):
+    """A factor with a real analytic density -- standard normal around the sum of its
+    conditioning (0 if unconditioned) -- so log_prob re-plug is checkable against the
+    log-prob stored at sample time."""
+
+    def __init__(self, name, conditioning=()):
+        self.parameters = [name]
+        self.conditioning = list(conditioning)
+        self._name = name
+
+    def _mean(self, given, n_rows):
+        if self.conditioning:
+            return sum(given[k] for k in self.conditioning)
+        return torch.zeros(n_rows)
+
+    def sample_and_log_prob(self, n, context, given=None):
+        n_rows = 1 if not self.conditioning else next(iter(given.values())).shape[0]
+        mean = self._mean(given, n_rows).repeat_interleave(n)
+        vals = mean + torch.randn(n_rows * n)
+        return {self._name: vals}, self.log_prob({self._name: vals}, context, given)
+
+    def log_prob(self, theta_i, context, given=None):
+        v = theta_i[self._name]
+        mean = self._mean(given, v.shape[0])
+        return -0.5 * (v - mean) ** 2 - 0.5 * math.log(2 * math.pi)
+
+
+def test_log_prob_replug_kind_aware():
+    # Re-plug: the kind-aware log_prob rebuilds consumed columns via the reparam
+    # inverses, sums factor densities and Jacobians, and skips target corrections.
+    torch.manual_seed(0)
+    comp = ChainComposer(
+        [
+            _GaussFactor("u"),
+            _MockReparam(shift=10.0, log_det_val=0.5),
+            _MockTargetCorrection(reads="v", emits="corr", consume=False),
+            _GaussFactor("w", conditioning=["v"]),
+        ]
+    )
+    out = comp.sample(64, context=None)
+    assert "u" not in out  # consumed by the reparam
+    samples = {k: v for k, v in out.items() if k != "log_prob"}
+    lp = comp.log_prob(samples, context=None)
+    # u is rebuilt as v - shift; the float round trip is not bit-exact, so compare
+    # with a tolerance.
+    assert torch.allclose(lp, out["log_prob"], atol=1e-5)
+    # And against a hand-built sum: gauss(u) + gauss(w | v) - log_det.
+    manual = (
+        comp.steps[0].log_prob({"u": samples["v"] - 10.0}, None)
+        + comp.steps[3].log_prob({"w": samples["w"]}, None, {"v": samples["v"]})
+        - 0.5
+    )
+    assert torch.allclose(lp, manual, atol=1e-6)
+
+
+def test_log_prob_raises_for_density_free_chain():
+    block = GibbsBlock(
+        init_factor=_ConstFactor("proxy"),
+        factors=[_ConstFactor("theta", conditioning=["proxy"])],
+        num_iterations=2,
+    )
+    comp = ChainComposer([block])
+    out = comp.sample(4, context=None)
+    with pytest.raises(ValueError, match="density-free"):
+        comp.log_prob(dict(out), context=None)
+
+
+class _InPlaceReparam(Reparametrization):
+    """A bijection that replaces its input under the same name (``x -> x + shift``),
+    so evaluation must restore the column state per step position."""
+
+    def __init__(self, shift=5.0):
+        self.conditioning = ["x"]
+        self.parameters = ["x"]
+        self._shift = shift
+
+    def forward(self, given, context):
+        return {"x": given["x"] + self._shift}
+
+    def inverse(self, params, context):
+        return {"x": params["x"] - self._shift}
+
+
+def test_log_prob_replug_in_place_reparam():
+    # An in-place reparam overwrites its own column; the reverse fold restores the
+    # pre-transform value before the factor is scored, and a downstream consumer
+    # (the correction) is evaluated against the post-transform value it saw.
+    torch.manual_seed(0)
+    comp = ChainComposer(
+        [
+            _GaussFactor("x"),
+            _InPlaceReparam(shift=5.0),
+            _MockTargetCorrection(reads="x", emits="corr", consume=False),
+        ]
+    )
+    out = comp.sample(64, context=None)
+    assert torch.equal(out["corr"], 2 * out["x"])  # correction saw the shifted x
+    samples = {k: v for k, v in out.items() if k != "log_prob"}
+    lp = comp.log_prob(samples, context=None)
+    assert torch.allclose(lp, out["log_prob"], atol=1e-5)
+
+
+def test_validation_rejects_column_overwrite():
+    # Only a Reparametrization may replace an existing column (it is invertible);
+    # a factor re-producing a name would destroy state log_prob cannot rebuild.
+    with pytest.raises(ValueError, match="overwrite"):
+        ChainComposer([_ConstFactor("a"), _ConstFactor("a")])
+    # The in-place reparam is the allowed exception.
+    ChainComposer([_ConstFactor("x"), _InPlaceReparam()])

From 8a2547e2c4c6d5bdebf9b51e25fd9a0fb2b45c20 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 5 Jul 2026 19:20:00 +0100
Subject: [PATCH 24/65] Add a chain device policy, cache the likelihood,
 invalidate stale contexts

- SamplerContext/GWSamplerContext gain device (the model device): steps that
  create fresh tensors (DeltaFactor) create them there; steps that transform
  existing rows (the Reparametrization.log_det default, TargetCorrection,
  GNPEKernelFactor.log_prob) follow their inputs, so a chain stays on one
  device end to end (the GPU pipe case). The kernel density converts each
  side to numpy before subtracting (times and proxies may live on different
  devices; the kernel is a bilby PriorDict -- the same object that samples
  the blur). SyntheticPhaseFactor returns its phase and log-prob on the
  chain device.
- GWSamplerContext.likelihood() keeps the most recently built likelihood
  with its arguments and rebuilds only when they change: the synthetic-phase
  factor requests one per chain chunk, and importance sampling will request
  its own configuration once. The arguments are captured via locals() so the
  comparison tracks the signature automatically.
- Remove the never-used phase_grid constructor argument from
  StationaryGaussianGWLikelihood, Result._build_likelihood, and
  GWSamplerContext.likelihood(); the grid is assigned as an attribute by its
  consumers (synthetic phase) or set internally by phase marginalization.
- Result.reset_event() drops the live sampler_context: after resetting to
  (possibly regenerated) event data it no longer describes self.context.
- Tests: device placement (meta unit test + an MPS end-to-end mixed chain),
  likelihood cache semantics, reset_event invalidation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py               | 26 ++++++---
 dingo/core/result.py                |  5 ++
 dingo/gw/inference/factors.py       | 71 ++++++++++++++++++------
 dingo/gw/likelihood.py              |  6 +--
 dingo/gw/result.py                  |  2 -
 tests/core/test_factors.py          | 49 +++++++++++++++++
 tests/core/test_result.py           | 26 +++++++--
 tests/gw/test_gw_sampler_context.py | 83 +++++++++++++++++++++++++++++
 8 files changed, 236 insertions(+), 32 deletions(-)
 create mode 100644 tests/gw/test_gw_sampler_context.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index fd8a36b04..eb3269bc8 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -69,9 +69,14 @@ class SamplerContext(Protocol):
     derived from it (the prepared-data representation, the likelihood, event metadata).
     Concrete implementations are domain-specific; see
     `dingo.gw.inference.factors.GWSamplerContext`.
+
+    `device` is the torch device the chain runs on: steps that create fresh tensors
+    (rather than transforming existing ones) create them here, so their outputs can
+    join a chain whose network factors run on a GPU.
     """
 
     event_metadata: Optional[dict]
+    device: Union[torch.device, str]
 
     def prepared_data(self) -> torch.Tensor:
         """The data representation the factors condition on
@@ -326,10 +331,15 @@ def __init__(self, values: dict[str, float]):
         self.conditioning = []
 
     def sample_and_log_prob(self, num_samples, context, given=None):
+        # A delta factor creates fresh tensors, so it places them on the chain's
+        # device (unlike steps that transform existing rows, which follow their
+        # inputs).
+        device = getattr(context, "device", None)
         samples = {
-            p: torch.full((num_samples,), float(v)) for p, v in self.values.items()
+            p: torch.full((num_samples,), float(v), device=device)
+            for p, v in self.values.items()
         }
-        return samples, torch.zeros(num_samples)
+        return samples, torch.zeros(num_samples, device=device)
 
     def log_prob(self, theta_i, context, given=None):
         """0 per row, matching sample time: the point mass is evaluated on its own
@@ -372,9 +382,10 @@ def inverse(
     def log_det(
         self, given: dict[str, torch.Tensor], context: "SamplerContext"
     ) -> torch.Tensor:
-        """`log|det J|` of `forward`, per row. Default 0 (measure-preserving)."""
-        n = next(iter(given.values())).shape[0]
-        return torch.zeros(n)
+        """`log|det J|` of `forward`, per row. Default 0 (measure-preserving), on
+        the device of the transformed rows."""
+        reference_column = next(iter(given.values()))
+        return torch.zeros_like(reference_column)
 
     def sample_and_log_prob(
         self,
@@ -429,8 +440,9 @@ def sample_and_log_prob(
         if num_samples != 1:
             raise ValueError("A target correction is 1:1; use fan_out=1.")
         out = self.correction(given, context)
-        n = next(iter(out.values())).shape[0]
-        return out, torch.zeros(n)
+        # 0 proposal contribution per row, on the device of the emitted column.
+        reference_column = next(iter(out.values()))
+        return out, torch.zeros_like(reference_column)
 
 
 class Step(Protocol):
diff --git a/dingo/core/result.py b/dingo/core/result.py
index 8f9696052..1490c63b0 100644
--- a/dingo/core/result.py
+++ b/dingo/core/result.py
@@ -164,6 +164,11 @@ def reset_event(self, event_dataset):
         event_dataset: EventDataset
             New event to be used for importance sampling.
         """
+        # The live sampler context describes the data the samples were drawn from;
+        # after resetting to (possibly regenerated) event data it no longer describes
+        # self.context, so drop it and fall back to self-building.
+        self.sampler_context = None
+
         context = event_dataset.data
         event_metadata = event_dataset.settings
 
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 97f6bfdfd..295b63bda 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -14,7 +14,7 @@
 import copy
 import logging
 import time
-from typing import Optional
+from typing import Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -88,6 +88,7 @@ def __init__(
         raw_context: dict,
         event_metadata: Optional[dict] = None,
         base_metadata: Optional[dict] = None,
+        device: Union[torch.device, str] = "cpu",
     ):
         """
         Parameters
@@ -110,6 +111,9 @@ def __init__(
         base_metadata : dict, optional
             The base (parameter / waveform / domain) metadata, used to build the
             likelihood.
+        device : torch.device or str, default "cpu"
+            The torch device the chain runs on (the model device); steps that create
+            fresh tensors (e.g. `DeltaFactor`) create them here.
         """
         self.domain = domain
         self.detectors = detectors
@@ -118,8 +122,11 @@ def __init__(
         self.base_metadata = base_metadata
         self.raw_context = raw_context
         self.event_metadata = event_metadata
+        self.device = device
         self._prepared: Optional[torch.Tensor] = None
         self._prior: Optional[PriorDict] = None
+        self._likelihood: Optional[StationaryGaussianGWLikelihood] = None
+        self._likelihood_settings: Optional[dict] = None
 
     @classmethod
     def from_model(
@@ -177,6 +184,7 @@ def from_model(
             raw_context=raw_context,
             event_metadata=event_metadata,
             base_metadata=meta,
+            device=model.device,
         )
 
     def prepared_data(self) -> torch.Tensor:
@@ -208,7 +216,6 @@ def likelihood(
         time_marginalization_kwargs: Optional[dict] = None,
         phase_marginalization_kwargs: Optional[dict] = None,
         calibration_marginalization_kwargs: Optional[dict] = None,
-        phase_grid: Optional[np.ndarray] = None,
         use_base_domain: bool = False,
     ) -> StationaryGaussianGWLikelihood:
         """
@@ -221,6 +228,13 @@ def likelihood(
         event time (the training-frame right-ascension correction is already applied to the
         samples), or the training reference time when no event time is set.
 
+        The most recently built likelihood is cached with its arguments: a repeated
+        call with the same arguments returns the shared instance (the synthetic-phase
+        factor requests one per chain chunk), and a call with different arguments
+        builds a replacement. In the exact synthetic-phase mode the consumer assigns
+        a `phase_grid` attribute on the shared instance, as
+        `Result.sample_synthetic_phase` does.
+
         Importance-sampling domain rebuilds (an updated `T` / `delta_f` or frequency range)
         are not wired here yet: the likelihood uses this context's domain.
 
@@ -233,8 +247,6 @@ def likelihood(
             Analytically marginalize over `phase`. Requires a uniform [0, 2 pi) phase prior.
         calibration_marginalization_kwargs : dict, optional
             Marginalize over detector calibration uncertainty.
-        phase_grid : np.ndarray, optional
-            Phase grid for the exact phase-marginalized / synthetic-phase likelihood.
         use_base_domain : bool, default False
             For a multibanded domain, evaluate on the base (undecimated) frequency domain
             rather than the decimated one.
@@ -243,6 +255,15 @@ def likelihood(
         -------
         StationaryGaussianGWLikelihood
         """
+        # Capture the call's arguments for the cache comparison; this must stay the
+        # first statement so locals() holds exactly the arguments. The comparison
+        # happens before the marginalization bounds are filled in below (which is
+        # deterministic given the same arguments).
+        settings = dict(locals())
+        settings.pop("self")
+        if settings == self._likelihood_settings:
+            return self._likelihood
+
         # Marginalizing over a parameter needs the (uniform) prior the network
         # marginalized over; use it to parameterize the requested marginalization.
         if time_marginalization_kwargs is not None:
@@ -284,7 +305,7 @@ def likelihood(
         else:
             t_ref = self.t_ref
 
-        return StationaryGaussianGWLikelihood(
+        likelihood = StationaryGaussianGWLikelihood(
             wfg_kwargs=dataset_settings["waveform_generator"],
             wfg_domain=wfg_domain,
             data_domain=self.domain,
@@ -293,7 +314,6 @@ def likelihood(
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
             calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-            phase_grid=phase_grid,
             use_base_domain=use_base_domain,
             frequency_update=dict(
                 minimum_frequency=self._frequency(
@@ -304,6 +324,9 @@ def likelihood(
                 ),
             ),
         )
+        self._likelihood = likelihood
+        self._likelihood_settings = settings
+        return likelihood
 
     def _frequency(self, key: str, default: float):
         """The event's frequency-range override for `key` (min/max), else `default`."""
@@ -386,7 +409,9 @@ def sample_and_log_prob(self, num_samples, context, given=None):
             raise ValueError(
                 "Synthetic phase is 1:1; draw one phase per sample (fan_out=1)."
             )
-        n = len(next(iter(given.values())))
+        reference = next(iter(given.values()))
+        device = reference.device if torch.is_tensor(reference) else None
+        n = len(reference)
         logger.info(f"Estimating synthetic phase for {n} samples.")
         t0 = time.time()
         phases, phase_posterior = self._phase_profile(given, context)
@@ -394,15 +419,20 @@ def sample_and_log_prob(self, num_samples, context, given=None):
             phases, phase_posterior, self.num_processes
         )
         logger.info(f"Done. This took {time.time() - t0:.2f} s.")
-        return {"phase": torch.as_tensor(new_phase)}, torch.as_tensor(log_prob)
+        return (
+            {"phase": torch.as_tensor(new_phase, device=device)},
+            torch.as_tensor(log_prob, device=device),
+        )
 
     def log_prob(self, theta_i, context, given=None):
         """Evaluate `log q(phase | theta_rest, d)` at the given phases (re-plug / IS)."""
+        reference = next(iter(given.values()))
+        device = reference.device if torch.is_tensor(reference) else None
         phases, phase_posterior = self._phase_profile(given, context)
         log_prob = interpolated_log_prob_multi(
             phases, phase_posterior, _to_numpy(theta_i["phase"]), self.num_processes
         )
-        return torch.as_tensor(log_prob)
+        return torch.as_tensor(log_prob, device=device)
 
     def _phase_profile(self, given, context):
         """The phase grid and the mass-covered (un-normalized) phase distribution, one row
@@ -557,14 +587,21 @@ def sample_and_log_prob(self, num_samples, context, given=None):
 
     def log_prob(self, theta_i, context, given=None):
         """`log p(theta_hat | theta)` from the kernel, at the proxies (`theta_i`) and
-        the detector times (`given`)."""
-        diffs = {}
-        for k in self.kernel.keys():
-            diff = given[k] - theta_i[f"{k}_proxy"]
-            if torch.is_tensor(diff):
-                diff = diff.detach().cpu().numpy()
-            diffs[k] = np.asarray(diff)
-        return torch.as_tensor(self.kernel.ln_prob(diffs, axis=0), dtype=torch.float32)
+        the detector times (`given`).
+
+        The kernel is a bilby `PriorDict` -- the same object that samples the blur --
+        so the density is evaluated in numpy (converting each side first: the times
+        and proxies may live on different devices) and returned on the detector
+        times' device."""
+        reference = next(iter(given.values()))
+        device = reference.device if torch.is_tensor(reference) else None
+        diffs = {
+            k: _to_numpy(given[k]) - _to_numpy(theta_i[f"{k}_proxy"])
+            for k in self.kernel.keys()
+        }
+        return torch.as_tensor(
+            self.kernel.ln_prob(diffs, axis=0), dtype=torch.float32, device=device
+        )
 
 
 class GNPEFlowFactor(Factor):
diff --git a/dingo/gw/likelihood.py b/dingo/gw/likelihood.py
index 49b29941f..43df8f80f 100644
--- a/dingo/gw/likelihood.py
+++ b/dingo/gw/likelihood.py
@@ -39,13 +39,11 @@ def __init__(
         time_marginalization_kwargs: Optional[dict] = None,
         phase_marginalization_kwargs: Optional[dict] = None,
         calibration_marginalization_kwargs: Optional[dict] = None,
-        phase_grid=None,
         use_base_domain: bool = False,
         frequency_update: Optional[
             dict[str, float | dict[str, float | list[float]]]
         ] = None,
     ):
-        # TODO: Does the phase_grid argument ever get used?
         """
         Parameters
         ----------
@@ -152,7 +150,9 @@ def __init__(
             ]
         )
         self.whiten = True
-        self.phase_grid = phase_grid
+        # Assigned by phase-grid consumers (synthetic phase) or set internally by
+        # phase marginalization; never a construction input.
+        self.phase_grid = None
 
         # optionally initialize time marginalization
         self.time_marginalization = False
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 6e3b3bf0f..0e36ba86f 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -303,7 +303,6 @@ def _build_likelihood(
         time_marginalization_kwargs: Optional[dict] = None,
         phase_marginalization_kwargs: Optional[dict] = None,
         calibration_marginalization_kwargs: Optional[dict] = None,
-        phase_grid: Optional[np.ndarray] = None,
     ):
         """
         Build the likelihood function based on model metadata. This is called at the
@@ -381,7 +380,6 @@ def _build_likelihood(
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
             calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-            phase_grid=phase_grid,
             use_base_domain=self.use_base_domain,
             frequency_update=dict(
                 minimum_frequency=self.minimum_frequency,
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 1a9d65101..99f02d072 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -15,6 +15,7 @@
 
 from dingo.core.factors import (
     ChainComposer,
+    DeltaFactor,
     Factor,
     GibbsBlock,
     Reparametrization,
@@ -475,3 +476,51 @@ def test_validation_rejects_column_overwrite():
         ChainComposer([_ConstFactor("a"), _ConstFactor("a")])
     # The in-place reparam is the allowed exception.
     ChainComposer([_ConstFactor("x"), _InPlaceReparam()])
+
+
+class _DeviceContext:
+    """Bare context stub carrying only a device (these steps ignore the rest)."""
+
+    def __init__(self, device):
+        self.device = device
+        self.event_metadata = None
+
+
+def test_delta_factor_creates_on_context_device():
+    # A DeltaFactor creates fresh tensors, so they land on the chain's device;
+    # the 'meta' device stands in for a GPU without needing one.
+    delta = DeltaFactor({"chi_1": 0.0, "chi_2": 0.3})
+    samples, lp = delta.sample_and_log_prob(4, _DeviceContext(torch.device("meta")))
+    assert all(v.device.type == "meta" for v in samples.values())
+    assert lp.device.type == "meta"
+    samples_cpu, lp_cpu = delta.sample_and_log_prob(4, context=None)
+    assert lp_cpu.device.type == "cpu"  # no context device -> torch default
+
+
+_HAS_MPS = getattr(torch.backends, "mps", None) and torch.backends.mps.is_available()
+
+
+@pytest.mark.skipif(not _HAS_MPS, reason="needs an accelerator (MPS) for mixed devices")
+def test_chain_device_consistency_on_accelerator():
+    # A chain mixing a device-emitting factor, an (in-place) reparam, a correction,
+    # and a DeltaFactor filler must stay on one device end to end -- the GPU pipe
+    # case: constants follow the context device, transforms follow their inputs.
+    device = torch.device("mps")
+
+    class _DeviceGauss(_GaussFactor):
+        def sample_and_log_prob(self, n, context, given=None):
+            block, lp = super().sample_and_log_prob(n, context, given)
+            return {k: v.to(device) for k, v in block.items()}, lp.to(device)
+
+    comp = ChainComposer(
+        [
+            _DeviceGauss("x"),
+            _InPlaceReparam(shift=5.0),
+            _MockTargetCorrection(reads="x", emits="corr", consume=False),
+            DeltaFactor({"c": 1.5}),
+        ]
+    )
+    out = comp.sample(8, context=_DeviceContext(device))
+    assert out["log_prob"].device.type == "mps"
+    assert out["c"].device.type == "mps"
+    assert out["x"].device.type == "mps"
diff --git a/tests/core/test_result.py b/tests/core/test_result.py
index 752fd0227..e9a796de7 100644
--- a/tests/core/test_result.py
+++ b/tests/core/test_result.py
@@ -150,9 +150,9 @@ def test_higher_max_increases_output():
     ]
     # Expected output grows monotonically with k (stochastically, but very reliably).
     for i in range(len(counts) - 1):
-        assert counts[i] <= counts[i + 1], (
-            f"k={i+1} gave {counts[i]} samples but k={i+2} gave {counts[i+1]}"
-        )
+        assert (
+            counts[i] <= counts[i + 1]
+        ), f"k={i+1} gave {counts[i]} samples but k={i+2} gave {counts[i+1]}"
 
 
 # ---------------------------------------------------------------------------
@@ -233,3 +233,23 @@ def test_clip_weights_reproducible():
     out1 = result.rejection_sample(clip_weights=True, random_state=42)
     out2 = result.rejection_sample(clip_weights=True, random_state=42)
     pd.testing.assert_frame_equal(out1, out2)
+
+
+# ---------------------------------------------------------------------------
+# sampler_context lifecycle
+# ---------------------------------------------------------------------------
+
+
+def test_reset_event_invalidates_sampler_context():
+    """A live sampler context describes the data the samples were drawn from;
+    resetting to new event data must drop it so builders fall back to
+    self-building rather than reading the stale context."""
+    from types import SimpleNamespace
+
+    samples = pd.DataFrame({"x": [1.0, 2.0]})
+    result = Result(dictionary={"samples": samples}, sampler_context=object())
+    assert result.sampler_context is not None
+    event = SimpleNamespace(data={"waveform": np.zeros(3)}, settings={"f_min": 20.0})
+    result.reset_event(event)
+    assert result.sampler_context is None
+    assert result.event_metadata == {"f_min": 20.0}
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
new file mode 100644
index 000000000..2f97a1cfe
--- /dev/null
+++ b/tests/gw/test_gw_sampler_context.py
@@ -0,0 +1,83 @@
+"""
+Unit tests for ``GWSamplerContext`` plumbing: the default-likelihood cache and the
+device attribute. The likelihood itself is stubbed (no waveform generation); the
+model-based construction parity lives in the local harnesses.
+"""
+
+import pytest
+
+import dingo.gw.inference.factors as factors_module
+from dingo.gw.domains import build_domain
+from dingo.gw.inference.factors import GWSamplerContext
+
+_DOMAIN_SETTINGS = {
+    "type": "FrequencyDomain",
+    "f_min": 20.0,
+    "f_max": 1024.0,
+    "delta_f": 0.25,
+}
+
+_BASE_METADATA = {
+    "dataset_settings": {
+        "domain": _DOMAIN_SETTINGS,
+        "waveform_generator": {"approximant": "IMRPhenomD", "f_ref": 20.0},
+    },
+    "train_settings": {"data": {"inference_parameters": ["chirp_mass"]}},
+}
+
+
+class _StubLikelihood:
+    """Records constructions; stands in for StationaryGaussianGWLikelihood."""
+
+    constructions = 0
+
+    def __init__(self, **kwargs):
+        type(self).constructions += 1
+        self.kwargs = kwargs
+
+
+@pytest.fixture
+def context(monkeypatch):
+    monkeypatch.setattr(
+        factors_module, "StationaryGaussianGWLikelihood", _StubLikelihood
+    )
+    _StubLikelihood.constructions = 0
+    return GWSamplerContext(
+        domain=build_domain(_DOMAIN_SETTINGS),
+        detectors=["H1"],
+        t_ref=1126259462.4,
+        data_prep=None,
+        raw_context={},
+        base_metadata=_BASE_METADATA,
+    )
+
+
+def test_likelihood_rebuilt_only_when_settings_change(context):
+    # Repeated calls with unchanged arguments share the last-built instance (the
+    # synthetic-phase factor requests one per chain chunk); a settings change
+    # builds a replacement.
+    default = context.likelihood()
+    assert context.likelihood() is default
+    assert _StubLikelihood.constructions == 1
+
+    base = context.likelihood(use_base_domain=True)
+    assert base is not default
+    assert base.kwargs["use_base_domain"] is True
+    assert context.likelihood(use_base_domain=True) is base
+    assert _StubLikelihood.constructions == 2
+
+
+def test_context_device_default_and_explicit():
+    ctx = GWSamplerContext(
+        domain=None,
+        detectors=["H1"],
+        t_ref=0.0,
+        data_prep=None,
+        raw_context={},
+        device="meta",
+    )
+    assert ctx.device == "meta"
+    ctx_default = GWSamplerContext(
+        domain=None, detectors=["H1"], t_ref=0.0, data_prep=None, raw_context={}
+    )
+    assert ctx_default.device == "cpu"

From 78f3e7327d0b4eadc45761818119eb61629ad4a2 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 5 Jul 2026 19:26:08 +0100
Subject: [PATCH 25/65] Rename GWSamplerContext.raw_context to event_data

The word "context" meant four things (the sampler context, the network's
conditioning parameters, the Result storage key, and the raw strain/ASD
dict). The new API now reserves it for the sampler context: the data dict is
event_data, matching StationaryGaussianGWLikelihood(event_data=...) and
EventDataset.data. The serialized Result key ("context") and the training
metadata name (context_parameters) are frozen by back-compat and unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py       | 40 ++++++++++++++---------------
 tests/gw/test_gw_sampler_context.py |  6 ++---
 tests/gw/test_synthetic_phase.py    |  2 +-
 3 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 295b63bda..75aa9589f 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -85,7 +85,7 @@ def __init__(
         detectors: list[str],
         t_ref: float,
         data_prep: Compose,
-        raw_context: dict,
+        event_data: dict,
         event_metadata: Optional[dict] = None,
         base_metadata: Optional[dict] = None,
         device: Union[torch.device, str] = "cpu",
@@ -102,7 +102,7 @@ def __init__(
         data_prep : Compose
             The one-time data-preprocessing transform chain (whiten / decimate /
             repackage).
-        raw_context : dict
+        event_data : dict
             The raw event data `d` (strain + ASDs per detector), i.e. `EventDataset.data`.
             Consumed lazily by `prepared_data()` and reused for the likelihood.
         event_metadata : dict, optional
@@ -120,7 +120,7 @@ def __init__(
         self.t_ref = t_ref
         self._data_prep = data_prep
         self.base_metadata = base_metadata
-        self.raw_context = raw_context
+        self.event_data = event_data
         self.event_metadata = event_metadata
         self.device = device
         self._prepared: Optional[torch.Tensor] = None
@@ -132,7 +132,7 @@ def __init__(
     def from_model(
         cls,
         model: BasePosteriorModel,
-        raw_context: dict,
+        event_data: dict,
         event_metadata: Optional[dict] = None,
     ) -> "GWSamplerContext":
         """Build the context (domain + one-time data-prep chain) from a model's metadata.
@@ -143,7 +143,7 @@ def from_model(
         ----------
         model : BasePosteriorModel
             The model whose metadata defines the domain and preprocessing.
-        raw_context : dict
+        event_data : dict
             The raw event data (strain + ASDs).
         event_metadata : dict, optional
             Per-event metadata.
@@ -181,7 +181,7 @@ def from_model(
             detectors=detectors,
             t_ref=data_settings["ref_time"],
             data_prep=Compose(transforms),
-            raw_context=raw_context,
+            event_data=event_data,
             event_metadata=event_metadata,
             base_metadata=meta,
             device=model.device,
@@ -190,7 +190,7 @@ def from_model(
     def prepared_data(self) -> torch.Tensor:
         """One-time data preprocessing, computed once and cached."""
         if self._prepared is None:
-            self._prepared = self._data_prep(self.raw_context)
+            self._prepared = self._data_prep(self.event_data)
         return self._prepared
 
     @property
@@ -222,7 +222,7 @@ def likelihood(
         Build the exact GW likelihood on this event's data, in physical parameter space.
 
         The network's standardized, decimated view of the data is `prepared_data()`; the
-        likelihood instead takes the raw event data (`raw_context`) and builds its own
+        likelihood instead takes the raw event data (`event_data`) and builds its own
         representation -- decimating to the multibanded domain unless `use_base_domain`,
         and masking the ASDs to the event's frequency range. Its reference time is the
         event time (the training-frame right-ascension correction is already applied to the
@@ -309,7 +309,7 @@ def likelihood(
             wfg_kwargs=dataset_settings["waveform_generator"],
             wfg_domain=wfg_domain,
             data_domain=self.domain,
-            event_data=self.raw_context,
+            event_data=self.event_data,
             t_ref=t_ref,
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
@@ -949,7 +949,7 @@ def __init__(
     def from_model(
         cls,
         model: BasePosteriorModel,
-        raw_context: dict,
+        event_data: dict,
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
         """Build a plain-NPE GW sampler from a model and event data: the flow exposes
@@ -959,7 +959,7 @@ def from_model(
         ----------
         model : BasePosteriorModel
             The NPE model.
-        raw_context : dict
+        event_data : dict
             The raw event data (strain + ASDs).
         event_metadata : dict, optional
             Per-event metadata.
@@ -968,7 +968,7 @@ def from_model(
         -------
         GWComposedSampler
         """
-        context = GWSamplerContext.from_model(model, raw_context, event_metadata)
+        context = GWSamplerContext.from_model(model, event_data, event_metadata)
         metadata = _base_model_metadata(model)
         inference_parameters = metadata["train_settings"]["data"][
             "inference_parameters"
@@ -991,7 +991,7 @@ def from_gnpe_models(
         cls,
         init_model: BasePosteriorModel,
         main_model: BasePosteriorModel,
-        raw_context: dict,
+        event_data: dict,
         event_metadata: Optional[dict] = None,
         num_iterations: int = 30,
     ) -> "GWComposedSampler":
@@ -1008,7 +1008,7 @@ def from_gnpe_models(
             preprocessing.
         main_model : BasePosteriorModel
             The GNPE main network.
-        raw_context : dict
+        event_data : dict
             The raw event data (strain + ASDs).
         event_metadata : dict, optional
             Per-event metadata.
@@ -1023,7 +1023,7 @@ def from_gnpe_models(
         # Build the context from the main model: it owns the analysis (likelihood,
         # prior, inference parameters). The init model shares the data domain and
         # preprocessing (asserted above), so prepared_data() is identical either way.
-        context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
+        context = GWSamplerContext.from_model(main_model, event_data, event_metadata)
         metadata = _base_model_metadata(main_model)
         inference_parameters = metadata["train_settings"]["data"][
             "inference_parameters"
@@ -1051,7 +1051,7 @@ def from_singlestep_gnpe(
         cls,
         main_model: BasePosteriorModel,
         proxy_source: Factor,
-        raw_context: dict,
+        event_data: dict,
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
         """Build a single-step (density-preserving) time-GNPE sampler: a `ChainComposer`
@@ -1066,7 +1066,7 @@ def from_singlestep_gnpe(
         proxy_source : Factor
             Supplies the detector-time proxies -- a `DeltaFactor` for prior conditioning
             (BNS), or an unconditional NDE for density recovery.
-        raw_context : dict
+        event_data : dict
             The raw event data (strain + ASDs).
         event_metadata : dict, optional
             Per-event metadata.
@@ -1075,7 +1075,7 @@ def from_singlestep_gnpe(
         -------
         GWComposedSampler
         """
-        context = GWSamplerContext.from_model(main_model, raw_context, event_metadata)
+        context = GWSamplerContext.from_model(main_model, event_data, event_metadata)
         metadata = _base_model_metadata(main_model)
         inference_parameters = metadata["train_settings"]["data"][
             "inference_parameters"
@@ -1096,7 +1096,7 @@ def to_result(self):
         existing post-processing pipeline -- synthetic phase, importance sampling,
         evidence, plotting -- runs on the factorized sampler's output unchanged.
 
-        The raw event-data dict (`GWSamplerContext.raw_context`) is stored as the
+        The raw event-data dict (`GWSamplerContext.event_data`) is stored as the
         `Result` context (serialized), and the live `GWSamplerContext` is passed as
         `sampler_context` so `Result` can pull the prior (and, later, the likelihood)
         from it rather than rebuilding them from metadata.
@@ -1105,7 +1105,7 @@ def to_result(self):
 
         data_dict = {
             "samples": self.samples,
-            "context": self.context.raw_context,
+            "context": self.context.event_data,
             "event_metadata": self.context.event_metadata,
             "importance_sampling_metadata": None,
             "log_evidence": None,
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index 2f97a1cfe..dc2f51e01 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -47,7 +47,7 @@ def context(monkeypatch):
         detectors=["H1"],
         t_ref=1126259462.4,
         data_prep=None,
-        raw_context={},
+        event_data={},
         base_metadata=_BASE_METADATA,
     )
 
@@ -73,11 +73,11 @@ def test_context_device_default_and_explicit():
         detectors=["H1"],
         t_ref=0.0,
         data_prep=None,
-        raw_context={},
+        event_data={},
         device="meta",
     )
     assert ctx.device == "meta"
     ctx_default = GWSamplerContext(
-        domain=None, detectors=["H1"], t_ref=0.0, data_prep=None, raw_context={}
+        domain=None, detectors=["H1"], t_ref=0.0, data_prep=None, event_data={}
     )
     assert ctx_default.device == "cpu"
diff --git a/tests/gw/test_synthetic_phase.py b/tests/gw/test_synthetic_phase.py
index cc45224aa..71bec0623 100644
--- a/tests/gw/test_synthetic_phase.py
+++ b/tests/gw/test_synthetic_phase.py
@@ -119,7 +119,7 @@ def context(event_metadata):
             detectors=[],
             t_ref=0.0,
             data_prep=None,
-            raw_context={},
+            event_data={},
             event_metadata=event_metadata,
         )
 

From 602e5e64fa9915902f2040dae66af0a815591449 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 5 Jul 2026 21:28:50 +0100
Subject: [PATCH 26/65] Support unconditional (density-recovery) models in
 FlowFactor

An unconditional NDE (unconditional: True in its training metadata) takes no
network input: FlowFactor now draws it with no arguments and never touches
the context, mirroring the old Sampler._run_sampler branch. This is the
proxy source for density recovery on the spine: prepare_log_prob's re-sample
half becomes from_singlestep_gnpe(main, FlowFactor.from_model(nde)).

Metadata transparency, restoring the old two-handle split: the network-bound
settings (standardization, inference_parameters, context_parameters) are
always the model's OWN -- an unconditional NDE carries its own proxy
standardization, distinct from the base model's under metadata["base"].
FlowFactor and GNPEFlowFactor now read model.metadata directly (previously
_base_model_metadata, which for an NDE would have silently applied the base
model's standardization); _base_model_metadata keeps the analysis-side role
(domain / dataset / detector settings) and its docstring now states the rule.

- Mock CI test with a decoy base standardization (using it fails loudly).
- Parity harness (verify_unconditional_flowfactor.py, local models): replays
  prepare_log_prob end to end -- Gibbs run, train a small NDE on the proxies,
  then (A) FlowFactor(nde) vs GWSampler(nde)._run_sampler bit-exact, and
  (B) from_singlestep_gnpe(main, FlowFactor(nde)) vs the old
  GWSamplerGNPE(num_iterations=1, init_sampler=GWSampler(nde)) -- all 18
  shared columns bit-exact. (The old full run_sampler() path is itself broken
  for unconditional models -- _post_process reads metadata[dataset_settings];
  production only ever calls _run_sampler, which is the reference used.)
- Note the embed-once gap for row-identical data with row-varying
  conditioning (intrinsic/extrinsic split) as a TODO at the conditioned path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py         | 67 ++++++++++++++++++++++------------
 dingo/gw/inference/factors.py |  3 +-
 tests/core/test_factors.py    | 68 +++++++++++++++++++++++++++++++++++
 3 files changed, 114 insertions(+), 24 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index eb3269bc8..42da66612 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -177,9 +177,11 @@ def log_prob(
 
 
 def _base_model_metadata(model: BasePosteriorModel) -> dict:
-    """The training metadata describing the parameter standardization / data settings.
-    For an unconditional (density-recovery) model this lives under `metadata["base"]`.
-    """
+    """The base *analysis* metadata (dataset / domain / detector / data settings): for
+    an unconditional (density-recovery) model this lives under `metadata["base"]`, the
+    metadata of the model whose samples it was trained on. The network-bound settings
+    (`standardization`, `inference_parameters`) are always the model's own and are read
+    from `model.metadata` directly."""
     metadata = model.metadata
     if metadata["train_settings"]["data"].get("unconditional", False):
         return metadata["base"]
@@ -188,12 +190,14 @@ def _base_model_metadata(model: BasePosteriorModel) -> dict:
 
 class FlowFactor(Factor):
     """
-    Factor wrapping a posterior model (NPE flow, FMPE, ...). Reaches data only through
-    `SamplerContext.prepared_data()` and encapsulates the network's standardization, so
-    its interface is in physical parameter space.
-
-    An unconditioned model draws from the shared data context; a model with
-    `context_parameters` conditions on the values in `given`.
+    Factor wrapping a posterior model (NPE flow, FMPE, ...). Encapsulates the network's
+    standardization, so its interface is in physical parameter space.
+
+    Three conditioning shapes: a data-conditional model draws from the shared
+    `SamplerContext.prepared_data()`; a model with `context_parameters` (GNPE proxies)
+    additionally conditions on the values in `given`; and an *unconditional* model
+    (`unconditional` in its training metadata -- a density-recovery NDE or a
+    model-as-prior) takes no input at all and never touches the context.
     """
 
     def __init__(
@@ -230,7 +234,11 @@ def __init__(
         self.conditioning = conditioning or []
         # Network conditioning inputs (GNPE proxies). Empty for plain NPE.
         self.context_parameters = context_parameters or []
-        std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
+        data_settings = model.metadata["train_settings"]["data"]
+        self.unconditional = data_settings.get("unconditional", False)
+        # The model's *own* standardization: an unconditional (density-recovery) NDE
+        # carries its own, distinct from the base model's under metadata["base"].
+        std = data_settings["standardization"]
         self.standardization = Standardization(std["mean"], std["std"])
 
     @classmethod
@@ -238,7 +246,8 @@ def from_model(
         cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
     ) -> "FlowFactor":
         """Build a factor from a model, reading `parameters` and `context_parameters`
-        from its training metadata.
+        from its own training metadata (for an unconditional NDE these are its own,
+        e.g. the GNPE proxies it was trained on).
 
         Parameters
         ----------
@@ -252,7 +261,7 @@ def from_model(
         -------
         FlowFactor
         """
-        data_settings = _base_model_metadata(model)["train_settings"]["data"]
+        data_settings = model.metadata["train_settings"]["data"]
         context_parameters = data_settings.get("context_parameters") or []
         return cls(
             model=model,
@@ -263,12 +272,16 @@ def from_model(
         )
 
     def sample_and_log_prob(self, num_samples, context, given=None):
-        """Draw `num_samples` samples per conditioning row. Unconditioned: `num_samples`
-        draws from the shared data context. Conditioned: `num_samples` per context row,
-        returning `n_rows * num_samples` rows in row-major order."""
-        data = context.prepared_data()
+        """Draw `num_samples` samples per conditioning row. Unconditional: draws with
+        no input (the context is not touched). Data-conditional: `num_samples` draws
+        from the shared data context. Parameter-conditioned: `num_samples` per context
+        row, returning `n_rows * num_samples` rows in row-major order."""
         self.model.network.eval()
-        if not self.context_parameters:
+        if self.unconditional:
+            with torch.no_grad():
+                z, log_prob = self.model.sample_and_log_prob(num_samples=num_samples)
+        elif not self.context_parameters:
+            data = context.prepared_data()
             with torch.no_grad():
                 z, log_prob = self.model.sample_and_log_prob(
                     data.unsqueeze(0), num_samples=num_samples
@@ -277,9 +290,14 @@ def sample_and_log_prob(self, num_samples, context, given=None):
             z = z.squeeze(0)
             log_prob = log_prob.squeeze(0)
         else:
+            data = context.prepared_data()
             # Standardize the (physical) conditioning the network was trained on, giving
             # B = N context rows; expand the shared data to match. The network draws
             # num_samples per row -> (N, num_samples, dim).
+            # TODO: the embedding runs once per row, so row-identical data with
+            # row-varying conditioning (the intrinsic/extrinsic split) recomputes N
+            # identical data embeddings. Fixing this needs an embed/fuse split on the
+            # model (FlowWrapper) so the cached embedding can be reused across rows.
             ctx = self.standardization.standardize(given, self.context_parameters)
             n_rows = ctx.shape[0]
             data = data.expand(n_rows, *data.shape)
@@ -302,14 +320,17 @@ def log_prob(self, theta_i, context, given=None):
         }
         num_samples = next(iter(theta_net.values())).shape[0]
         z = self.standardization.standardize(theta_net, self._net_parameters)
-        data = context.prepared_data()
-        data = data.expand(num_samples, *data.shape)
         net_context: tuple[torch.Tensor, ...]
-        if not self.context_parameters:
-            net_context = (data,)
+        if self.unconditional:
+            net_context = ()
         else:
-            ctx = self.standardization.standardize(given, self.context_parameters)
-            net_context = (data, ctx)
+            data = context.prepared_data()
+            data = data.expand(num_samples, *data.shape)
+            if not self.context_parameters:
+                net_context = (data,)
+            else:
+                ctx = self.standardization.standardize(given, self.context_parameters)
+                net_context = (data, ctx)
         self.model.network.eval()
         with torch.no_grad():
             log_prob = self.model.log_prob(z, *net_context)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 75aa9589f..53ed63fb4 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -657,7 +657,8 @@ def __init__(
         self.conditioning = list(self.proxy_parameters)
         # For log_prob: the sampling path de-standardizes (and corrects the log-prob)
         # inside transform_post, but evaluating at a point needs the forward map too.
-        std = _base_model_metadata(model)["train_settings"]["data"]["standardization"]
+        # The model's own standardization (network-bound, like FlowFactor's).
+        std = model.metadata["train_settings"]["data"]["standardization"]
         self.standardization = Standardization(std["mean"], std["std"])
 
     @property
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 99f02d072..8e926a72b 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -17,6 +17,7 @@
     ChainComposer,
     DeltaFactor,
     Factor,
+    FlowFactor,
     GibbsBlock,
     Reparametrization,
     Stage,
@@ -524,3 +525,70 @@ def sample_and_log_prob(self, n, context, given=None):
     assert out["log_prob"].device.type == "mps"
     assert out["c"].device.type == "mps"
     assert out["x"].device.type == "mps"
+
+
+class _FakeUncondModel:
+    """A fake unconditional (density-recovery) model: a deterministic 'flow' whose
+    draws are arange-based, with its OWN standardization distinct from the decoy
+    base-model standardization under metadata["base"] -- the factor must read the
+    model's own metadata."""
+
+    def __init__(self):
+        from types import SimpleNamespace
+
+        self.metadata = {
+            "train_settings": {
+                "data": {
+                    "unconditional": True,
+                    "inference_parameters": ["H1_time_proxy"],
+                    "standardization": {
+                        "mean": {"H1_time_proxy": 2.0},
+                        "std": {"H1_time_proxy": 0.5},
+                    },
+                }
+            },
+            # Decoy: using the base model's settings here would be a bug.
+            "base": {
+                "train_settings": {
+                    "data": {
+                        "inference_parameters": ["x"],
+                        "standardization": {"mean": {"x": 0.0}, "std": {"x": 1.0}},
+                    }
+                }
+            },
+        }
+        self.network = SimpleNamespace(eval=lambda: None)
+
+    def sample_and_log_prob(self, num_samples=None):
+        z = torch.arange(num_samples, dtype=torch.float32).unsqueeze(-1)
+        lp = torch.full((num_samples,), -1.0)
+        return z, lp
+
+    def log_prob(self, z):
+        return torch.full((z.shape[0],), -1.0)
+
+
+def test_unconditional_flow_factor():
+    # A density-recovery NDE takes no input: the factor never touches the context
+    # (context=None), reads the model's own standardization and parameters, and
+    # serves as a chain root (e.g. the single-step GNPE proxy source).
+    factor = FlowFactor.from_model(_FakeUncondModel())
+    assert factor.unconditional
+    assert factor.parameters == ["H1_time_proxy"]
+    assert factor.conditioning == []
+
+    samples, lp = factor.sample_and_log_prob(4, context=None)
+    # z = arange(4), destandardized with the NDE's own mean/std: z * 0.5 + 2.0.
+    expected = torch.arange(4, dtype=torch.float32) * 0.5 + 2.0
+    assert torch.equal(samples["H1_time_proxy"], expected)
+    # Physical-space density: base lp plus the NDE's own -sum(log std).
+    assert torch.allclose(lp, torch.full((4,), -1.0 - math.log(0.5)))
+
+    # Re-plug through the no-context path.
+    lp2 = factor.log_prob(samples, context=None)
+    assert torch.allclose(lp2, lp)
+
+    # Composes as a chain root with a conditioned factor downstream.
+    comp = ChainComposer([factor, _ConstFactor("y", conditioning=["H1_time_proxy"])])
+    out = comp.sample(3, context=None)
+    assert out["y"].shape == (3,)

From 334fae1183d7672a4f257fedb5e81c2e9ab4e849 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 5 Jul 2026 23:52:11 +0100
Subject: [PATCH 27/65] Wire composed NPE sampling into dingo_pipe behind an
 implementation switch

Add --sampler-implementation {legacy, composed} (default legacy) to the
sampler options. The sampling stage builds a GWComposedSampler for plain
NPE when composed is selected (GNPE not yet supported on this path), and
GWComposedSampler gains the to_hdf5 shim mirroring Sampler.to_hdf5.

A/B parity vs legacy on GW200129: sampling-stage Result HDF5 bit-exact
(all columns, settings, context; ra stored float32), and the unmodified
importance-sampling stage consumes the composed file with identical
output up to ra-float32 precision (~1e-8 in log evidence).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py             |  7 ++++++
 dingo/pipe/parser.py                      | 27 +++++++++++++++++++----
 dingo/pipe/sampling.py                    | 14 ++++++++++++
 tests/pipe/test_sampler_implementation.py | 27 +++++++++++++++++++++++
 4 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 tests/pipe/test_sampler_implementation.py

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 53ed63fb4..5ebee3240 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -14,6 +14,7 @@
 import copy
 import logging
 import time
+from pathlib import Path
 from typing import Optional, Union
 
 import numpy as np
@@ -1114,3 +1115,9 @@ def to_result(self):
             "settings": copy.deepcopy(self.metadata),
         }
         return Result(dictionary=data_dict, sampler_context=self.context)
+
+    def to_hdf5(self, label="result", outdir="."):
+        """Export via `to_result` and save to `<outdir>/<label>.hdf5`."""
+        result = self.to_result()
+        Path(outdir).mkdir(parents=True, exist_ok=True)
+        result.to_file(file_name=Path(outdir, label + ".hdf5"))
diff --git a/dingo/pipe/parser.py b/dingo/pipe/parser.py
index ed65da6cb..3f081dd39 100644
--- a/dingo/pipe/parser.py
+++ b/dingo/pipe/parser.py
@@ -407,7 +407,10 @@ def create_parser(top_level=True, usage=None):
     # Mutually exclusive: --psd-dict / --asd-dataset
     if _BILBY_PIPE_V1_8:
         det_parser.add(
-            "--psd-dict", type=nonestr, default=None, help="Dictionary of PSD files to use"
+            "--psd-dict",
+            type=nonestr,
+            default=None,
+            help="Dictionary of PSD files to use",
         )
         det_parser.add(
             "--asd-dataset",
@@ -419,7 +422,10 @@ def create_parser(top_level=True, usage=None):
     else:
         psd_dict_parser = det_parser.add_mutually_exclusive_group()
         psd_dict_parser.add(
-            "--psd-dict", type=nonestr, default=None, help="Dictionary of PSD files to use"
+            "--psd-dict",
+            type=nonestr,
+            default=None,
+            help="Dictionary of PSD files to use",
         )
         psd_dict_parser.add(
             "--asd-dataset",
@@ -1253,7 +1259,9 @@ def create_parser(top_level=True, usage=None):
     )
     # Mutually exclusive: --prior-file / --prior-dict
     if _BILBY_PIPE_V1_8:
-        prior_parser.add("--prior-file", type=nonestr, default=None, help="The prior file")
+        prior_parser.add(
+            "--prior-file", type=nonestr, default=None, help="The prior file"
+        )
         prior_parser.add(
             "--prior-dict",
             type=nonestr,
@@ -1268,7 +1276,9 @@ def create_parser(top_level=True, usage=None):
         )
     else:
         prior_parser_main = prior_parser.add_mutually_exclusive_group()
-        prior_parser_main.add("--prior-file", type=nonestr, default=None, help="The prior file")
+        prior_parser_main.add(
+            "--prior-file", type=nonestr, default=None, help="The prior file"
+        )
         prior_parser_main.add(
             "--prior-dict",
             type=nonestr,
@@ -1522,6 +1532,15 @@ def create_parser(top_level=True, usage=None):
         help="Reference time for neural network. Do not add this manually as it is "
         "specified by the network.",
     )
+    sampler_parser.add(
+        "--sampler-implementation",
+        type=str,
+        default="legacy",
+        choices=["legacy", "composed"],
+        help="Sampler implementation: 'legacy' (GWSampler / GWSamplerGNPE) or "
+        "'composed' (the factorized GWComposedSampler; experimental, currently "
+        "plain NPE only). Default 'legacy'.",
+    )
     sampler_parser.add(
         "--recover-log-prob",
         action=StoreBoolean,
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index 9138066bc..8965d18f3 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -14,6 +14,7 @@
 
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
 from dingo.gw.data.event_dataset import EventDataset
+from dingo.gw.inference.factors import GWComposedSampler
 from dingo.gw.inference.gw_samplers import GWSampler, GWSamplerGNPE
 from dingo.gw.inference.inference_utils import prepare_log_prob
 from dingo.pipe.default_settings import DENSITY_RECOVERY_SETTINGS
@@ -52,6 +53,7 @@ def __init__(self, args, unknown_args):
         self.model_init = args.model_init and (
             resolve_filename_with_transfer_fallback(args.model_init) or args.model_init
         )
+        self.sampler_implementation = args.sampler_implementation
         self.recover_log_prob = args.recover_log_prob
         self.device = args.device
         self.num_gnpe_iterations = args.num_gnpe_iterations
@@ -119,6 +121,18 @@ def _load_sampler(self):
             filename=self.model, device=self.device, load_training_info=False
         )
 
+        if self.sampler_implementation == "composed":
+            if self.model_init is not None:
+                raise NotImplementedError(
+                    "sampler-implementation = composed does not support GNPE models "
+                    "yet. Use sampler-implementation = legacy."
+                )
+            self.gnpe = False
+            self.dingo_sampler = GWComposedSampler.from_model(
+                model, self.context, event_metadata=self.event_metadata
+            )
+            return
+
         if self.model_init is not None:
             self.gnpe = True
             init_model = build_model_from_kwargs(
diff --git a/tests/pipe/test_sampler_implementation.py b/tests/pipe/test_sampler_implementation.py
new file mode 100644
index 000000000..6ecda5453
--- /dev/null
+++ b/tests/pipe/test_sampler_implementation.py
@@ -0,0 +1,27 @@
+"""Tests for the sampler-implementation switch in the dingo_pipe parser."""
+
+import pytest
+from bilby_pipe.utils import parse_args
+
+from dingo.pipe.parser import create_parser
+
+
+def _parse(tmp_path, extra=()):
+    ini = tmp_path / "test.ini"
+    ini.write_text("model = model.pt\n")
+    args, _ = parse_args([str(ini), *extra], create_parser(top_level=False))
+    return args
+
+
+def test_default_is_legacy(tmp_path):
+    assert _parse(tmp_path).sampler_implementation == "legacy"
+
+
+def test_composed_parses(tmp_path):
+    args = _parse(tmp_path, ["--sampler-implementation", "composed"])
+    assert args.sampler_implementation == "composed"
+
+
+def test_invalid_value_rejected(tmp_path):
+    with pytest.raises(SystemExit):
+        _parse(tmp_path, ["--sampler-implementation", "bogus"])

From bb2575a5775efd9a99a0e3267d1323f5dd28bd47 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 5 Jul 2026 23:52:11 +0100
Subject: [PATCH 28/65] Accept Path as well as str in DingoDataset file APIs

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/dataset.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/dingo/core/dataset.py b/dingo/core/dataset.py
index 49f9718b8..ef61fa6fc 100644
--- a/dingo/core/dataset.py
+++ b/dingo/core/dataset.py
@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import List, Optional, Union
 import ast
 import h5py
@@ -93,7 +94,7 @@ class DingoDataset:
 
     def __init__(
         self,
-        file_name: Optional[str] = None,
+        file_name: Optional[Union[str, Path]] = None,
         dictionary: Optional[dict] = None,
         data_keys: Optional[List] = None,
         leave_on_disk_keys: Optional[list] = None,
@@ -104,7 +105,7 @@ def __init__(
 
         Parameters
         ----------
-        file_name : str
+        file_name : str or Path
             HDF5 file containing a dataset
         dictionary : dict
             Contains settings and data entries. The data keys should be the same as
@@ -135,7 +136,7 @@ def __init__(
         elif dictionary is not None:
             self.from_dictionary(dictionary)
 
-    def to_file(self, file_name: str, mode: str = "w"):
+    def to_file(self, file_name: Union[str, Path], mode: str = "w"):
         print("Saving dataset to " + str(file_name))
         save_dict = {
             k: v
@@ -149,7 +150,7 @@ def to_file(self, file_name: str, mode: str = "w"):
             if self.dataset_type:
                 f.attrs["dataset_type"] = self.dataset_type
 
-    def from_file(self, file_name: str):
+    def from_file(self, file_name: Union[str, Path]):
         print(f"Loading dataset from {str(file_name)}.")
         if self._leave_on_disk_keys:
             print(f"Omitting data keys {self._leave_on_disk_keys}.")

From 7e83f8d877ec53a0d0744be20bf0e78fe5e5f5e0 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 00:14:46 +0100
Subject: [PATCH 29/65] Wire composed GNPE with density recovery into
 dingo_pipe

With sampler-implementation = composed and a model-init, the sampling
stage builds the Gibbs chain via from_gnpe_models. Log-prob recovery
mirrors prepare_log_prob on the spine: run the Gibbs chain, train the
unconditional proxy NDE (train_unconditional_flow, unchanged), then swap
in the density-preserving single-step sampler with the NDE as proxy
source (from_singlestep_gnpe).

A/B parity vs legacy on the GW200129 toy GNPE pair: one seed carries
bit-exact through Gibbs, NDE training, and the final single-step run
(all shared columns including log_prob and delta_log_prob_target), and
the unmodified importance-sampling stage (approximation-22 synthetic
phase) matches to ra-float32 precision. Legacy-only provenance keys in
the Result settings (init_model, num_iterations) are accepted as a gap
for now; proper provenance comes with the serialization rework.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/pipe/parser.py   |  4 +--
 dingo/pipe/sampling.py | 67 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/dingo/pipe/parser.py b/dingo/pipe/parser.py
index 3f081dd39..de066d493 100644
--- a/dingo/pipe/parser.py
+++ b/dingo/pipe/parser.py
@@ -1538,8 +1538,8 @@ def create_parser(top_level=True, usage=None):
         default="legacy",
         choices=["legacy", "composed"],
         help="Sampler implementation: 'legacy' (GWSampler / GWSamplerGNPE) or "
-        "'composed' (the factorized GWComposedSampler; experimental, currently "
-        "plain NPE only). Default 'legacy'.",
+        "'composed' (the factorized GWComposedSampler; experimental). "
+        "Default 'legacy'.",
     )
     sampler_parser.add(
         "--recover-log-prob",
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index 8965d18f3..cb5233a2a 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 """Script to sample from a Dingo model. Based on bilby_pipe data analysis script."""
+import copy
 import sys
 from pathlib import Path
 import os
@@ -12,9 +13,10 @@
     resolve_filename_with_transfer_fallback,
 )
 
+from dingo.core.factors import FlowFactor
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
 from dingo.gw.data.event_dataset import EventDataset
-from dingo.gw.inference.factors import GWComposedSampler
+from dingo.gw.inference.factors import GNPEKernelFactor, GWComposedSampler
 from dingo.gw.inference.gw_samplers import GWSampler, GWSamplerGNPE
 from dingo.gw.inference.inference_utils import prepare_log_prob
 from dingo.pipe.default_settings import DENSITY_RECOVERY_SETTINGS
@@ -123,14 +125,25 @@ def _load_sampler(self):
 
         if self.sampler_implementation == "composed":
             if self.model_init is not None:
-                raise NotImplementedError(
-                    "sampler-implementation = composed does not support GNPE models "
-                    "yet. Use sampler-implementation = legacy."
+                self.gnpe = True
+                self._main_model = model
+                init_model = build_model_from_kwargs(
+                    filename=self.model_init,
+                    device=self.device,
+                    load_training_info=False,
+                )
+                self.dingo_sampler = GWComposedSampler.from_gnpe_models(
+                    init_model,
+                    model,
+                    self.context,
+                    event_metadata=self.event_metadata,
+                    num_iterations=self.num_gnpe_iterations,
+                )
+            else:
+                self.gnpe = False
+                self.dingo_sampler = GWComposedSampler.from_model(
+                    model, self.context, event_metadata=self.event_metadata
                 )
-            self.gnpe = False
-            self.dingo_sampler = GWComposedSampler.from_model(
-                model, self.context, event_metadata=self.event_metadata
-            )
             return
 
         if self.model_init is not None:
@@ -194,6 +207,31 @@ def density_recovery_settings(self, settings):
     #     result_dir = os.path.join(self.outdir, "result")
     #     return os.path.relpath(result_dir)
 
+    def _recover_log_prob_composed(self):
+        """Density recovery on the composed sampler, mirroring `prepare_log_prob`:
+        run the Gibbs chain, train an unconditional NDE on the GNPE proxies, then
+        swap in a single-step (density-preserving) sampler with the NDE as the
+        proxy source."""
+        settings = copy.deepcopy(self.density_recovery_settings)
+        self.dingo_sampler.run_sampler(
+            settings["num_samples"], batch_size=self.batch_size
+        )
+        result = self.dingo_sampler.to_result()
+        nde_settings = settings["nde_settings"]
+        nde_settings["training"]["device"] = str(self._main_model.device)
+        proxy_parameters = GNPEKernelFactor.from_model(self._main_model).parameters
+        unconditional_model = result.train_unconditional_flow(
+            proxy_parameters,
+            nde_settings,
+            threshold_std=settings.get("threshold_std", float("inf")),
+        )
+        self.dingo_sampler = GWComposedSampler.from_singlestep_gnpe(
+            self._main_model,
+            FlowFactor.from_model(unconditional_model),
+            self.context,
+            event_metadata=self.event_metadata,
+        )
+
     def run_sampler(self):
         if self.gnpe and self.recover_log_prob:
             logger.info(
@@ -202,11 +240,14 @@ def run_sampler(self):
             )
 
             # Note that this will not save any low latency samples at present.
-            prepare_log_prob(
-                self.dingo_sampler,
-                batch_size=self.batch_size,
-                **self.density_recovery_settings,
-            )
+            if self.sampler_implementation == "composed":
+                self._recover_log_prob_composed()
+            else:
+                prepare_log_prob(
+                    self.dingo_sampler,
+                    batch_size=self.batch_size,
+                    **self.density_recovery_settings,
+                )
 
         self.dingo_sampler.run_sampler(self.num_samples, batch_size=self.batch_size)
         self.dingo_sampler.to_hdf5(label=self.label, outdir=self.result_directory)

From 759fbb0b1d52f5b30b01244458bbe1b4e5442951 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 10:20:05 +0100
Subject: [PATCH 30/65] Reconstruct the sampler context from serialized results

Add GWSamplerContext.from_model_metadata, which builds the context from
a model-metadata dict (e.g. a saved Result.settings) -- no model needed:
the domain, data-prep chain, prior, and likelihood are all defined by
the metadata, and the device only sets where prepared_data() lives.
from_model becomes a thin wrapper that reads the model's own metadata
(data preparation is network-bound) and rejects unconditional models,
which prepare no data.

Result now reconstructs the context from its serialized payload when it
is not handed a live one -- on loading from file, and around the new
event payload after reset_event() -- via the _build_context hook. Prior
construction (and later the likelihood) thereby delegates to the context
no matter how the Result was born; the self-build paths remain as a
fallback for non-model payloads such as parameter subsets.

Parity: pipe NPE and GNPE sampling + importance-sampling stages, prior
delegation, and integrated-sampler harnesses all bit-exact, with IS
evidence identical to previously recorded values.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/result.py                | 25 +++++++----
 dingo/gw/inference/factors.py       | 68 +++++++++++++++++++++++------
 dingo/gw/result.py                  | 20 +++++++++
 tests/gw/test_gw_sampler_context.py | 58 ++++++++++++++++++++++--
 4 files changed, 146 insertions(+), 25 deletions(-)

diff --git a/dingo/core/result.py b/dingo/core/result.py
index 1490c63b0..2a2fea1c8 100644
--- a/dingo/core/result.py
+++ b/dingo/core/result.py
@@ -87,10 +87,10 @@ def __init__(self, file_name=None, dictionary=None, sampler_context=None):
         self.context = None
         self.samples = None
         self.log_noise_evidence = None
-        # Live sampler context (a GWSamplerContext), set when a sampler builds the
-        # Result rather than loading it from file. When present, _build_prior (and,
-        # later, _build_likelihood) delegate to it; it is not serialized, so a Result
-        # loaded from file has sampler_context=None and self-builds as before.
+        # Sampler context (a GWSamplerContext): passed in live when a sampler builds
+        # the Result, and otherwise reconstructed from the serialized payload by
+        # _build_context(), so that prior (and, later, likelihood) construction
+        # delegates to it no matter how the Result was born.
         self.sampler_context = sampler_context
         super().__init__(
             file_name=file_name,
@@ -102,6 +102,8 @@ def __init__(self, file_name=None, dictionary=None, sampler_context=None):
         if self.importance_sampling_metadata is None:
             self.importance_sampling_metadata = {}
 
+        if self.sampler_context is None:
+            self.sampler_context = self._build_context()
         self._build_prior()
         self._build_domain()
         if self.importance_sampling_metadata.get("updates"):
@@ -150,6 +152,12 @@ def _build_prior(self):
     def _build_likelihood(self, **likelihood_kwargs):
         self.likelihood = None
 
+    def _build_context(self):
+        """Reconstruct the sampler context from the serialized payload; overridden
+        by domain-specific subclasses. Called when no live context was passed in,
+        and again after reset_event()."""
+        return None
+
     def reset_event(self, event_dataset):
         """
         Set the Result context and event_metadata based on an EventDataset.
@@ -164,11 +172,6 @@ def reset_event(self, event_dataset):
         event_dataset: EventDataset
             New event to be used for importance sampling.
         """
-        # The live sampler context describes the data the samples were drawn from;
-        # after resetting to (possibly regenerated) event data it no longer describes
-        # self.context, so drop it and fall back to self-building.
-        self.sampler_context = None
-
         context = event_dataset.data
         event_metadata = event_dataset.settings
 
@@ -199,6 +202,10 @@ def reset_event(self, event_dataset):
             self._rebuild_domain(verbose=True)
         self.event_metadata = event_metadata
 
+        # The old context described the data the samples were drawn from; rebuild it
+        # around the new (possibly regenerated) event payload.
+        self.sampler_context = self._build_context()
+
     def _rebuild_domain(self, verbose=False):
         pass
 
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 5ebee3240..86952345b 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -130,33 +130,37 @@ def __init__(
         self._likelihood_settings: Optional[dict] = None
 
     @classmethod
-    def from_model(
+    def from_model_metadata(
         cls,
-        model: BasePosteriorModel,
+        metadata: dict,
         event_data: dict,
         event_metadata: Optional[dict] = None,
+        device: Union[torch.device, str] = "cpu",
     ) -> "GWSamplerContext":
-        """Build the context (domain + one-time data-prep chain) from a model's metadata.
-        The data-prep chain reproduces `GWSampler` preprocessing for the plain-NPE case
-        (parameter-dependent transforms are handled per-factor).
+        """Build the context from a model-metadata dict -- no model required. The
+        domain, one-time data-prep chain, prior, and likelihood are all defined by
+        the metadata (e.g. a saved `Result.settings`); `device` only sets where
+        `prepared_data()` and chain-created tensors live.
 
         Parameters
         ----------
-        model : BasePosteriorModel
-            The model whose metadata defines the domain and preprocessing.
+        metadata : dict
+            Conditional-model metadata (`dataset_settings` + `train_settings`),
+            e.g. the `settings` of a saved `Result`.
         event_data : dict
             The raw event data (strain + ASDs).
         event_metadata : dict, optional
             Per-event metadata.
+        device : torch.device or str, default "cpu"
+            Device for `prepared_data()` and chain-created tensors.
 
         Returns
         -------
         GWSamplerContext
         """
-        meta = _base_model_metadata(model)
-        data_settings = meta["train_settings"]["data"]
+        data_settings = metadata["train_settings"]["data"]
 
-        domain = build_domain(meta["dataset_settings"]["domain"])
+        domain = build_domain(metadata["dataset_settings"]["domain"])
         if "domain_update" in data_settings:
             domain.update(data_settings["domain_update"])
         detectors = data_settings["detectors"]
@@ -173,7 +177,7 @@ def from_model(
         # TODO: frequency-range cropping (MaskDataForFrequencyRangeUpdate) -- follow-up.
         transforms += [
             RepackageStrainsAndASDS(ifos=detectors, first_index=domain.min_idx),
-            ToTorch(device=model.device),
+            ToTorch(device=device),
             GetItem("waveform"),
         ]
 
@@ -184,8 +188,46 @@ def from_model(
             data_prep=Compose(transforms),
             event_data=event_data,
             event_metadata=event_metadata,
-            base_metadata=meta,
-            device=model.device,
+            base_metadata=metadata,
+            device=device,
+        )
+
+    @classmethod
+    def from_model(
+        cls,
+        model: BasePosteriorModel,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+    ) -> "GWSamplerContext":
+        """Build the context from a model: its own metadata and its device. Data
+        preparation is network-bound, so the settings come from `model.metadata`
+        (for a conditional model this equals the base analysis metadata). An
+        unconditional model prepares no data, so no context can be built from one;
+        for the prior/likelihood views alone, use
+        `from_model_metadata(_base_model_metadata(model), ...)`.
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The (conditional) model whose metadata defines the domain and
+            preprocessing.
+        event_data : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+
+        Returns
+        -------
+        GWSamplerContext
+        """
+        if model.metadata["train_settings"]["data"].get("unconditional", False):
+            raise ValueError(
+                "An unconditional model has no data preparation, so a context "
+                "cannot be built from it. For the prior/likelihood views, use "
+                "GWSamplerContext.from_model_metadata(_base_model_metadata(model), ...)."
+            )
+        return cls.from_model_metadata(
+            model.metadata, event_data, event_metadata, device=model.device
         )
 
     def prepared_data(self) -> torch.Tensor:
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 0e36ba86f..bb79f06f6 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -212,6 +212,26 @@ def _rebuild_domain(self, verbose=False):
             if verbose:
                 print("No domain updates found; domain not rebuilt.")
 
+    def _build_context(self):
+        """Reconstruct the per-event sampler context from the serialized payload
+        (settings + event data + event metadata), so that prior (and, later,
+        likelihood) construction delegates to `GWSamplerContext` no matter how the
+        Result was born -- live from a sampler or loaded from file."""
+        if self.context is None or self.settings is None:
+            return None
+        from dingo.gw.inference.factors import GWSamplerContext
+
+        try:
+            # base_metadata resolves the unconditional ("base") indirection, so
+            # density-recovery results reconstruct from the analysis metadata.
+            return GWSamplerContext.from_model_metadata(
+                self.base_metadata, self.context, self.event_metadata
+            )
+        except (KeyError, TypeError):
+            # Settings are not a full model metadata (e.g. parameter subsets); the
+            # self-build paths below still apply.
+            return None
+
     def _build_prior(self):
         """Build the prior based on model metadata. Called by __init__()."""
         if self.sampler_context is not None:
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index dc2f51e01..04a67291d 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -1,9 +1,12 @@
 """
-Unit tests for ``GWSamplerContext`` plumbing: the default-likelihood cache and the
-device attribute. The likelihood itself is stubbed (no waveform generation); the
-model-based construction parity lives in the local harnesses.
+Unit tests for ``GWSamplerContext`` plumbing: the default-likelihood cache, the
+device attribute, and the metadata constructors. The likelihood itself is stubbed
+(no waveform generation); the model-based construction parity lives in the local
+harnesses.
 """
 
+import copy
+
 import pytest
 
 import dingo.gw.inference.factors as factors_module
@@ -81,3 +84,52 @@ def test_context_device_default_and_explicit():
         domain=None, detectors=["H1"], t_ref=0.0, data_prep=None, event_data={}
     )
     assert ctx_default.device == "cpu"
+
+
+# Full conditional-model metadata, as serialized in Result.settings.
+_MODEL_METADATA = {
+    "dataset_settings": {
+        "domain": _DOMAIN_SETTINGS,
+        "waveform_generator": {"approximant": "IMRPhenomD", "f_ref": 20.0},
+    },
+    "train_settings": {
+        "data": {
+            "detectors": ["H1", "L1"],
+            "ref_time": 1126259462.4,
+            "inference_parameters": ["chirp_mass"],
+        }
+    },
+}
+
+
+class _StubModel:
+    metadata = _MODEL_METADATA
+    device = "cpu"
+
+
+def test_from_model_metadata_matches_from_model():
+    ctx_meta = GWSamplerContext.from_model_metadata(_MODEL_METADATA, event_data={})
+    ctx_model = GWSamplerContext.from_model(_StubModel(), event_data={})
+    for ctx in (ctx_meta, ctx_model):
+        assert ctx.detectors == ["H1", "L1"]
+        assert ctx.t_ref == 1126259462.4
+        assert ctx.domain.domain_dict == ctx_meta.domain.domain_dict
+        assert ctx.base_metadata is _MODEL_METADATA
+        assert ctx.device == "cpu"
+        assert ctx._data_prep is not None
+
+
+def test_from_model_rejects_unconditional():
+    nde_metadata = copy.deepcopy(_MODEL_METADATA)
+    nde_metadata["train_settings"]["data"]["unconditional"] = True
+    nde_metadata["base"] = _MODEL_METADATA
+
+    class _StubNDE:
+        metadata = nde_metadata
+        device = "cpu"
+
+    with pytest.raises(ValueError, match="unconditional"):
+        GWSamplerContext.from_model(_StubNDE(), event_data={})
+    # The analysis views remain available from the base metadata.
+    ctx = GWSamplerContext.from_model_metadata(nde_metadata["base"], event_data={})
+    assert ctx.detectors == ["H1", "L1"]

From c01871596814a24b566cf39638489d09d874dc1c Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 11:24:36 +0100
Subject: [PATCH 31/65] Delegate likelihood construction to the sampler context

Result._build_likelihood keeps marginalization validation (against its
evolving prior) and provenance, then delegates construction to
context.likelihood(). Importance-sampling state enters the context as
arguments -- the (possibly rebuilt) data domain, an updated waveform-
generator delta_f (a changed T cannot be handled by domain projection),
and the event frequency range -- with defaults reproducing the context's
own views, so the context stays IS-state-free. Time-marginalization
bounds are filled from the network prior only when the caller does not
provide them; bounds from an updated prior therefore win.

prepared_data() now raises for an event frequency range strictly
narrower than the data domain: the composed data-prep chain does not
implement input masking (MaskDataForFrequencyRangeUpdate) yet, and
failing loudly beats silently feeding the network unmasked data. Bounds
merely wider than the domain are fine and expected -- generation writes
base-domain bounds, which can exceed a multibanded band edge (1099.0 vs
1098.875), a case the legacy path covers with an identity mask.

Parity: full harness battery bit-exact, importance-sampling evidence
identical to previously recorded values; 85 CI tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py       | 89 +++++++++++++++++++++++------
 dingo/gw/result.py                  | 54 +++++++++++------
 tests/gw/test_gw_sampler_context.py | 50 ++++++++++++++++
 3 files changed, 160 insertions(+), 33 deletions(-)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 86952345b..18456cada 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -233,9 +233,37 @@ def from_model(
     def prepared_data(self) -> torch.Tensor:
         """One-time data preprocessing, computed once and cached."""
         if self._prepared is None:
+            self._require_no_frequency_cropping()
             self._prepared = self._data_prep(self.event_data)
         return self._prepared
 
+    def _require_no_frequency_cropping(self):
+        """Frequency-range cropping -- masking the network input to an event range
+        narrower than the network domain (legacy `MaskDataForFrequencyRangeUpdate`)
+        -- is not implemented in the composed data-prep chain yet; fail loudly
+        rather than silently feed the network unmasked data. The likelihood view is
+        unaffected (it applies the range via ASD masking). Event bounds merely
+        *wider* than the data domain are fine (nothing to remove): generation
+        writes base-domain bounds, which can exceed a multibanded domain's
+        quantized band edge (e.g. 1099.0 vs 1098.875)."""
+        if self.event_metadata is None:
+            return
+        for key, bound, narrower in (
+            ("minimum_frequency", self.domain.f_min, lambda v, b: v > b),
+            ("maximum_frequency", self.domain.f_max, lambda v, b: v < b),
+        ):
+            value = self.event_metadata.get(key)
+            if value is None:
+                continue
+            values = value.values() if isinstance(value, dict) else [value]
+            if any(narrower(v, bound) for v in values):
+                raise NotImplementedError(
+                    f"Event {key} = {value} is narrower than the network domain "
+                    f"bound ({bound}): frequency-range cropping is not implemented "
+                    "in the composed data-prep chain yet. "
+                    "Use sampler-implementation = legacy."
+                )
+
     @property
     def prior(self) -> PriorDict:
         """The static prior over all parameters, built once from the model metadata
@@ -260,6 +288,9 @@ def likelihood(
         phase_marginalization_kwargs: Optional[dict] = None,
         calibration_marginalization_kwargs: Optional[dict] = None,
         use_base_domain: bool = False,
+        data_domain=None,
+        wfg_delta_f: Optional[float] = None,
+        frequency_update: Optional[dict] = None,
     ) -> StationaryGaussianGWLikelihood:
         """
         Build the exact GW likelihood on this event's data, in physical parameter space.
@@ -278,14 +309,18 @@ def likelihood(
         a `phase_grid` attribute on the shared instance, as
         `Result.sample_synthetic_phase` does.
 
-        Importance-sampling domain rebuilds (an updated `T` / `delta_f` or frequency range)
-        are not wired here yet: the likelihood uses this context's domain.
+        The context itself is importance-sampling-state-free: IS-time settings
+        updates (a rebuilt data domain, an updated `T` / `delta_f`, an explicit
+        frequency range) enter as arguments, supplied by the IS layer
+        (`Result._build_likelihood`).
 
         Parameters
         ----------
         time_marginalization_kwargs : dict, optional
-            Analytically marginalize over `geocent_time`; `t_lower` / `t_upper` are filled
-            from the network's (uniform) time prior. Requires a time-marginalized network.
+            Analytically marginalize over `geocent_time`. `t_lower` / `t_upper` are
+            filled from the network's (uniform) time prior when not already provided
+            (a caller with an updated prior passes its own bounds). Requires a
+            time-marginalized network.
         phase_marginalization_kwargs : dict, optional
             Analytically marginalize over `phase`. Requires a uniform [0, 2 pi) phase prior.
         calibration_marginalization_kwargs : dict, optional
@@ -293,6 +328,15 @@ def likelihood(
         use_base_domain : bool, default False
             For a multibanded domain, evaluate on the base (undecimated) frequency domain
             rather than the decimated one.
+        data_domain : Domain, optional
+            Override for the data domain (e.g. rebuilt with importance-sampling
+            updates); defaults to this context's domain.
+        wfg_delta_f : float, optional
+            Override for the waveform-generator domain `delta_f` (an updated
+            `T = 1/delta_f` cannot be handled by domain projection).
+        frequency_update : dict, optional
+            `minimum_frequency` / `maximum_frequency` for ASD masking; defaults to
+            the event metadata's values with this context's domain as fallback.
 
         Returns
         -------
@@ -309,7 +353,11 @@ def likelihood(
 
         # Marginalizing over a parameter needs the (uniform) prior the network
         # marginalized over; use it to parameterize the requested marginalization.
-        if time_marginalization_kwargs is not None:
+        # Bounds already provided by the caller (e.g. from an updated prior) win.
+        if time_marginalization_kwargs is not None and not (
+            "t_lower" in time_marginalization_kwargs
+            and "t_upper" in time_marginalization_kwargs
+        ):
             time_prior = self._marginalized_prior("geocent_time")
             if time_prior is None:
                 raise NotImplementedError(
@@ -338,8 +386,12 @@ def likelihood(
 
         dataset_settings = self.base_metadata["dataset_settings"]
         # WaveformGenerator domain -- generally the dataset domain (it may be wider than
-        # the data domain for generation). delta_f / T importance-sampling updates: TODO.
-        wfg_domain = build_domain(dataset_settings["domain"])
+        # the data domain for generation). An updated T = 1/delta_f enters here (it
+        # cannot be handled by domain projection).
+        wfg_domain_dict = dataset_settings["domain"]
+        if wfg_delta_f is not None:
+            wfg_domain_dict = {**wfg_domain_dict, "delta_f": wfg_delta_f}
+        wfg_domain = build_domain(wfg_domain_dict)
 
         # Likelihood reference time: the event time (the training-frame RA correction has
         # already been applied to the samples), falling back to the training reference.
@@ -348,24 +400,29 @@ def likelihood(
         else:
             t_ref = self.t_ref
 
+        if data_domain is None:
+            data_domain = self.domain
+        if frequency_update is None:
+            frequency_update = dict(
+                minimum_frequency=self._frequency(
+                    "minimum_frequency", self.domain.f_min
+                ),
+                maximum_frequency=self._frequency(
+                    "maximum_frequency", self.domain.f_max
+                ),
+            )
+
         likelihood = StationaryGaussianGWLikelihood(
             wfg_kwargs=dataset_settings["waveform_generator"],
             wfg_domain=wfg_domain,
-            data_domain=self.domain,
+            data_domain=data_domain,
             event_data=self.event_data,
             t_ref=t_ref,
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
             calibration_marginalization_kwargs=calibration_marginalization_kwargs,
             use_base_domain=use_base_domain,
-            frequency_update=dict(
-                minimum_frequency=self._frequency(
-                    "minimum_frequency", self.domain.f_min
-                ),
-                maximum_frequency=self._frequency(
-                    "maximum_frequency", self.domain.f_max
-                ),
-            ),
+            frequency_update=frequency_update,
         )
         self._likelihood = likelihood
         self._likelihood_settings = settings
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index bb79f06f6..8940d5f0d 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -369,26 +369,46 @@ def _build_likelihood(
         self.phase_marginalization_kwargs = phase_marginalization_kwargs
         self.calibration_marginalization_kwargs = calibration_marginalization_kwargs
 
-        # Choose the WaveformGenerator domain. This is ultimately projected to the data domain, but generally we
-        # allow it to be different, e.g., to start integrating from lower frequencies than the lower bound of the
-        # likelihood integral. Usually we use the same domain as that used for the WaveformDataset. However,
-        # if we make a change to certain settings during importance sampling, we need to ensure the projection is
-        # still compatible:
+        # The WaveformGenerator domain is ultimately projected to the data domain, but
+        # generally we allow it to be different, e.g., to start integrating from lower
+        # frequencies than the lower bound of the likelihood integral. An updated
+        # T = 1/delta_f cannot be handled by the domain projection, so it enters at
+        # the level of the WaveformGenerator domain.
         #
-        # * delta_f (=1/T): cannot be changed in a domain projection, so update at the level of the
-        #   WaveformGenerator.
-        #
-        # TODO: Add functionality to update other waveform settings, i.e., approximant, generation minimum and
-        #  maximum frequencies, reference frequency, and starting frequency.
+        # TODO: Add functionality to update other waveform settings, i.e., approximant,
+        #  generation minimum and maximum frequencies, reference frequency, and
+        #  starting frequency.
+        wfg_delta_f = None
+        if "T" in (self.importance_sampling_metadata.get("updates") or {}):
+            wfg_delta_f = 1 / self.importance_sampling_metadata["updates"]["T"]
+            print(
+                f"Updating waveform generation delta_f from "
+                f'{self.base_metadata["dataset_settings"]["domain"]["delta_f"]} '
+                f"to {wfg_delta_f}."
+            )
+
+        if self.sampler_context is not None:
+            # Delegate construction to the sampler context (the single owner). The
+            # importance-sampling state -- validated marginalization bounds, the
+            # (possibly rebuilt) data domain, delta_f / frequency updates -- enters
+            # as arguments; the context itself stays IS-state-free.
+            self.likelihood = self.sampler_context.likelihood(
+                time_marginalization_kwargs=time_marginalization_kwargs,
+                phase_marginalization_kwargs=phase_marginalization_kwargs,
+                calibration_marginalization_kwargs=calibration_marginalization_kwargs,
+                use_base_domain=self.use_base_domain,
+                data_domain=self.domain,
+                wfg_delta_f=wfg_delta_f,
+                frequency_update=dict(
+                    minimum_frequency=self.minimum_frequency,
+                    maximum_frequency=self.maximum_frequency,
+                ),
+            )
+            return
 
         wfg_domain_dict = self.base_metadata["dataset_settings"]["domain"].copy()
-        if "updates" in self.importance_sampling_metadata:
-            if "T" in self.importance_sampling_metadata["updates"]:
-                delta_f_new = 1 / self.importance_sampling_metadata["updates"]["T"]
-                print(
-                    f'Updating waveform generation delta_f from {wfg_domain_dict["delta_f"]} to {delta_f_new}.'
-                )
-                wfg_domain_dict["delta_f"] = delta_f_new
+        if wfg_delta_f is not None:
+            wfg_domain_dict["delta_f"] = wfg_delta_f
         wfg_domain = build_domain(wfg_domain_dict)
 
         self.likelihood = StationaryGaussianGWLikelihood(
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index 04a67291d..a14132314 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -70,6 +70,56 @@ def test_likelihood_rebuilt_only_when_settings_change(context):
     assert _StubLikelihood.constructions == 2
 
 
+def test_likelihood_importance_sampling_overrides(context):
+    # IS-state overrides pass through: an explicit data domain, an updated
+    # waveform-generator delta_f, and an explicit frequency range.
+    likelihood = context.likelihood(
+        data_domain="rebuilt-domain",
+        wfg_delta_f=0.5,
+        frequency_update={"minimum_frequency": 21.0, "maximum_frequency": 512.0},
+    )
+    assert likelihood.kwargs["data_domain"] == "rebuilt-domain"
+    assert likelihood.kwargs["wfg_domain"].delta_f == 0.5
+    assert likelihood.kwargs["frequency_update"]["maximum_frequency"] == 512.0
+    # Defaults reproduce the context's own views.
+    default = context.likelihood()
+    assert default.kwargs["data_domain"] is context.domain
+    assert default.kwargs["wfg_domain"].delta_f == _DOMAIN_SETTINGS["delta_f"]
+
+
+def test_prepared_data_rejects_frequency_cropping():
+    # A narrower event frequency range needs input masking, which the composed
+    # data-prep chain does not implement yet -> loud failure. Bounds equal to or
+    # wider than the domain (generation writes base-domain bounds) are fine.
+    def make(event_metadata):
+        return GWSamplerContext(
+            domain=build_domain(_DOMAIN_SETTINGS),
+            detectors=["H1"],
+            t_ref=0.0,
+            data_prep=lambda data: "prepared",
+            event_data={},
+            event_metadata=event_metadata,
+        )
+
+    with pytest.raises(NotImplementedError, match="cropping"):
+        make({"minimum_frequency": 25.0}).prepared_data()
+    with pytest.raises(NotImplementedError, match="cropping"):
+        make({"maximum_frequency": {"H1": 512.0}}).prepared_data()
+    assert make({"minimum_frequency": 20.0}).prepared_data() == "prepared"
+    assert make({"maximum_frequency": 1099.0}).prepared_data() == "prepared"
+    assert make(None).prepared_data() == "prepared"
+
+
+def test_likelihood_caller_marginalization_bounds_win(context):
+    # Bounds provided by the caller (e.g. from an updated prior at the IS layer)
+    # are used as-is; the network-prior fill only runs when they are missing.
+    likelihood = context.likelihood(
+        time_marginalization_kwargs={"n_fft": 2, "t_lower": 1.0, "t_upper": 2.0}
+    )
+    kwargs = likelihood.kwargs["time_marginalization_kwargs"]
+    assert (kwargs["t_lower"], kwargs["t_upper"]) == (1.0, 2.0)
+
+
 def test_context_device_default_and_explicit():
     ctx = GWSamplerContext(
         domain=None,

From 9f6ff752efca16980ea8e58f4cd6e63d7dd238ff Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 13:23:46 +0100
Subject: [PATCH 32/65] Run synthetic phase as a chain over the sampler context

Add SampleTableFactor, a chain root that emits a fixed table of existing
samples together with their stored proposal log-prob: a chain can
thereby continue from previously drawn samples, and the composer's
ordinary log-prob fold yields the joint proposal density -- the legacy
samples["log_prob"] += delta was chain composition all along.

Result.sample_synthetic_phase (forward mode, with a sampler context) now
assembles [SampleTableFactor, SyntheticPhaseFactor] and runs it over the
context; the sample-aware parts stay on Result (phase-prior validation,
the within-prior mask and its RNG order, the out-of-prior scatter, prior
bookkeeping). SyntheticPhaseFactor gains likelihood_kwargs so the
importance-sampling layer selects the likelihood view (use_base_domain,
rebuilt domain, delta_f / frequency updates) through the new shared
Result._likelihood_view_kwargs(); the factor thereby shares the cached
likelihood instance built by _build_likelihood, closing the deferred
use_base_domain gap. The legacy body remains for context-less Results
and the inverse mode.

Parity: synthetic-phase profile bit-exact in both modes; the pipe
importance-sampling harnesses reproduce previously recorded values
exactly; 90 CI tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py            | 48 +++++++++++++++++++
 dingo/gw/inference/factors.py    |  8 +++-
 dingo/gw/result.py               | 79 +++++++++++++++++++++++++++++---
 tests/core/test_factors.py       | 45 ++++++++++++++++++
 tests/gw/test_synthetic_phase.py | 21 +++++++++
 5 files changed, 193 insertions(+), 8 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 42da66612..c7f2e908f 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -372,6 +372,54 @@ def log_prob(self, theta_i, context, given=None):
         return torch.zeros_like(reference_column)
 
 
+class SampleTableFactor(Factor):
+    """
+    Chain root emitting a fixed table of existing samples (with their stored proposal
+    log-prob, if available) instead of drawing new ones.
+
+    This is how a chain continues from previously drawn samples: the
+    importance-sampling side runs its post-sampling steps (e.g. synthetic phase) as a
+    chain rooted in the proposal sample table, and the composer's ordinary log-prob
+    fold then yields the joint proposal density `log q(theta) + log q(extra | theta)`
+    with no special-casing. Without a stored log-prob the chain is density-free (as
+    with a `GibbsBlock`).
+    """
+
+    def __init__(self, table: dict, log_prob=None):
+        """
+        Parameters
+        ----------
+        table : dict
+            The existing samples, one array-like entry per parameter column.
+        log_prob : array-like, optional
+            The stored proposal log-prob of the table rows. If omitted, the chain
+            has no tractable density.
+        """
+        self.table = {k: torch.as_tensor(v) for k, v in table.items()}
+        self.table_log_prob = (
+            torch.as_tensor(log_prob) if log_prob is not None else None
+        )
+        self.parameters = list(self.table)
+        self.conditioning: list[str] = []
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        """Emit the table; `num_samples` must equal the table length (a fixed table
+        cannot be chunked, so run the chain with `batch_size=None`)."""
+        n = len(next(iter(self.table.values())))
+        if num_samples != n:
+            raise ValueError(
+                f"A sample table is fixed: num_samples must equal the table length "
+                f"({n}), got {num_samples}. Run the chain with batch_size=None."
+            )
+        return dict(self.table), self.table_log_prob
+
+    def log_prob(self, theta_i, context, given=None):
+        raise NotImplementedError(
+            "A sample table is not a density; its rows carry their stored log-prob. "
+            "Evaluate log_prob through the chain that produced the samples instead."
+        )
+
+
 class Reparametrization(ABC):
     """
     A deterministic bijection `Step`: it transforms existing parameters (no sampling)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 18456cada..3515d1fb5 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -479,6 +479,7 @@ def __init__(
         approximation_22_mode: bool = False,
         uniform_weight: float = 0.01,
         num_processes: int = 1,
+        likelihood_kwargs: Optional[dict] = None,
     ):
         """
         Parameters
@@ -494,6 +495,10 @@ def __init__(
             Weight of the uniform floor added to the phase distribution for mass coverage.
         num_processes : int, default 1
             Parallel processes for the per-sample likelihood evaluation and phase sampling.
+        likelihood_kwargs : dict, optional
+            Arguments for `context.likelihood()` selecting the likelihood view (e.g.
+            `use_base_domain`, a rebuilt `data_domain`, frequency updates), supplied
+            by the importance-sampling layer. Defaults to the context's own views.
         """
         self.parameters = ["phase"]
         self.conditioning = list(conditioning)
@@ -501,6 +506,7 @@ def __init__(
         self.approximation_22_mode = approximation_22_mode
         self.uniform_weight = uniform_weight
         self.num_processes = num_processes
+        self.likelihood_kwargs = likelihood_kwargs
 
     def sample_and_log_prob(self, num_samples, context, given=None):
         """Draw one phase per `theta_rest` row (`num_samples` must be 1); return the phases
@@ -539,7 +545,7 @@ def _phase_profile(self, given, context):
         per sample: evaluate `log L` on the grid, exponentiate (shifted by the per-row
         max), and add the uniform floor."""
         theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
-        likelihood = context.likelihood()
+        likelihood = context.likelihood(**(self.likelihood_kwargs or {}))
         phases = np.linspace(0, 2 * np.pi, self.n_grid)
         if self.approximation_22_mode:
             # Assume the waveform is (2, 2)-dominated (transforms as exp(2i phase)), so the
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 8940d5f0d..ccc51fdbe 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -318,6 +318,23 @@ def update_prior(self, prior_update):
             # Recalculate the importance-sampling weights and log evidence.
             self._calculate_evidence()
 
+    def _likelihood_view_kwargs(self) -> dict:
+        """The `context.likelihood()` arguments encoding this Result's
+        importance-sampling view: the (possibly rebuilt) data domain, the
+        base-domain choice, and delta_f / frequency updates."""
+        wfg_delta_f = None
+        if "T" in (self.importance_sampling_metadata.get("updates") or {}):
+            wfg_delta_f = 1 / self.importance_sampling_metadata["updates"]["T"]
+        return dict(
+            use_base_domain=self.use_base_domain,
+            data_domain=self.domain,
+            wfg_delta_f=wfg_delta_f,
+            frequency_update=dict(
+                minimum_frequency=self.minimum_frequency,
+                maximum_frequency=self.maximum_frequency,
+            ),
+        )
+
     def _build_likelihood(
         self,
         time_marginalization_kwargs: Optional[dict] = None,
@@ -396,13 +413,7 @@ def _build_likelihood(
                 time_marginalization_kwargs=time_marginalization_kwargs,
                 phase_marginalization_kwargs=phase_marginalization_kwargs,
                 calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-                use_base_domain=self.use_base_domain,
-                data_domain=self.domain,
-                wfg_delta_f=wfg_delta_f,
-                frequency_update=dict(
-                    minimum_frequency=self.minimum_frequency,
-                    maximum_frequency=self.maximum_frequency,
-                ),
+                **self._likelihood_view_kwargs(),
             )
             return
 
@@ -536,6 +547,52 @@ def sample_calibration_parameters(self, calibration_sampling_kwargs: dict):
         # Rebuild the prior (which will include calibration priors from prior_update)
         self._build_prior()
 
+    def _sample_synthetic_phase_chain(
+        self, theta, within_prior, approximation_22_mode, num_processes
+    ):
+        """Synthetic phase as a chain over the sampler context: the root emits the
+        within-prior proposal samples with their stored log-prob, and
+        `SyntheticPhaseFactor` draws `phase` -- the composer's ordinary log-prob fold
+        returns the joint proposal density `log q(theta) + log q(phase | theta, d)`,
+        which replaces the samples' log_prob (the legacy `log_prob += delta`). The
+        Result's likelihood view (base domain, rebuilt domain, frequency updates) is
+        passed to the factor, so it shares the cached likelihood instance built by
+        `_build_likelihood`. Out-of-prior rows keep the legacy convention:
+        `phase = 0`, `log_prob = nan`."""
+        from dingo.core.factors import ChainComposer, SampleTableFactor
+        from dingo.gw.inference.factors import SyntheticPhaseFactor
+
+        theta_within = theta.iloc[np.flatnonzero(within_prior)]
+        table = SampleTableFactor(
+            {k: theta_within[k].to_numpy() for k in theta_within.columns},
+            log_prob=self.samples["log_prob"].to_numpy()[within_prior],
+        )
+        factor = SyntheticPhaseFactor(
+            conditioning=list(theta_within.columns),
+            n_grid=self.synthetic_phase_kwargs["n_grid"],
+            approximation_22_mode=approximation_22_mode,
+            uniform_weight=self.synthetic_phase_kwargs.get("uniform_weight", 0.01),
+            num_processes=num_processes,
+            likelihood_kwargs=self._likelihood_view_kwargs(),
+        )
+        chain = ChainComposer([table, factor])
+        out, log_prob = chain.sample_and_log_prob(
+            len(theta_within), self.sampler_context
+        )
+
+        phase_array = np.full(len(theta), 0.0)
+        phase_array[within_prior] = out["phase"].numpy()
+        log_prob_array = np.full(len(theta), np.nan)
+        log_prob_array[within_prior] = log_prob.numpy()
+        self.samples["phase"] = phase_array
+        self.samples["log_prob"] = log_prob_array
+
+        # Insert the phase prior in the prior, since now the phase is present.
+        self.prior["phase"] = self.phase_prior
+        self.phase_prior = None
+        # reset likelihood for safety (matches the legacy path)
+        self.likelihood = None
+
     def sample_synthetic_phase(
         self,
         synthetic_phase_kwargs,
@@ -627,6 +684,14 @@ def sample_synthetic_phase(
         if not inverse:
             self._build_likelihood()
 
+        if not inverse and self.sampler_context is not None:
+            # Run the synthetic-phase step as a chain over the sampler context.
+            self._sample_synthetic_phase_chain(
+                theta, within_prior, approximation_22_mode, num_processes
+            )
+            print(f"Done. This took {time.time() - t0:.2f} s.")
+            return
+
         if inverse:
             # We estimate the log_prob for given phases, so first save the evaluation
             # points.
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 8e926a72b..485402cda 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -20,6 +20,7 @@
     FlowFactor,
     GibbsBlock,
     Reparametrization,
+    SampleTableFactor,
     Stage,
     TargetCorrection,
     chunk_and_concat,
@@ -204,6 +205,50 @@ def test_gibbs_block_runs_as_a_density_free_step():
     assert torch.equal(out["theta"], torch.arange(6, dtype=torch.float32))
 
 
+def test_sample_table_root_feeds_chain_and_sums_log_prob():
+    # A chain continuing from existing samples: the table root emits the rows with
+    # their stored log-prob, a 1:1 factor adds a block, and the ordinary fold yields
+    # the joint proposal density (stored + delta) -- the importance-sampling shape.
+    stored = torch.tensor([0.5, 0.6, 0.7], dtype=torch.float64)
+    table = SampleTableFactor(
+        {"a": torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)}, log_prob=stored
+    )
+
+    class _PlusOne(Factor):
+        parameters = ["b"]
+        conditioning = ["a"]
+
+        def sample_and_log_prob(self, n, context, given=None):
+            assert n == 1  # 1:1
+            return {"b": given["a"] + 1.0}, torch.full_like(given["a"], 0.25)
+
+        def log_prob(self, theta_i, context, given=None):
+            return torch.full_like(given["a"], 0.25)
+
+    out, lp = ChainComposer([table, _PlusOne()]).sample_and_log_prob(3, None)
+    assert torch.equal(out["a"], torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64))
+    assert torch.equal(out["b"], torch.tensor([2.0, 3.0, 4.0], dtype=torch.float64))
+    assert torch.equal(lp, stored + 0.25)
+
+
+def test_sample_table_requires_full_length():
+    table = SampleTableFactor({"a": torch.arange(3.0)})
+    with pytest.raises(ValueError, match="table length"):
+        table.sample_and_log_prob(2, None)
+
+
+def test_sample_table_without_log_prob_is_density_free():
+    table = SampleTableFactor({"a": torch.arange(3.0)})
+    _, lp = ChainComposer([table]).sample_and_log_prob(3, None)
+    assert lp is None
+
+
+def test_sample_table_log_prob_raises():
+    table = SampleTableFactor({"a": torch.arange(3.0)})
+    with pytest.raises(NotImplementedError, match="stored log-prob"):
+        table.log_prob({"a": torch.zeros(3)}, None)
+
+
 class _MockReparam(Reparametrization):
     """A mock bijection ``u -> v = u + shift`` with a constant nonzero ``log|det J|`` so the
     density contribution (``-log_det``) is checkable."""
diff --git a/tests/gw/test_synthetic_phase.py b/tests/gw/test_synthetic_phase.py
index 71bec0623..e0968db17 100644
--- a/tests/gw/test_synthetic_phase.py
+++ b/tests/gw/test_synthetic_phase.py
@@ -127,3 +127,24 @@ def context(event_metadata):
     assert with_override._frequency("minimum_frequency", 20.0) == 25.0
     assert with_override._frequency("maximum_frequency", 1024.0) == 1024.0
     assert context(None)._frequency("minimum_frequency", 20.0) == 20.0
+
+
+def test_likelihood_kwargs_forwarded_to_context():
+    # The IS layer selects the likelihood view (base domain, rebuilt domain, ...);
+    # the factor forwards it verbatim to context.likelihood().
+    recorded = {}
+
+    class _RecordingContext(_MockContext):
+        def likelihood(self, **kwargs):
+            recorded.update(kwargs)
+            return super().likelihood()
+
+    factor = SyntheticPhaseFactor(
+        conditioning=["chirp_mass"],
+        n_grid=11,
+        approximation_22_mode=True,
+        likelihood_kwargs={"use_base_domain": True},
+    )
+    _seed()
+    factor.sample_and_log_prob(1, _RecordingContext(), _given())
+    assert recorded == {"use_base_domain": True}

From 6aa580838f7b036d615870a066c3dce70e1ca968 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 14:24:03 +0100
Subject: [PATCH 33/65] Fix two latent bugs caught by end-to-end testing

Domain.__eq__ now follows the equality protocol (NotImplemented for
non-Domain operands) instead of assuming other has a domain_dict:
generic container comparisons -- e.g. the likelihood cache comparing a
settings dict holding a Domain against one holding None -- invoke
__eq__ with foreign types and must evaluate unequal, not raise.

The log-probs pipeline plot has been silently missing since bilby_pipe
began normalizing --log_probs to --log-probs at parse time: the
underscore-defined option no longer matched, and the flag was dropped
as an unknown argument without an error. Define the option dash-style;
args.log_probs binds as before, and the plot node's underscore form
normalizes onto it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/domains/base.py            |  7 +++----
 dingo/pipe/plot.py                  |  5 ++++-
 tests/gw/test_gw_sampler_context.py | 16 ++++++++++++++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/dingo/gw/domains/base.py b/dingo/gw/domains/base.py
index fe8fb1e60..77b4f2cc9 100644
--- a/dingo/gw/domains/base.py
+++ b/dingo/gw/domains/base.py
@@ -70,7 +70,6 @@ def domain_dict(self):
         pass
 
     def __eq__(self, other):
-        if self.domain_dict == other.domain_dict:
-            return True
-        else:
-            return False
+        if not isinstance(other, Domain):
+            return NotImplemented
+        return self.domain_dict == other.domain_dict
diff --git a/dingo/pipe/plot.py b/dingo/pipe/plot.py
index 7dc8af756..1eec0690d 100644
--- a/dingo/pipe/plot.py
+++ b/dingo/pipe/plot.py
@@ -23,8 +23,11 @@ def create_parser():
     parser.add(
         "--weights", action="store_true", help="Generate plot of importance " "weights"
     )
+    # Dash-style: bilby_pipe's parse_args normalizes --log_probs to --log-probs
+    # before parsing, so an underscore-defined option is silently dropped as
+    # unknown. argparse still binds the value to args.log_probs.
     parser.add(
-        "--log_probs",
+        "--log-probs",
         action="store_true",
         help="Generate plot of target" "versus proposal log probability",
     )
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index a14132314..660bd2fb1 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -70,6 +70,22 @@ def test_likelihood_rebuilt_only_when_settings_change(context):
     assert _StubLikelihood.constructions == 2
 
 
+def test_likelihood_cache_handles_domain_and_none_settings(context):
+    # A bare call (data_domain=None) followed by an explicit domain (either order)
+    # must be a clean cache miss, not a Domain == None comparison error.
+    bare = context.likelihood()
+    explicit = context.likelihood(data_domain=build_domain(_DOMAIN_SETTINGS))
+    assert explicit is not bare
+    assert context.likelihood() is not explicit
+
+
+def test_domain_eq_is_none_safe():
+    domain = build_domain(_DOMAIN_SETTINGS)
+    assert domain == build_domain(_DOMAIN_SETTINGS)
+    assert domain != None  # noqa: E711
+    assert (domain == 3) is False
+
+
 def test_likelihood_importance_sampling_overrides(context):
     # IS-state overrides pass through: an explicit data domain, an updated
     # waveform-generator delta_f, and an explicit frequency range.

From 0fadbd648b6f5ffee6a372ad7849c426db963a02 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 14:55:39 +0100
Subject: [PATCH 34/65] Move frequency-update validation to
 dingo.gw.frequency_updates

check_frequency_updates and the min/max validators move verbatim out of
gw_samplers.py to a neutral module: dingo_pipe imports them at INI
time, the legacy sampler setters keep using them, and the composed
data-prep cropping port will call them at construction -- they outlive
the legacy samplers. The module header documents the licensing rule the
code embodies: narrowing the frequency range is only permitted when the
network was trained with random strain cropping covering the requested
bounds.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/frequency_updates.py     | 188 ++++++++++++++++++++++++++++++
 dingo/gw/inference/gw_samplers.py | 182 +----------------------------
 dingo/pipe/main.py                |   2 +-
 3 files changed, 194 insertions(+), 178 deletions(-)
 create mode 100644 dingo/gw/frequency_updates.py

diff --git a/dingo/gw/frequency_updates.py b/dingo/gw/frequency_updates.py
new file mode 100644
index 000000000..d51e1000f
--- /dev/null
+++ b/dingo/gw/frequency_updates.py
@@ -0,0 +1,188 @@
+"""
+Validation of event- and importance-sampling-time frequency-range updates against a
+model's frequency domain.
+
+A frequency range narrower than the network's domain is only permitted when the
+network was trained with random strain cropping (`random_strain_cropping` in the
+training data settings), within the bounds that the cropping covered. These
+functions are called at INI-parse time (`dingo_pipe`, on the model metadata) and by
+the samplers when event metadata carries `minimum_frequency` / `maximum_frequency`.
+"""
+
+import numpy as np
+
+from dingo.gw.domains import (
+    MultibandedFrequencyDomain,
+    UniformFrequencyDomain,
+    build_domain_from_model_metadata,
+)
+
+
+def _validate_maximum_frequency(
+    f_max: dict[str, float] | float,
+    detectors: list[str],
+    domain: UniformFrequencyDomain | MultibandedFrequencyDomain,
+    crop_settings: dict | None,
+):
+    if isinstance(f_max, float):
+        f_max = {d: f_max for d in detectors}
+    if set(f_max) != set(detectors):
+        raise ValueError(
+            f"f_max must have exactly detectors {detectors}, got " f"{list(f_max)}."
+        )
+    f_max_vals = np.array([f_max[d] for d in detectors])
+
+    # Hard upper bound
+    if np.any(f_max_vals > domain.f_max):
+        raise ValueError(f"f_max {f_max} > domain.f_max = {domain.f_max}.")
+
+    # Nothing changed
+    if np.all(f_max_vals == domain.f_max):
+        return
+
+    # Cropping must be on
+    if not crop_settings or crop_settings.get("cropping_probability", 0.0) == 0.0:
+        raise ValueError(
+            f"Cropping disabled; cannot lower maximum frequency to {f_max}."
+        )
+
+    # Extract lower bounds
+    floors = crop_settings.get("f_max_lower")
+    if floors is None:
+        floors = domain.f_max
+    if not isinstance(floors, dict):
+        floors = {d: floors for d in detectors}
+
+    # Check lower bound.
+    if not crop_settings.get("independent_detectors", True):
+        if len(set(f_max_vals)) > 1:
+            raise ValueError(
+                f"Independent max frequencies per detector not enabled. "
+                f"All frequencies must match, got f_max = {f_max}."
+            )
+        # TODO: Risk of non-constant floors with non-independent detectors.
+        assert len(set(floors.values())) == 1
+    for d in detectors:
+        if f_max[d] < floors[d]:
+            raise ValueError(
+                f"Maximum frequency requested for {d} ({f_max[d]} Hz) "
+                f"less than lower bound of {floors[d]} Hz."
+            )
+
+
+def _validate_minimum_frequency(
+    f_min: dict[str, float] | float,
+    detectors: list[str],
+    domain: UniformFrequencyDomain | MultibandedFrequencyDomain,
+    crop_settings: dict | None,
+):
+    if isinstance(f_min, float):
+        f_min = {d: f_min for d in detectors}
+    if set(f_min) != set(detectors):
+        raise ValueError(
+            f"f_min must have exactly detectors {detectors}, got {list(f_min)}."
+        )
+    f_min_vals = np.array([f_min[d] for d in detectors])
+
+    # Hard lower bound
+    if np.any(f_min_vals < domain.f_min):
+        raise ValueError(f"f_min {f_min} < domain.f_min = {domain.f_min}.")
+
+    # Nothing changed
+    if np.all(f_min_vals == domain.f_min):
+        return
+
+    # Cropping must be on
+    if not crop_settings or crop_settings.get("cropping_probability", 0.0) == 0.0:
+        raise ValueError(
+            f"Cropping disabled; cannot raise minimum frequency to {f_min}."
+        )
+
+    # Extract upper bounds
+    caps = crop_settings.get("f_min_upper")
+    if caps is None:
+        caps = domain.f_min
+    if not isinstance(caps, dict):
+        caps = {d: caps for d in detectors}
+
+    # Check upper bound.
+    if not crop_settings.get("independent_detectors", True):
+        if len(set(f_min_vals)) > 1:
+            raise ValueError(
+                f"Independent min frequencies per detector not enabled. "
+                f"All frequencies must match, got f_min = {f_min}."
+            )
+        # TODO: Risk of non-constant caps with non-independent detectors.
+        assert len(set(caps.values())) == 1
+    for d in detectors:
+        if f_min[d] > caps[d]:
+            raise ValueError(
+                f"Minimum frequency requested for {d} ({f_min[d]} Hz) "
+                f"greater than upper bound of {caps[d]} Hz."
+            )
+
+
+def check_frequency_updates(
+    model_metadata: dict,
+    f_min: dict[str, float] | float | None = None,
+    f_max: dict[str, float] | float | None = None,
+):
+    """
+    Validate and apply optional minimum and maximum frequency constraints
+    for a model’s frequency domain.
+
+    This function checks that any provided per-detector minimum (`f_min`)
+    or maximum (`f_max`) frequencies—either as a single float applied to
+    all detectors or as a dict mapping each detector to its own value—:
+      - Match exactly the set of detectors in the model metadata.
+      - Respect the hard bounds defined by the domain (`domain.f_min` /
+        `domain.f_max`).
+      - Comply with optional random-strain-cropping settings (probability,
+        independent vs. joint detectors, and per-detector caps/floors).
+
+    Parameters
+    ----------
+    model_metadata : dict
+        Dictionary containing the model’s training settings and data.
+        Must include:
+          - `["train_settings"]["data"]["detectors"]`: list of detector names.
+          - `["train_settings"]["data"]["random_strain_cropping"]`: optional
+            dict of cropping parameters.
+    f_min : dict[str, float], float, or None, optional
+        Single float or per-detector dict of minimum frequencies to enforce.
+        If a float is provided, it is applied to all detectors. Each value
+        must be ≥ `domain.f_min`. If `None`, no minimum-frequency
+        validation is performed.
+    f_max : dict[str, float], float, or None, optional
+        Single float or per-detector dict of maximum frequencies to enforce.
+        If a float is provided, it is applied to all detectors. Each value
+        must be ≤ `domain.f_max`. If `None`, no maximum-frequency
+        validation is performed.
+
+    Raises
+    ------
+    ValueError
+        - If `model_metadata` does not describe a `UniformFrequencyDomain`
+          or `MultibandedFrequencyDomain`.
+        - If `f_min`/`f_max` keys don’t exactly match the detector list.
+        - If any requested frequency lies outside the hard domain bounds.
+        - If cropping is disabled but a change in frequency is requested.
+        - If per-detector constraints (independent vs. joint) or
+          cropping caps/floors are violated.
+
+    Returns
+    -------
+    None
+    """
+    crop_settings = model_metadata["train_settings"]["data"].get(
+        "random_strain_cropping"
+    )
+    detectors = model_metadata["train_settings"]["data"]["detectors"]
+    domain = build_domain_from_model_metadata(model_metadata, base=True)
+    if not isinstance(domain, (UniformFrequencyDomain, MultibandedFrequencyDomain)):
+        raise ValueError("Frequency updates only possible for frequency domains.")
+
+    if f_min is not None:
+        _validate_minimum_frequency(f_min, detectors, domain, crop_settings)
+    if f_max is not None:
+        _validate_maximum_frequency(f_max, detectors, domain, crop_settings)
diff --git a/dingo/gw/inference/gw_samplers.py b/dingo/gw/inference/gw_samplers.py
index 2a1b63d89..52189e44a 100644
--- a/dingo/gw/inference/gw_samplers.py
+++ b/dingo/gw/inference/gw_samplers.py
@@ -11,11 +11,14 @@
 from dingo.core.transforms import GetItem, RenameKey
 from dingo.gw.domains import (
     MultibandedFrequencyDomain,
-    build_domain_from_model_metadata,
     UniformFrequencyDomain,
     Domain,
 )
 from dingo.gw.domains import build_domain
+from dingo.gw.frequency_updates import (
+    _validate_maximum_frequency,
+    _validate_minimum_frequency,
+)
 from dingo.gw.gwutils import get_extrinsic_prior_dict
 from dingo.gw.prior import build_prior_with_defaults
 from dingo.gw.result import Result
@@ -38,8 +41,7 @@
 class SamplerProtocol(Protocol):
     base_model_metadata: dict
 
-    def _initialize_transforms(self) -> None:
-        ...
+    def _initialize_transforms(self) -> None: ...
 
 
 class _GWMixinProtocol(SamplerProtocol):
@@ -173,7 +175,6 @@ def _build_domain(self: Sampler):
         if "domain_update" in data_settings:
             self.domain.update(data_settings["domain_update"])
 
-
     def _correct_reference_time(
         self: Sampler, samples: Union[dict, pd.DataFrame], inverse: bool = False
     ):
@@ -507,176 +508,3 @@ def _kernel_log_prob(self, samples):
             for k in self.gnpe_kernel.keys()
         }
         return self.gnpe_kernel.ln_prob(gnpe_proxies_diff, axis=0)
-
-
-# Functions for frequency cropping. Used by Sampler classes and dingo-pipe.
-
-
-def _validate_maximum_frequency(
-    f_max: dict[str, float] | float,
-    detectors: list[str],
-    domain: UniformFrequencyDomain | MultibandedFrequencyDomain,
-    crop_settings: dict | None,
-):
-    if isinstance(f_max, float):
-        f_max = {d: f_max for d in detectors}
-    if set(f_max) != set(detectors):
-        raise ValueError(
-            f"f_max must have exactly detectors {detectors}, got " f"{list(f_max)}."
-        )
-    f_max_vals = np.array([f_max[d] for d in detectors])
-
-    # Hard upper bound
-    if np.any(f_max_vals > domain.f_max):
-        raise ValueError(f"f_max {f_max} > domain.f_max = {domain.f_max}.")
-
-    # Nothing changed
-    if np.all(f_max_vals == domain.f_max):
-        return
-
-    # Cropping must be on
-    if not crop_settings or crop_settings.get("cropping_probability", 0.0) == 0.0:
-        raise ValueError(
-            f"Cropping disabled; cannot lower maximum frequency to {f_max}."
-        )
-
-    # Extract lower bounds
-    floors = crop_settings.get("f_max_lower")
-    if floors is None:
-        floors = domain.f_max
-    if not isinstance(floors, dict):
-        floors = {d: floors for d in detectors}
-
-    # Check lower bound.
-    if not crop_settings.get("independent_detectors", True):
-        if len(set(f_max_vals)) > 1:
-            raise ValueError(
-                f"Independent max frequencies per detector not enabled. "
-                f"All frequencies must match, got f_max = {f_max}."
-            )
-        # TODO: Risk of non-constant floors with non-independent detectors.
-        assert len(set(floors.values())) == 1
-    for d in detectors:
-        if f_max[d] < floors[d]:
-            raise ValueError(
-                f"Maximum frequency requested for {d} ({f_max[d]} Hz) "
-                f"less than lower bound of {floors[d]} Hz."
-            )
-
-
-def _validate_minimum_frequency(
-    f_min: dict[str, float] | float,
-    detectors: list[str],
-    domain: UniformFrequencyDomain | MultibandedFrequencyDomain,
-    crop_settings: dict | None,
-):
-    if isinstance(f_min, float):
-        f_min = {d: f_min for d in detectors}
-    if set(f_min) != set(detectors):
-        raise ValueError(
-            f"f_min must have exactly detectors {detectors}, got {list(f_min)}."
-        )
-    f_min_vals = np.array([f_min[d] for d in detectors])
-
-    # Hard lower bound
-    if np.any(f_min_vals < domain.f_min):
-        raise ValueError(f"f_min {f_min} < domain.f_min = {domain.f_min}.")
-
-    # Nothing changed
-    if np.all(f_min_vals == domain.f_min):
-        return
-
-    # Cropping must be on
-    if not crop_settings or crop_settings.get("cropping_probability", 0.0) == 0.0:
-        raise ValueError(
-            f"Cropping disabled; cannot raise minimum frequency to {f_min}."
-        )
-
-    # Extract upper bounds
-    caps = crop_settings.get("f_min_upper")
-    if caps is None:
-        caps = domain.f_min
-    if not isinstance(caps, dict):
-        caps = {d: caps for d in detectors}
-
-    # Check upper bound.
-    if not crop_settings.get("independent_detectors", True):
-        if len(set(f_min_vals)) > 1:
-            raise ValueError(
-                f"Independent min frequencies per detector not enabled. "
-                f"All frequencies must match, got f_min = {f_min}."
-            )
-        # TODO: Risk of non-constant caps with non-independent detectors.
-        assert len(set(caps.values())) == 1
-    for d in detectors:
-        if f_min[d] > caps[d]:
-            raise ValueError(
-                f"Minimum frequency requested for {d} ({f_min[d]} Hz) "
-                f"greater than upper bound of {caps[d]} Hz."
-            )
-
-
-def check_frequency_updates(
-    model_metadata: dict,
-    f_min: dict[str, float] | float | None = None,
-    f_max: dict[str, float] | float | None = None,
-):
-    """
-    Validate and apply optional minimum and maximum frequency constraints
-    for a model’s frequency domain.
-
-    This function checks that any provided per-detector minimum (`f_min`)
-    or maximum (`f_max`) frequencies—either as a single float applied to
-    all detectors or as a dict mapping each detector to its own value—:
-      - Match exactly the set of detectors in the model metadata.
-      - Respect the hard bounds defined by the domain (`domain.f_min` /
-        `domain.f_max`).
-      - Comply with optional random-strain-cropping settings (probability,
-        independent vs. joint detectors, and per-detector caps/floors).
-
-    Parameters
-    ----------
-    model_metadata : dict
-        Dictionary containing the model’s training settings and data.
-        Must include:
-          - `["train_settings"]["data"]["detectors"]`: list of detector names.
-          - `["train_settings"]["data"]["random_strain_cropping"]`: optional
-            dict of cropping parameters.
-    f_min : dict[str, float], float, or None, optional
-        Single float or per-detector dict of minimum frequencies to enforce.
-        If a float is provided, it is applied to all detectors. Each value
-        must be ≥ `domain.f_min`. If `None`, no minimum-frequency
-        validation is performed.
-    f_max : dict[str, float], float, or None, optional
-        Single float or per-detector dict of maximum frequencies to enforce.
-        If a float is provided, it is applied to all detectors. Each value
-        must be ≤ `domain.f_max`. If `None`, no maximum-frequency
-        validation is performed.
-
-    Raises
-    ------
-    ValueError
-        - If `model_metadata` does not describe a `UniformFrequencyDomain`
-          or `MultibandedFrequencyDomain`.
-        - If `f_min`/`f_max` keys don’t exactly match the detector list.
-        - If any requested frequency lies outside the hard domain bounds.
-        - If cropping is disabled but a change in frequency is requested.
-        - If per-detector constraints (independent vs. joint) or
-          cropping caps/floors are violated.
-
-    Returns
-    -------
-    None
-    """
-    crop_settings = model_metadata["train_settings"]["data"].get(
-        "random_strain_cropping"
-    )
-    detectors = model_metadata["train_settings"]["data"]["detectors"]
-    domain = build_domain_from_model_metadata(model_metadata, base=True)
-    if not isinstance(domain, (UniformFrequencyDomain, MultibandedFrequencyDomain)):
-        raise ValueError("Frequency updates only possible for frequency domains.")
-
-    if f_min is not None:
-        _validate_minimum_frequency(f_min, detectors, domain, crop_settings)
-    if f_max is not None:
-        _validate_maximum_frequency(f_max, detectors, domain, crop_settings)
diff --git a/dingo/pipe/main.py b/dingo/pipe/main.py
index 1c52cedcf..52117ec9b 100644
--- a/dingo/pipe/main.py
+++ b/dingo/pipe/main.py
@@ -28,7 +28,7 @@
 
 from ..gw.domains.build_domain import build_domain_from_model_metadata
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
-from ..gw.inference.gw_samplers import check_frequency_updates
+from ..gw.frequency_updates import check_frequency_updates
 from ..gw.injection import Injection
 from ..gw.noise.asd_dataset import ASDDataset
 

From 7b0671350d943c3264033c9ec42d2928a4dc5c8e Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 15:12:01 +0100
Subject: [PATCH 35/65] Port event frequency-range cropping to the composed
 data-prep chain

from_model_metadata now splices MaskDataForFrequencyRangeUpdate into
the chain (after whitening, before repackaging -- the legacy slot)
whenever the event range differs from the data-domain bounds. The
condition is the legacy frequency_updates property verbatim, including
the identity-mask case where generation wrote base-domain bounds that
exceed a multibanded band edge. prepared_data() validates the request
exactly as the legacy setters did -- hard bounds against the (base)
domain, narrowing only under the training random-strain-cropping
license -- via the rehomed validators, replacing the interim tripwire.
The likelihood view is unaffected (it applies the range via ASD
masking).

Tests cover equal-bounds no-op, unlicensed narrowing, out-of-domain
rejection, and value-level masking of the prepared tensor. The toy
harnesses now exercise the identity mask on the composed path and
remain bit-exact.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py       | 84 +++++++++++++++++++----------
 tests/gw/test_gw_sampler_context.py | 71 ++++++++++++++++--------
 2 files changed, 106 insertions(+), 49 deletions(-)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 3515d1fb5..6f3b67f3a 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -45,6 +45,10 @@
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
 from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
+from dingo.gw.frequency_updates import (
+    _validate_maximum_frequency,
+    _validate_minimum_frequency,
+)
 from dingo.gw.gwutils import get_extrinsic_prior_dict
 from dingo.gw.likelihood import StationaryGaussianGWLikelihood
 from dingo.gw.prior import build_prior_with_defaults
@@ -54,6 +58,7 @@
     GetDetectorTimes,
     GNPEBase,
     GNPECoalescenceTimes,
+    MaskDataForFrequencyRangeUpdate,
     PostCorrectGeocentTime,
     RepackageStrainsAndASDS,
     SelectStandardizeRepackageParameters,
@@ -65,6 +70,26 @@
 logger = logging.getLogger(__name__)
 
 
+def _frequency_range_update(domain, event_metadata) -> Optional[dict]:
+    """The event's requested frequency range when it differs from the data-domain
+    bounds (the legacy `frequency_updates` condition), else `None`. Values may be
+    floats or per-detector dicts; defaults are the domain bounds. Note a request
+    merely *wider* than the domain also triggers (generation writes base-domain
+    bounds, e.g. 1099.0 vs a multibanded band edge of 1098.875) -- the resulting
+    mask is then an identity, matching legacy behavior."""
+    if event_metadata is None:
+        return None
+    minimum = event_metadata.get("minimum_frequency", domain.f_min)
+    maximum = event_metadata.get("maximum_frequency", domain.f_max)
+
+    def normalize(value):
+        return set(value.values()) if isinstance(value, dict) else {value}
+
+    if normalize(minimum) == {domain.f_min} and normalize(maximum) == {domain.f_max}:
+        return None
+    return {"minimum_frequency": minimum, "maximum_frequency": maximum}
+
+
 class GWSamplerContext:
     """
     Per-event shared GW state: the data `d` and everything derived from it. Referenced
@@ -173,8 +198,16 @@ def from_model_metadata(
             )
         # Whiten and scale (the network expects standardized data).
         transforms.append(WhitenAndScaleStrain(domain.noise_std))
+        # Event frequency-range update: mask the whitened strain/ASDs outside the
+        # requested range. Must precede repackaging (ranges may be per-detector).
+        # The request is validated against the training crop license the first time
+        # prepared_data() runs.
+        range_update = _frequency_range_update(domain, event_metadata)
+        if range_update is not None:
+            transforms.append(
+                MaskDataForFrequencyRangeUpdate(domain=domain, **range_update)
+            )
         # Repackage strains/ASDs into an array, move to torch, extract the waveform.
-        # TODO: frequency-range cropping (MaskDataForFrequencyRangeUpdate) -- follow-up.
         transforms += [
             RepackageStrainsAndASDS(ifos=detectors, first_index=domain.min_idx),
             ToTorch(device=device),
@@ -231,38 +264,33 @@ def from_model(
         )
 
     def prepared_data(self) -> torch.Tensor:
-        """One-time data preprocessing, computed once and cached."""
+        """One-time data preprocessing, computed once and cached. An event
+        frequency-range update is validated against the training crop license
+        before the first preparation."""
         if self._prepared is None:
-            self._require_no_frequency_cropping()
+            self._validate_frequency_range()
             self._prepared = self._data_prep(self.event_data)
         return self._prepared
 
-    def _require_no_frequency_cropping(self):
-        """Frequency-range cropping -- masking the network input to an event range
-        narrower than the network domain (legacy `MaskDataForFrequencyRangeUpdate`)
-        -- is not implemented in the composed data-prep chain yet; fail loudly
-        rather than silently feed the network unmasked data. The likelihood view is
-        unaffected (it applies the range via ASD masking). Event bounds merely
-        *wider* than the data domain are fine (nothing to remove): generation
-        writes base-domain bounds, which can exceed a multibanded domain's
-        quantized band edge (e.g. 1099.0 vs 1098.875)."""
-        if self.event_metadata is None:
+    def _validate_frequency_range(self):
+        """Validate an event frequency-range update, exactly as the legacy sampler
+        setters did: hard bounds against the (base) domain, and narrowing only when
+        the network was trained with random strain cropping covering the requested
+        range. Applies to the network-input view only -- the likelihood view
+        applies the range independently via ASD masking."""
+        update = _frequency_range_update(self.domain, self.event_metadata)
+        if update is None:
             return
-        for key, bound, narrower in (
-            ("minimum_frequency", self.domain.f_min, lambda v, b: v > b),
-            ("maximum_frequency", self.domain.f_max, lambda v, b: v < b),
-        ):
-            value = self.event_metadata.get(key)
-            if value is None:
-                continue
-            values = value.values() if isinstance(value, dict) else [value]
-            if any(narrower(v, bound) for v in values):
-                raise NotImplementedError(
-                    f"Event {key} = {value} is narrower than the network domain "
-                    f"bound ({bound}): frequency-range cropping is not implemented "
-                    "in the composed data-prep chain yet. "
-                    "Use sampler-implementation = legacy."
-                )
+        domain = getattr(self.domain, "base_domain", self.domain)
+        crop_settings = self.base_metadata["train_settings"]["data"].get(
+            "random_strain_cropping"
+        )
+        _validate_minimum_frequency(
+            update["minimum_frequency"], self.detectors, domain, crop_settings
+        )
+        _validate_maximum_frequency(
+            update["maximum_frequency"], self.detectors, domain, crop_settings
+        )
 
     @property
     def prior(self) -> PriorDict:
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index 660bd2fb1..952d20635 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -7,7 +7,9 @@
 
 import copy
 
+import numpy as np
 import pytest
+import torch
 
 import dingo.gw.inference.factors as factors_module
 from dingo.gw.domains import build_domain
@@ -103,27 +105,54 @@ def test_likelihood_importance_sampling_overrides(context):
     assert default.kwargs["wfg_domain"].delta_f == _DOMAIN_SETTINGS["delta_f"]
 
 
-def test_prepared_data_rejects_frequency_cropping():
-    # A narrower event frequency range needs input masking, which the composed
-    # data-prep chain does not implement yet -> loud failure. Bounds equal to or
-    # wider than the domain (generation writes base-domain bounds) are fine.
-    def make(event_metadata):
-        return GWSamplerContext(
-            domain=build_domain(_DOMAIN_SETTINGS),
-            detectors=["H1"],
-            t_ref=0.0,
-            data_prep=lambda data: "prepared",
-            event_data={},
-            event_metadata=event_metadata,
-        )
-
-    with pytest.raises(NotImplementedError, match="cropping"):
-        make({"minimum_frequency": 25.0}).prepared_data()
-    with pytest.raises(NotImplementedError, match="cropping"):
-        make({"maximum_frequency": {"H1": 512.0}}).prepared_data()
-    assert make({"minimum_frequency": 20.0}).prepared_data() == "prepared"
-    assert make({"maximum_frequency": 1099.0}).prepared_data() == "prepared"
-    assert make(None).prepared_data() == "prepared"
+def _event_data(n_bins):
+    strain = np.ones(n_bins, dtype=complex)
+    asd = np.ones(n_bins)
+    return {
+        "waveform": {"H1": strain.copy(), "L1": strain.copy()},
+        "asds": {"H1": asd.copy(), "L1": asd.copy()},
+    }
+
+
+def _crop_context(event_metadata, crop_settings=None):
+    meta = copy.deepcopy(_MODEL_METADATA)
+    if crop_settings is not None:
+        meta["train_settings"]["data"]["random_strain_cropping"] = crop_settings
+    n_bins = int(_DOMAIN_SETTINGS["f_max"] / _DOMAIN_SETTINGS["delta_f"]) + 1
+    return GWSamplerContext.from_model_metadata(
+        meta, _event_data(n_bins), event_metadata=event_metadata
+    )
+
+
+def test_frequency_range_equal_bounds_is_a_no_op():
+    reference = _crop_context(None).prepared_data()
+    same = _crop_context(
+        {"minimum_frequency": 20.0, "maximum_frequency": 1024.0}
+    ).prepared_data()
+    assert torch.equal(reference, same)
+
+
+def test_frequency_range_narrowing_requires_crop_license():
+    with pytest.raises(ValueError, match="Cropping disabled"):
+        _crop_context({"minimum_frequency": 25.0}).prepared_data()
+
+
+def test_frequency_range_beyond_domain_rejected():
+    # The O1-file-reuse scenario: event bounds wider than the network domain.
+    with pytest.raises(ValueError, match="domain.f_max"):
+        _crop_context({"maximum_frequency": 1099.0}).prepared_data()
+
+
+def test_frequency_range_cropping_masks_network_input():
+    ctx = _crop_context(
+        {"minimum_frequency": 25.0},
+        crop_settings={"cropping_probability": 0.5, "f_min_upper": 30.0},
+    )
+    out = ctx.prepared_data()  # (n_det, 3 channels, n_bins), real strain first
+    frequencies = ctx.domain.sample_frequencies[ctx.domain.min_idx :]
+    strain_real = out[0, 0].numpy()
+    assert (strain_real[frequencies < 25.0] == 0).all()
+    assert (strain_real[frequencies >= 25.0] != 0).all()
 
 
 def test_likelihood_caller_marginalization_bounds_win(context):

From 106d7f94ffa97a2ec2ef8d910a5e0d640249cce5 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 16:02:11 +0100
Subject: [PATCH 36/65] Record structured sampler provenance in composed
 results

Every chain step gains a describe() method producing a literal-only
descriptor (class, parameter block, conditioning; salient configuration
where it exists: DeltaFactor values, the FlowFactor unconditional flag,
nested GibbsBlock structure with its iteration count, synthetic-phase
grid settings). GWComposedSampler.to_result() assembles them into
settings["sampler"] -- {version, implementation, chain, ...} -- and the
pipe adds model checkpoint paths and, for density recovery, the full
recipe: the Gibbs chain descriptor that generated the NDE's training
samples plus the trainer-enriched NDE settings, including the fitted
proxy standardization of the ephemeral flow.

This replaces the legacy provenance -- ad-hoc num_iterations/init_model
keys written by setter side effects, describing only the final sampler
state -- with a deliberate record of how the samples were made. Nothing
consumes the block yet; the version field is insurance for future
consumers such as chain reconstruction from file.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py               | 43 +++++++++++++++++++++++++++++
 dingo/gw/inference/factors.py       | 33 +++++++++++++++++++++-
 dingo/pipe/sampling.py              | 20 ++++++++++++++
 tests/core/test_factors.py          | 21 ++++++++++++++
 tests/gw/test_gw_sampler_context.py | 17 ++++++++++++
 5 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index c7f2e908f..75732ea78 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -133,6 +133,17 @@ def chunk_and_concat(
     return samples, log_prob
 
 
+def _describe_default(step) -> dict:
+    """Default provenance descriptor for a chain step: class name, the parameter
+    block it emits, and what it conditions on. Literal-only (round-trips through
+    `str`/`ast.literal_eval` in saved settings)."""
+    return {
+        "step": type(step).__name__,
+        "parameters": list(step.parameters),
+        "conditioning": list(step.conditioning),
+    }
+
+
 class Factor(ABC):
     """
     A conditional distribution `q_i(theta_i | f_i(theta_<i, d))` over one parameter
@@ -175,6 +186,11 @@ def log_prob(
         """Evaluate `log q_i` at given physical `theta_i` (one conditioning row per
         row)."""
 
+    def describe(self) -> dict:
+        """Provenance descriptor for saved-result metadata; steps with salient
+        configuration override this."""
+        return _describe_default(self)
+
 
 def _base_model_metadata(model: BasePosteriorModel) -> dict:
     """The base *analysis* metadata (dataset / domain / detector / data settings): for
@@ -336,6 +352,9 @@ def log_prob(self, theta_i, context, given=None):
             log_prob = self.model.log_prob(z, *net_context)
         return log_prob + self.standardization.log_det(self._net_parameters)
 
+    def describe(self) -> dict:
+        return {**_describe_default(self), "unconditional": self.unconditional}
+
 
 class DeltaFactor(Factor):
     """`q_i = delta(theta_i - c)`: a point mass pinning parameters to fixed values,
@@ -371,6 +390,12 @@ def log_prob(self, theta_i, context, given=None):
         reference_column = next(iter(theta_i.values()))
         return torch.zeros_like(reference_column)
 
+    def describe(self) -> dict:
+        return {
+            **_describe_default(self),
+            "values": {k: float(v) for k, v in self.values.items()},
+        }
+
 
 class SampleTableFactor(Factor):
     """
@@ -475,6 +500,10 @@ def consumes(self) -> list[str]:
         `ChainComposer.log_prob` rebuilds them via `inverse`."""
         return [c for c in self.conditioning if c not in self.parameters]
 
+    def describe(self) -> dict:
+        """Provenance descriptor for saved-result metadata."""
+        return _describe_default(self)
+
 
 class TargetCorrection(ABC):
     """
@@ -513,6 +542,10 @@ def sample_and_log_prob(
         reference_column = next(iter(out.values()))
         return out, torch.zeros_like(reference_column)
 
+    def describe(self) -> dict:
+        """Provenance descriptor for saved-result metadata."""
+        return _describe_default(self)
+
 
 class Step(Protocol):
     """
@@ -763,6 +796,16 @@ def sample_and_log_prob(
         `num_samples` is the walker (root) count -- Gibbs does not fan out."""
         return self._run_once(num_samples, context), None
 
+    def describe(self) -> dict:
+        """Provenance descriptor: the Gibbs structure, with nested descriptors for
+        the init factor and the swept factors."""
+        return {
+            "step": type(self).__name__,
+            "num_iterations": self.num_iterations,
+            "init": self.init_factor.describe(),
+            "factors": [factor.describe() for factor in self.factors],
+        }
+
     def _run_once(
         self, num_samples: int, context: SamplerContext
     ) -> dict[str, torch.Tensor]:
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 6f3b67f3a..c5db1ef1a 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -568,6 +568,16 @@ def log_prob(self, theta_i, context, given=None):
         )
         return torch.as_tensor(log_prob, device=device)
 
+    def describe(self) -> dict:
+        return {
+            "step": type(self).__name__,
+            "parameters": list(self.parameters),
+            "conditioning": list(self.conditioning),
+            "n_grid": self.n_grid,
+            "approximation_22_mode": self.approximation_22_mode,
+            "uniform_weight": self.uniform_weight,
+        }
+
     def _phase_profile(self, given, context):
         """The phase grid and the mass-covered (un-normalized) phase distribution, one row
         per sample: evaluate `log L` on the grid, exponentiate (shifted by the per-row
@@ -1079,6 +1089,25 @@ def __init__(
         super().__init__(composer, context)
         self.metadata = metadata
         self.inference_parameters = inference_parameters
+        # Extra provenance merged into settings["sampler"] by to_result -- e.g. the
+        # pipe records model checkpoint paths and the density-recovery recipe.
+        # Literal-only values (the settings dict round-trips through str/literal_eval).
+        self.provenance_extra: dict = {}
+
+    def sampler_provenance(self) -> dict:
+        """Provenance of how the samples were made, stored as `settings["sampler"]`
+        in the exported `Result`: the executed chain in order (one descriptor per
+        step, via `Step.describe()`), plus anything in `provenance_extra`. Nothing
+        consumes this yet -- it is deliberate, structured provenance (the legacy
+        samplers recorded only ad-hoc `num_iterations` / `init_model` keys via
+        setter side effects); the `version` field is insurance for future
+        consumers such as chain reconstruction from file."""
+        return {
+            "version": 1,
+            "implementation": "composed",
+            "chain": [step.describe() for step in self.composer.steps],
+            **copy.deepcopy(self.provenance_extra),
+        }
 
     @classmethod
     def from_model(
@@ -1238,6 +1267,8 @@ def to_result(self):
         """
         from dingo.gw.result import Result
 
+        settings = copy.deepcopy(self.metadata)
+        settings["sampler"] = self.sampler_provenance()
         data_dict = {
             "samples": self.samples,
             "context": self.context.event_data,
@@ -1245,7 +1276,7 @@ def to_result(self):
             "importance_sampling_metadata": None,
             "log_evidence": None,
             "log_noise_evidence": None,
-            "settings": copy.deepcopy(self.metadata),
+            "settings": settings,
         }
         return Result(dictionary=data_dict, sampler_context=self.context)
 
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index cb5233a2a..e3c63f341 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -144,6 +144,7 @@ def _load_sampler(self):
                 self.dingo_sampler = GWComposedSampler.from_model(
                     model, self.context, event_metadata=self.event_metadata
                 )
+            self.dingo_sampler.provenance_extra["models"] = self._model_paths()
             return
 
         if self.model_init is not None:
@@ -207,6 +208,12 @@ def density_recovery_settings(self, settings):
     #     result_dir = os.path.join(self.outdir, "result")
     #     return os.path.relpath(result_dir)
 
+    def _model_paths(self) -> dict:
+        paths = {"model": self.model}
+        if self.model_init is not None:
+            paths["model_init"] = self.model_init
+        return paths
+
     def _recover_log_prob_composed(self):
         """Density recovery on the composed sampler, mirroring `prepare_log_prob`:
         run the Gibbs chain, train an unconditional NDE on the GNPE proxies, then
@@ -216,6 +223,10 @@ def _recover_log_prob_composed(self):
         self.dingo_sampler.run_sampler(
             settings["num_samples"], batch_size=self.batch_size
         )
+        # The Gibbs chain that generates the NDE's training samples, recorded in the
+        # recovery provenance below (the final sampler's own chain is the recovered
+        # single-step one).
+        gibbs_chain = self.dingo_sampler.sampler_provenance()["chain"]
         result = self.dingo_sampler.to_result()
         nde_settings = settings["nde_settings"]
         nde_settings["training"]["device"] = str(self._main_model.device)
@@ -231,6 +242,15 @@ def _recover_log_prob_composed(self):
             self.context,
             event_metadata=self.event_metadata,
         )
+        self.dingo_sampler.provenance_extra["models"] = self._model_paths()
+        # The recovery recipe: the Gibbs chain that generated the NDE's training
+        # samples (with its iteration count inside the GibbsBlock descriptor) plus
+        # the NDE training settings. The legacy metadata recorded only the final
+        # single-step state.
+        self.dingo_sampler.provenance_extra["density_recovery"] = {
+            "chain": gibbs_chain,
+            **settings,
+        }
 
     def run_sampler(self):
         if self.gnpe and self.recover_log_prob:
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 485402cda..dafd733fe 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -249,6 +249,27 @@ def test_sample_table_log_prob_raises():
         table.log_prob({"a": torch.zeros(3)}, None)
 
 
+def test_describe_descriptors_are_structured_and_literal():
+    import ast
+
+    delta = DeltaFactor({"chi_1": 0.1})
+    d = delta.describe()
+    assert d["step"] == "DeltaFactor" and d["values"] == {"chi_1": 0.1}
+    assert isinstance(d["values"]["chi_1"], float)
+
+    gibbs = GibbsBlock(
+        _ConstFactor("proxy"), [_ConstFactor("theta", ["proxy"])], num_iterations=7
+    )
+    g = gibbs.describe()
+    assert g["num_iterations"] == 7
+    assert g["init"]["step"] == "_ConstFactor" and g["init"]["parameters"] == ["proxy"]
+    assert g["factors"][0]["conditioning"] == ["proxy"]
+
+    # Descriptors must survive the settings round-trip (str / literal_eval).
+    for block in (d, g):
+        assert ast.literal_eval(str(block)) == block
+
+
 class _MockReparam(Reparametrization):
     """A mock bijection ``u -> v = u + shift`` with a constant nonzero ``log|det J|`` so the
     density contribution (``-log_det``) is checkable."""
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index 952d20635..9f15fb8b1 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -143,6 +143,23 @@ def test_frequency_range_beyond_domain_rejected():
         _crop_context({"maximum_frequency": 1099.0}).prepared_data()
 
 
+def test_sampler_provenance_block():
+    import ast
+
+    from dingo.core.factors import ChainComposer, DeltaFactor
+    from dingo.gw.inference.factors import GWComposedSampler
+
+    sampler = GWComposedSampler(
+        ChainComposer([DeltaFactor({"a": 1.0})]), None, {}, ["a"]
+    )
+    sampler.provenance_extra["models"] = {"model": "model.pt"}
+    block = sampler.sampler_provenance()
+    assert block["version"] == 1 and block["implementation"] == "composed"
+    assert block["chain"][0]["step"] == "DeltaFactor"
+    assert block["models"] == {"model": "model.pt"}
+    assert ast.literal_eval(str(block)) == block
+
+
 def test_frequency_range_cropping_masks_network_input():
     ctx = _crop_context(
         {"minimum_frequency": 25.0},

From 4a9cc24e4dad80db16364bafcc5053e46cf3c403 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 16:24:18 +0100
Subject: [PATCH 37/65] Delete the legacy samplers

The composed sampler is the sole implementation: GWSampler,
GWSamplerGNPE, the core Sampler/GNPESampler base classes, and
prepare_log_prob are removed, and the pipe's sampler-implementation
option defaults to composed (legacy is retained only to give a clear
error). The last commit containing the legacy implementation is tagged
legacy-samplers-final; A/B parity runs against it remain possible from
a git worktree at that tag.

The two standalone consumers migrate to the composed API:
importance_weights.py resamples an unconditional NDE through a
one-factor chain over the result's sampler context (the previous code
path was broken -- the full run_sampler on an NDE died in
post-processing), and the diagnostics slice plots take a Result plus an
explicit proposal log-prob callable. Documentation is rewritten off the
legacy API, including two autoclass directives that pointed at deleted
classes.

Deliberately dropped as unused: remove_init_outliers (never set by the
pipe) and the low-latency samples save in density recovery.

Functional parity with the deleted implementation was established
bit-exactly by the stage and pipeline harnesses before deletion; the
post-deletion end-to-end dingo_pipe run (generation through plots)
passes on this tree.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/samplers.py                        | 573 ------------------
 dingo/gw/importance_sampling/diagnostics.py   |  53 +-
 .../importance_sampling/importance_weights.py |  17 +-
 dingo/gw/inference/factors.py                 |  34 +-
 dingo/gw/inference/gw_samplers.py             | 510 ----------------
 dingo/gw/inference/inference_utils.py         |  68 ---
 dingo/gw/result.py                            |  13 +-
 dingo/pipe/parser.py                          |   9 +-
 dingo/pipe/sampling.py                        |  78 +--
 docs/source/dingo_pipe.md                     |   2 +-
 docs/source/example_injection.md              |   8 +-
 docs/source/gnpe.md                           |  10 +-
 docs/source/inference.md                      |  19 +-
 docs/source/result.md                         |   2 +-
 tests/pipe/test_sampler_implementation.py     |  12 +-
 15 files changed, 137 insertions(+), 1271 deletions(-)
 delete mode 100644 dingo/core/samplers.py
 delete mode 100644 dingo/gw/inference/gw_samplers.py
 delete mode 100644 dingo/gw/inference/inference_utils.py

diff --git a/dingo/core/samplers.py b/dingo/core/samplers.py
deleted file mode 100644
index efc4fc1fa..000000000
--- a/dingo/core/samplers.py
+++ /dev/null
@@ -1,573 +0,0 @@
-import copy
-import math
-import time
-from pathlib import Path
-from typing import Optional, Union
-import sys
-
-import numpy as np
-import pandas as pd
-import torch
-from torchvision.transforms import Compose
-
-from dingo.core.posterior_models import BasePosteriorModel
-from dingo.core.result import Result
-from dingo.core.result import DATA_KEYS as RESULT_DATA_KEYS
-from dingo.core.utils import torch_detach_to_cpu, IterationTracker
-
-# FIXME: transform below should be in core
-from dingo.gw.transforms import SelectStandardizeRepackageParameters
-
-#
-# Sampler classes are motivated by the approach of Bilby.
-#
-
-
-class Sampler(object):
-    """
-    Sampler class that wraps a PosteriorModel. Allows for conditional and unconditional
-    models.
-
-    Draws samples from the model based on (optional) context data.
-
-    This is intended for use either as a standalone sampler, or as a sampler producing
-    initial sample points for a GNPE sampler.
-
-    Methods
-    -------
-    run_sampler
-    log_prob
-    to_result
-    to_hdf5
-
-    Attributes
-    ----------
-    model : BasePosteriorModel
-    inference_parameters : list
-    samples : DataFrame
-        Samples produced from the model by run_sampler().
-    context : dict
-    metadata : dict
-    event_metadata : dict
-    unconditional_model : bool
-        Whether the model is unconditional, in which case it is not provided context
-        information.
-    transform_pre, transform_post : Transform
-        Transforms to be applied to data and parameters during inference. These are
-        typically implemented in a subclass.
-    """
-
-    def __init__(
-        self,
-        model: BasePosteriorModel,
-    ):
-        """
-        Parameters
-        ----------
-        model : BasePosteriorModel
-        """
-        self.model = model
-        self.event_metadata = None
-
-        self.metadata = self.model.metadata.copy()
-        if self.metadata["train_settings"]["data"].get("unconditional", False):
-            self.unconditional_model = True
-            # For unconditional models, the context will be stored with the model. It
-            # is needed for calculating the likelihood for importance sampling.
-            # However, it will not be used when sampling from the model, since it is
-            # unconditional.
-            self._context = self.model.context
-            self.base_model_metadata = self.metadata["base"]
-        else:
-            self.unconditional_model = False
-            self._context = None
-            self.base_model_metadata = self.metadata
-
-        self.inference_parameters = self.metadata["train_settings"]["data"][
-            "inference_parameters"
-        ]
-
-        self.samples = None
-        self._build_domain()
-
-        if self.unconditional_model:
-            # Must be after _build_domain(), since setter might access domain.
-            self.event_metadata = self.model.event_metadata
-
-        # Must be after _build_domain() since transforms can depend on domain.
-        self._initialize_transforms()
-
-        self._pesummary_package = "core"
-        self._result_class = Result
-
-    @property
-    def context(self):
-        """Data on which to condition the sampler. For injections, there should be a
-        'parameters' key with truth values."""
-        return self._context
-
-    @context.setter
-    def context(self, value):
-        if value is not None and "parameters" in value:
-            if self.event_metadata is None:
-                self.event_metadata = {}
-            self.event_metadata["injection_parameters"] = value.pop("parameters")
-        self._context = value
-
-    @property
-    def event_metadata(self):
-        """Metadata for data analyzed. Can in principle influence any post-sampling
-        parameter transformations (e.g., sky position correction), as well as the
-        likelihood detector positions."""
-        return self._event_metadata
-
-    @event_metadata.setter
-    def event_metadata(self, value):
-        self._event_metadata = value
-
-    def _initialize_transforms(self):
-        self.transform_pre = Compose([])
-
-        # De-standardize data and extract inference parameters. This needs to be here
-        # (and not just in subclasses) in order for run_sampler() to properly execute.
-        self.transform_post = SelectStandardizeRepackageParameters(
-            {"inference_parameters": self.inference_parameters},
-            self.metadata["train_settings"]["data"]["standardization"],
-            inverse=True,
-            as_type="dict",
-        )
-
-    def _run_sampler(
-        self,
-        num_samples: int,
-        context: Optional[dict] = None,
-    ) -> dict:
-        if not self.unconditional_model:
-            if context is None:
-                raise ValueError("Context required to run sampler.")
-            x = context.copy()
-            x["parameters"] = {}
-            x["extrinsic_parameters"] = {}
-
-            # transforms_pre are expected to transform the data in the same way for each
-            # requested sample. We therefore apply pre-processing only once.
-            x = self.transform_pre(context)
-            # Require a batch dimension for the embedding network.
-            x = x.unsqueeze(0)
-            x = [x]
-        else:
-            if context is not None:
-                print("Unconditional model. Ignoring context.")
-            x = []
-
-        # For a normalizing flow, we get the log_prob for "free" when sampling,
-        # so we always include this. For other architectures, it may make sense to
-        # have a flag for whether to calculate the log_prob.
-        self.model.network.eval()
-        with torch.no_grad():
-            y, log_prob = self.model.sample_and_log_prob(*x, num_samples=num_samples)
-
-        if not self.unconditional_model:
-            # Squeeze the batch dimension added earlier.
-            y = y.squeeze(0)
-            log_prob = log_prob.squeeze(0)
-
-        samples = self.transform_post({"parameters": y, "log_prob": log_prob})
-        result = samples["parameters"]
-        result["log_prob"] = samples["log_prob"]
-        return result
-
-    def run_sampler(
-        self,
-        num_samples: int,
-        batch_size: Optional[int] = None,
-    ):
-        """
-        Generates samples and stores them in self.samples. Conditions the model on
-        self.context if appropriate (i.e., if the model is not unconditional).
-
-        If possible, it also calculates the log_prob and saves it as a column in
-        self.samples. When using GNPE it is not possible to obtain the log_prob due to
-        the many Gibbs iterations. However, in the case of just one iteration, and when
-        starting from a sampler for the proxy, the GNPESampler does calculate the
-        log_prob.
-
-        Allows for batched sampling, e.g., if limited by GPU memory. Actual sampling
-        for each batch is performed by _run_sampler(), which will differ for Sampler
-        and GNPESampler.
-
-        Parameters
-        ----------
-        num_samples : int
-            Number of samples requested.
-        batch_size : int, optional
-            Batch size for sampler.
-        """
-        self.samples = None
-
-        print(f"Running sampler to generate {num_samples} samples.")
-        t0 = time.time()
-        if not self.unconditional_model:
-            if self.context is None:
-                raise ValueError("Context must be set in order to run sampler.")
-            context = self.context
-        else:
-            context = None
-
-        # Carry out batched sampling by calling _run_sample() on each batch and
-        # consolidating the results.
-        if batch_size is None:
-            batch_size = num_samples
-        full_batches, remainder = divmod(num_samples, batch_size)
-        samples = [self._run_sampler(batch_size, context) for _ in range(full_batches)]
-        if remainder > 0:
-            samples.append(self._run_sampler(remainder, context))
-        samples = {p: torch.cat([s[p] for s in samples]) for p in samples[0].keys()}
-        samples = {k: v.cpu().numpy() for k, v in samples.items()}
-
-        # Apply any post-sampling transformation to sampled parameters (e.g.,
-        # correction for t_ref) and represent as DataFrame.
-        self._post_process(samples)
-        self.samples = pd.DataFrame(samples)
-        print(f"Done. This took {time.time() - t0:.1f} s.")
-        sys.stdout.flush()
-
-    def log_prob(self, samples: pd.DataFrame | dict) -> np.ndarray:
-        """
-        Calculate the model log probability at specific sample points.
-
-        Parameters
-        ----------
-        samples : pd.DataFrame | dict
-            Sample points at which to calculate the log probability.
-
-        Returns
-        -------
-        np.array of log probabilities.
-        """
-        if self.context is None and not self.unconditional_model:
-            raise ValueError("Context must be set in order to calculate log_prob.")
-
-        if isinstance(samples, dict):
-            if np.isscalar(list(samples.values())[0]):
-                samples = pd.DataFrame([samples])
-            else:
-                samples = pd.DataFrame(samples)
-
-        # This undoes any post-correction that would have been done to the samples,
-        # before evaluating the log_prob. E.g., the t_ref / sky position correction.
-        samples = samples.copy()
-        self._post_process(samples, inverse=True)
-
-        # Standardize the sample parameters and place on device.
-        y = samples[self.inference_parameters].to_numpy()
-        standardization = self.metadata["train_settings"]["data"]["standardization"]
-        mean = np.array([standardization["mean"][p] for p in self.inference_parameters])
-        std = np.array([standardization["std"][p] for p in self.inference_parameters])
-        y = (y - mean) / std
-        y = torch.from_numpy(y).to(device=self.model.device, dtype=torch.float32)
-
-        if not self.unconditional_model:
-            x = self.context.copy()
-            x["parameters"] = {}
-            x["extrinsic_parameters"] = {}
-
-            # Context is the same for each sample. Expand across batch dimension after
-            # pre-processing.
-            x = self.transform_pre(self.context)
-            x = x.expand(len(samples), *x.shape)  # TODO: Make this more efficient.
-            x = [x]
-        else:
-            x = []
-
-        self.model.network.eval()
-        with torch.no_grad():
-            log_prob = self.model.log_prob(y, *x)
-
-        log_prob = log_prob.cpu().numpy()
-        log_prob -= np.sum(np.log(std))
-
-        # Pre-processing step may have included a log_prob with the samples.
-        if "log_prob" in samples:
-            log_prob += samples["log_prob"].to_numpy()
-
-        return log_prob
-
-    def _post_process(self, samples: Union[dict, pd.DataFrame], inverse: bool = False):
-        pass
-
-    def _build_domain(self):
-        self.domain = None
-
-    def to_result(self) -> Result:
-        """
-        Export samples, metadata, and context information to a Result instance,
-        which can be used for saving or, e.g., importance sampling, training an
-        unconditional flow, etc.
-
-        Returns
-        -------
-        Result
-        """
-        data_dict = {k: getattr(self, k, None) for k in RESULT_DATA_KEYS}
-        # *COPY* the metadata to avoid recursion errors when creating new objects.
-        data_dict["settings"] = copy.deepcopy(self.metadata)
-        return self._result_class(dictionary=data_dict)
-
-    def to_hdf5(self, label="result", outdir="."):
-        dataset = self.to_result()
-        file_name = label + ".hdf5"
-        Path(outdir).mkdir(parents=True, exist_ok=True)
-        dataset.to_file(file_name=Path(outdir, file_name))
-
-
-class GNPESampler(Sampler):
-    """
-    Base class for GNPE sampler. It wraps a PosteriorModel *and* a standard Sampler for
-    initialization. The former is used to generate initial samples for Gibbs sampling.
-
-    A GNPE network is conditioned on additional "proxy" context theta^, i.e.,
-
-    p(theta | theta^, d)
-
-    The theta^ depend on theta via a fixed kernel p(theta^ | theta). Combining these
-    known distributions, this class uses Gibbs sampling to draw samples from the joint
-    distribution,
-
-    p(theta, theta^ | d)
-
-    The advantage of this approach is that we are allowed to perform any transformation of
-    d that depends on theta^. In particular, we can use this freedom to simplify the
-    data, e.g., by aligning data to have merger times = 0 in each detector. The merger
-    times are unknown quantities that must be inferred jointly with all other
-    parameters, and GNPE provides a means to do this iteratively. See
-    https://arxiv.org/abs/2111.13139 for additional details.
-
-    Gibbs sampling breaks access to the probability density, so this must be recovered
-    through other means. One way is to train an unconditional flow to represent p(theta^
-    | d) for fixed d based on the samples produced through the GNPE Gibbs sampling.
-    Starting from these, a single Gibbs iteration gives theta from the GNPE network,
-    along with the probability density in the joint space. This is implemented in
-    GNPESampler provided the init_sampler provides proxies directly and num_iterations
-    = 1.
-
-    Attributes (beyond those of Sampler)
-    ----------
-    init_sampler : Sampler
-        Used for providing initial samples for Gibbs sampling.
-    num_iterations : int
-        Number of Gibbs iterations to perform.
-    iteration_tracker : IterationTracker  **not set up**
-    remove_init_outliers : float  **not set up**
-    """
-
-    def __init__(
-        self,
-        model: BasePosteriorModel,
-        init_sampler: Sampler,
-        num_iterations: int = 1,
-    ):
-        """
-        Parameters
-        ----------
-        model : BasePosteriorModel
-        init_sampler : Sampler
-            Used for generating initial samples
-        num_iterations : int
-            Number of GNPE iterations to be performed by sampler.
-        """
-        self.gnpe_parameters = []  # Should be set in subclass _initialize_transform()
-
-        super().__init__(model)
-        self.init_sampler = init_sampler
-        self.num_iterations = num_iterations
-        self.iteration_tracker = None
-        # remove self.remove_init_outliers of lowest log_prob init samples before gnpe
-        self.remove_init_outliers = 0.0
-
-    @property
-    def init_sampler(self):
-        return self._init_sampler
-
-    @init_sampler.setter
-    def init_sampler(self, value):
-        self._init_sampler = value
-        # Copy this so it persists if we delete the init model.
-        self.metadata["init_model"] = self._init_sampler.model.metadata.copy()
-        if self._init_sampler.unconditional_model:
-            self.context = self._init_sampler.context
-            self.event_metadata = self._init_sampler.event_metadata
-
-    @property
-    def num_iterations(self):
-        """The number of GNPE iterations to perform when sampling."""
-        return self._num_iterations
-
-    @num_iterations.setter
-    def num_iterations(self, value):
-        self._num_iterations = value
-        self.metadata["num_iterations"] = self._num_iterations
-
-    @property
-    def gnpe_proxy_parameters(self):
-        return [p + "_proxy" for p in self.gnpe_parameters]
-
-    def _kernel_log_prob(self, samples):
-        raise NotImplementedError("To be implemented in subclass.")
-
-    def _run_sampler(
-        self,
-        num_samples: int,
-        context: Optional[dict] = None,
-    ) -> dict:
-        if context is None:
-            raise ValueError("self.context must be set to run sampler.")
-
-        data_ = self.init_sampler.transform_pre(context)
-
-        # TODO: Reimplement outlier removal in IterationTracker? Save setting somewhere.
-        if self.remove_init_outliers == 0.0:
-            init_samples = self.init_sampler._run_sampler(num_samples, context)
-        else:
-            if self.num_iterations == 1:
-                print(
-                    f"Warning: Removing initial outliers, but only carrying out "
-                    f"{self.num_iterations} GNPE iteration. This risks biasing "
-                    f"results."
-                )
-            init_samples = self.init_sampler._run_sampler(
-                math.ceil(num_samples / (1 - self.remove_init_outliers)), context
-            )
-            thr = torch.quantile(init_samples["log_prob"], self.remove_init_outliers)
-            inds = torch.where(init_samples["log_prob"] >= thr)[0][:num_samples]
-            init_samples = {k: v[inds] for k, v in init_samples.items()}
-
-        # We could be starting with either the GNPE parameters *or* their proxies,
-        # depending on the nature of the initialization network.
-
-        start_with_proxies = False
-        proxy_log_prob = None
-        proxies = {}
-
-        if {p + "_proxy" for p in self.gnpe_parameters}.issubset(init_samples.keys()):
-            start_with_proxies = True
-            proxy_log_prob = init_samples["log_prob"]
-            proxies = {k: init_samples[k] for k in self.gnpe_proxy_parameters}
-            # proxies is a dict of torch.Tensors, since it came from _run_sampler(),
-            # not run_sampler(). Clone it for a later assertion check.
-            init_proxies = {k: v.clone() for k, v in proxies.items()}
-
-        x = {"extrinsic_parameters": init_samples, "parameters": {}}
-
-        #
-        # Gibbs sample.
-        #
-
-        # Saving as an attribute so we can later access the intermediate Gibbs samples.
-        self.iteration_tracker = IterationTracker(store_data=True)
-
-        for i in range(self.num_iterations):
-            start_time = time.time()
-
-            if start_with_proxies and i == 0:
-                x["extrinsic_parameters"] = proxies.copy()
-            else:
-                x["extrinsic_parameters"] = {
-                    k: x["extrinsic_parameters"][k] for k in self.gnpe_parameters
-                }
-
-            # TODO: Depending on whether start_with_proxies is True, this might end up
-            #  comparing proxies vs gnpe_parameters for the first iteration.
-            self.iteration_tracker.update(
-                {k: v.cpu().numpy() for k, v in x["extrinsic_parameters"].items()}
-            )
-
-            d = data_.clone()
-            x["data"] = d.expand(num_samples, *d.shape)
-
-            x = self.transform_pre(x)
-
-            time_sample_start = time.time()
-            self.model.network.eval()
-            with torch.no_grad():
-                if "context_parameters" in x:
-                    y, log_prob = self.model.sample_and_log_prob(
-                        x["data"], x["context_parameters"]
-                    )
-                else:
-                    y, log_prob = self.model.sample_and_log_prob(x["data"])
-
-            # Squeeze the extra dimension added by sample_and_log_prob(num_samples=1).
-            y = y.squeeze(1)
-            log_prob = log_prob.squeeze(1)
-
-            time_sample_end = time.time()
-
-            x["parameters"] = y
-            x["log_prob"] = log_prob
-            x = self.transform_post(x)
-
-            # Extract the proxy parameters from x["extrinsic_parameters"]. These have
-            # not been standardized. They are persistent from before sampling took place,
-            # since this is when they were placed here and their values should not have
-            # changed.
-            proxies = {
-                p: x["extrinsic_parameters"][p] for p in self.gnpe_proxy_parameters
-            }
-
-            print(
-                f"it {i}.\tmin pvalue: {self.iteration_tracker.pvalue_min:.3f}"
-                f"\tproxy mean: ",
-                *[f"{torch.mean(v).item():.5f}" for v in proxies.values()],
-                "\tproxy std:",
-                *[f"{torch.std(v).item():.5f}" for v in proxies.values()],
-                "\ttimes:",
-                time_sample_start - start_time,
-                time_sample_end - time_sample_start,
-                time.time() - time_sample_end,
-            )
-
-        #
-        # Prepare final result.
-        #
-
-        if start_with_proxies and self.num_iterations == 1:
-            # In this case it makes sense to save the log_prob and the proxy parameters.
-
-            samples = x["parameters"]
-            samples["log_prob"] = x["log_prob"] + proxy_log_prob
-
-            # The log_prob returned by gnpe is not just the log_prob over parameters
-            # theta, but instead the log_prob in the *joint* space q(theta,theta^|x),
-            # including the proxies theta^. For importance sampling this means,
-            # that the target density is
-            #
-            #       p(theta,theta^|x) = p(theta^|theta) * p(theta|x).
-            #
-            # We compute log[p(theta^|theta)] below and store it as
-            # samples["delta_log_prob_target"], such that for importance sampling we
-            # only need to evaluate log[p(theta|x)] and add this correction.
-
-            # Proxies should be sitting in extrinsic_parameters.
-            all_params = {**x["extrinsic_parameters"], **samples}
-            all_params = {k: torch_detach_to_cpu(v) for k, v in all_params.items()}
-            kernel_log_prob = self._kernel_log_prob(all_params)
-            samples["delta_log_prob_target"] = torch.Tensor(kernel_log_prob)
-
-        else:
-            # Otherwise we only save the inference parameters, and no log_prob.
-            # Alternatively we could save the entire chain and the log_prob, but this
-            # is not useful for our purposes.
-
-            samples = x["parameters"]
-
-        # Include the proxies along with the inference parameters. The variable proxies
-        # gets set at the end of each Gibbs loop.
-        samples.update(proxies)
-
-        # Safety check for unconditional flows. Make sure the proxies haven't changed.
-        if start_with_proxies and self.num_iterations == 1:
-            for k in proxies:
-                assert torch.equal(proxies[k], init_proxies[k])
-
-        return samples
diff --git a/dingo/gw/importance_sampling/diagnostics.py b/dingo/gw/importance_sampling/diagnostics.py
index 39cc95bb5..0ae8cb15d 100644
--- a/dingo/gw/importance_sampling/diagnostics.py
+++ b/dingo/gw/importance_sampling/diagnostics.py
@@ -8,8 +8,19 @@
 
 
 def plot_posterior_slice2d(
-    sampler, theta, theta_range, n_grid=100, num_processes=1, outname=None
+    result, theta, theta_range, n_grid=100, num_processes=1, outname=None
 ):
+    """Plot a 2D slice of the (unnormalized) target posterior through `theta`.
+
+    Parameters
+    ----------
+    result : Result
+        An importance-sampled result; provides the likelihood and log evidence.
+    theta : pd.Series
+        The parameter point the slice passes through.
+    theta_range : dict
+        `{parameter: (min, max)}` for the two scanned parameters.
+    """
     # prepare theta grid
     keys = list(theta_range.keys())
     axis0 = np.linspace(theta_range[keys[0]][0], theta_range[keys[0]][1], n_grid)
@@ -26,9 +37,9 @@ def plot_posterior_slice2d(
 
     # compute log_prob
     log_probs_target = (
-        sampler.likelihood.log_likelihood_multi(theta_grid, num_processes)
-        # + sampler.prior.ln_prob(theta_grid, axis=0)
-        - sampler.log_evidence
+        result.likelihood.log_likelihood_multi(theta_grid, num_processes)
+        # + result.prior.ln_prob(theta_grid, axis=0)
+        - result.log_evidence
     )
     log_probs_target -= np.max(log_probs_target)
     plt.clf()
@@ -47,7 +58,8 @@ def plot_posterior_slice2d(
 
 
 def plot_posterior_slice(
-    sampler,
+    result,
+    log_prob_proposal_fn,
     theta,
     theta_range,
     outname=None,
@@ -56,6 +68,21 @@ def plot_posterior_slice(
     parameters=None,
     normalize_conditionals=False,
 ):
+    """Plot 1D slices of the target posterior and the proposal through `theta`.
+
+    Parameters
+    ----------
+    result : Result
+        An importance-sampled result; provides the likelihood, prior, and log
+        evidence.
+    log_prob_proposal_fn : callable
+        Evaluates the proposal log-prob on a DataFrame of parameter points, e.g.
+        `FlowFactor.from_model(nde).log_prob` wrapped for DataFrames.
+    theta : pd.Series
+        The parameter point the slices pass through.
+    theta_range : dict
+        `{parameter: (min, max)}` for each scanned parameter.
+    """
     # Put a cap on number of processes to avoid overhead.
     # num_processes = min(num_processes, n_grid // 50)
 
@@ -81,13 +108,12 @@ def plot_posterior_slice(
         theta_param[param] = param_axis
         # evaluate the posterior at theta_grid
         log_probs_target = (
-            sampler.likelihood.log_likelihood_multi(theta_param, num_processes)
-            + sampler.prior.ln_prob(theta_param, axis=0)
-            - sampler.log_evidence
+            result.likelihood.log_likelihood_multi(theta_param, num_processes)
+            + result.prior.ln_prob(theta_param, axis=0)
+            - result.log_evidence
         )
-        # evaluate nde at theta_grid
-        # log_probs_proposal = get_log_probs_from_proposal(nde, theta_param)
-        log_probs_proposal = sampler.log_prob(theta_param)
+        # evaluate the proposal at theta_grid
+        log_probs_proposal = log_prob_proposal_fn(theta_param)
 
         # plot
         target = np.exp(log_probs_target)
@@ -181,7 +207,8 @@ def plot_diagnostics(
     #     for idx, (_, theta_idx) in enumerate(theta_slice_plots.iterrows()):
     #         # 1d slice plots
     #         plot_posterior_slice(
-    #             sampler,
+    #             result,
+    #             log_prob_proposal_fn,
     #             theta_idx,
     #             theta_range_1d,
     #             num_processes=num_processes,
@@ -201,7 +228,7 @@ def plot_diagnostics(
     #             theta_range_2d["phase"] = (0, 2 * np.pi)
     #             for param_pair in params_slice2d:
     #                 plot_posterior_slice2d(
-    #                     sampler,
+    #                     result,
     #                     theta_idx,
     #                     {k: theta_range_2d[k] for k in param_pair},
     #                     num_processes=num_processes,
diff --git a/dingo/gw/importance_sampling/importance_weights.py b/dingo/gw/importance_sampling/importance_weights.py
index 7f88fa67e..eecd35d92 100644
--- a/dingo/gw/importance_sampling/importance_weights.py
+++ b/dingo/gw/importance_sampling/importance_weights.py
@@ -2,6 +2,7 @@
 Step 1: Train unconditional nde
 Step 2: Set up likelihood and prior
 """
+
 from pathlib import Path
 
 import yaml
@@ -11,7 +12,8 @@
 
 from dingo.core.posterior_models import NormalizingFlowPosteriorModel
 from dingo.gw.result import Result
-from dingo.gw.inference.gw_samplers import GWSampler
+from dingo.core.factors import ChainComposer, FlowFactor
+from dingo.gw.inference.factors import GWComposedSampler
 from dingo.gw.importance_sampling.diagnostics import plot_diagnostics
 
 
@@ -116,9 +118,16 @@ def main():
             print(f"Renaming trained nde model to {nde_name}.")
             rename(join(args.outdir, "model_latest.pt"), nde_name)
 
-        # Step 1a: Sample from proposal.
-        nde_sampler = GWSampler(model=nde)
-        nde_sampler.run_sampler(num_samples=settings["num_samples"])
+        # Step 1a: Sample from proposal. The unconditional NDE is the sole
+        # chain factor; the event payload and analysis views come from the
+        # original result's sampler context.
+        nde_sampler = GWComposedSampler(
+            ChainComposer([FlowFactor.from_model(nde)]),
+            result.sampler_context,
+            nde.metadata,
+            nde.metadata["train_settings"]["data"]["inference_parameters"],
+        )
+        nde_sampler.run_sampler(settings["num_samples"])
         result = nde_sampler.to_result()
 
     # else:
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index c5db1ef1a..b8ae74127 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -72,11 +72,11 @@
 
 def _frequency_range_update(domain, event_metadata) -> Optional[dict]:
     """The event's requested frequency range when it differs from the data-domain
-    bounds (the legacy `frequency_updates` condition), else `None`. Values may be
-    floats or per-detector dicts; defaults are the domain bounds. Note a request
-    merely *wider* than the domain also triggers (generation writes base-domain
-    bounds, e.g. 1099.0 vs a multibanded band edge of 1098.875) -- the resulting
-    mask is then an identity, matching legacy behavior."""
+    bounds, else `None`. Values may be floats or per-detector dicts; defaults are
+    the domain bounds. A request merely *wider* than the domain also triggers --
+    data generation writes base-domain bounds, which can exceed a multibanded
+    domain's quantized band edge (e.g. 1099.0 vs 1098.875) -- and the resulting
+    mask is then an identity."""
     if event_metadata is None:
         return None
     minimum = event_metadata.get("minimum_frequency", domain.f_min)
@@ -273,11 +273,11 @@ def prepared_data(self) -> torch.Tensor:
         return self._prepared
 
     def _validate_frequency_range(self):
-        """Validate an event frequency-range update, exactly as the legacy sampler
-        setters did: hard bounds against the (base) domain, and narrowing only when
-        the network was trained with random strain cropping covering the requested
-        range. Applies to the network-input view only -- the likelihood view
-        applies the range independently via ASD masking."""
+        """Validate an event frequency-range update: hard bounds against the (base)
+        domain, and narrowing only when the network was trained with random strain
+        cropping covering the requested range. Applies to the network-input view
+        only -- the likelihood view applies the range independently via ASD
+        masking."""
         update = _frequency_range_update(self.domain, self.event_metadata)
         if update is None:
             return
@@ -611,8 +611,9 @@ def _phase_profile(self, given, context):
 
 
 def _build_gnpe_transforms(model: BasePosteriorModel):
-    """Build the time-shift GNPE per-step transforms from a model's metadata (the same
-    chains as `GWSamplerGNPE._initialize_transforms`).
+    """Build the time-shift GNPE per-step transforms from a model's metadata: the
+    proxy blur, the per-row time-shift alignment applied before the network, and the
+    post-network geocent-time correction.
 
     Returns
     -------
@@ -1097,11 +1098,10 @@ def __init__(
     def sampler_provenance(self) -> dict:
         """Provenance of how the samples were made, stored as `settings["sampler"]`
         in the exported `Result`: the executed chain in order (one descriptor per
-        step, via `Step.describe()`), plus anything in `provenance_extra`. Nothing
-        consumes this yet -- it is deliberate, structured provenance (the legacy
-        samplers recorded only ad-hoc `num_iterations` / `init_model` keys via
-        setter side effects); the `version` field is insurance for future
-        consumers such as chain reconstruction from file."""
+        step, via `Step.describe()`), plus anything in `provenance_extra`. The
+        block is purely a record -- nothing consumes it at load time -- and the
+        `version` field allows future consumers (e.g. chain reconstruction from
+        file) to evolve the format safely."""
         return {
             "version": 1,
             "implementation": "composed",
diff --git a/dingo/gw/inference/gw_samplers.py b/dingo/gw/inference/gw_samplers.py
deleted file mode 100644
index 52189e44a..000000000
--- a/dingo/gw/inference/gw_samplers.py
+++ /dev/null
@@ -1,510 +0,0 @@
-from typing import Union, Protocol
-
-import numpy as np
-import pandas as pd
-from astropy.time import Time
-from bilby.core.prior import PriorDict, DeltaFunction, Constraint
-from bilby.gw.detector import InterferometerList
-from torchvision.transforms import Compose
-
-from dingo.core.samplers import Sampler, GNPESampler
-from dingo.core.transforms import GetItem, RenameKey
-from dingo.gw.domains import (
-    MultibandedFrequencyDomain,
-    UniformFrequencyDomain,
-    Domain,
-)
-from dingo.gw.domains import build_domain
-from dingo.gw.frequency_updates import (
-    _validate_maximum_frequency,
-    _validate_minimum_frequency,
-)
-from dingo.gw.gwutils import get_extrinsic_prior_dict
-from dingo.gw.prior import build_prior_with_defaults
-from dingo.gw.result import Result
-from dingo.gw.transforms import (
-    WhitenAndScaleStrain,
-    RepackageStrainsAndASDS,
-    ToTorch,
-    SelectStandardizeRepackageParameters,
-    GNPECoalescenceTimes,
-    TimeShiftStrain,
-    GNPEBase,
-    PostCorrectGeocentTime,
-    CopyToExtrinsicParameters,
-    GetDetectorTimes,
-    DecimateWaveformsAndASDS,
-    MaskDataForFrequencyRangeUpdate,
-)
-
-
-class SamplerProtocol(Protocol):
-    base_model_metadata: dict
-
-    def _initialize_transforms(self) -> None: ...
-
-
-class _GWMixinProtocol(SamplerProtocol):
-    detectors: list[str]
-    domain: Domain
-    random_strain_cropping: dict
-
-
-class GWSamplerMixin(object):
-    """
-    Mixin class designed to add gravitational wave functionality to Sampler classes:
-        * builder for data domain
-        * correction for fixed detector locations during training (t_ref)
-    """
-
-    def __init__(self: SamplerProtocol, **kwargs):
-        """
-        Parameters
-        ----------
-        kwargs
-            Keyword arguments that are forwarded to the superclass.
-        """
-        # Has to be specified before init, because the information is required in _initialize_transforms()
-        self._minimum_frequency = None
-        self._maximum_frequency = None
-        super().__init__(**kwargs)
-        self.t_ref = self.base_model_metadata["train_settings"]["data"]["ref_time"]
-        self._pesummary_package = "gw"
-        self._result_class = Result
-
-    @property
-    def detectors(self: SamplerProtocol):
-        return self.base_model_metadata["train_settings"]["data"]["detectors"]
-
-    @property
-    def random_strain_cropping(self: SamplerProtocol):
-        return self.base_model_metadata["train_settings"]["data"].get(
-            "random_strain_cropping"
-        )
-
-    @property
-    def minimum_frequency(self) -> float | dict[str, float]:
-        if self._minimum_frequency is not None:
-            return self._minimum_frequency
-        else:
-            return self.domain.f_min
-
-    @minimum_frequency.setter
-    def minimum_frequency(self: _GWMixinProtocol, value: dict[str, float] | float):
-        if isinstance(self.domain, MultibandedFrequencyDomain):
-            domain = self.domain.base_domain
-        elif isinstance(self.domain, UniformFrequencyDomain):
-            domain = self.domain
-        else:
-            raise ValueError("Frequency updates only possible for frequency domains.")
-        _validate_minimum_frequency(
-            value,
-            self.detectors,
-            domain,
-            self.random_strain_cropping,
-        )  # TODO: Ensure minimum frequency is a dict?
-        self._minimum_frequency = value
-        self._initialize_transforms()
-
-    @property
-    def maximum_frequency(self) -> float | dict[str, float]:
-        if self._maximum_frequency is not None:
-            return self._maximum_frequency
-        else:
-            return self.domain.f_max
-
-    @maximum_frequency.setter
-    def maximum_frequency(self: _GWMixinProtocol, value: Union[float, dict]):
-        if isinstance(self.domain, MultibandedFrequencyDomain):
-            domain = self.domain.base_domain
-        elif isinstance(self.domain, UniformFrequencyDomain):
-            domain = self.domain
-        else:
-            raise ValueError("Frequency updates only possible for frequency domains.")
-        _validate_maximum_frequency(
-            value,
-            self.detectors,
-            domain,
-            self.random_strain_cropping,
-        )
-        self._maximum_frequency = value
-        self._initialize_transforms()
-
-    @property
-    def frequency_updates(self) -> bool:
-        def normalize(val):
-            if isinstance(val, dict):
-                return set(val.values())
-            return {val}
-
-        return normalize(self.minimum_frequency) != {self.domain.f_min} or normalize(
-            self.maximum_frequency
-        ) != {self.domain.f_max}
-
-    @property
-    def event_metadata(self):
-        if self._event_metadata is not None:
-            metadata = self._event_metadata.copy()
-        else:
-            metadata = {}
-        metadata["minimum_frequency"] = self.minimum_frequency
-        metadata["maximum_frequency"] = self.maximum_frequency
-        return metadata
-
-    @event_metadata.setter
-    def event_metadata(self, value):
-        if value is not None:
-            value = value.copy()
-            if "minimum_frequency" in value:
-                self.minimum_frequency = value.pop("minimum_frequency")
-            if "maximum_frequency" in value:
-                self.maximum_frequency = value.pop("maximum_frequency")
-        self._event_metadata = value
-
-    def _build_domain(self: Sampler):
-        """
-        Construct the domain object based on model metadata.
-
-        Called by __init__() immediately after _build_prior().
-        """
-        self.domain = build_domain(
-            self.base_model_metadata["dataset_settings"]["domain"]
-        )
-
-        data_settings = self.base_model_metadata["train_settings"]["data"]
-        if "domain_update" in data_settings:
-            self.domain.update(data_settings["domain_update"])
-
-    def _correct_reference_time(
-        self: Sampler, samples: Union[dict, pd.DataFrame], inverse: bool = False
-    ):
-        """
-        Correct the sky position of an event based on the reference time of the model.
-        This is necessary since the model was trained with with fixed detector (reference)
-        positions. This transforms the right ascension based on the e difference between
-        the time of the event and t_ref.
-
-        The correction is only applied if the event time can be found in self.metadata[
-        'event'].
-
-        This method modifies the samples in place.
-
-        Parameters
-        ----------
-        samples : dict or pd.DataFrame
-        inverse : bool, default True
-            Whether to apply instead the inverse transformation. This is used prior to
-            calculating the log_prob.
-        """
-        if self.event_metadata is not None:
-            t_event = self.event_metadata.get("time_event")
-            if t_event is not None and t_event != self.t_ref and "ra" in samples:
-                ra = samples["ra"]
-                time_reference = Time(self.t_ref, format="gps", scale="utc")
-                time_event = Time(t_event, format="gps", scale="utc")
-                longitude_event = time_event.sidereal_time("apparent", "greenwich")
-                longitude_reference = time_reference.sidereal_time(
-                    "apparent", "greenwich"
-                )
-                delta_longitude = longitude_event - longitude_reference
-                ra_correction = delta_longitude.rad
-                if not inverse:
-                    samples["ra"] = (ra + ra_correction) % (2 * np.pi)
-                else:
-                    samples["ra"] = (ra - ra_correction) % (2 * np.pi)
-
-    def _post_process(self, samples: Union[dict, pd.DataFrame], inverse: bool = False):
-        """
-        Post-processing of parameter samples.
-        * Add any fixed parameters from the prior.
-        * Correct the sky position for a potentially fixed reference time.
-          (see self._correct_reference_time)
-
-        This method modifies the samples in place.
-
-        Parameters
-        ----------
-        samples : dict or pd.DataFrame
-        inverse : bool, default True
-            Whether to apply instead the inverse transformation. This is used prior to
-            calculating the log_prob.
-        """
-        intrinsic_prior = self.metadata["dataset_settings"]["intrinsic_prior"]
-        extrinsic_prior = get_extrinsic_prior_dict(
-            self.metadata["train_settings"]["data"]["extrinsic_prior"]
-        )
-        prior = build_prior_with_defaults({**intrinsic_prior, **extrinsic_prior})
-
-        if not inverse:
-            # Add fixed parameters from prior.
-            num_samples = len(samples[list(samples.keys())[0]])
-            for k, p in prior.items():
-                if isinstance(p, DeltaFunction) and k not in samples:
-                    v = p.peak
-                    print(f"Adding fixed parameter {k} = {v} from prior.")
-                    samples[k] = p.peak * np.ones(num_samples)
-        else:
-            # Drop non-inference parameters from samples.
-            # NOTE: Important to drop "log_prob" in particular before running
-            # Sampler.log_prob(), otherwise log probabilities are added.
-            drop_parameters = [
-                k for k in samples.keys() if k not in self.inference_parameters
-            ]
-            if isinstance(samples, pd.DataFrame):
-                samples.drop(columns=drop_parameters, inplace=True, errors="ignore")
-            elif isinstance(samples, dict):
-                for k in drop_parameters:
-                    samples.pop(k, None)
-
-        if not self.unconditional_model:
-            self._correct_reference_time(samples, inverse)
-
-
-class GWSampler(GWSamplerMixin, Sampler):
-    """
-    Sampler for gravitational-wave inference using neural posterior estimation.
-    Augments the base class by defining transform_pre and transform_post to prepare
-    data for the inference network.
-
-    transform_pre :
-        * Decimates data (if necessary and using MultibandedFrequencyDomain).
-        * Whitens strain.
-        * Repackages strain data and the inverse ASDs (suitably scaled) into a torch
-          tensor.
-
-    transform_post :
-        * Extract the desired inference parameters from the network output (
-          array-like), de-standardize them, and repackage as a dict.
-
-    Also mixes in GW functionality for building the domain and correcting the reference
-    time.
-
-    Allows for conditional and unconditional models, and draws samples from the model
-    based on (optional) context data.
-
-    This is intended for use either as a standalone sampler, or as a sampler producing
-    initial sample points for a GNPE sampler.
-    """
-
-    def _initialize_transforms(self):
-        # preprocessing transforms:
-        transform_pre = []
-        #   * in case of MultibandedFrequencyDomain, decimate data from base domain
-        if isinstance(self.domain, MultibandedFrequencyDomain):
-            transform_pre.append(
-                DecimateWaveformsAndASDS(self.domain, decimation_mode="whitened")
-            )
-
-        #   * whiten and scale strain (since the inference network expects standardized
-        #   data)
-        transform_pre.append(WhitenAndScaleStrain(self.domain.noise_std))
-        if self.frequency_updates:
-            # * update frequency range
-            # Needs to happen before RepackageStrainsAndASDs since we might need to apply
-            # detectors specific frequency updates.
-            transform_pre.append(
-                MaskDataForFrequencyRangeUpdate(
-                    domain=self.domain,
-                    minimum_frequency=self.minimum_frequency,
-                    maximum_frequency=self.maximum_frequency,
-                )
-            )
-        #   * repackage strains and asds from dicts to an array
-        #   * convert array to torch tensor on the correct device
-        #   * extract only strain/waveform from the sample
-        transform_pre += [
-            # Use base metadata so that unconditional samplers still know how to
-            # transform data, since this transform is used by the GNPE sampler as
-            # well.
-            RepackageStrainsAndASDS(
-                ifos=self.detectors,
-                first_index=self.domain.min_idx,
-            ),
-            ToTorch(device=self.model.device),
-            GetItem("waveform"),
-        ]
-        self.transform_pre = Compose(transform_pre)
-
-        # postprocessing transforms:
-        #   * de-standardize data and extract inference parameters
-        self.transform_post = SelectStandardizeRepackageParameters(
-            {"inference_parameters": self.inference_parameters},
-            self.metadata["train_settings"]["data"]["standardization"],
-            inverse=True,
-            as_type="dict",
-        )
-
-
-class GWSamplerGNPE(GWSamplerMixin, GNPESampler):
-    """
-    Gravitational-wave GNPE sampler. It wraps a PosteriorModel and a standard Sampler for
-    initialization. The former is used to generate initial samples for Gibbs sampling.
-
-    Compared to the base class, this class implements the required transforms for
-    preparing data and parameters for the network. This includes GNPE transforms,
-    data processing transforms, and standardization/de-standardization of parameters.
-
-    A GNPE network is conditioned on additional "proxy" context theta^, i.e.,
-
-    p(theta | theta^, d)
-
-    The theta^ depend on theta via a fixed kernel p(theta^ | theta). Combining these
-    known distributions, this class uses Gibbs sampling to draw samples from the joint
-    distribution,
-
-    p(theta, theta^ | d)
-
-    The advantage of this approach is that we are allowed to perform any transformation of
-    d that depends on theta^. In particular, we can use this freedom to simplify the
-    data, e.g., by aligning data to have merger times = 0 in each detector. The merger
-    times are unknown quantities that must be inferred jointly with all other
-    parameters, and GNPE provides a means to do this iteratively. See
-    https://arxiv.org/abs/2111.13139 for additional details.
-
-    Gibbs sampling breaks access to the probability density, so this must be recovered
-    through other means. One way is to train an unconditional flow to represent p(theta^
-    | d) for fixed d based on the samples produced through the GNPE Gibbs sampling.
-    Starting from these, a single Gibbs iteration gives theta from the GNPE network,
-    along with the probability density in the joint space. This is implemented in
-    GNPESampler provided the init_sampler provides proxies directly and num_iterations
-    = 1.
-
-    Attributes (beyond those of Sampler)
-    ------------------------------------
-    init_sampler : Sampler
-        Used for providing initial samples for Gibbs sampling.
-    num_iterations : int
-        Number of Gibbs iterations to perform.
-    iteration_tracker : IterationTracker
-        **not set up**
-    remove_init_outliers : float
-        **not set up**
-    """
-
-    @property
-    def minimum_frequency(self) -> float | dict[str, float]:
-        if self.init_sampler is not None:
-            return self.init_sampler.minimum_frequency
-        else:
-            raise AttributeError(
-                "init_sampler not set. Cannot access minimum frequency."
-            )
-
-    @minimum_frequency.setter
-    def minimum_frequency(self, value):
-        if self.init_sampler is not None:
-            self.init_sampler.minimum_frequency = value
-        else:
-            raise AttributeError(
-                "init_sampler not set. Cannot update minimum frequency."
-            )
-
-    @property
-    def maximum_frequency(self) -> float | dict[str, float]:
-        if self.init_sampler is not None:
-            return self.init_sampler.maximum_frequency
-        else:
-            raise AttributeError(
-                "init_sampler not set. Cannot access maximum frequency."
-            )
-
-    @maximum_frequency.setter
-    def maximum_frequency(self, value):
-        if self.init_sampler is not None:
-            self.init_sampler.maximum_frequency = value
-        else:
-            raise AttributeError(
-                "init_sampler not set. Cannot update maximum frequency."
-            )
-
-    def _initialize_transforms(self):
-        """
-        Builds the transforms that are used in the GNPE loop.
-        """
-        data_settings = self.metadata["train_settings"]["data"]
-        ifo_list = InterferometerList(data_settings["detectors"])
-
-        gnpe_time_settings = data_settings.get("gnpe_time_shifts")
-        gnpe_chirp_settings = data_settings.get("gnpe_chirp")
-        gnpe_phase_settings = data_settings.get("gnpe_phase")
-        if (
-            not gnpe_time_settings
-            and not gnpe_chirp_settings
-            and not gnpe_phase_settings
-        ):
-            raise KeyError(
-                "GNPE inference requires network trained for either chirp mass, "
-                "coalescence time, or phase GNPE."
-            )
-
-        # transforms for gnpe loop, to be applied prior to sampling step:
-        #   * reset the sample (e.g., clone non-gnpe transformed waveform)
-        #   * blurring detector times to obtain gnpe proxies
-        #   * shifting the strain by - gnpe proxies
-        #   * repackaging & standardizing proxies to sample['context_parameters']
-        #     for conditioning of the inference network
-        transform_pre = []
-        transform_pre.append(RenameKey("data", "waveform"))
-        if gnpe_time_settings:
-            transform_pre.append(
-                GNPECoalescenceTimes(
-                    ifo_list,
-                    gnpe_time_settings["kernel"],
-                    gnpe_time_settings["exact_equiv"],
-                    inference=True,
-                )
-            )
-            transform_pre.append(TimeShiftStrain(ifo_list, self.domain))
-        transform_pre.append(
-            SelectStandardizeRepackageParameters(
-                {"context_parameters": data_settings["context_parameters"]},
-                data_settings["standardization"],
-                device=self.model.device,
-            )
-        )
-        transform_pre.append(RenameKey("waveform", "data"))
-
-        # Extract GNPE information (list of parameters, dict of kernels) from the
-        # transforms.
-        self.gnpe_parameters = []
-        self.gnpe_kernel = PriorDict()
-        for transform in transform_pre:
-            if isinstance(transform, GNPEBase):
-                self.gnpe_parameters += transform.input_parameter_names
-                for k, v in transform.kernel.items():
-                    self.gnpe_kernel[k] = v
-        print("GNPE parameters: ", self.gnpe_parameters)
-        print("GNPE kernel: ", self.gnpe_kernel)
-
-        self.transform_pre = Compose(transform_pre)
-
-        # transforms for gnpe loop, to be applied after sampling step:
-        #   * de-standardization of parameters
-        #   * post correction for geocent time (required for gnpe with exact equivariance)
-        #   * computation of detectortimes from parameters (required for next gnpe
-        #       iteration)
-        self.transform_post = Compose(
-            [
-                SelectStandardizeRepackageParameters(
-                    {"inference_parameters": self.inference_parameters},
-                    data_settings["standardization"],
-                    inverse=True,
-                    as_type="dict",
-                ),
-                PostCorrectGeocentTime(),
-                CopyToExtrinsicParameters(
-                    "ra", "dec", "geocent_time", "chirp_mass", "mass_ratio", "phase"
-                ),
-                GetDetectorTimes(ifo_list, data_settings["ref_time"]),
-            ]
-        )
-
-    def _kernel_log_prob(self, samples):
-        # TODO: Reimplement as a method of GNPEBase.
-        if len({"chirp_mass", "mass_ratio", "phase"} & self.gnpe_kernel.keys()) > 0:
-            raise NotImplementedError("kernel log_prob only implemented for time gnpe.")
-        gnpe_proxies_diff = {
-            k: np.array(samples[k] - samples[f"{k}_proxy"])
-            for k in self.gnpe_kernel.keys()
-        }
-        return self.gnpe_kernel.ln_prob(gnpe_proxies_diff, axis=0)
diff --git a/dingo/gw/inference/inference_utils.py b/dingo/gw/inference/inference_utils.py
deleted file mode 100644
index 1ed544d8d..000000000
--- a/dingo/gw/inference/inference_utils.py
+++ /dev/null
@@ -1,68 +0,0 @@
-from typing import Optional
-
-import numpy as np
-
-from dingo.gw.inference.gw_samplers import GWSampler
-
-
-def prepare_log_prob(
-    sampler,
-    num_samples: int,
-    nde_settings: dict,
-    batch_size: Optional[int] = None,
-    threshold_std: Optional[float] = np.inf,
-    remove_init_outliers: Optional[float] = 0.0,
-    low_latency_label: str = None,
-    outdir: str = None,
-):
-    """
-    Prepare gnpe sampling with log_prob. This is required, since in its vanilla
-    form gnpe does not provide the density for its samples.
-
-    Specifically, we train an unconditional neural density estimator (nde) for the
-    gnpe proxies. This requires running the gnpe sampler till convergence, and
-    extracting the gnpe proxies after the final gnpe iteration. The nde is trained
-    to match the distribution over gnpe proxies, which provides a way of rapidly
-    sampling (converged!) gnpe proxies *and* evaluating the log_prob.
-
-    After this preparation step, self.run_sampler can leverage
-    self.gnpe_proxy_sampler (which is based on the aforementioned trained nde) to
-    sample gnpe proxies, such that one gnpe iteration is sufficient. The
-    log_prob of
-    the samples in the *joint* space (inference parameters + gnpe proxies) is then
-    simply given by the sum of the corresponding log_probs (from self.model and
-    self.gnpe_proxy_sampler.model).
-
-    Parameters
-    ----------
-    num_samples: int
-        number of samples for training of nde
-    batch_size: int = None
-        batch size for sampler
-    threshold_std: float = np.inf
-        gnpe proxies deviating by more then threshold_std standard deviations from
-        the proxy mean (along any axis) are discarded.
-    low_latency_label: str = None
-        File label for low latency samples (= samples used for training nde).
-        If None, these samples are not saved.
-    outdir: str = None
-        Directory in which low latency samples are saved. Needs to be set if
-        low_latency_label is not None.
-    """
-    sampler.remove_init_outliers = remove_init_outliers
-    sampler.run_sampler(num_samples, batch_size)
-    if low_latency_label is not None:
-        sampler.to_hdf5(label=low_latency_label, outdir=outdir)
-    result = sampler.to_result()
-    nde_settings["training"]["device"] = str(sampler.model.device)
-    unconditional_model = result.train_unconditional_flow(
-        sampler.gnpe_proxy_parameters,
-        nde_settings,
-        threshold_std=threshold_std,
-    )
-
-    # Prepare sampler with unconditional model as initialization. This should only use
-    # one iteration and also not remove any outliers.
-    sampler.init_sampler = GWSampler(model=unconditional_model)
-    sampler.num_iterations = 1
-    sampler.remove_init_outliers = 0.0  # Turn off for final sampler.
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index ccc51fdbe..b7bc49c35 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -554,11 +554,11 @@ def _sample_synthetic_phase_chain(
         within-prior proposal samples with their stored log-prob, and
         `SyntheticPhaseFactor` draws `phase` -- the composer's ordinary log-prob fold
         returns the joint proposal density `log q(theta) + log q(phase | theta, d)`,
-        which replaces the samples' log_prob (the legacy `log_prob += delta`). The
-        Result's likelihood view (base domain, rebuilt domain, frequency updates) is
-        passed to the factor, so it shares the cached likelihood instance built by
-        `_build_likelihood`. Out-of-prior rows keep the legacy convention:
-        `phase = 0`, `log_prob = nan`."""
+        which becomes the samples' log_prob. The Result's likelihood view (base
+        domain, rebuilt domain, frequency updates) is passed to the factor, so it
+        shares the cached likelihood instance built by `_build_likelihood`.
+        Out-of-prior rows get `phase = 0` and `log_prob = nan` (they receive zero
+        weight in importance sampling regardless)."""
         from dingo.core.factors import ChainComposer, SampleTableFactor
         from dingo.gw.inference.factors import SyntheticPhaseFactor
 
@@ -590,7 +590,8 @@ def _sample_synthetic_phase_chain(
         # Insert the phase prior in the prior, since now the phase is present.
         self.prior["phase"] = self.phase_prior
         self.phase_prior = None
-        # reset likelihood for safety (matches the legacy path)
+        # Reset so importance sampling rebuilds the likelihood with its own
+        # marginalization settings rather than reusing the phase-full one.
         self.likelihood = None
 
     def sample_synthetic_phase(
diff --git a/dingo/pipe/parser.py b/dingo/pipe/parser.py
index de066d493..86691bf17 100644
--- a/dingo/pipe/parser.py
+++ b/dingo/pipe/parser.py
@@ -1535,11 +1535,12 @@ def create_parser(top_level=True, usage=None):
     sampler_parser.add(
         "--sampler-implementation",
         type=str,
-        default="legacy",
+        default="composed",
         choices=["legacy", "composed"],
-        help="Sampler implementation: 'legacy' (GWSampler / GWSamplerGNPE) or "
-        "'composed' (the factorized GWComposedSampler; experimental). "
-        "Default 'legacy'.",
+        help="Sampler implementation. 'composed' (the factorized "
+        "GWComposedSampler) is the only implementation; 'legacy' is accepted "
+        "for a clear error message (the last commit containing the legacy "
+        "samplers is tagged legacy-samplers-final). Default 'composed'.",
     )
     sampler_parser.add(
         "--recover-log-prob",
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index e3c63f341..2df334f94 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -17,8 +17,6 @@
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
 from dingo.gw.data.event_dataset import EventDataset
 from dingo.gw.inference.factors import GNPEKernelFactor, GWComposedSampler
-from dingo.gw.inference.gw_samplers import GWSampler, GWSamplerGNPE
-from dingo.gw.inference.inference_utils import prepare_log_prob
 from dingo.pipe.default_settings import DENSITY_RECOVERY_SETTINGS
 from dingo.pipe.parser import create_parser
 
@@ -118,53 +116,38 @@ def _load_event(self):
         self.event_metadata = event_dataset.settings
 
     def _load_sampler(self):
-        """Load the sampler and set its context based on event data."""
+        """Build the sampler from the model(s) and the event data."""
+        if self.sampler_implementation == "legacy":
+            raise ValueError(
+                "The legacy samplers (GWSampler / GWSamplerGNPE) have been removed; "
+                "use sampler-implementation = composed. The last commit containing "
+                "them is tagged legacy-samplers-final."
+            )
         model = build_model_from_kwargs(
             filename=self.model, device=self.device, load_training_info=False
         )
 
-        if self.sampler_implementation == "composed":
-            if self.model_init is not None:
-                self.gnpe = True
-                self._main_model = model
-                init_model = build_model_from_kwargs(
-                    filename=self.model_init,
-                    device=self.device,
-                    load_training_info=False,
-                )
-                self.dingo_sampler = GWComposedSampler.from_gnpe_models(
-                    init_model,
-                    model,
-                    self.context,
-                    event_metadata=self.event_metadata,
-                    num_iterations=self.num_gnpe_iterations,
-                )
-            else:
-                self.gnpe = False
-                self.dingo_sampler = GWComposedSampler.from_model(
-                    model, self.context, event_metadata=self.event_metadata
-                )
-            self.dingo_sampler.provenance_extra["models"] = self._model_paths()
-            return
-
         if self.model_init is not None:
             self.gnpe = True
+            self._main_model = model
             init_model = build_model_from_kwargs(
-                filename=self.model_init, device=self.device, load_training_info=False
+                filename=self.model_init,
+                device=self.device,
+                load_training_info=False,
             )
-            init_sampler = GWSampler(model=init_model)
-            self.dingo_sampler = GWSamplerGNPE(
-                model=model,
-                init_sampler=init_sampler,
+            self.dingo_sampler = GWComposedSampler.from_gnpe_models(
+                init_model,
+                model,
+                self.context,
+                event_metadata=self.event_metadata,
                 num_iterations=self.num_gnpe_iterations,
             )
-
         else:
             self.gnpe = False
-            self.dingo_sampler = GWSampler(model=model)
-
-        self.dingo_sampler.context = self.context
-        self.dingo_sampler.event_metadata = self.event_metadata
+            self.dingo_sampler = GWComposedSampler.from_model(
+                model, self.context, event_metadata=self.event_metadata
+            )
+        self.dingo_sampler.provenance_extra["models"] = self._model_paths()
 
     @property
     def density_recovery_settings(self):
@@ -215,10 +198,10 @@ def _model_paths(self) -> dict:
         return paths
 
     def _recover_log_prob_composed(self):
-        """Density recovery on the composed sampler, mirroring `prepare_log_prob`:
-        run the Gibbs chain, train an unconditional NDE on the GNPE proxies, then
-        swap in a single-step (density-preserving) sampler with the NDE as the
-        proxy source."""
+        """Recover the sample density for a GNPE model: run the Gibbs chain, train
+        an unconditional NDE on the GNPE proxies, then swap in a single-step
+        (density-preserving) sampler with the NDE as the proxy source. The final
+        samples then carry a log_prob, as required for importance sampling."""
         settings = copy.deepcopy(self.density_recovery_settings)
         self.dingo_sampler.run_sampler(
             settings["num_samples"], batch_size=self.batch_size
@@ -245,8 +228,7 @@ def _recover_log_prob_composed(self):
         self.dingo_sampler.provenance_extra["models"] = self._model_paths()
         # The recovery recipe: the Gibbs chain that generated the NDE's training
         # samples (with its iteration count inside the GibbsBlock descriptor) plus
-        # the NDE training settings. The legacy metadata recorded only the final
-        # single-step state.
+        # the NDE training settings.
         self.dingo_sampler.provenance_extra["density_recovery"] = {
             "chain": gibbs_chain,
             **settings,
@@ -259,15 +241,7 @@ def run_sampler(self):
                 "samples and training a new network to recover it."
             )
 
-            # Note that this will not save any low latency samples at present.
-            if self.sampler_implementation == "composed":
-                self._recover_log_prob_composed()
-            else:
-                prepare_log_prob(
-                    self.dingo_sampler,
-                    batch_size=self.batch_size,
-                    **self.density_recovery_settings,
-                )
+            self._recover_log_prob_composed()
 
         self.dingo_sampler.run_sampler(self.num_samples, batch_size=self.batch_size)
         self.dingo_sampler.to_hdf5(label=self.label, outdir=self.result_directory)
diff --git a/docs/source/dingo_pipe.md b/docs/source/dingo_pipe.md
index 9d9b8478a..fce9f4eb5 100644
--- a/docs/source/dingo_pipe.md
+++ b/docs/source/dingo_pipe.md
@@ -94,7 +94,7 @@ Injections mirror the [approach of `bilby_pipe`](https://lscsoft.docs.ligo.org/b
 
 ## Sampling
 
-The next step is sampling from the Dingo model. The model is loaded into a [GWSampler](dingo.gw.inference.gw_samplers.GWSampler) or [GWSamplerGNPE](dingo.gw.inference.gw_samplers.GWSamplerGNPE) object. (If using [GNPE](gnpe) it is necessary to specify a `model-init`.) The Sampler `context` is then set from the EventDataset prepared in the previous step. `num-samples` samples are then generated in batches of size `batch-size`. The samples (and context) are stored in a [Result](dingo.gw.result.Result) object and saved in HDF5 format.
+The next step is sampling from the Dingo model. The model is loaded into a [GWComposedSampler](dingo.gw.inference.factors.GWComposedSampler), built from the EventDataset prepared in the previous step. (If using [GNPE](gnpe) it is necessary to specify a `model-init`.) `num-samples` samples are then generated in batches of size `batch-size`. The samples (and context) are stored in a [Result](dingo.gw.result.Result) object and saved in HDF5 format.
 
 If using GNPE, one can optionally specify `num-gnpe-iterations` (it defaults to 30). Importantly, obtaining the log probability when using GNPE requires an [extra step of training an unconditional flow](result.md#density-recovery). This is done using the `recover-log-prob` flag, which defaults to `True`. The default density recovery settings can be overwritten by providing a `density-recovery-settings` dictionary in the `.ini` file.
 
diff --git a/docs/source/example_injection.md b/docs/source/example_injection.md
index 1f0ec73d8..f92fd25c5 100644
--- a/docs/source/example_injection.md
+++ b/docs/source/example_injection.md
@@ -56,11 +56,11 @@ strain_data = injection_generator.injection(theta)
 Then one can create a injections and do inference on them.
 
 ```
-from dingo.gw.inference.gw_samplers import GWSamplerGNPE, GWSampler
+from dingo.gw.inference.factors import GWComposedSampler
 
-init_sampler = GWSampler(model=init_pm)
-sampler = GWSamplerGNPE(model=main_pm, init_sampler=init_sampler, num_iterations=30)
-sampler.context = strain_data
+sampler = GWComposedSampler.from_gnpe_models(
+    init_pm, main_pm, strain_data, num_iterations=30
+)
 sampler.run_sampler(num_samples=50_000, batch_size=10_000)
 result = sampler.to_result()
 result.plot_corner()
diff --git a/docs/source/gnpe.md b/docs/source/gnpe.md
index 189b8009e..e82ae51b7 100644
--- a/docs/source/gnpe.md
+++ b/docs/source/gnpe.md
@@ -105,16 +105,16 @@ The number of Gibbs iterations is also specified here (typically 30 is appropria
 
 
 
-## The `GNPESampler` class
+## The `GibbsBlock` step
 
-The inference script above uses the `GWSamplerGNPE` class, which is based on `GNPESampler`,
+GNPE inference uses `GWComposedSampler.from_gnpe_models(init_model, main_model, ...)`,
+which wraps the cyclic GNPE loop in a `GibbsBlock` chain step,
 ```{eval-rst}
-.. autoclass:: dingo.core.samplers.GNPESampler
+.. autoclass:: dingo.core.factors.GibbsBlock
    :members:
-   :inherited-members:
    :show-inheritance:
 ```
-In addition to storing a `PosteriorModel`, a `GNPESampler` also stores a second `Sampler` instance, which is based on the initialization network.  When `run_sampler()` is called, it first generates samples from the initialization network, perturbs them with the kernel to obtain proxy samples, and then performs `num_iterations` Gibbs steps to obtain the final samples.
+The block seeds the chain with samples from the initialization network, perturbs them with the kernel to obtain proxy samples, and then performs `num_iterations` Gibbs sweeps (kernel and main-network factors in turn) to obtain the final samples. Because Gibbs sampling breaks density access, the chain yields no log probability; `dingo_pipe` recovers it by training an unconditional flow for the proxies and re-sampling with a single density-preserving step (see [density recovery](result.md#density-recovery)).
 
 ```{eval-rst}
 .. footbibliography::
diff --git a/docs/source/inference.md b/docs/source/inference.md
index 767137d80..c1103c4a8 100644
--- a/docs/source/inference.md
+++ b/docs/source/inference.md
@@ -4,19 +4,22 @@ With a trained network, inference can be performed on injections or real data. F
 injections, see the [discussion in the examples](example_injection.md). For real data, we 
 recommend to use [dingo_pipe](dingo_pipe.md).
 
-## The `Sampler` class
+## The `GWComposedSampler` class
 
-Inference uses the `Sampler` class, or more specifically, the `GWSampler` class,
-which inherits from it.
+Inference uses the `GWComposedSampler` class, which represents the posterior as a
+chain of conditional factors (the network flow, deterministic reparametrizations
+such as the sky-frame rotation, and fillers for fixed parameters).
 
 ```{eval-rst}
-.. autoclass:: dingo.gw.inference.gw_samplers.GWSampler
+.. autoclass:: dingo.gw.inference.factors.GWComposedSampler
     :members:
     :inherited-members:
     :show-inheritance:
 ```
 
-This is instantiated based on a `PosteriorModel`. To draw samples, the `context` property must first be set to the data to be analyzed. For gravitational waves this should be a dictionary with the following keys:
+A plain-NPE sampler is built with `GWComposedSampler.from_model(model, event_data,
+event_metadata)`, where `model` is a `PosteriorModel` and `event_data` is the data
+to be analyzed---a dictionary with the following keys:
 
 waveform
 : (unwhitened) strain data in each detector
@@ -27,14 +30,14 @@ asds
 parameters (optional)
 : for injections, the true parameters of the signal (for saving; ignored for sampling)
 
-Once this is set, the `run_sampler()` method draws the requested samples from the posterior conditioned on the context. It applies some post-processing (to de-standardize the data, and to correct for the rotation of the Earth between the network reference time and the event time), and then stores the result as a DataFrame in `GWSampler.samples`. The DataFrame contains columns for each inference parameter, as well as the log probability of the sample under the posterior model.
+The `run_sampler()` method then draws the requested samples from the posterior conditioned on the data. All processing---de-standardization inside the network factor, and the correction for the rotation of the Earth between the network reference time and the event time---is expressed as chain steps, and the samples are stored as a DataFrame in `GWComposedSampler.samples`. The DataFrame contains columns for each inference parameter, as well as the log probability of the sample under the posterior model.
 
-The `GWSampler.metadata` attribute contains all settings that went into producing the samples, including training datasets, network training settings, event metadata (for real events) and possible injection parameters. Finally, the `to_samples_dataset()` method returns a `SamplesDataset` containing all results, including the samples, settings, and context. This can be saved easily as HDF5.
+The `GWComposedSampler.metadata` attribute contains all settings that went into producing the samples. The `to_result()` method returns a [Result](dingo.gw.result.Result) containing the samples, settings (including structured sampler provenance under `settings["sampler"]`), and data; `to_hdf5()` saves it directly.
 
 
 ## Injections
 
-Injections (i.e., simulated data) are produced using the `Injection` class. It includes options for fixed or random parameters (drawn from a prior), and it returns injections in a format that can be directly set as `GWSampler.context`.
+Injections (i.e., simulated data) are produced using the `Injection` class. It includes options for fixed or random parameters (drawn from a prior), and it returns injections in a format that can be passed directly as the `event_data` of a `GWComposedSampler`.
 
 ```{eval-rst}
 .. autoclass:: dingo.gw.injection.Injection
diff --git a/docs/source/result.md b/docs/source/result.md
index 261c2ab6a..083c1a8e1 100644
--- a/docs/source/result.md
+++ b/docs/source/result.md
@@ -19,7 +19,7 @@ When sampling with GNPE, there is no direct access to the probability density $q
 The `Result` class contains methods to enable *recovery* of the probability density for a collection of samples. The approach is as follows:
 1. Start from the samples $\{(\theta_i, \hat\theta_i)\}_{i=1}^N$ from the final Gibbs iteration, including parameters $\theta$ and proxy parameters $\hat\theta$. By default these are included in the `samples` attribute generated by the `Sampler`.
 2. Train an *unconditional* density estimator $q(\hat\theta)$ to model the proxy parameters. This is done by (1) using `parameter_subset()` to produce a new `Result` containing just the proxies, and (2) using `train_unconditional_flow()` on this subset.
-3. Generate new samples $(\theta, \hat\theta) \sim q(\theta, \hat\theta | d) = q(\theta | d, \hat\theta) q(\hat\theta)$. This can be accomplished using `GNPESampler.sample()` with `num_iterations = 1` and setting the initial sampler to be the unconditional flow trained in the previous step. Since this does not involve multiple iterations, the density is obtained as well, so importance sampling can be performed.
+3. Generate new samples $(\theta, \hat\theta) \sim q(\theta, \hat\theta | d) = q(\theta | d, \hat\theta) q(\hat\theta)$. This is accomplished using `GWComposedSampler.from_singlestep_gnpe()`, with the unconditional flow trained in the previous step as the proxy source. The resulting chain is autoregressive rather than iterative, so the density is obtained as well, and importance sampling can be performed.
 
 ```{note}
 Density recovery can also be achieved using an unconditional density estimator for $\theta$ (trained on samples $\{\theta_i\}_{i=1}^N$ from GNPE). Since $\theta$ typically comprises 14 parameters (versus 2 or 3 for $\hat\theta$) it is usually more straightforward to learn the proxies.
diff --git a/tests/pipe/test_sampler_implementation.py b/tests/pipe/test_sampler_implementation.py
index 6ecda5453..a2b7e31cb 100644
--- a/tests/pipe/test_sampler_implementation.py
+++ b/tests/pipe/test_sampler_implementation.py
@@ -13,13 +13,15 @@ def _parse(tmp_path, extra=()):
     return args
 
 
-def test_default_is_legacy(tmp_path):
-    assert _parse(tmp_path).sampler_implementation == "legacy"
+def test_default_is_composed(tmp_path):
+    assert _parse(tmp_path).sampler_implementation == "composed"
 
 
-def test_composed_parses(tmp_path):
-    args = _parse(tmp_path, ["--sampler-implementation", "composed"])
-    assert args.sampler_implementation == "composed"
+def test_legacy_still_parses(tmp_path):
+    # Accepted at parse time so that old INI files fail at sampler construction
+    # with a clear pointer to the legacy-samplers-final tag, not an argparse error.
+    args = _parse(tmp_path, ["--sampler-implementation", "legacy"])
+    assert args.sampler_implementation == "legacy"
 
 
 def test_invalid_value_rejected(tmp_path):

From 07681d0017b912d15eee4ba63e19e095ca744d2a Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 6 Jul 2026 19:58:05 +0100
Subject: [PATCH 38/65] Consolidate Result onto the sampler context

The remaining self-build paths are removed: the prior and domain are
taken from the sampler context (the single construction owner), the
likelihood is delegated unconditionally, and sample_synthetic_phase
runs only as a chain over the context. The unreachable inverse mode is
removed with it (its validation raised whenever phase was already
present, and its only intended consumer was the removed density-recovery
workflow); the capability remains available as
SyntheticPhaseFactor.log_prob.

Context reconstruction now requires only the model metadata: a result
whose strain payload was stripped or left on disk keeps a working prior
and domain, and only the likelihood demands event data. Payloads without
full metadata load in a transport-only mode (prior/domain None, with a
visible message) instead of failing inside construction, and
_rebuild_domain tolerates that mode.

Result retains what is genuinely its own: the importance-sampling view
assembly, marginalization validation against the evolving prior, weight
and evidence computation, prior_update state, serialization, provenance,
and exports.

Verified: 131 tests; the importance-sampling harnesses reproduce
previously recorded evidence to the last digit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py       |  13 +-
 dingo/gw/result.py                  | 309 ++++++++--------------------
 tests/gw/test_gw_sampler_context.py |  10 +
 3 files changed, 110 insertions(+), 222 deletions(-)

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index b8ae74127..55427ce4e 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -379,6 +379,12 @@ def likelihood(
         if settings == self._likelihood_settings:
             return self._likelihood
 
+        if self.event_data is None:
+            raise ValueError(
+                "Building the likelihood requires event data (strain + ASDs), "
+                "which this context does not carry."
+            )
+
         # Marginalizing over a parameter needs the (uniform) prior the network
         # marginalized over; use it to parameterize the requested marginalization.
         # Bounds already provided by the caller (e.g. from an updated prior) win.
@@ -496,8 +502,11 @@ class SyntheticPhaseFactor(Factor):
 
     Two grid modes: `approximation_22_mode=True` assumes a (2, 2)-dominated signal (the
     whole waveform transforms as `exp(2i phase)`), giving `log L` from the complex overlap
-    `Re[(d | h(phase=0)) exp(2i phase)]`; `False` (the production default) sums the modes
-    exactly and requires the waveform generator's `spin_conversion_phase = 0`.
+    `Re[(d | h(phase=0)) exp(2i phase)]`; `False` sums the modes exactly and requires the
+    waveform generator's `spin_conversion_phase = 0`. Note the entry points differ on the
+    default: this factor and `dingo_pipe`'s `PhaseRecoveryDefault` use the exact mode,
+    while `Result.sample_synthetic_phase` defaults to the (2, 2) approximation when the
+    key is omitted.
     """
 
     def __init__(
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index b7bc49c35..6f5542a6d 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -8,18 +8,10 @@
 from bilby.gw.prior import CalibrationPriorDict
 from bilby_pipe.utils import CALIBRATION_CORRECTION_TYPE_LOOKUP
 
-from dingo.core.density import (
-    interpolated_sample_and_log_prob_multi,
-    interpolated_log_prob_multi,
-)
-from dingo.core.multiprocessing import apply_func_with_multiprocessing
 from dingo.core.result import Result as CoreResult
 from dingo.gw.conversion import change_spin_conversion_phase
 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.likelihood import StationaryGaussianGWLikelihood
-from dingo.gw.prior import build_prior_with_defaults
 from dingo.core.utils.backward_compatibility import check_minimum_version
 
 
@@ -49,8 +41,8 @@ class Result(CoreResult):
         context : dict
             Context data from which the samples were produced (e.g., strain data, ASDs).
         metadata : dict
-            Metadata inherited from the Sampler object. This describes the neural
-            networks and sampling settings used.
+            Metadata describing the neural networks and sampling settings used,
+            including structured sampler provenance under `settings["sampler"]`.
         event_metadata : dict
             Metadata for the event analyzed, including time, data conditioning, channel,
             and detector information.
@@ -161,17 +153,16 @@ def maximum_frequency(self, value: dict[str, float] | float):
         self.event_metadata["maximum_frequency"] = value
 
     def _build_domain(self):
-        """
-        Construct the domain object based on model metadata
-
-        Called by __init__() immediately after _build_prior().
-        """
-        self.domain = build_domain(self.base_metadata["dataset_settings"]["domain"])
+        """Take the data domain from the sampler context. After importance-sampling
+        updates, `_rebuild_domain` rebinds `self.domain` to a new object, and the
+        updated domain reaches the likelihood through `_likelihood_view_kwargs`;
+        the context keeps the original training-settings domain. Called by
+        __init__()."""
         check_minimum_version(self.version, raise_exception=False)
-
-        data_settings = self.base_metadata["train_settings"]["data"]
-        if "domain_update" in data_settings:
-            self.domain.update(data_settings["domain_update"])
+        if self.sampler_context is None:
+            self.domain = None
+            return
+        self.domain = self.sampler_context.domain
 
     def _rebuild_domain(self, verbose=False):
         """Rebuild the domain based on settings updated for importance sampling.
@@ -186,6 +177,10 @@ def _rebuild_domain(self, verbose=False):
 
         domain_keys = ["minimum_frequency", "maximum_frequency", "T"]
         if any(k in updates for k in domain_keys):
+            if self.domain is None:
+                # Transport-only result (no reconstructable context): nothing to
+                # rebuild; domain-dependent operations fail at use.
+                return
             # TODO: Make compatible with MultibandedFrequencyDomain.
             if isinstance(self.domain, MultibandedFrequencyDomain):
                 raise NotImplementedError()
@@ -217,7 +212,11 @@ def _build_context(self):
         (settings + event data + event metadata), so that prior (and, later,
         likelihood) construction delegates to `GWSamplerContext` no matter how the
         Result was born -- live from a sampler or loaded from file."""
-        if self.context is None or self.settings is None:
+        # Only the metadata is required: the prior and domain views are defined
+        # without event data (a result whose strain payload was stripped or left
+        # on disk still has a working prior); the likelihood view checks for the
+        # event data itself.
+        if self.settings is None:
             return None
         from dingo.gw.inference.factors import GWSamplerContext
 
@@ -227,25 +226,30 @@ def _build_context(self):
             return GWSamplerContext.from_model_metadata(
                 self.base_metadata, self.context, self.event_metadata
             )
-        except (KeyError, TypeError):
-            # Settings are not a full model metadata (e.g. parameter subsets); the
-            # self-build paths below still apply.
+        except (KeyError, TypeError) as e:
+            # Settings are not a full model metadata (e.g. minimal test payloads);
+            # the result is then transport-only. Make the degradation visible
+            # rather than silent.
+            print(
+                f"Could not reconstruct a sampler context from the result settings "
+                f"({type(e).__name__}: {e}); prior, domain, and likelihood are "
+                f"unavailable."
+            )
             return None
 
     def _build_prior(self):
-        """Build the prior based on model metadata. Called by __init__()."""
-        if self.sampler_context is not None:
-            # Delegate the static prior to the sampler context (the single owner).
-            # Deepcopy because the marginalization split-off below mutates it.
-            self.prior = copy.deepcopy(self.sampler_context.prior)
-        else:
-            intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
-            extrinsic_prior = get_extrinsic_prior_dict(
-                self.base_metadata["train_settings"]["data"]["extrinsic_prior"]
-            )
-            self.prior = build_prior_with_defaults(
-                {**intrinsic_prior, **extrinsic_prior}
-            )
+        """Take the static prior from the sampler context (its single owner), then
+        apply the evolving analysis state: any importance-sampling prior update,
+        and the split-off of time / phase priors for marginalized networks. Called
+        by __init__(). Without a reconstructable context (a payload without full
+        model metadata) the result is transport-only and the prior is `None`."""
+        if self.sampler_context is None:
+            self.prior = None
+            self.geocent_time_prior = None
+            self.phase_prior = None
+            return
+        # Deepcopy because the marginalization split-off below mutates it.
+        self.prior = copy.deepcopy(self.sampler_context.prior)
 
         prior_update = self.importance_sampling_metadata.get("prior_update")
         if prior_update is not None:
@@ -386,56 +390,34 @@ def _build_likelihood(
         self.phase_marginalization_kwargs = phase_marginalization_kwargs
         self.calibration_marginalization_kwargs = calibration_marginalization_kwargs
 
-        # The WaveformGenerator domain is ultimately projected to the data domain, but
-        # generally we allow it to be different, e.g., to start integrating from lower
-        # frequencies than the lower bound of the likelihood integral. An updated
-        # T = 1/delta_f cannot be handled by the domain projection, so it enters at
-        # the level of the WaveformGenerator domain.
+        if self.sampler_context is None:
+            raise ValueError(
+                "Building the likelihood requires a sampler context; this result "
+                "does not carry full model metadata."
+            )
+
+        # Construction is owned by the sampler context. The importance-sampling
+        # state -- validated marginalization bounds, the (possibly rebuilt) data
+        # domain, delta_f / frequency updates -- enters as arguments; the context
+        # itself stays IS-state-free. An updated T = 1/delta_f enters at the level
+        # of the WaveformGenerator domain (it cannot be handled by the domain
+        # projection).
         #
-        # TODO: Add functionality to update other waveform settings, i.e., approximant,
-        #  generation minimum and maximum frequencies, reference frequency, and
-        #  starting frequency.
-        wfg_delta_f = None
-        if "T" in (self.importance_sampling_metadata.get("updates") or {}):
-            wfg_delta_f = 1 / self.importance_sampling_metadata["updates"]["T"]
+        # TODO: Add functionality to update other waveform settings, i.e.,
+        #  approximant, generation minimum and maximum frequencies, reference
+        #  frequency, and starting frequency.
+        view = self._likelihood_view_kwargs()
+        if view["wfg_delta_f"] is not None:
             print(
                 f"Updating waveform generation delta_f from "
                 f'{self.base_metadata["dataset_settings"]["domain"]["delta_f"]} '
-                f"to {wfg_delta_f}."
+                f"to {view['wfg_delta_f']}."
             )
-
-        if self.sampler_context is not None:
-            # Delegate construction to the sampler context (the single owner). The
-            # importance-sampling state -- validated marginalization bounds, the
-            # (possibly rebuilt) data domain, delta_f / frequency updates -- enters
-            # as arguments; the context itself stays IS-state-free.
-            self.likelihood = self.sampler_context.likelihood(
-                time_marginalization_kwargs=time_marginalization_kwargs,
-                phase_marginalization_kwargs=phase_marginalization_kwargs,
-                calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-                **self._likelihood_view_kwargs(),
-            )
-            return
-
-        wfg_domain_dict = self.base_metadata["dataset_settings"]["domain"].copy()
-        if wfg_delta_f is not None:
-            wfg_domain_dict["delta_f"] = wfg_delta_f
-        wfg_domain = build_domain(wfg_domain_dict)
-
-        self.likelihood = StationaryGaussianGWLikelihood(
-            wfg_kwargs=self.base_metadata["dataset_settings"]["waveform_generator"],
-            wfg_domain=wfg_domain,
-            data_domain=self.domain,
-            event_data=self.context,
-            t_ref=self.t_ref,
+        self.likelihood = self.sampler_context.likelihood(
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
             calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-            use_base_domain=self.use_base_domain,
-            frequency_update=dict(
-                minimum_frequency=self.minimum_frequency,
-                maximum_frequency=self.maximum_frequency,
-            ),
+            **view,
         )
 
     def sample_calibration_parameters(self, calibration_sampling_kwargs: dict):
@@ -594,59 +576,39 @@ def _sample_synthetic_phase_chain(
         # marginalization settings rather than reusing the phase-full one.
         self.likelihood = None
 
-    def sample_synthetic_phase(
-        self,
-        synthetic_phase_kwargs,
-        inverse: bool = False,
-    ):
+    def sample_synthetic_phase(self, synthetic_phase_kwargs):
         """
         Sample a synthetic phase for the waveform. This is a post-processing step
         applicable to samples theta in the full parameter space, except for the phase
-        parameter (i.e., 14D samples). This step adds a phase parameter to the samples
-        based on likelihood evaluations.
-
-        A synthetic phase is sampled as follows.
-
-            * Compute and cache the modes for the waveform mu(theta, phase=0) for
-              phase 0, organize them such that each contribution m transforms as
-              exp(-i * m * phase).
-            * Compute the likelihood on a phase grid, by computing mu(theta, phase) from
-              the cached modes. In principle this likelihood is exact, however, it can
-              deviate slightly from the likelihood computed without cached modes for
-              various technical reasons (e.g., slightly different windowing of cached
-              modes compared to full waveform when transforming TD waveform to FD).
-              These small deviations can be fully accounted for by importance sampling.
-              *Note*: when approximation_22_mode=True, the entire waveform is assumed
-              to transform as exp(2i*phase), in which case the likelihood is only exact
-              if the waveform is fully dominated by the (2, 2) mode.
-            * Build a synthetic conditional phase distribution based on this grid. We
-              use an interpolated prior distribution bilby.core.prior.Interped,
-              such that we can sample and also evaluate the log_prob. We add a constant
-              background with weight self.synthetic_phase_kwargs to the kde to make
-              sure that we keep a mass-covering property. With this, the importance
-              sampling will yield exact results even when the synthetic phase conditional
-              is just an approximation.
-
-        Besides adding phase samples to self.samples['phase'], this method also modifies
-        self.samples['log_prob'] by adding the log_prob of the synthetic phase
-        conditional.
-
-        This method modifies self.samples in place.
+        parameter (i.e., 14D samples). It adds a `phase` column to the samples and
+        replaces `log_prob` with the joint proposal density
+        `log q(theta) + log q(phase | theta, d)`.
+
+        The phase distribution `q(phase | theta, d)` is constructed per sample by
+        `SyntheticPhaseFactor` from the likelihood on a phase grid (with a uniform
+        floor for mass coverage, so importance sampling remains exact even where the
+        conditional is approximate); the step runs as a chain rooted in the
+        within-prior proposal samples. Out-of-prior samples receive `phase = 0` and
+        `log_prob = nan`, and carry zero weight in importance sampling.
+
+        This method modifies self.samples in place. Afterwards the phase prior
+        rejoins `self.prior`.
 
         Parameters
         ----------
         synthetic_phase_kwargs : dict
-            This should consist of the kwargs
-                approximation_22_mode (optional)
-                num_processes (optional)
-                n_grid
-                uniform_weight (optional)
-        inverse : bool, default False
-            Whether to apply instead the inverse transformation. This is used prior to
-            calculating the log_prob. In inverse mode, the posterior probability over
-            phase is calculated for given samples. It is stored in self.samples[
-            'log_prob'].
+            Keys: `n_grid` (required), `approximation_22_mode` (optional; default
+            True assumes a (2, 2)-dominated waveform, otherwise the exact mode sum
+            is used, which requires the waveform generator's
+            `spin_conversion_phase = 0`), `uniform_weight` (optional),
+            `num_processes` (optional).
         """
+        if self.sampler_context is None:
+            raise ValueError(
+                "Synthetic phase requires a sampler context; this result does not "
+                "carry full model metadata."
+            )
+
         self.synthetic_phase_kwargs = synthetic_phase_kwargs
 
         approximation_22_mode = self.synthetic_phase_kwargs.get(
@@ -662,9 +624,6 @@ def sample_synthetic_phase(
                 f" However, the prior is {self.phase_prior}."
             )
 
-        # This builds on the Bilby approach to sampling the phase when using a
-        # phase-marginalized likelihood.
-
         # Restrict to samples that are within the prior.
         param_keys = [k for k, v in self.prior.items() if not isinstance(v, Constraint)]
         theta = self.samples[param_keys]
@@ -682,100 +641,10 @@ def sample_synthetic_phase(
         print(f"Estimating synthetic phase for {num_valid_samples} samples.")
         t0 = time.time()
 
-        if not inverse:
-            self._build_likelihood()
-
-        if not inverse and self.sampler_context is not None:
-            # Run the synthetic-phase step as a chain over the sampler context.
-            self._sample_synthetic_phase_chain(
-                theta, within_prior, approximation_22_mode, num_processes
-            )
-            print(f"Done. This took {time.time() - t0:.2f} s.")
-            return
-
-        if inverse:
-            # We estimate the log_prob for given phases, so first save the evaluation
-            # points.
-            sample_phase = theta["phase"].to_numpy(copy=True)
-
-        # For each sample, build the posterior over phase given the remaining parameters.
-
-        phases = np.linspace(0, 2 * np.pi, self.synthetic_phase_kwargs["n_grid"])
-        if approximation_22_mode:
-            # For each sample, the un-normalized posterior depends only on (d | h(phase)):
-            # The prior p(phase), and the inner products (h | h), and (d | d) only contribute
-            # to the normalization. (We check above that p(phase) is constant.)
-            theta["phase"] = 0.0
-            d_inner_h_complex = self.likelihood.d_inner_h_complex_multi(
-                theta.iloc[within_prior],
-                num_processes,
-            )
-
-            # Evaluate the log posterior over the phase across the grid.
-            phasor = np.exp(2j * phases)
-            phase_log_posterior = np.outer(d_inner_h_complex, phasor).real
-        else:
-            self.likelihood.phase_grid = phases
-
-            phase_log_posterior = apply_func_with_multiprocessing(
-                self.likelihood.log_likelihood_phase_grid,
-                theta.iloc[within_prior],
-                num_processes=num_processes,
-            )
-
-        phase_posterior = np.exp(
-            phase_log_posterior - np.amax(phase_log_posterior, axis=1, keepdims=True)
+        self._build_likelihood()
+        self._sample_synthetic_phase_chain(
+            theta, within_prior, approximation_22_mode, num_processes
         )
-        # Include a floor value to maintain mass coverage.
-        phase_posterior += phase_posterior.mean(
-            axis=-1, keepdims=True
-        ) * self.synthetic_phase_kwargs.get("uniform_weight", 0.01)
-
-        if not inverse:
-            # Forward direction:
-            #   (1) Sample a new phase according to the synthetic posterior.
-            #   (2) Add the log_prob to the existing log_prob.
-
-            new_phase, delta_log_prob = interpolated_sample_and_log_prob_multi(
-                phases,
-                phase_posterior,
-                num_processes,
-            )
-
-            phase_array = np.full(len(theta), 0.0)
-            phase_array[within_prior] = new_phase
-            delta_log_prob_array = np.full(len(theta), -np.nan)
-            delta_log_prob_array[within_prior] = delta_log_prob
-
-            self.samples["phase"] = phase_array
-            self.samples["log_prob"] += delta_log_prob_array
-
-            # Insert the phase prior in the prior, since now the phase is present.
-            self.prior["phase"] = self.phase_prior
-            self.phase_prior = None
-
-            # reset likelihood for safety
-            # TODO: Can this be removed?
-            self.likelihood = None
-
-        else:
-            # TODO: Possibly remove.
-            # Inverse direction:
-            #   (1) Evaluate the synthetic log prob for given phase points, and save it.
-
-            log_prob = interpolated_log_prob_multi(
-                phases,
-                phase_posterior,
-                sample_phase[within_prior],
-                num_processes,
-            )
-
-            # Outside of prior, set log_prob to -np.nan.
-            log_prob_array = np.full(len(theta), -np.nan)
-            log_prob_array[within_prior] = log_prob
-            self.samples["log_prob"] = log_prob_array
-            del self.samples["phase"]
-
         print(f"Done. This took {time.time() - t0:.2f} s.")
 
     def get_samples_bilby_phase(self, num_processes=1):
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index 9f15fb8b1..b4d0733c1 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -143,6 +143,16 @@ def test_frequency_range_beyond_domain_rejected():
         _crop_context({"maximum_frequency": 1099.0}).prepared_data()
 
 
+def test_context_without_event_data_serves_metadata_views_only():
+    # A stripped payload (settings but no strain data) still yields the
+    # metadata-derived views; only the likelihood demands event data.
+    ctx = GWSamplerContext.from_model_metadata(_MODEL_METADATA, event_data=None)
+    assert ctx.detectors == ["H1", "L1"]
+    assert ctx.domain.f_max == _DOMAIN_SETTINGS["f_max"]
+    with pytest.raises(ValueError, match="event data"):
+        ctx.likelihood()
+
+
 def test_sampler_provenance_block():
     import ast
 

From 53c5ab22bee9b6e66ed681943fc05889eb9ab535 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 10:04:40 +0100
Subject: [PATCH 39/65] Derive importance-sampling contexts instead of
 threading likelihood views

GWSamplerContext.derive() bakes importance-sampling settings updates into
an immutable derived context: an updated duration rebuilds the data domain
at delta_f = 1/T and enters waveform generation as wfg_delta_f, and
use_base_domain switches a multibanded likelihood to the base domain.
likelihood() reads the representation from context state and takes only
marginalization kwargs; data representation is context identity,
marginalization is per-request.

Result reconstructs the derived context from its serialized metadata
(_build_context), so from-file results carry the full view. The threaded
likelihood view is gone: _likelihood_view_kwargs and _rebuild_domain are
removed, the use_base_domain setter re-derives, SyntheticPhaseFactor builds
its likelihood from the context directly, and the vestigial pre-build in
sample_synthetic_phase is dropped. Result.domain is a pure alias of the
context's domain.

On the context, detectors and t_ref become properties over the model
metadata, and base_metadata is renamed model_metadata (the base
indirection is resolved by Result, not here).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/result.py                |  10 +-
 dingo/gw/inference/factors.py       | 234 +++++++++++++++++++---------
 dingo/gw/result.py                  | 129 +++++----------
 tests/gw/test_gw_sampler_context.py | 120 +++++++++-----
 tests/gw/test_synthetic_phase.py    |  16 +-
 5 files changed, 288 insertions(+), 221 deletions(-)

diff --git a/dingo/core/result.py b/dingo/core/result.py
index 2a2fea1c8..f2ada5ed4 100644
--- a/dingo/core/result.py
+++ b/dingo/core/result.py
@@ -106,8 +106,6 @@ def __init__(self, file_name=None, dictionary=None, sampler_context=None):
             self.sampler_context = self._build_context()
         self._build_prior()
         self._build_domain()
-        if self.importance_sampling_metadata.get("updates"):
-            self._rebuild_domain()
 
     @property
     def metadata(self):
@@ -199,15 +197,13 @@ def reset_event(self, event_dataset):
                 print(f"  {k}:  {event_metadata[k]}")
                 self.importance_sampling_metadata["updates"][k] = event_metadata[k]
 
-            self._rebuild_domain(verbose=True)
         self.event_metadata = event_metadata
 
         # The old context described the data the samples were drawn from; rebuild it
-        # around the new (possibly regenerated) event payload.
+        # around the new (possibly regenerated) event payload -- recorded settings
+        # updates enter as a derived representation -- and re-alias the domain.
         self.sampler_context = self._build_context()
-
-    def _rebuild_domain(self, verbose=False):
-        pass
+        self._build_domain()
 
     @property
     def num_samples(self):
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 55427ce4e..97e0ff40c 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -20,6 +20,7 @@
 import numpy as np
 import pandas as pd
 import torch
+import yaml
 from astropy.time import Time
 from bilby.core.prior import DeltaFunction, PriorDict, Uniform
 from bilby.gw.detector import InterferometerList
@@ -103,28 +104,35 @@ class GWSamplerContext:
 
     Event metadata lives here (not on individual factors): it is a property of the data,
     and it drives frequency cropping, the t_ref/RA correction, and the likelihood.
+
+    The data representation is part of a context's identity: importance-sampling
+    settings updates produce a *derived* context (`derive`) rather than arguments to
+    `likelihood()`. Contexts are treated as immutable -- to change the representation,
+    derive a new one.
+
+    Note that the representation vocabulary here (frequency domains, multibanded
+    decimation, the base-domain likelihood view, frequency-range masking) is specific
+    to this domain family. A new domain family should get its own context class
+    implementing the same interface (`prepared_data` / `likelihood` / `prior` /
+    `derive`) rather than extending this one.
     """
 
     def __init__(
         self,
         domain,
-        detectors: list[str],
-        t_ref: float,
         data_prep: Compose,
         event_data: dict,
         event_metadata: Optional[dict] = None,
-        base_metadata: Optional[dict] = None,
+        model_metadata: Optional[dict] = None,
         device: Union[torch.device, str] = "cpu",
+        use_base_domain: bool = False,
+        wfg_delta_f: Optional[float] = None,
     ):
         """
         Parameters
         ----------
         domain : Domain
             The data frequency domain.
-        detectors : list[str]
-            Detector names.
-        t_ref : float
-            Training reference GPS time.
         data_prep : Compose
             The one-time data-preprocessing transform chain (whiten / decimate /
             repackage).
@@ -134,26 +142,44 @@ def __init__(
         event_metadata : dict, optional
             Per-event metadata; drives frequency cropping, the RA correction, and the
             likelihood reference time.
-        base_metadata : dict, optional
-            The base (parameter / waveform / domain) metadata, used to build the
-            likelihood.
+        model_metadata : dict, optional
+            The metadata of the model defining this analysis (dataset + train
+            settings); the source for the prior, the likelihood, the detector
+            names, and the reference time.
         device : torch.device or str, default "cpu"
             The torch device the chain runs on (the model device); steps that create
             fresh tensors (e.g. `DeltaFactor`) create them here.
+        use_base_domain : bool, default False
+            For a multibanded domain, evaluate the likelihood on the base
+            (undecimated) frequency domain rather than the decimated one. Set via
+            `derive`.
+        wfg_delta_f : float, optional
+            Override for the waveform-generator domain `delta_f` (an updated
+            `T = 1/delta_f` cannot be handled by domain projection). Set via `derive`.
         """
         self.domain = domain
-        self.detectors = detectors
-        self.t_ref = t_ref
         self._data_prep = data_prep
-        self.base_metadata = base_metadata
+        self.model_metadata = model_metadata
         self.event_data = event_data
         self.event_metadata = event_metadata
         self.device = device
+        self.use_base_domain = use_base_domain
+        self.wfg_delta_f = wfg_delta_f
         self._prepared: Optional[torch.Tensor] = None
         self._prior: Optional[PriorDict] = None
         self._likelihood: Optional[StationaryGaussianGWLikelihood] = None
         self._likelihood_settings: Optional[dict] = None
 
+    @property
+    def detectors(self) -> list[str]:
+        """Detector names, read from the model metadata."""
+        return self.model_metadata["train_settings"]["data"]["detectors"]
+
+    @property
+    def t_ref(self) -> float:
+        """Training reference GPS time, read from the model metadata."""
+        return self.model_metadata["train_settings"]["data"]["ref_time"]
+
     @classmethod
     def from_model_metadata(
         cls,
@@ -201,7 +227,10 @@ def from_model_metadata(
         # Event frequency-range update: mask the whitened strain/ASDs outside the
         # requested range. Must precede repackaging (ranges may be per-detector).
         # The request is validated against the training crop license the first time
-        # prepared_data() runs.
+        # prepared_data() runs -- deliberately not here: the license governs the
+        # network-input view only, and contexts are also reconstructed for
+        # likelihood-only use (e.g. from saved importance-sampling results), where
+        # a range that is illegal as network input is legal for ASD masking.
         range_update = _frequency_range_update(domain, event_metadata)
         if range_update is not None:
             transforms.append(
@@ -216,12 +245,10 @@ def from_model_metadata(
 
         return cls(
             domain=domain,
-            detectors=detectors,
-            t_ref=data_settings["ref_time"],
             data_prep=Compose(transforms),
             event_data=event_data,
             event_metadata=event_metadata,
-            base_metadata=metadata,
+            model_metadata=metadata,
             device=device,
         )
 
@@ -263,6 +290,83 @@ def from_model(
             model.metadata, event_data, event_metadata, device=model.device
         )
 
+    def derive(
+        self,
+        updates: Optional[dict] = None,
+        use_base_domain: Optional[bool] = None,
+    ) -> "GWSamplerContext":
+        """Derive a context for the same event under a different data representation.
+
+        The derived context shares the event payload (`event_data`, `event_metadata`),
+        the analysis metadata, and the reference time by construction -- parameter
+        meaning is preserved, so importance weights between the two representations
+        remain well-defined. Only the representation changes: an updated duration `T`
+        rebuilds the data domain at `delta_f = 1/T` (and enters waveform generation as
+        `wfg_delta_f`), and `use_base_domain` switches a multibanded likelihood to the
+        undecimated domain. Frequency-range updates act through `event_metadata` (ASD
+        masking in the likelihood) and need no derived state. Caches start fresh.
+
+        Parameters
+        ----------
+        updates : dict, optional
+            Importance-sampling settings updates, as recorded in
+            `Result.importance_sampling_metadata["updates"]`. May contain `T`,
+            `minimum_frequency`, `maximum_frequency` (and non-domain keys, which are
+            ignored here), but no other quantities that define a new domain.
+        use_base_domain : bool, optional
+            For a multibanded domain, evaluate the likelihood on the base
+            (undecimated) frequency domain. `None` keeps this context's setting.
+
+        Returns
+        -------
+        GWSamplerContext
+        """
+        domain = self.domain
+        wfg_delta_f = self.wfg_delta_f
+        if updates and any(
+            k in updates for k in ("minimum_frequency", "maximum_frequency", "T")
+        ):
+            # TODO: Make compatible with MultibandedFrequencyDomain.
+            if isinstance(domain, MultibandedFrequencyDomain):
+                raise NotImplementedError()
+
+            updates = updates.copy()
+            # A duration update is generation-level on both sides: the pipe
+            # regenerates the event data at the new T, and the waveform must be
+            # generated natively at delta_f = 1/T (a resolution change is not a
+            # projection of existing samples). wfg_delta_f persists on the derived
+            # context, reaching every downstream likelihood build and further
+            # derivations.
+            if "T" in updates:
+                updates["delta_f"] = 1.0 / updates["T"]
+                wfg_delta_f = updates["delta_f"]
+                print(
+                    f"Updating waveform generation delta_f from "
+                    f'{self.model_metadata["dataset_settings"]["domain"]["delta_f"]} '
+                    f"to {wfg_delta_f}."
+                )
+
+            domain_dict = domain.domain_dict  # Existing settings
+            domain_dict.update(
+                (k, updates[k]) for k in set(domain_dict).intersection(updates)
+            )
+            print("Rebuilding domain as follows:")
+            print(yaml.dump(domain_dict, default_flow_style=False, sort_keys=False))
+            domain = build_domain(domain_dict)
+
+        return type(self)(
+            domain=domain,
+            data_prep=self._data_prep,
+            event_data=self.event_data,
+            event_metadata=self.event_metadata,
+            model_metadata=self.model_metadata,
+            device=self.device,
+            use_base_domain=(
+                self.use_base_domain if use_base_domain is None else use_base_domain
+            ),
+            wfg_delta_f=wfg_delta_f,
+        )
+
     def prepared_data(self) -> torch.Tensor:
         """One-time data preprocessing, computed once and cached. An event
         frequency-range update is validated against the training crop license
@@ -282,7 +386,7 @@ def _validate_frequency_range(self):
         if update is None:
             return
         domain = getattr(self.domain, "base_domain", self.domain)
-        crop_settings = self.base_metadata["train_settings"]["data"].get(
+        crop_settings = self.model_metadata["train_settings"]["data"].get(
             "random_strain_cropping"
         )
         _validate_minimum_frequency(
@@ -302,8 +406,8 @@ def prior(self) -> PriorDict:
         applied downstream (they depend on the evolving analysis state), not here.
         """
         if self._prior is None:
-            data_settings = self.base_metadata["train_settings"]["data"]
-            intrinsic_prior = self.base_metadata["dataset_settings"]["intrinsic_prior"]
+            data_settings = self.model_metadata["train_settings"]["data"]
+            intrinsic_prior = self.model_metadata["dataset_settings"]["intrinsic_prior"]
             extrinsic_prior = get_extrinsic_prior_dict(data_settings["extrinsic_prior"])
             self._prior = build_prior_with_defaults(
                 {**intrinsic_prior, **extrinsic_prior}
@@ -315,10 +419,6 @@ def likelihood(
         time_marginalization_kwargs: Optional[dict] = None,
         phase_marginalization_kwargs: Optional[dict] = None,
         calibration_marginalization_kwargs: Optional[dict] = None,
-        use_base_domain: bool = False,
-        data_domain=None,
-        wfg_delta_f: Optional[float] = None,
-        frequency_update: Optional[dict] = None,
     ) -> StationaryGaussianGWLikelihood:
         """
         Build the exact GW likelihood on this event's data, in physical parameter space.
@@ -330,17 +430,16 @@ def likelihood(
         event time (the training-frame right-ascension correction is already applied to the
         samples), or the training reference time when no event time is set.
 
+        The data representation (the domain, `use_base_domain`, `wfg_delta_f`, the
+        event's frequency range) is context state -- importance-sampling settings
+        updates enter by deriving a new context (`derive`), not as arguments here.
+        Marginalization is per-request and enters as arguments.
+
         The most recently built likelihood is cached with its arguments: a repeated
         call with the same arguments returns the shared instance (the synthetic-phase
         factor requests one per chain chunk), and a call with different arguments
         builds a replacement. In the exact synthetic-phase mode the consumer assigns
-        a `phase_grid` attribute on the shared instance, as
-        `Result.sample_synthetic_phase` does.
-
-        The context itself is importance-sampling-state-free: IS-time settings
-        updates (a rebuilt data domain, an updated `T` / `delta_f`, an explicit
-        frequency range) enter as arguments, supplied by the IS layer
-        (`Result._build_likelihood`).
+        a `phase_grid` attribute on the shared instance.
 
         Parameters
         ----------
@@ -353,18 +452,6 @@ def likelihood(
             Analytically marginalize over `phase`. Requires a uniform [0, 2 pi) phase prior.
         calibration_marginalization_kwargs : dict, optional
             Marginalize over detector calibration uncertainty.
-        use_base_domain : bool, default False
-            For a multibanded domain, evaluate on the base (undecimated) frequency domain
-            rather than the decimated one.
-        data_domain : Domain, optional
-            Override for the data domain (e.g. rebuilt with importance-sampling
-            updates); defaults to this context's domain.
-        wfg_delta_f : float, optional
-            Override for the waveform-generator domain `delta_f` (an updated
-            `T = 1/delta_f` cannot be handled by domain projection).
-        frequency_update : dict, optional
-            `minimum_frequency` / `maximum_frequency` for ASD masking; defaults to
-            the event metadata's values with this context's domain as fallback.
 
         Returns
         -------
@@ -373,9 +460,12 @@ def likelihood(
         # Capture the call's arguments for the cache comparison; this must stay the
         # first statement so locals() holds exactly the arguments. The comparison
         # happens before the marginalization bounds are filled in below (which is
-        # deterministic given the same arguments).
+        # deterministic given the same arguments). The representation fields join the
+        # key as a guard against off-contract in-place mutation.
         settings = dict(locals())
         settings.pop("self")
+        settings["use_base_domain"] = self.use_base_domain
+        settings["wfg_delta_f"] = self.wfg_delta_f
         if settings == self._likelihood_settings:
             return self._likelihood
 
@@ -387,7 +477,11 @@ def likelihood(
 
         # Marginalizing over a parameter needs the (uniform) prior the network
         # marginalized over; use it to parameterize the requested marginalization.
-        # Bounds already provided by the caller (e.g. from an updated prior) win.
+        # Bounds already provided by the caller win: the importance-sampling layer
+        # validates against its evolved prior (prior updates, time/phase
+        # split-offs), which this sample-free context cannot see, and passes
+        # explicit bounds. The fill below covers standalone (chain) use from the
+        # network's static prior.
         if time_marginalization_kwargs is not None and not (
             "t_lower" in time_marginalization_kwargs
             and "t_upper" in time_marginalization_kwargs
@@ -418,13 +512,17 @@ def likelihood(
                     f"but is {phase_prior}."
                 )
 
-        dataset_settings = self.base_metadata["dataset_settings"]
-        # WaveformGenerator domain -- generally the dataset domain (it may be wider than
-        # the data domain for generation). An updated T = 1/delta_f enters here (it
-        # cannot be handled by domain projection).
+        dataset_settings = self.model_metadata["dataset_settings"]
+        # WaveformGenerator domain -- deliberately the dataset domain WITHOUT any
+        # domain_update, so waveform generation mirrors how the training set was
+        # generated (possibly wider than the network's data domain); the likelihood
+        # projects the generated waveform onto the data domain. An updated
+        # T = 1/delta_f enters here: unlike range masking or decimation, a
+        # resolution change needs samples that do not exist on the old grid, so it
+        # cannot be handled by domain projection and must apply at generation.
         wfg_domain_dict = dataset_settings["domain"]
-        if wfg_delta_f is not None:
-            wfg_domain_dict = {**wfg_domain_dict, "delta_f": wfg_delta_f}
+        if self.wfg_delta_f is not None:
+            wfg_domain_dict = {**wfg_domain_dict, "delta_f": self.wfg_delta_f}
         wfg_domain = build_domain(wfg_domain_dict)
 
         # Likelihood reference time: the event time (the training-frame RA correction has
@@ -434,28 +532,21 @@ def likelihood(
         else:
             t_ref = self.t_ref
 
-        if data_domain is None:
-            data_domain = self.domain
-        if frequency_update is None:
-            frequency_update = dict(
-                minimum_frequency=self._frequency(
-                    "minimum_frequency", self.domain.f_min
-                ),
-                maximum_frequency=self._frequency(
-                    "maximum_frequency", self.domain.f_max
-                ),
-            )
+        frequency_update = dict(
+            minimum_frequency=self._frequency("minimum_frequency", self.domain.f_min),
+            maximum_frequency=self._frequency("maximum_frequency", self.domain.f_max),
+        )
 
         likelihood = StationaryGaussianGWLikelihood(
             wfg_kwargs=dataset_settings["waveform_generator"],
             wfg_domain=wfg_domain,
-            data_domain=data_domain,
+            data_domain=self.domain,
             event_data=self.event_data,
             t_ref=t_ref,
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
             calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-            use_base_domain=use_base_domain,
+            use_base_domain=self.use_base_domain,
             frequency_update=frequency_update,
         )
         self._likelihood = likelihood
@@ -472,7 +563,10 @@ def _marginalized_prior(self, name: str):
         """The prior over `name` if the network marginalized it (i.e. `name` is not an
         inference parameter), else `None` -- used to parameterize likelihood
         marginalization over time / phase."""
-        if name in self.base_metadata["train_settings"]["data"]["inference_parameters"]:
+        if (
+            name
+            in self.model_metadata["train_settings"]["data"]["inference_parameters"]
+        ):
             return None
         return self.prior.get(name)
 
@@ -516,7 +610,6 @@ def __init__(
         approximation_22_mode: bool = False,
         uniform_weight: float = 0.01,
         num_processes: int = 1,
-        likelihood_kwargs: Optional[dict] = None,
     ):
         """
         Parameters
@@ -532,10 +625,6 @@ def __init__(
             Weight of the uniform floor added to the phase distribution for mass coverage.
         num_processes : int, default 1
             Parallel processes for the per-sample likelihood evaluation and phase sampling.
-        likelihood_kwargs : dict, optional
-            Arguments for `context.likelihood()` selecting the likelihood view (e.g.
-            `use_base_domain`, a rebuilt `data_domain`, frequency updates), supplied
-            by the importance-sampling layer. Defaults to the context's own views.
         """
         self.parameters = ["phase"]
         self.conditioning = list(conditioning)
@@ -543,7 +632,6 @@ def __init__(
         self.approximation_22_mode = approximation_22_mode
         self.uniform_weight = uniform_weight
         self.num_processes = num_processes
-        self.likelihood_kwargs = likelihood_kwargs
 
     def sample_and_log_prob(self, num_samples, context, given=None):
         """Draw one phase per `theta_rest` row (`num_samples` must be 1); return the phases
@@ -592,7 +680,9 @@ def _phase_profile(self, given, context):
         per sample: evaluate `log L` on the grid, exponentiate (shifted by the per-row
         max), and add the uniform floor."""
         theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
-        likelihood = context.likelihood(**(self.likelihood_kwargs or {}))
+        # The context (possibly derived for importance sampling) carries the
+        # representation; the phase-full likelihood needs no arguments.
+        likelihood = context.likelihood()
         phases = np.linspace(0, 2 * np.pi, self.n_grid)
         if self.approximation_22_mode:
             # Assume the waveform is (2, 2)-dominated (transforms as exp(2i phase)), so the
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 6f5542a6d..6593b3027 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -3,15 +3,12 @@
 from typing import Optional
 
 import numpy as np
-import yaml
 from bilby.core.prior import Uniform, Constraint, PriorDict, DeltaFunction
 from bilby.gw.prior import CalibrationPriorDict
 from bilby_pipe.utils import CALIBRATION_CORRECTION_TYPE_LOOKUP
 
 from dingo.core.result import Result as CoreResult
 from dingo.gw.conversion import change_spin_conversion_phase
-from dingo.gw.domains import MultibandedFrequencyDomain
-from dingo.gw.domains import build_domain
 from dingo.core.utils.backward_compatibility import check_minimum_version
 
 
@@ -111,6 +108,16 @@ def use_base_domain(self) -> bool:
     def use_base_domain(self, value: bool):
         if hasattr(self.domain, "base_domain"):
             self.importance_sampling_metadata["use_base_domain"] = value
+            # The representation is context state: re-derive rather than mutate, so
+            # the fresh context starts with an empty likelihood cache and a
+            # previously built likelihood cannot leak across the change.
+            if (
+                self.sampler_context is not None
+                and self.sampler_context.use_base_domain != value
+            ):
+                self.sampler_context = self.sampler_context.derive(
+                    use_base_domain=value
+                )
 
     @property
     def f_ref(self):
@@ -153,60 +160,15 @@ def maximum_frequency(self, value: dict[str, float] | float):
         self.event_metadata["maximum_frequency"] = value
 
     def _build_domain(self):
-        """Take the data domain from the sampler context. After importance-sampling
-        updates, `_rebuild_domain` rebinds `self.domain` to a new object, and the
-        updated domain reaches the likelihood through `_likelihood_view_kwargs`;
-        the context keeps the original training-settings domain. Called by
-        __init__()."""
+        """Take the data domain from the sampler context -- its single owner. A
+        context derived with importance-sampling updates already carries the
+        rebuilt domain. Called by __init__() and after reset_event()."""
         check_minimum_version(self.version, raise_exception=False)
         if self.sampler_context is None:
             self.domain = None
             return
         self.domain = self.sampler_context.domain
 
-    def _rebuild_domain(self, verbose=False):
-        """Rebuild the domain based on settings updated for importance sampling.
-
-        These settings should be saved in self.importance_sampling_metadata["updates"],
-        which is expected to be populated by reset_event()."""
-        updates = self.importance_sampling_metadata["updates"].copy()
-
-        # Assume that updates can contain T, f_s, roll_off, f_min, f_max, but no other
-        # quantities that define a new domain (e.g., delta_f). Typical event metadata
-        # will be constructed in this way.
-
-        domain_keys = ["minimum_frequency", "maximum_frequency", "T"]
-        if any(k in updates for k in domain_keys):
-            if self.domain is None:
-                # Transport-only result (no reconstructable context): nothing to
-                # rebuild; domain-dependent operations fail at use.
-                return
-            # TODO: Make compatible with MultibandedFrequencyDomain.
-            if isinstance(self.domain, MultibandedFrequencyDomain):
-                raise NotImplementedError()
-
-            if "T" in updates:
-                updates["delta_f"] = 1.0 / updates["T"]
-
-            domain_dict = self.domain.domain_dict  # Existing settings
-            domain_dict.update(
-                (k, updates[k]) for k in set(domain_dict).intersection(updates)
-            )
-
-            if verbose:
-                print("Rebuilding domain as follows:")
-                print(
-                    yaml.dump(
-                        domain_dict,
-                        default_flow_style=False,
-                        sort_keys=False,
-                    )
-                )
-            self.domain = build_domain(domain_dict)
-        else:
-            if verbose:
-                print("No domain updates found; domain not rebuilt.")
-
     def _build_context(self):
         """Reconstruct the per-event sampler context from the serialized payload
         (settings + event data + event metadata), so that prior (and, later,
@@ -223,7 +185,7 @@ def _build_context(self):
         try:
             # base_metadata resolves the unconditional ("base") indirection, so
             # density-recovery results reconstruct from the analysis metadata.
-            return GWSamplerContext.from_model_metadata(
+            context = GWSamplerContext.from_model_metadata(
                 self.base_metadata, self.context, self.event_metadata
             )
         except (KeyError, TypeError) as e:
@@ -236,6 +198,14 @@ def _build_context(self):
                 f"unavailable."
             )
             return None
+        # Importance-sampling settings updates change the data representation, which
+        # lives on a derived context (same event, different representation).
+        metadata = self.importance_sampling_metadata or {}
+        updates = metadata.get("updates")
+        use_base_domain = metadata.get("use_base_domain", False)
+        if updates or use_base_domain:
+            context = context.derive(updates=updates, use_base_domain=use_base_domain)
+        return context
 
     def _build_prior(self):
         """Take the static prior from the sampler context (its single owner), then
@@ -322,23 +292,6 @@ def update_prior(self, prior_update):
             # Recalculate the importance-sampling weights and log evidence.
             self._calculate_evidence()
 
-    def _likelihood_view_kwargs(self) -> dict:
-        """The `context.likelihood()` arguments encoding this Result's
-        importance-sampling view: the (possibly rebuilt) data domain, the
-        base-domain choice, and delta_f / frequency updates."""
-        wfg_delta_f = None
-        if "T" in (self.importance_sampling_metadata.get("updates") or {}):
-            wfg_delta_f = 1 / self.importance_sampling_metadata["updates"]["T"]
-        return dict(
-            use_base_domain=self.use_base_domain,
-            data_domain=self.domain,
-            wfg_delta_f=wfg_delta_f,
-            frequency_update=dict(
-                minimum_frequency=self.minimum_frequency,
-                maximum_frequency=self.maximum_frequency,
-            ),
-        )
-
     def _build_likelihood(
         self,
         time_marginalization_kwargs: Optional[dict] = None,
@@ -360,6 +313,10 @@ def _build_likelihood(
         calibration_marginalization_kwargs: dict
             Calibration marginalization parameters. If None, no calibration marginalization is used.
         """
+        # Marginalization is validated here against the *evolved* prior (any
+        # importance-sampling prior update, the time / phase split-offs), which
+        # the sample-free context cannot see; the explicit bounds set below
+        # therefore take precedence over the context's static-prior fill.
         if time_marginalization_kwargs is not None:
             if self.geocent_time_prior is None:
                 raise NotImplementedError(
@@ -396,28 +353,17 @@ def _build_likelihood(
                 "does not carry full model metadata."
             )
 
-        # Construction is owned by the sampler context. The importance-sampling
-        # state -- validated marginalization bounds, the (possibly rebuilt) data
-        # domain, delta_f / frequency updates -- enters as arguments; the context
-        # itself stays IS-state-free. An updated T = 1/delta_f enters at the level
-        # of the WaveformGenerator domain (it cannot be handled by the domain
-        # projection).
+        # Construction is owned by the sampler context; its (possibly derived)
+        # representation already encodes the importance-sampling settings updates.
+        # Validated marginalization bounds enter as arguments.
         #
         # TODO: Add functionality to update other waveform settings, i.e.,
         #  approximant, generation minimum and maximum frequencies, reference
         #  frequency, and starting frequency.
-        view = self._likelihood_view_kwargs()
-        if view["wfg_delta_f"] is not None:
-            print(
-                f"Updating waveform generation delta_f from "
-                f'{self.base_metadata["dataset_settings"]["domain"]["delta_f"]} '
-                f"to {view['wfg_delta_f']}."
-            )
         self.likelihood = self.sampler_context.likelihood(
             time_marginalization_kwargs=time_marginalization_kwargs,
             phase_marginalization_kwargs=phase_marginalization_kwargs,
             calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-            **view,
         )
 
     def sample_calibration_parameters(self, calibration_sampling_kwargs: dict):
@@ -536,11 +482,11 @@ def _sample_synthetic_phase_chain(
         within-prior proposal samples with their stored log-prob, and
         `SyntheticPhaseFactor` draws `phase` -- the composer's ordinary log-prob fold
         returns the joint proposal density `log q(theta) + log q(phase | theta, d)`,
-        which becomes the samples' log_prob. The Result's likelihood view (base
-        domain, rebuilt domain, frequency updates) is passed to the factor, so it
-        shares the cached likelihood instance built by `_build_likelihood`.
-        Out-of-prior rows get `phase = 0` and `log_prob = nan` (they receive zero
-        weight in importance sampling regardless)."""
+        which becomes the samples' log_prob. The factor builds the phase-full
+        likelihood from the sampler context, whose (possibly derived) representation
+        encodes the importance-sampling view (base domain, rebuilt domain, frequency
+        updates). Out-of-prior rows get `phase = 0` and `log_prob = nan` (they
+        receive zero weight in importance sampling regardless)."""
         from dingo.core.factors import ChainComposer, SampleTableFactor
         from dingo.gw.inference.factors import SyntheticPhaseFactor
 
@@ -555,7 +501,6 @@ def _sample_synthetic_phase_chain(
             approximation_22_mode=approximation_22_mode,
             uniform_weight=self.synthetic_phase_kwargs.get("uniform_weight", 0.01),
             num_processes=num_processes,
-            likelihood_kwargs=self._likelihood_view_kwargs(),
         )
         chain = ChainComposer([table, factor])
         out, log_prob = chain.sample_and_log_prob(
@@ -572,8 +517,9 @@ def _sample_synthetic_phase_chain(
         # Insert the phase prior in the prior, since now the phase is present.
         self.prior["phase"] = self.phase_prior
         self.phase_prior = None
-        # Reset so importance sampling rebuilds the likelihood with its own
-        # marginalization settings rather than reusing the phase-full one.
+        # Any previously built likelihood does not describe the now-phase-full
+        # samples; importance sampling rebuilds with its own marginalization
+        # settings.
         self.likelihood = None
 
     def sample_synthetic_phase(self, synthetic_phase_kwargs):
@@ -641,7 +587,6 @@ def sample_synthetic_phase(self, synthetic_phase_kwargs):
         print(f"Estimating synthetic phase for {num_valid_samples} samples.")
         t0 = time.time()
 
-        self._build_likelihood()
         self._sample_synthetic_phase_chain(
             theta, within_prior, approximation_22_mode, num_processes
         )
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index b4d0733c1..6d6572170 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -1,8 +1,8 @@
 """
-Unit tests for ``GWSamplerContext`` plumbing: the default-likelihood cache, the
-device attribute, and the metadata constructors. The likelihood itself is stubbed
-(no waveform generation); the model-based construction parity lives in the local
-harnesses.
+Unit tests for ``GWSamplerContext`` plumbing: the likelihood cache, derived
+contexts (importance-sampling representations), the device attribute, and the
+metadata constructors. The likelihood itself is stubbed (no waveform generation);
+the model-based construction parity lives in the local harnesses.
 """
 
 import copy
@@ -27,7 +27,13 @@
         "domain": _DOMAIN_SETTINGS,
         "waveform_generator": {"approximant": "IMRPhenomD", "f_ref": 20.0},
     },
-    "train_settings": {"data": {"inference_parameters": ["chirp_mass"]}},
+    "train_settings": {
+        "data": {
+            "inference_parameters": ["chirp_mass"],
+            "detectors": ["H1"],
+            "ref_time": 1126259462.4,
+        }
+    },
 }
 
 
@@ -49,11 +55,9 @@ def context(monkeypatch):
     _StubLikelihood.constructions = 0
     return GWSamplerContext(
         domain=build_domain(_DOMAIN_SETTINGS),
-        detectors=["H1"],
-        t_ref=1126259462.4,
         data_prep=None,
         event_data={},
-        base_metadata=_BASE_METADATA,
+        model_metadata=_BASE_METADATA,
     )
 
 
@@ -65,22 +69,15 @@ def test_likelihood_rebuilt_only_when_settings_change(context):
     assert context.likelihood() is default
     assert _StubLikelihood.constructions == 1
 
-    base = context.likelihood(use_base_domain=True)
-    assert base is not default
-    assert base.kwargs["use_base_domain"] is True
-    assert context.likelihood(use_base_domain=True) is base
+    marg_kwargs = {"num_calibration_curves": 100}
+    marginalized = context.likelihood(calibration_marginalization_kwargs=marg_kwargs)
+    assert marginalized is not default
+    assert context.likelihood(calibration_marginalization_kwargs=marg_kwargs) is (
+        marginalized
+    )
     assert _StubLikelihood.constructions == 2
 
 
-def test_likelihood_cache_handles_domain_and_none_settings(context):
-    # A bare call (data_domain=None) followed by an explicit domain (either order)
-    # must be a clean cache miss, not a Domain == None comparison error.
-    bare = context.likelihood()
-    explicit = context.likelihood(data_domain=build_domain(_DOMAIN_SETTINGS))
-    assert explicit is not bare
-    assert context.likelihood() is not explicit
-
-
 def test_domain_eq_is_none_safe():
     domain = build_domain(_DOMAIN_SETTINGS)
     assert domain == build_domain(_DOMAIN_SETTINGS)
@@ -88,21 +85,66 @@ def test_domain_eq_is_none_safe():
     assert (domain == 3) is False
 
 
-def test_likelihood_importance_sampling_overrides(context):
-    # IS-state overrides pass through: an explicit data domain, an updated
-    # waveform-generator delta_f, and an explicit frequency range.
-    likelihood = context.likelihood(
-        data_domain="rebuilt-domain",
-        wfg_delta_f=0.5,
-        frequency_update={"minimum_frequency": 21.0, "maximum_frequency": 512.0},
-    )
-    assert likelihood.kwargs["data_domain"] == "rebuilt-domain"
-    assert likelihood.kwargs["wfg_domain"].delta_f == 0.5
-    assert likelihood.kwargs["frequency_update"]["maximum_frequency"] == 512.0
-    # Defaults reproduce the context's own views.
+def test_derived_context_has_independent_likelihood_cache(context):
+    # The representation is context state: a derived context builds its own
+    # likelihood (fresh cache) without disturbing the original's.
     default = context.likelihood()
-    assert default.kwargs["data_domain"] is context.domain
-    assert default.kwargs["wfg_domain"].delta_f == _DOMAIN_SETTINGS["delta_f"]
+    derived = context.derive(use_base_domain=True)
+    base = derived.likelihood()
+    assert base is not default
+    assert base.kwargs["use_base_domain"] is True
+    assert derived.likelihood() is base
+    assert context.likelihood() is default
+    assert context.use_base_domain is False
+
+
+def test_derive_with_updated_duration_rebuilds_domain(context):
+    # An updated duration T rebuilds the data domain at delta_f = 1/T and enters
+    # waveform generation as wfg_delta_f; the event payload and metadata are
+    # shared with the original context.
+    derived = context.derive(updates={"T": 2.0})
+    assert derived.domain.delta_f == 0.5
+    assert derived.wfg_delta_f == 0.5
+    assert derived.event_data is context.event_data
+    assert derived.model_metadata is context.model_metadata
+    assert derived.t_ref == context.t_ref
+    likelihood = derived.likelihood()
+    assert likelihood.kwargs["data_domain"] is derived.domain
+    assert likelihood.kwargs["wfg_domain"].delta_f == 0.5
+    # The original context is untouched.
+    assert context.domain.delta_f == _DOMAIN_SETTINGS["delta_f"]
+    assert context.wfg_delta_f is None
+    assert context.likelihood().kwargs["wfg_domain"].delta_f == (
+        _DOMAIN_SETTINGS["delta_f"]
+    )
+
+
+def test_derive_with_frequency_updates_keeps_domain_grid(context):
+    # min/max frequency updates apply via ASD masking (through the event
+    # metadata), not the domain grid: the rebuild triggers but reproduces the
+    # same domain settings.
+    derived = context.derive(updates={"minimum_frequency": 25.0})
+    assert derived.domain is not context.domain
+    assert derived.domain.domain_dict == context.domain.domain_dict
+
+
+def test_likelihood_frequency_range_follows_event_metadata(context):
+    # The likelihood masks the ASDs to the event's frequency range; without an
+    # event override the range defaults to the domain bounds.
+    default_update = context.likelihood().kwargs["frequency_update"]
+    assert default_update == {
+        "minimum_frequency": _DOMAIN_SETTINGS["f_min"],
+        "maximum_frequency": _DOMAIN_SETTINGS["f_max"],
+    }
+    with_event = GWSamplerContext(
+        domain=build_domain(_DOMAIN_SETTINGS),
+        data_prep=None,
+        event_data={},
+        event_metadata={"minimum_frequency": 21.0, "maximum_frequency": 512.0},
+        model_metadata=_BASE_METADATA,
+    )
+    update = with_event.likelihood().kwargs["frequency_update"]
+    assert update == {"minimum_frequency": 21.0, "maximum_frequency": 512.0}
 
 
 def _event_data(n_bins):
@@ -195,16 +237,12 @@ def test_likelihood_caller_marginalization_bounds_win(context):
 def test_context_device_default_and_explicit():
     ctx = GWSamplerContext(
         domain=None,
-        detectors=["H1"],
-        t_ref=0.0,
         data_prep=None,
         event_data={},
         device="meta",
     )
     assert ctx.device == "meta"
-    ctx_default = GWSamplerContext(
-        domain=None, detectors=["H1"], t_ref=0.0, data_prep=None, event_data={}
-    )
+    ctx_default = GWSamplerContext(domain=None, data_prep=None, event_data={})
     assert ctx_default.device == "cpu"
 
 
@@ -236,7 +274,7 @@ def test_from_model_metadata_matches_from_model():
         assert ctx.detectors == ["H1", "L1"]
         assert ctx.t_ref == 1126259462.4
         assert ctx.domain.domain_dict == ctx_meta.domain.domain_dict
-        assert ctx.base_metadata is _MODEL_METADATA
+        assert ctx.model_metadata is _MODEL_METADATA
         assert ctx.device == "cpu"
         assert ctx._data_prep is not None
 
diff --git a/tests/gw/test_synthetic_phase.py b/tests/gw/test_synthetic_phase.py
index e0968db17..c9f92eee3 100644
--- a/tests/gw/test_synthetic_phase.py
+++ b/tests/gw/test_synthetic_phase.py
@@ -116,8 +116,6 @@ def test_context_frequency_override():
     def context(event_metadata):
         return GWSamplerContext(
             domain=None,
-            detectors=[],
-            t_ref=0.0,
             data_prep=None,
             event_data={},
             event_metadata=event_metadata,
@@ -129,22 +127,22 @@ def context(event_metadata):
     assert context(None)._frequency("minimum_frequency", 20.0) == 20.0
 
 
-def test_likelihood_kwargs_forwarded_to_context():
-    # The IS layer selects the likelihood view (base domain, rebuilt domain, ...);
-    # the factor forwards it verbatim to context.likelihood().
-    recorded = {}
+def test_factor_builds_likelihood_from_context():
+    # The factor takes the likelihood from the context with no arguments -- the
+    # data representation lives on the (possibly derived) context, not on the
+    # factor.
+    recorded = []
 
     class _RecordingContext(_MockContext):
         def likelihood(self, **kwargs):
-            recorded.update(kwargs)
+            recorded.append(kwargs)
             return super().likelihood()
 
     factor = SyntheticPhaseFactor(
         conditioning=["chirp_mass"],
         n_grid=11,
         approximation_22_mode=True,
-        likelihood_kwargs={"use_base_domain": True},
     )
     _seed()
     factor.sample_and_log_prob(1, _RecordingContext(), _given())
-    assert recorded == {"use_base_domain": True}
+    assert recorded == [{}]

From f23b76991e8ee00e7fb154a7aaa7a1ec88fb1cd1 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 11:34:23 +0100
Subject: [PATCH 40/65] Add the spin-convention reparametrization

SpinConventionReparam relabels theta_jn / phi_jl between Dingo's internal
spin-phase convention (spin_conversion_phase = 0) and the physical (Bilby)
one, as a framework Reparametrization: the export methods
(get_samples_bilby_phase, get_pesummary_samples) route through its
to_physical DataFrame API, and to_network provides the reverse direction
for ingesting external posteriors. Stored samples keep the network
convention throughout the density / likelihood / synthetic-phase pipeline;
no marked parameter names are needed because the two conventions never
coexist in one table.

The bijection preserves the spherical measure but not the flat coordinate
measure: log_det = log sin(theta_jn) - log sin(theta_jn'), verified
numerically against finite differences through the LAL conversion
(agreement ~1e-9, permanent test). Forward chain sampling is
density-correct and dtype-preserving; chain re-plug raises pending a
conditioning-aware inverse. Export output is bit-exact against the
previous direct calls.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/inference/factors.py            | 135 +++++++++++++++++
 dingo/gw/result.py                       |  30 +---
 tests/gw/test_spin_convention_reparam.py | 179 +++++++++++++++++++++++
 3 files changed, 322 insertions(+), 22 deletions(-)
 create mode 100644 tests/gw/test_spin_convention_reparam.py

diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 97e0ff40c..ffc2ab00b 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -45,6 +45,7 @@
 from dingo.core.multiprocessing import apply_func_with_multiprocessing
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.core.transforms import GetItem, RenameKey
+from dingo.gw.conversion import change_spin_conversion_phase
 from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
 from dingo.gw.frequency_updates import (
     _validate_maximum_frequency,
@@ -1056,6 +1057,140 @@ def inverse(self, params, context):
         return {"ra@t_ref": ra_tref.float()}
 
 
+class SpinConventionReparam(Reparametrization):
+    """
+    Relabel the precessing-spin angles between Dingo's internal spin-phase
+    convention and the physical (Bilby) one.
+
+    Dingo fixes the spin-conversion phase (usually to 0) so that the Cartesian
+    spins decouple from the coalescence phase; the whole density / likelihood /
+    synthetic-phase pipeline requires that convention, so stored samples keep the
+    plain names `theta_jn` / `phi_jl` in the *network* convention throughout. The
+    physical convention (spin conversion at the sample's own phase) is what Bilby
+    and PESummary mean by the same names, so the relabel happens at the export
+    boundary -- this class is its single home. Only `theta_jn` and `phi_jl`
+    change; the conversion phase and reference frequency are read from the model
+    metadata, and a model trained without a fixed conversion phase (`None`)
+    relabels to the identity.
+
+    Unlike `RAReparam` no marked intermediate name is needed: the two conventions
+    never coexist in one table -- each world's plain names denote its own
+    convention, and this bijection is the boundary crossing.
+
+    Exporting a finished weighted sample set needs no Jacobian (proposal, prior,
+    and likelihood transform together), which is how `to_physical` is used. As a
+    chain `Step` the bijection is *not* measure-preserving in the flat
+    `(theta_jn, phi_jl)` coordinates: it rotates the line of sight rigidly about
+    the orbital angular momentum, preserving the spherical measure
+    `sin(theta_jn) dtheta dphi`, so
+    `log_det = log sin(theta_jn) - log sin(theta_jn')` (verified numerically
+    against finite differences through the LAL conversion). Forward chain
+    sampling is therefore density-correct; chain re-plug
+    (`ChainComposer.log_prob`) additionally needs the conditioning-aware
+    `inverse`, which the chain protocol does not yet provide.
+    """
+
+    def __init__(self, num_processes: int = 1):
+        """
+        Parameters
+        ----------
+        num_processes : int, default 1
+            Parallel processes for the per-sample LAL spin conversion.
+        """
+        self.parameters = ["theta_jn", "phi_jl"]
+        self.conditioning = [
+            "theta_jn",
+            "phi_jl",
+            "phase",
+            "chirp_mass",
+            "mass_ratio",
+            "a_1",
+            "a_2",
+            "tilt_1",
+            "tilt_2",
+            "phi_12",
+        ]
+        self.num_processes = num_processes
+
+    @property
+    def consumes(self) -> list[str]:
+        # The bijection overwrites theta_jn / phi_jl in place; the remaining
+        # conditioning (phase, masses, tilts, ...) is read-only and must stay in
+        # the chain -- the default conditioning-minus-parameters would drop it.
+        return []
+
+    def log_det(self, given, context):
+        """`log|det J|` of `forward`, per row. The map preserves the spherical
+        measure, so the flat-coordinate Jacobian is
+        `sin(theta_jn) / sin(theta_jn')` -- verified numerically against finite
+        differences through the LAL conversion in
+        `test_jacobian_matches_sin_ratio` (agreement ~1e-9)."""
+        converted = self.forward(given, context)
+        return self._log_det(given["theta_jn"], converted["theta_jn"])
+
+    @staticmethod
+    def _log_det(theta_jn_in, theta_jn_out):
+        # Compute in double (the LAL conversion is double precision anyway),
+        # return in the input dtype: a reparametrization preserves the chain's
+        # dtype rather than promoting the summed log_prob.
+        log_det = torch.log(torch.sin(theta_jn_in.double())) - torch.log(
+            torch.sin(theta_jn_out.double())
+        )
+        return log_det.to(theta_jn_in.dtype)
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        """Apply `forward`; contribute `-log|det J|`. Overridden to share the
+        single LAL conversion between the transform and its Jacobian (the base
+        implementation would convert twice)."""
+        if num_samples != 1:
+            raise ValueError("A reparametrization is 1:1; use fan_out=1.")
+        out = self.forward(given, context)
+        return out, -self._log_det(given["theta_jn"], out["theta_jn"])
+
+    @staticmethod
+    def _model_convention(model_metadata: dict) -> tuple[float, Optional[float]]:
+        """The reference frequency and spin-conversion phase the model trained with."""
+        wfg_settings = model_metadata["dataset_settings"]["waveform_generator"]
+        return wfg_settings["f_ref"], wfg_settings.get("spin_conversion_phase")
+
+    def to_physical(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFrame:
+        """Relabel samples from the model's convention to the physical (Bilby) one."""
+        f_ref, sc_phase = self._model_convention(model_metadata)
+        return change_spin_conversion_phase(
+            samples, f_ref, sc_phase, None, num_processes=self.num_processes
+        )
+
+    def to_network(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFrame:
+        """Relabel samples from the physical (Bilby) convention to the model's,
+        e.g. to ingest external posteriors for comparison."""
+        f_ref, sc_phase = self._model_convention(model_metadata)
+        return change_spin_conversion_phase(
+            samples, f_ref, None, sc_phase, num_processes=self.num_processes
+        )
+
+    def forward(self, given, context):
+        # The conversion runs in double; the outputs return in the input dtype
+        # and device (cf. RAReparam: compute in float64, store the chain dtype).
+        reference = given["theta_jn"]
+        theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
+        converted = self.to_physical(theta, context.model_metadata)
+        return {
+            k: torch.as_tensor(converted[k].to_numpy()).to(
+                dtype=reference.dtype, device=reference.device
+            )
+            for k in self.parameters
+        }
+
+    def inverse(self, params, context):
+        # The chain protocol hands inverse only the parameters block, but this
+        # bijection also needs the invariant conditioning (phase, masses, tilts).
+        raise NotImplementedError(
+            "The spin-convention inverse needs the conditioning block, which the "
+            "chain re-plug protocol does not provide; convert DataFrames with "
+            "to_network() instead."
+        )
+
+
 class GNPEKernelCorrection(TargetCorrection):
     """
     The single-step GNPE kernel correction as a target-side chain step.
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 6593b3027..be8972336 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -8,7 +8,6 @@
 from bilby_pipe.utils import CALIBRATION_CORRECTION_TYPE_LOOKUP
 
 from dingo.core.result import Result as CoreResult
-from dingo.gw.conversion import change_spin_conversion_phase
 from dingo.core.utils.backward_compatibility import check_minimum_version
 
 
@@ -607,17 +606,10 @@ def get_samples_bilby_phase(self, num_processes=1):
         pd.DataFrame
             Samples
         """
-        spin_conversion_phase_old = self.base_metadata["dataset_settings"][
-            "waveform_generator"
-        ].get("spin_conversion_phase")
-
-        # Redefine phase parameter to be consistent with Bilby.
-        return change_spin_conversion_phase(
-            self.samples,
-            self.f_ref,
-            spin_conversion_phase_old,
-            None,
-            num_processes=num_processes,
+        from dingo.gw.inference.factors import SpinConventionReparam
+
+        return SpinConventionReparam(num_processes=num_processes).to_physical(
+            self.samples, self.base_metadata
         )
 
     def get_pesummary_samples(
@@ -679,17 +671,11 @@ def get_pesummary_samples(
             if "time" in col:
                 samples.loc[:, col] += self.t_ref
 
-        spin_conversion_phase_old = self.base_metadata["dataset_settings"][
-            "waveform_generator"
-        ].get("spin_conversion_phase")
+        # Redefine the spin angles to the physical (Bilby) convention.
+        from dingo.gw.inference.factors import SpinConventionReparam
 
-        # Redefine phase parameter to be consistent with Bilby. COMMENTED BECAUSE SLOW
-        samples = change_spin_conversion_phase(
-            samples,
-            self.f_ref,
-            spin_conversion_phase_old,
-            None,
-            num_processes=num_processes,
+        samples = SpinConventionReparam(num_processes=num_processes).to_physical(
+            samples, self.base_metadata
         )
 
         self._pesummary_samples = samples
diff --git a/tests/gw/test_spin_convention_reparam.py b/tests/gw/test_spin_convention_reparam.py
new file mode 100644
index 000000000..9278b7854
--- /dev/null
+++ b/tests/gw/test_spin_convention_reparam.py
@@ -0,0 +1,179 @@
+"""Unit tests for ``SpinConventionReparam``: the step contract (in-place bijection,
+guarded log_det / inverse), the metadata-driven convention including the
+``None``-convention identity, and a real LAL roundtrip between the network and
+physical conventions."""
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from dingo.gw.inference.factors import SpinConventionReparam
+
+_MODEL_METADATA_SC0 = {
+    "dataset_settings": {
+        "waveform_generator": {
+            "approximant": "IMRPhenomXPHM",
+            "f_ref": 20.0,
+            "spin_conversion_phase": 0.0,
+        }
+    }
+}
+
+_MODEL_METADATA_NONE = {
+    "dataset_settings": {
+        "waveform_generator": {"approximant": "IMRPhenomXPHM", "f_ref": 20.0}
+    }
+}
+
+
+def _samples(n=5, seed=0):
+    rng = np.random.default_rng(seed)
+    return pd.DataFrame(
+        {
+            "theta_jn": rng.uniform(0.1, np.pi - 0.1, n),
+            "phi_jl": rng.uniform(0, 2 * np.pi, n),
+            "phase": rng.uniform(0, 2 * np.pi, n),
+            "chirp_mass": rng.uniform(20.0, 40.0, n),
+            "mass_ratio": rng.uniform(0.5, 1.0, n),
+            "a_1": rng.uniform(0.1, 0.9, n),
+            "a_2": rng.uniform(0.1, 0.9, n),
+            "tilt_1": rng.uniform(0.1, np.pi - 0.1, n),
+            "tilt_2": rng.uniform(0.1, np.pi - 0.1, n),
+            "phi_12": rng.uniform(0, 2 * np.pi, n),
+        }
+    )
+
+
+def test_contract():
+    reparam = SpinConventionReparam()
+    assert reparam.parameters == ["theta_jn", "phi_jl"]
+    # The bijection overwrites its parameters in place; the read-only
+    # conditioning (phase, masses, tilts, ...) must not be consumed.
+    assert reparam.consumes == []
+    assert set(reparam.parameters) < set(reparam.conditioning)
+
+
+def test_inverse_is_guarded():
+    # The chain re-plug protocol hands inverse only the parameters block, which
+    # is insufficient for this bijection; the DataFrame API is the working
+    # physical -> network direction.
+    with pytest.raises(NotImplementedError, match="conditioning"):
+        SpinConventionReparam().inverse({}, None)
+
+
+def test_none_convention_is_identity():
+    # A model trained without a fixed spin-conversion phase already outputs the
+    # physical convention; the relabel is the identity.
+    samples = _samples()
+    out = SpinConventionReparam().to_physical(samples, _MODEL_METADATA_NONE)
+    pd.testing.assert_frame_equal(out, samples.astype(float))
+
+
+def test_roundtrip_network_physical_network():
+    reparam = SpinConventionReparam()
+    samples = _samples()
+    physical = reparam.to_physical(samples, _MODEL_METADATA_SC0)
+    back = reparam.to_network(physical, _MODEL_METADATA_SC0)
+    # Only the two convention-dependent angles change (the others are invariants
+    # of the common rotation, recomputed through LAL trig, so equal to roundoff);
+    # the roundtrip closes; phase is explicitly restored, so it is exact.
+    assert not np.allclose(physical["theta_jn"], samples["theta_jn"])
+    for col in samples.columns:
+        assert np.allclose(back[col], samples[col], atol=1e-10), col
+    for col in ["tilt_1", "tilt_2", "phi_12", "a_1", "a_2"]:
+        assert np.allclose(physical[col], samples[col], atol=1e-12), col
+    assert np.array_equal(physical["phase"], samples["phase"].astype(float))
+
+
+def test_jacobian_matches_sin_ratio():
+    # The convention change rotates the line of sight rigidly about L, so it
+    # preserves the spherical measure sin(theta_jn) dtheta dphi; in the flat
+    # (theta_jn, phi_jl) coordinates the Jacobian determinant is then
+    # sin(theta_jn) / sin(theta_jn'). Verify numerically by central finite
+    # differences through the actual LAL conversion, and check that log_det /
+    # sample_and_log_prob implement exactly this.
+    reparam = SpinConventionReparam()
+    base = _samples(n=20, seed=3)
+    # Keep clear of the coordinate poles, where finite differences degrade.
+    base["theta_jn"] = np.clip(base["theta_jn"], 0.3, np.pi - 0.3)
+    h = 1e-6
+
+    # One batched conversion: for each base row, rows [base, th+h, th-h, ph+h, ph-h].
+    blocks = []
+    for _, row in base.iterrows():
+        for col, delta in [
+            (None, 0.0),
+            ("theta_jn", h),
+            ("theta_jn", -h),
+            ("phi_jl", h),
+            ("phi_jl", -h),
+        ]:
+            r = row.copy()
+            if col is not None:
+                r[col] += delta
+            blocks.append(r)
+    stacked = pd.DataFrame(blocks).reset_index(drop=True)
+    converted = reparam.to_physical(stacked, _MODEL_METADATA_SC0)
+
+    def wrap(d):
+        return (d + np.pi) % (2 * np.pi) - np.pi
+
+    theta_out = converted["theta_jn"].to_numpy().reshape(-1, 5)
+    phi_out = converted["phi_jl"].to_numpy().reshape(-1, 5)
+    dtheta_dtheta = (theta_out[:, 1] - theta_out[:, 2]) / (2 * h)
+    dphi_dtheta = wrap(phi_out[:, 1] - phi_out[:, 2]) / (2 * h)
+    dtheta_dphi = (theta_out[:, 3] - theta_out[:, 4]) / (2 * h)
+    dphi_dphi = wrap(phi_out[:, 3] - phi_out[:, 4]) / (2 * h)
+    det = dtheta_dtheta * dphi_dphi - dtheta_dphi * dphi_dtheta
+
+    expected = np.sin(base["theta_jn"].to_numpy()) / np.sin(theta_out[:, 0])
+    assert np.max(np.abs(det / expected - 1)) < 1e-5
+    # The map genuinely distorts the flat coordinate measure (order-one effect):
+    # a zero log_det here would corrupt chain densities per-row.
+    assert np.max(np.abs(expected - 1)) > 0.01
+
+    # The class implements exactly this Jacobian, and the sampling path
+    # contributes its negative alongside the forward transform.
+    class _Context:
+        model_metadata = _MODEL_METADATA_SC0
+
+    given = {k: torch.as_tensor(base[k].to_numpy()) for k in reparam.conditioning}
+    log_det = reparam.log_det(given, _Context())
+    assert np.allclose(log_det.numpy(), np.log(expected), atol=1e-12)
+    out, log_prob_contribution = reparam.sample_and_log_prob(1, _Context(), given)
+    assert np.allclose(out["theta_jn"].numpy(), theta_out[:, 0])
+    assert np.allclose(log_prob_contribution.numpy(), -np.log(expected), atol=1e-12)
+
+
+def test_chain_dtype_preserved():
+    # A reparametrization must not promote the chain's dtype: float32 columns in,
+    # float32 transform and log_det contribution out (computation is double
+    # internally).
+    class _Context:
+        model_metadata = _MODEL_METADATA_SC0
+
+    reparam = SpinConventionReparam()
+    samples = _samples()
+    given = {
+        k: torch.as_tensor(samples[k].to_numpy(), dtype=torch.float32)
+        for k in reparam.conditioning
+    }
+    out, log_prob_contribution = reparam.sample_and_log_prob(1, _Context(), given)
+    assert out["theta_jn"].dtype == torch.float32
+    assert out["phi_jl"].dtype == torch.float32
+    assert log_prob_contribution.dtype == torch.float32
+
+
+def test_forward_matches_to_physical():
+    class _Context:
+        model_metadata = _MODEL_METADATA_SC0
+
+    reparam = SpinConventionReparam()
+    samples = _samples()
+    given = {k: torch.as_tensor(samples[k].to_numpy()) for k in reparam.conditioning}
+    out = reparam.forward(given, _Context())
+    expected = reparam.to_physical(samples, _MODEL_METADATA_SC0)
+    assert set(out) == {"theta_jn", "phi_jl"}
+    for k in out:
+        assert np.allclose(out[k].numpy(), expected[k].to_numpy())

From 4d04bb90d41d9ea7426676aef7b6dcdc42301377 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 13:40:42 +0100
Subject: [PATCH 41/65] Port the BNS heterodyning transforms and likelihood
 decomposition

From PR #355: factor_fiducial_waveform and HeterodynePhase (leading-order
chirp heterodyning with an optional 1PN correction; a fixed-parameters
mode for inference-time use) and the GNPEChirp proxy transform, with
their tests -- including a slow integration test on real NRTidal
waveforms (roundtrip mismatch < 1e-14, an order of magnitude fewer zero
crossings after heterodyning).

From PR #354: d_inner_h_complex_multi gains return_rho2opt, returning
(h|h) alongside (d|h) so the full likelihood decomposition
log L = log_Zn + Re[(d|h) e^(2i phase)] - (h|h)/2 is available from a
single waveform generation (the basis for the chirp-mass sweep), with
its tests. The synthetic-phase profile path is unchanged (re-verified
bit-exact).

PR #354's sampler extensions (fixed_context_parameters,
_preprocess_data) are deliberately not ported: the composed spine
already expresses them as a DeltaFactor pin and the context's one-time
prepared_data. The remaining #355 changes (the sample-dict interface for
the dataset-compression transforms) are training-side and come later as
a package.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/likelihood.py                     |  38 ++-
 dingo/gw/transforms/gnpe_transforms.py     | 111 +++++++
 dingo/gw/transforms/waveform_transforms.py | 186 ++++++++++++
 tests/gw/test_likelihood.py                | 148 ++++++++++
 tests/gw/transforms/test_heterodyne.py     | 324 +++++++++++++++++++++
 5 files changed, 794 insertions(+), 13 deletions(-)
 create mode 100644 tests/gw/test_likelihood.py
 create mode 100644 tests/gw/transforms/test_heterodyne.py

diff --git a/dingo/gw/likelihood.py b/dingo/gw/likelihood.py
index 43df8f80f..ea9927322 100644
--- a/dingo/gw/likelihood.py
+++ b/dingo/gw/likelihood.py
@@ -640,7 +640,10 @@ def _log_likelihood_calibration_marginalized(self, theta):
         return logsumexp(likelihoods) - np.log(len(likelihoods))
 
     def d_inner_h_complex_multi(
-        self, theta: pd.DataFrame, num_processes: int = 1
+        self,
+        theta: pd.DataFrame,
+        num_processes: int = 1,
+        return_rho2opt: bool = False,
     ) -> np.ndarray:
         """
         Calculate the complex inner product (d | h(theta)) between the stored data d
@@ -652,26 +655,31 @@ def d_inner_h_complex_multi(
             Parameters at which to evaluate h.
         num_processes : int
             Number of parallel processes to use.
+        return_rho2opt : bool
+            If True, also return rho2opt = (h|h) for each sample. Useful for
+            synthetic phase sampling where the full likelihood decomposition
+            log L = log_Zn + Re[(d|h)*exp(2i*phase)] - 0.5*(h|h) is needed.
 
         Returns
         -------
-        complex : Inner product
+        np.ndarray or tuple of np.ndarray
+            Complex inner products. If return_rho2opt is True, returns
+            (d_inner_h_complex, rho2opt).
         """
         with threadpool_limits(limits=1, user_api="blas"):
-            # Generator object for theta rows. For idx this yields row idx of
-            # theta dataframe, converted to dict, ready to be passed to
-            # self.log_likelihood.
             theta_generator = (d[1].to_dict() for d in theta.iterrows())
 
             if num_processes > 1:
                 with Pool(processes=num_processes) as pool:
-                    d_inner_h_complex = pool.map(
-                        self.d_inner_h_complex, theta_generator
-                    )
+                    results = pool.map(self._d_inner_h_complex, theta_generator)
             else:
-                d_inner_h_complex = list(map(self.d_inner_h_complex, theta_generator))
+                results = list(map(self._d_inner_h_complex, theta_generator))
 
-        return np.array(d_inner_h_complex)
+        d_inner_h = np.array([r[0] for r in results])
+        if return_rho2opt:
+            rho2opt = np.array([r[1] for r in results]).real
+            return d_inner_h, rho2opt
+        return d_inner_h
 
     def d_inner_h_complex(self, theta):
         """
@@ -685,20 +693,24 @@ def d_inner_h_complex(self, theta):
 
         Returns
         -------
-        complex : Inner product
+        complex : Inner product (d|h)
         """
         # TODO: Implement for time marginalization.
-        return self._d_inner_h_complex(theta)
+        return self._d_inner_h_complex(theta)[0]
 
     def _d_inner_h_complex(self, theta):
+        """Return (d|h) and (h|h) as a tuple. Computing both is essentially
+        free since the waveform mu is already generated."""
         mu = self.signal(theta)["waveform"]
         d = self.whitened_strains
-        return sum(
+        d_inner_h = sum(
             [
                 inner_product_complex(d_ifo, mu_ifo)
                 for d_ifo, mu_ifo in zip(d.values(), mu.values())
             ]
         )
+        rho2opt = sum([inner_product(mu_ifo, mu_ifo) for mu_ifo in mu.values()])
+        return d_inner_h, rho2opt
 
 
 def inner_product(a, b, min_idx=0, delta_f=None, psd=None):
diff --git a/dingo/gw/transforms/gnpe_transforms.py b/dingo/gw/transforms/gnpe_transforms.py
index eb4db2ba4..55b52aee0 100644
--- a/dingo/gw/transforms/gnpe_transforms.py
+++ b/dingo/gw/transforms/gnpe_transforms.py
@@ -3,6 +3,8 @@
 from bilby.core.prior import PriorDict
 from abc import ABC, abstractmethod
 
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
+
 
 class GNPEBase(ABC):
     """
@@ -227,3 +229,112 @@ def __call__(self, input_sample):
         extrinsic_parameters.update(new_parameters)
         sample["extrinsic_parameters"] = extrinsic_parameters
         return sample
+
+
+class GNPEChirp(GNPEBase):
+    """
+    Relative binning / heterodyning GNPE transform, which factors out the overall chirp
+    from the waveform. This is done based on the proxy parameters chirp_mass_proxy and
+    optionally mass_ratio_proxy. These are defined as blurred version of the parameters
+    chirp_mass and mass_ratio.
+
+    At leading order, the data are transformed by dividing by a fiducial waveform of the
+    form
+
+        exp( - 1j * (3/128) * (pi G chirp_mass_proxy f / c**3)**(-5/3) ) ;
+
+    see 2001.11412, eq. (7.2). This is the leading order chirp due to the emission of
+    quadrupole radiation.
+
+    At next to leading order, this transform also optionally implements 1PN corrections
+    involving the mass ratio. We do not include any amplitude in the fiducial waveform,
+    since at inference time this transform will be applied to noisy data. Multiplying
+    the frequency-domain noise by a complex number of unit norm is allowed because it
+    only changes the phase, not the overall amplitude, which would change the noise PSD.
+    """
+
+    def __init__(self, kernel, domain, order: int = 0, inference: bool = False):
+        """
+        Parameters
+        ----------
+        kernel : dict or str
+            Defines a Bilby prior. If a dict, keys should include chirp_mass,
+            and (possibly) mass_ratio.
+        domain : Domain
+            Only works for a FrequencyDomain at present.
+        order : int
+            Twice the post-Newtonian order for the expansion. Valid orders are 0 and 2.
+        inference : bool = False
+            Whether to use inference or training mode.
+        """
+        # We copy the kernel because the PriorDict constructor modifies the argument.
+        kernel = kernel.copy()
+
+        if order == 0:
+            if "chirp_mass" not in kernel:
+                raise KeyError("Kernel must include chirp_mass key.")
+            if "mass_ratio" in kernel:
+                print(
+                    "Warning: mass_ratio kernel provided, but will be ignored for "
+                    "order 0 GNPE."
+                )
+                kernel.pop("mass_ratio")
+        elif order == 2:
+            if "chirp_mass" not in kernel or "mass_ratio" not in kernel:
+                raise KeyError("Kernel must include chirp_mass and mass_ratio keys.")
+            raise NotImplementedError(
+                "Mass ratio conditioning for 2nd order not set up."
+            )
+        else:
+            raise ValueError(f"Order {order} invalid. Acceptable values are 0 and 2.")
+
+        operators = {"chirp_mass": "+", "mass_ratio": "+"}
+        super().__init__(kernel, operators)
+
+        self.inference = inference
+        self.phase_heterodyning_transform = HeterodynePhase(
+            domain, order, inverse=False
+        )
+
+    def __call__(self, input_sample):
+        sample = input_sample.copy()
+
+        extrinsic_parameters = sample["extrinsic_parameters"].copy()
+
+        # If proxies already exist, use them. Otherwise, sample them. Proxies may
+        # already exist if provided by an unconditional initialization network when
+        # attempting to recover the density from GNPE samples, or when using fixed
+        # initialization parameters.
+        if set(self.proxy_list).issubset(extrinsic_parameters.keys()):
+            proxies = {p: extrinsic_parameters[p] for p in self.proxy_list}
+        else:
+            # The relevant parameters could be in either the intrinsic or extrinsic
+            # parameters list. At inference time, we put all GNPE parameters into the
+            # extrinsic parameters list.
+            parameters = {**sample.get("parameters", {}), **extrinsic_parameters}
+            proxies = self.sample_proxies(parameters)
+            extrinsic_parameters.update(proxies)
+            delta_parameters = {
+                "delta_"
+                + p[: -len("_proxy")]: parameters[p[: -len("_proxy")]]
+                - proxies[p]
+                for p in self.proxy_list
+            }
+            extrinsic_parameters.update(delta_parameters)
+            sample["extrinsic_parameters"] = extrinsic_parameters
+
+        # The only situation where we would expect to not have a waveform to transform
+        # would be when calculating parameter standardizations, since we just want to
+        # draw samples of the parameters at that point, and not prepare any data.
+        if "waveform" in sample:
+            sample["waveform"] = self.phase_heterodyning_transform(
+                {
+                    "waveform": sample["waveform"],
+                    "parameters": {
+                        "chirp_mass": proxies["chirp_mass_proxy"],
+                        "mass_ratio": proxies.get("mass_ratio_proxy"),
+                    },
+                }
+            )["waveform"]
+
+        return sample
diff --git a/dingo/gw/transforms/waveform_transforms.py b/dingo/gw/transforms/waveform_transforms.py
index f6d022435..31349b321 100644
--- a/dingo/gw/transforms/waveform_transforms.py
+++ b/dingo/gw/transforms/waveform_transforms.py
@@ -1,5 +1,8 @@
 from typing import Optional
+
+import lal
 import numpy as np
+import torch
 
 from dingo.gw.domains import MultibandedFrequencyDomain, UniformFrequencyDomain
 
@@ -555,3 +558,186 @@ def create_mask_based_on_frequency_update(
             frequency_masks[d] = frequency_masks[d] & mask_max
 
     return frequency_masks
+
+
+def factor_fiducial_waveform(
+    data, domain, chirp_mass, mass_ratio=None, order=0, inverse=False
+):
+    """
+    Relative binning / heterodyning. Divides the data by a fiducial waveform defined by
+    the chirp mass and (optionally) mass ratio. Allows for batching.
+
+    At leading order, this factors out the overall chirp by dividing the data by a
+    fiducial waveform of the form
+        exp( - 1j * (3/128) * (pi G chirp_mass f / c**3)**(-5/3) ) ;
+    see 2001.11412, eq. (7.2). This is the leading order chirp due to the emission of
+    quadrupole radiation.
+
+    At next-to-leading order, this implements 1PN corrections involving the mass ratio.
+
+    We do not include any amplitude in the fiducial waveform, since at inference time
+    this transformation will be applied to noisy data. Multiplying the frequency-domain
+    noise by a complex number of unit norm is allowed because it only changes the
+    phase, not the overall amplitude, which would change the noise PSD.
+
+    Parameters
+    ----------
+    data : Union[dict, torch.Tensor]
+        If a dict, the keys would correspond to different detectors or
+        polarizations. For a Tensor, these would be within different components.
+        This method uses the same fiducial waveform for each detector.
+    domain : Domain
+        Only works for a FrequencyDomain or MultibandedFrequencyDomain at present.
+    chirp_mass : Union[np.array, torch.Tensor, float]
+        Chirp mass parameter(s).
+    mass_ratio : Union[np.array, torch.Tensor, float]
+        Mass ratio parameter(s).
+    order : int
+        Twice the post-Newtonian order for the expansion. Valid orders are 0 and 2.
+    inverse : bool
+        If True, apply the inverse transformation (restore chirp phase).
+
+    Returns
+    -------
+    data: Union[dict, torch.Tensor]
+        Transformed data, of the same form as data.
+    """
+    if order not in (0, 2):
+        raise ValueError(f"Order {order} invalid. Acceptable values are 0 and 2.")
+    if order == 2:
+        if mass_ratio is None:
+            raise ValueError("Mass ratio required for 2nd order heterodyning.")
+        elif (
+            isinstance(chirp_mass, (np.ndarray, torch.Tensor))
+            and chirp_mass.shape != mass_ratio.shape
+        ):
+            raise ValueError(
+                f"Shape of chirp_mass ({chirp_mass.shape}) and mass_ratio "
+                f"({mass_ratio.shape}) don't match"
+            )
+    if not np.isfinite(chirp_mass).all():
+        raise ValueError("Got nan or inf elements in chirp_mass.")
+
+    if type(data) == dict:
+        f = domain.get_sample_frequencies_astype(list(data.values())[0])
+    else:
+        f = domain.get_sample_frequencies_astype(data)
+
+    # Expand across possible batch dimension.
+    if type(chirp_mass) == np.float64 or type(chirp_mass) == float:
+        # Use np.outer.squeeze instead of chirp_mass * f, as the latter has a
+        # different precision, and we don't want the behaviour of this function to
+        # depend on whether chirp_mass is an array or float.
+        mc_f = np.outer(chirp_mass, f).squeeze()
+        if f[0] == 0.0:
+            mc_f[0] = 1.0
+    elif type(chirp_mass) == np.ndarray:
+        mc_f = np.outer(chirp_mass, f)
+        if mass_ratio is not None:
+            mass_ratio = mass_ratio[:, None]
+    elif type(chirp_mass) == torch.Tensor:
+        if chirp_mass.dim() == 0:
+            mc_f = chirp_mass * f  # scalar tensor, no need for outer
+        else:
+            mc_f = torch.outer(chirp_mass, f)
+            if mass_ratio is not None:
+                mass_ratio = mass_ratio[:, None]
+    else:
+        raise TypeError(
+            f"Invalid type {type(chirp_mass)}. "
+            f"Only implemented for floats, arrays and tensors"
+        )
+
+    # Avoid taking a negative power of 0 in the first index. This will get
+    # chopped off or multiplied by 0 later anyway.
+    if (f[..., 0] == 0.0).all():
+        mc_f[..., 0] = 1.0
+
+    # Leading (0PN) phase
+    pi_mc_f_SI = np.pi * (lal.GMSUN_SI / lal.C_SI**3) * mc_f
+    fiducial_phase = (3 / 128) * (pi_mc_f_SI) ** (-5 / 3)
+
+    # 1PN correction
+    if order >= 2:
+        assert mass_ratio is not None
+        symmetric_mass_ratio = mass_ratio / (1 + mass_ratio) ** 2
+        pi_m_f_SI = pi_mc_f_SI / symmetric_mass_ratio ** (3 / 5)
+        correction = 1 + (55 * symmetric_mass_ratio / 9 + 3715 / 756) * (pi_m_f_SI) ** (
+            2 / 3
+        )
+
+        fiducial_phase *= correction
+
+    if inverse:
+        fiducial_phase *= -1
+
+    if type(data) == dict:
+        result = {}
+        for k, v in data.items():
+            result[k] = domain.add_phase(v, -fiducial_phase)
+    else:
+        result = domain.add_phase(data, -fiducial_phase)
+
+    return result
+
+
+class HeterodynePhase(object):
+    """
+    Transform operator for applying phase heterodyning. See docstring of
+    factor_fiducial_waveform for details.
+    """
+
+    def __init__(
+        self,
+        domain,
+        order: int = 0,
+        inverse: bool = False,
+        fixed_parameters: dict = None,
+    ):
+        """
+        Parameters
+        ----------
+        domain: Domain
+            Domain of the waveform data.
+        order : int
+            Twice the post-Newtonian order for the expansion. Either 0 or 2.
+        inverse : bool
+            Whether to apply the forward or inverse transform. Default: False.
+        fixed_parameters: dict = None
+            If set, extract chirp_mass and mass_ratio parameters from this dict instead
+            from the input sample.
+        """
+        self.domain = domain
+        self.order = order
+        self.inverse = inverse
+        self.fixed_parameters = fixed_parameters
+
+    def __call__(self, input_sample):
+        """
+        Parameters
+        ----------
+        input_sample : dict
+            sample["waveform"]: Dict with values being arrays containing waveforms,
+            or torch Tensor with the waveform.
+            sample["parameters"] contains parameters of the binary system. Required for
+            chirp mass (and mass ratio if self.order == 2).
+
+        Returns
+        -------
+        dict of the same form as the input, but with transformed (heterodyned) waveforms.
+        """
+        sample = input_sample.copy()
+        waveform = sample["waveform"]
+        if self.fixed_parameters:
+            parameters = self.fixed_parameters
+        else:
+            parameters = sample["parameters"]
+        sample["waveform"] = factor_fiducial_waveform(
+            waveform,
+            self.domain,
+            parameters["chirp_mass"],
+            mass_ratio=parameters.get("mass_ratio"),
+            order=self.order,
+            inverse=self.inverse,
+        )
+        return sample
diff --git a/tests/gw/test_likelihood.py b/tests/gw/test_likelihood.py
new file mode 100644
index 000000000..20da4fa28
--- /dev/null
+++ b/tests/gw/test_likelihood.py
@@ -0,0 +1,148 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from dingo.gw.domains import UniformFrequencyDomain
+from dingo.gw.likelihood import (
+    StationaryGaussianGWLikelihood,
+    inner_product,
+    inner_product_complex,
+)
+
+
+@pytest.fixture
+def likelihood_and_theta():
+    """Set up a simple likelihood for testing."""
+    f_min = 20.0
+    f_max = 1024.0
+    T = 8.0
+    domain = UniformFrequencyDomain(f_min, f_max, delta_f=1 / T)
+
+    waveform = {
+        "H1": np.where(domain.frequency_mask, (1.0 + 1j) * 1e-20, 0.0),
+        "L1": np.where(domain.frequency_mask, (1.0 + 1j) * 1e-20, 0.0),
+    }
+    asds = {
+        "H1": np.where(domain.frequency_mask, 1e-20, 1.0),
+        "L1": np.where(domain.frequency_mask, 1e-20, 1.0),
+    }
+
+    sample = {"waveform": waveform, "asds": asds}
+
+    likelihood = StationaryGaussianGWLikelihood(
+        wfg_kwargs={
+            "approximant": "IMRPhenomXPHM",
+            "f_ref": 20.0,
+            "spin_conversion_phase": 0.0,
+        },
+        wfg_domain=domain,
+        data_domain=domain,
+        event_data=sample,
+        t_ref=1248242632.0,
+    )
+
+    theta = {
+        "chirp_mass": 50.0,
+        "mass_ratio": 0.8,
+        "a_1": 0.3,
+        "a_2": 0.4,
+        "tilt_1": 1.3,
+        "tilt_2": 1.4,
+        "phi_12": 0.3,
+        "phi_jl": 3.3,
+        "theta_jn": 1.2,
+        "luminosity_distance": 4000.0,
+        "geocent_time": -0.03,
+        "dec": -0.3,
+        "ra": 1.5,
+        "psi": 1.2,
+        "phase": 1.4,
+    }
+
+    return likelihood, theta
+
+
+class TestLikelihoodDecomposition:
+    """Test the likelihood decomposition log L = log_Zn + kappa2 - 0.5*rho2opt."""
+
+    def test_log_likelihood_decomposition(self, likelihood_and_theta):
+        """log L = log_Zn + Re(d|h) - 0.5*(h|h)"""
+        likelihood, theta = likelihood_and_theta
+        log_l = likelihood.log_likelihood(theta)
+        d_inner_h, rho2opt = likelihood._d_inner_h_complex(theta)
+        kappa2 = d_inner_h.real
+        reconstructed = likelihood.log_Zn + kappa2 - 0.5 * rho2opt
+        assert np.isclose(log_l, reconstructed)
+
+    def test_rho2opt_is_positive(self, likelihood_and_theta):
+        """rho2opt = (h|h) should be non-negative."""
+        likelihood, theta = likelihood_and_theta
+        _, rho2opt = likelihood._d_inner_h_complex(theta)
+        assert rho2opt >= 0
+
+    def test_rho2opt_matches_direct_computation(self, likelihood_and_theta):
+        """rho2opt from _d_inner_h_complex should match direct (h|h)."""
+        likelihood, theta = likelihood_and_theta
+        _, rho2opt = likelihood._d_inner_h_complex(theta)
+
+        mu = likelihood.signal(theta)["waveform"]
+        rho2opt_direct = sum([inner_product(mu_ifo, mu_ifo) for mu_ifo in mu.values()])
+        assert np.isclose(rho2opt, rho2opt_direct)
+
+
+class TestDInnerHComplex:
+    """Test the complex inner product methods."""
+
+    def test_d_inner_h_complex_returns_complex(self, likelihood_and_theta):
+        """d_inner_h_complex returns a complex scalar (backward compat)."""
+        likelihood, theta = likelihood_and_theta
+        result = likelihood.d_inner_h_complex(theta)
+        assert np.isscalar(result) or isinstance(result, (complex, np.complexfloating))
+
+    def test_private_returns_tuple(self, likelihood_and_theta):
+        """_d_inner_h_complex returns (d_inner_h, rho2opt) tuple."""
+        likelihood, theta = likelihood_and_theta
+        result = likelihood._d_inner_h_complex(theta)
+        assert isinstance(result, tuple)
+        assert len(result) == 2
+
+    def test_public_matches_private(self, likelihood_and_theta):
+        """d_inner_h_complex matches _d_inner_h_complex[0]."""
+        likelihood, theta = likelihood_and_theta
+        public = likelihood.d_inner_h_complex(theta)
+        private_tuple = likelihood._d_inner_h_complex(theta)
+        assert np.isclose(public, private_tuple[0])
+
+    def test_multi_without_rho2opt(self, likelihood_and_theta):
+        """d_inner_h_complex_multi without return_rho2opt returns array only."""
+        likelihood, theta = likelihood_and_theta
+        theta_df = pd.DataFrame([theta, theta])
+        result = likelihood.d_inner_h_complex_multi(theta_df)
+        assert isinstance(result, np.ndarray)
+        assert result.shape == (2,)
+        assert np.iscomplexobj(result)
+
+    def test_multi_with_rho2opt(self, likelihood_and_theta):
+        """d_inner_h_complex_multi with return_rho2opt returns tuple."""
+        likelihood, theta = likelihood_and_theta
+        theta_df = pd.DataFrame([theta, theta, theta])
+        d_inner_h, rho2opt = likelihood.d_inner_h_complex_multi(
+            theta_df, return_rho2opt=True
+        )
+        assert d_inner_h.shape == (3,)
+        assert rho2opt.shape == (3,)
+        assert np.iscomplexobj(d_inner_h)
+        assert not np.iscomplexobj(rho2opt)
+        assert np.all(rho2opt >= 0)
+
+    def test_multi_rho2opt_consistency(self, likelihood_and_theta):
+        """rho2opt from multi matches single-sample computation."""
+        likelihood, theta = likelihood_and_theta
+        d_inner_h_single, rho2opt_single = likelihood._d_inner_h_complex(theta)
+
+        theta_df = pd.DataFrame([theta])
+        d_inner_h_multi, rho2opt_multi = likelihood.d_inner_h_complex_multi(
+            theta_df, return_rho2opt=True
+        )
+        assert np.isclose(d_inner_h_single, d_inner_h_multi[0])
+        assert np.isclose(rho2opt_single, rho2opt_multi[0])
diff --git a/tests/gw/transforms/test_heterodyne.py b/tests/gw/transforms/test_heterodyne.py
new file mode 100644
index 000000000..e8d427ab5
--- /dev/null
+++ b/tests/gw/transforms/test_heterodyne.py
@@ -0,0 +1,324 @@
+import pytest
+import numpy as np
+import pandas as pd
+import torch
+
+from dingo.gw.domains import UniformFrequencyDomain, build_domain
+from dingo.gw.transforms.waveform_transforms import (
+    factor_fiducial_waveform,
+    HeterodynePhase,
+)
+from dingo.gw.transforms.gnpe_transforms import GNPEChirp
+from dingo.gw.gwutils import get_mismatch
+
+
+@pytest.fixture
+def domain():
+    return UniformFrequencyDomain(f_min=20.0, f_max=1024.0, delta_f=0.125)
+
+
+@pytest.fixture
+def complex_waveform_dict(domain):
+    """Dict of complex numpy waveforms keyed by detector name."""
+    n = len(domain)
+    rng = np.random.default_rng(42)
+    return {
+        "H1": rng.standard_normal(n) + 1j * rng.standard_normal(n),
+        "L1": rng.standard_normal(n) + 1j * rng.standard_normal(n),
+    }
+
+
+@pytest.fixture
+def complex_waveform_tensor(domain):
+    """Complex torch tensor waveform, batched."""
+    torch.manual_seed(42)
+    batch = 8
+    n = len(domain)
+    return torch.randn(batch, n) + 1j * torch.randn(batch, n)
+
+
+# --- Tests for factor_fiducial_waveform ---
+
+
+def test_factor_fiducial_waveform_roundtrip_dict(domain, complex_waveform_dict):
+    """Heterodyne then inverse heterodyne recovers original data (dict input)."""
+    mc = 1.2
+    hetero = factor_fiducial_waveform(complex_waveform_dict, domain, mc)
+    recovered = factor_fiducial_waveform(hetero, domain, mc, inverse=True)
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(recovered[k], complex_waveform_dict[k], atol=1e-10)
+
+
+def test_factor_fiducial_waveform_roundtrip_tensor(domain, complex_waveform_tensor):
+    """Heterodyne then inverse heterodyne recovers original data (tensor input)."""
+    mc = torch.tensor(1.2)
+    hetero = factor_fiducial_waveform(complex_waveform_tensor, domain, mc)
+    recovered = factor_fiducial_waveform(hetero, domain, mc, inverse=True)
+    torch.testing.assert_close(recovered, complex_waveform_tensor, atol=1e-5, rtol=1e-5)
+
+
+def test_factor_fiducial_waveform_roundtrip_batched_tensor(domain):
+    """Roundtrip with batched chirp mass tensor."""
+    torch.manual_seed(0)
+    batch = 4
+    n = len(domain)
+    data = torch.randn(batch, n) + 1j * torch.randn(batch, n)
+    mc = torch.tensor([1.0, 1.2, 1.4, 1.6])
+    hetero = factor_fiducial_waveform(data, domain, mc)
+    recovered = factor_fiducial_waveform(hetero, domain, mc, inverse=True)
+    torch.testing.assert_close(recovered, data, atol=1e-4, rtol=1e-4)
+
+
+def test_factor_fiducial_waveform_unit_norm(domain, complex_waveform_dict):
+    """Verify the phase factor has unit norm (doesn't change amplitude)."""
+    mc = 1.2
+    hetero = factor_fiducial_waveform(complex_waveform_dict, domain, mc)
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(
+            np.abs(hetero[k]), np.abs(complex_waveform_dict[k]), atol=1e-12
+        )
+
+
+def test_factor_fiducial_waveform_order_2(domain, complex_waveform_dict):
+    """Verify order=2 (1PN correction) runs and differs from order=0 result."""
+    mc = 1.2
+    q = 0.8
+    result_0 = factor_fiducial_waveform(complex_waveform_dict, domain, mc, order=0)
+    result_2 = factor_fiducial_waveform(
+        complex_waveform_dict, domain, mc, mass_ratio=q, order=2
+    )
+    # They should differ
+    for k in complex_waveform_dict:
+        assert not np.allclose(
+            result_0[k], result_2[k]
+        ), "order=0 and order=2 results should differ"
+    # But amplitudes should still be preserved
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(
+            np.abs(result_2[k]), np.abs(complex_waveform_dict[k]), atol=1e-12
+        )
+
+
+def test_factor_fiducial_waveform_invalid_order(domain, complex_waveform_dict):
+    """Verify invalid order raises ValueError."""
+    with pytest.raises(ValueError, match="Order 1 invalid"):
+        factor_fiducial_waveform(complex_waveform_dict, domain, 1.2, order=1)
+
+
+# --- Tests for HeterodynePhase ---
+
+
+def test_heterodyne_phase_transform(domain, complex_waveform_dict):
+    """Test HeterodynePhase.__call__ with a sample dict."""
+    transform = HeterodynePhase(domain, order=0)
+    sample = {
+        "waveform": complex_waveform_dict,
+        "parameters": {"chirp_mass": 1.2},
+    }
+    result = transform(sample)
+    # Waveform should be modified
+    for k in complex_waveform_dict:
+        assert not np.allclose(result["waveform"][k], complex_waveform_dict[k])
+    # Amplitudes preserved
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(
+            np.abs(result["waveform"][k]),
+            np.abs(complex_waveform_dict[k]),
+            atol=1e-12,
+        )
+
+
+def test_heterodyne_phase_fixed_parameters(domain, complex_waveform_dict):
+    """Test HeterodynePhase with fixed_parameters overriding sample parameters."""
+    mc_fixed = 1.4
+    transform = HeterodynePhase(
+        domain, order=0, fixed_parameters={"chirp_mass": mc_fixed}
+    )
+    sample = {
+        "waveform": complex_waveform_dict,
+        "parameters": {"chirp_mass": 999.0},  # should be ignored
+    }
+    result = transform(sample)
+
+    # Compare with direct call using fixed chirp mass
+    expected = factor_fiducial_waveform(complex_waveform_dict, domain, mc_fixed)
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(result["waveform"][k], expected[k], atol=1e-12)
+
+
+# --- Tests for GNPEChirp ---
+
+
+def test_gnpe_chirp_training(domain, complex_waveform_dict):
+    """Test GNPEChirp in training mode: proxies sampled, delta stored, waveform modified."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0, inference=False)
+
+    mc_true = 1.2
+    sample = {
+        "waveform": complex_waveform_dict,
+        "parameters": {"chirp_mass": mc_true},
+        "extrinsic_parameters": {"chirp_mass": mc_true},
+    }
+    result = transform(sample)
+    ep = result["extrinsic_parameters"]
+
+    # Proxy should exist and be close to true value
+    assert "chirp_mass_proxy" in ep
+    assert abs(ep["chirp_mass_proxy"] - mc_true) <= 0.01
+
+    # Delta should be chirp_mass - proxy
+    assert "delta_chirp_mass" in ep
+    np.testing.assert_allclose(ep["delta_chirp_mass"], mc_true - ep["chirp_mass_proxy"])
+
+    # Waveform should be modified
+    for k in complex_waveform_dict:
+        assert not np.allclose(result["waveform"][k], complex_waveform_dict[k])
+
+
+def test_gnpe_chirp_inference(domain):
+    """Test GNPEChirp in inference mode with pre-existing proxies."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0, inference=True)
+
+    batch = 4
+    n = len(domain)
+    torch.manual_seed(0)
+    waveform = torch.randn(batch, n) + 1j * torch.randn(batch, n)
+
+    mc_true = torch.full((batch,), 1.2)
+    mc_proxy = torch.full((batch,), 1.19)
+
+    sample = {
+        "waveform": waveform,
+        "extrinsic_parameters": {
+            "chirp_mass": mc_true,
+            "chirp_mass_proxy": mc_proxy,
+        },
+    }
+    result = transform(sample)
+    ep = result["extrinsic_parameters"]
+
+    # Should reuse existing proxy
+    torch.testing.assert_close(ep["chirp_mass_proxy"], mc_proxy)
+
+
+def test_gnpe_chirp_context_parameters(domain):
+    """Verify context_parameters list is correct (should be proxy names from GNPEBase)."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0)
+    assert transform.context_parameters == ["chirp_mass_proxy"]
+
+
+def test_gnpe_chirp_no_waveform(domain):
+    """Verify GNPEChirp works when no 'waveform' key (standardization case)."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0)
+
+    sample = {
+        "parameters": {"chirp_mass": 1.2},
+        "extrinsic_parameters": {"chirp_mass": 1.2},
+    }
+    result = transform(sample)
+    ep = result["extrinsic_parameters"]
+
+    assert "chirp_mass_proxy" in ep
+    assert "delta_chirp_mass" in ep
+    assert "waveform" not in result
+
+
+# --- Integration test with real BNS waveforms (from Max's branch) ---
+
+
+def _number_of_zero_crossings(x):
+    """Count zero crossings along the last axis."""
+    return np.sum(np.diff(np.sign(x), axis=-1) != 0, axis=-1)
+
+
+def _max_zero_crossings(polarizations, min_idx=0):
+    return np.max(
+        [
+            np.max(_number_of_zero_crossings(v[..., min_idx:].real))
+            for v in polarizations.values()
+        ]
+    )
+
+
+def _max_mismatch(a, b, domain):
+    if a.keys() != b.keys():
+        raise ValueError()
+    return np.max([np.max(get_mismatch(a[k], b[k], domain)) for k in a.keys()])
+
+
+@pytest.mark.slow
+def test_heterodyning_bns_waveforms():
+    """
+    Integration test: generate BNS waveforms, heterodyne them, and verify that
+    (1) heterodyning is invertible (zero mismatch on roundtrip),
+    (2) heterodyned waveforms have significantly fewer zero crossings.
+
+    Adapted from Max Dax's test on the bns_add_dingo_pipe_max branch.
+    """
+    from dingo.gw.prior import build_prior_with_defaults
+    from dingo.gw.waveform_generator import (
+        WaveformGenerator,
+        generate_waveforms_parallel,
+    )
+
+    domain_settings = {
+        "type": "UniformFrequencyDomain",
+        "f_min": 20.0,
+        "f_max": 1024.0,
+        "delta_f": 0.00390625,
+    }
+    prior_settings = {
+        "mass_1": "bilby.core.prior.Constraint(minimum=1.0, maximum=2.5)",
+        "mass_2": "bilby.core.prior.Constraint(minimum=1.0, maximum=2.5)",
+        "chirp_mass": "bilby.gw.prior.UniformInComponentsChirpMass(minimum=1.0, maximum=2.0)",
+        "mass_ratio": "bilby.gw.prior.UniformInComponentsMassRatio(minimum=0.125, maximum=1.0)",
+        "phase": "default",
+        "a_1": "bilby.core.prior.Uniform(minimum=0.0, maximum=0.99)",
+        "a_2": "bilby.core.prior.Uniform(minimum=0.0, maximum=0.99)",
+        "tilt_1": "default",
+        "tilt_2": "default",
+        "phi_12": "default",
+        "phi_jl": "default",
+        "theta_jn": "default",
+        "lambda_1": "bilby.core.prior.Uniform(minimum=0.0, maximum=5000)",
+        "luminosity_distance": 100.0,
+        "geocent_time": 0.0,
+    }
+    wfg_settings = {"approximant": "IMRPhenomPv2_NRTidal", "f_ref": 10.0}
+
+    ufd = build_domain(domain_settings)
+    prior = build_prior_with_defaults(prior_settings)
+    wfg = WaveformGenerator(domain=ufd, **wfg_settings)
+
+    parameters = prior.sample(10)
+    polarizations = generate_waveforms_parallel(wfg, pd.DataFrame(parameters))
+
+    # Heterodyne and inverse heterodyne
+    polarizations_het = HeterodynePhase(ufd)(
+        {"waveform": polarizations, "parameters": parameters}
+    )["waveform"]
+    polarizations_roundtrip = HeterodynePhase(ufd, inverse=True)(
+        {"waveform": polarizations_het, "parameters": parameters}
+    )["waveform"]
+
+    # Roundtrip should recover original waveforms to machine precision.
+    assert _max_mismatch(polarizations, polarizations_roundtrip, ufd) < 1e-14
+    # Heterodyned waveforms should differ from originals.
+    assert _max_mismatch(polarizations, polarizations_het, ufd) > 1e-2
+
+    # Heterodyned waveforms should have far fewer zero crossings (the whole point).
+    n_roots = _max_zero_crossings(polarizations, ufd.min_idx)
+    n_roots_het = _max_zero_crossings(polarizations_het, ufd.min_idx)
+    assert n_roots > 10 * n_roots_het

From f649e9b8a0bba1fc676caa77daf53f51eb1d35ec Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 14:59:08 +0100
Subject: [PATCH 42/65] Fix the window-factor compatibility script for model
 files

The .pt branch referenced a nonexistent args.checkpoint (the parser
defines file_name), so it could never have run, and both branches missed
a window_factor nested in a multibanded domain's base_domain -- which is
exactly where pre-fix models carry it. Both file types now strip the key
at either level.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 compatibility/remove_domain_window_factor.py | 28 +++++++++++++-------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/compatibility/remove_domain_window_factor.py b/compatibility/remove_domain_window_factor.py
index 5b7dbac32..878afef15 100644
--- a/compatibility/remove_domain_window_factor.py
+++ b/compatibility/remove_domain_window_factor.py
@@ -27,24 +27,32 @@ def main():
     if path.suffix == ".hdf5":
         with h5py.File(path, "r+") as f:
             settings = ast.literal_eval(f.attrs["settings"])
-            try:
-                if "domain" in settings:
-                    del settings["domain"]["window_factor"]
-                elif "domain_dict" in settings:
-                    del settings["domain_dict"]["window_factor"]
+            domain_settings = settings.get("domain", settings.get("domain_dict"))
+            if domain_settings is not None and _remove_window_factor(domain_settings):
                 f.attrs["settings"] = str(settings)
-            except KeyError:
+            else:
                 print("Dataset is already in correct format.")
 
     elif path.suffix == ".pt":
-        d, _ = torch_load_with_fallback(args.checkpoint)
+        d, _ = torch_load_with_fallback(path)
 
-        try:
-            del d["metadata"]["dataset_settings"]["domain"]["window_factor"]
+        domain_settings = d["metadata"]["dataset_settings"]["domain"]
+        if _remove_window_factor(domain_settings):
             torch.save(d, path)
-        except KeyError:
+        else:
             print("Dataset is already in correct format.")
 
 
+def _remove_window_factor(domain_settings: dict) -> bool:
+    """Delete `window_factor` from the domain settings, including a multibanded
+    domain's nested base domain. Returns whether anything was removed."""
+    removed = False
+    for settings in (domain_settings, domain_settings.get("base_domain", {})):
+        if "window_factor" in settings:
+            del settings["window_factor"]
+            removed = True
+    return removed
+
+
 if __name__ == "__main__":
     main()

From 9c32a7878dac3750c3d07a853d2ef538e4904174 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 14:59:08 +0100
Subject: [PATCH 43/65] Add the proxy-offset reparametrization with
 conditioning-aware inverses

ProxyOffsetReparam reconstructs a physical parameter from a
proxy-conditioned network's offset output, X = delta_X + X_proxy (the
DINGO-BNS chirp-mass pattern): a pure shift at fixed proxy (log_det = 0)
that consumes the offset column while keeping the proxy recorded with
the samples, which is precisely what makes the chain output
self-sufficient for inversion.

Inverting it needs the proxy, so Reparametrization.inverse gains a
"given" argument carrying the non-consumed conditioning still present in
the chain -- the third bijection to need its parameterization at
inversion (RA needs none: its rotation is a per-event constant from the
context; the spin-convention change needed the invariant conditioning
and had to raise). The reverse fold supplies it, and
SpinConventionReparam.inverse is now implemented through it
(physical -> network, roundtrip-tested), closing the documented
protocol gap.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py                    | 16 ++++-
 dingo/gw/inference/factors.py            | 89 ++++++++++++++++++++----
 tests/core/test_factors.py               |  4 +-
 tests/gw/test_proxy_offset_reparam.py    | 61 ++++++++++++++++
 tests/gw/test_spin_convention_reparam.py | 23 ++++--
 5 files changed, 169 insertions(+), 24 deletions(-)
 create mode 100644 tests/gw/test_proxy_offset_reparam.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 75732ea78..d5a673e7d 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -469,9 +469,16 @@ def forward(
 
     @abstractmethod
     def inverse(
-        self, params: dict[str, torch.Tensor], context: "SamplerContext"
+        self,
+        params: dict[str, torch.Tensor],
+        context: "SamplerContext",
+        given: Optional[dict[str, torch.Tensor]] = None,
     ) -> dict[str, torch.Tensor]:
-        """Map the `parameters` block back to the conditioning block."""
+        """Rebuild the consumed inputs from the `parameters` block. `given` holds
+        the non-consumed conditioning columns still present in the chain (e.g. a
+        proxy the bijection shifts by, or invariant parameters a coordinate
+        change needs); bijections that depend only on their own outputs may
+        ignore it."""
 
     def log_det(
         self, given: dict[str, torch.Tensor], context: "SamplerContext"
@@ -727,7 +734,10 @@ def log_prob(
                 continue
             if isinstance(step, Reparametrization):
                 params = {k: values[k] for k in step.parameters}
-                values.update(step.inverse(params, context))
+                # The non-consumed conditioning is still present and may be
+                # needed to invert (e.g. a proxy the bijection shifts by).
+                available = {k: values[k] for k in step.conditioning if k in values}
+                values.update(step.inverse(params, context, available))
                 given = {k: values[k] for k in step.conditioning}
                 total = total - step.log_det(given, context)
             else:
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index ffc2ab00b..cf9a9cea1 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -1049,7 +1049,7 @@ def forward(self, given, context):
         ra = (given["ra@t_ref"].double() + correction) % (2 * np.pi)
         return {"ra": ra.float()}
 
-    def inverse(self, params, context):
+    def inverse(self, params, context, given=None):
         correction = self._correction(context)
         if correction == 0.0:
             return {"ra@t_ref": params["ra"]}
@@ -1084,10 +1084,9 @@ class SpinConventionReparam(Reparametrization):
     the orbital angular momentum, preserving the spherical measure
     `sin(theta_jn) dtheta dphi`, so
     `log_det = log sin(theta_jn) - log sin(theta_jn')` (verified numerically
-    against finite differences through the LAL conversion). Forward chain
-    sampling is therefore density-correct; chain re-plug
-    (`ChainComposer.log_prob`) additionally needs the conditioning-aware
-    `inverse`, which the chain protocol does not yet provide.
+    against finite differences through the LAL conversion), and `inverse`
+    rebuilds the network convention from the physical one using the invariant
+    conditioning the reverse fold supplies.
     """
 
     def __init__(self, num_processes: int = 1):
@@ -1181,14 +1180,78 @@ def forward(self, given, context):
             for k in self.parameters
         }
 
-    def inverse(self, params, context):
-        # The chain protocol hands inverse only the parameters block, but this
-        # bijection also needs the invariant conditioning (phase, masses, tilts).
-        raise NotImplementedError(
-            "The spin-convention inverse needs the conditioning block, which the "
-            "chain re-plug protocol does not provide; convert DataFrames with "
-            "to_network() instead."
-        )
+    def inverse(self, params, context, given=None):
+        # The physical -> network direction also needs the invariant
+        # conditioning (phase, masses, tilts), which the reverse fold supplies
+        # as `given`.
+        if given is None:
+            raise ValueError(
+                "The spin-convention inverse needs the conditioning block "
+                "(phase, masses, tilts); pass it as `given`, or convert "
+                "DataFrames with to_network()."
+            )
+        reference = params["theta_jn"]
+        rows = {**given, **params}
+        theta = pd.DataFrame({k: _to_numpy(v) for k, v in rows.items()})
+        converted = self.to_network(theta, context.model_metadata)
+        return {
+            k: torch.as_tensor(converted[k].to_numpy()).to(
+                dtype=reference.dtype, device=reference.device
+            )
+            for k in self.parameters
+        }
+
+
+class ProxyOffsetReparam(Reparametrization):
+    """
+    Reconstruct a physical parameter from a network's offset output and its proxy:
+    `X = delta_X + X_proxy`.
+
+    A proxy-conditioned network (e.g. the chirp-mass prior conditioning of
+    DINGO-BNS) infers the offset `delta_X = X - X_proxy` rather than `X` itself.
+    This step rebuilds `X`, consuming the offset column while keeping the proxy in
+    the chain (it is recorded with the samples, like the GNPE time proxies). A pure
+    shift at fixed proxy, so `log_det = 0`; `inverse` recovers the offset from the
+    proxy the reverse fold supplies.
+    """
+
+    def __init__(self, parameter_name: str):
+        """
+        Parameters
+        ----------
+        parameter_name : str
+            The physical parameter name `X`; the step reads `delta_X` and
+            `X_proxy` and produces `X`.
+        """
+        self.parameter_name = parameter_name
+        self.delta_name = f"delta_{parameter_name}"
+        self.proxy_name = f"{parameter_name}_proxy"
+        self.parameters = [parameter_name]
+        self.conditioning = [self.delta_name, self.proxy_name]
+
+    @property
+    def consumes(self) -> list[str]:
+        # The offset is replaced by the physical parameter; the proxy stays in
+        # the chain (recorded with the samples).
+        return [self.delta_name]
+
+    def forward(self, given, context):
+        return {self.parameter_name: given[self.delta_name] + given[self.proxy_name]}
+
+    def inverse(self, params, context, given=None):
+        if given is None or self.proxy_name not in given:
+            raise ValueError(
+                f"Inverting {self.parameter_name} = {self.delta_name} + "
+                f"{self.proxy_name} requires the proxy in `given`."
+            )
+        return {self.delta_name: params[self.parameter_name] - given[self.proxy_name]}
+
+    def describe(self) -> dict:
+        return {
+            "step": type(self).__name__,
+            "parameters": list(self.parameters),
+            "conditioning": list(self.conditioning),
+        }
 
 
 class GNPEKernelCorrection(TargetCorrection):
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index dafd733fe..2c0dfedec 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -283,7 +283,7 @@ def __init__(self, shift=10.0, log_det_val=0.5):
     def forward(self, given, context):
         return {"v": given["u"] + self._shift}
 
-    def inverse(self, params, context):
+    def inverse(self, params, context, given=None):
         return {"u": params["v"] - self._shift}
 
     def log_det(self, given, context):
@@ -513,7 +513,7 @@ def __init__(self, shift=5.0):
     def forward(self, given, context):
         return {"x": given["x"] + self._shift}
 
-    def inverse(self, params, context):
+    def inverse(self, params, context, given=None):
         return {"x": params["x"] - self._shift}
 
 
diff --git a/tests/gw/test_proxy_offset_reparam.py b/tests/gw/test_proxy_offset_reparam.py
new file mode 100644
index 000000000..66f23bc75
--- /dev/null
+++ b/tests/gw/test_proxy_offset_reparam.py
@@ -0,0 +1,61 @@
+"""Unit tests for ``ProxyOffsetReparam``: the offset reconstruction
+`X = delta_X + X_proxy` used by proxy-conditioned networks (e.g. DINGO-BNS chirp
+mass), including the reverse fold through a chain."""
+
+import pytest
+import torch
+
+from dingo.core.factors import ChainComposer, SampleTableFactor
+from dingo.gw.inference.factors import ProxyOffsetReparam
+
+
+def test_contract():
+    reparam = ProxyOffsetReparam("chirp_mass")
+    assert reparam.parameters == ["chirp_mass"]
+    assert reparam.conditioning == ["delta_chirp_mass", "chirp_mass_proxy"]
+    # The offset is consumed; the proxy stays in the chain (recorded with the
+    # samples, like the GNPE time proxies).
+    assert reparam.consumes == ["delta_chirp_mass"]
+
+
+def test_forward_inverse_and_log_det():
+    reparam = ProxyOffsetReparam("chirp_mass")
+    given = {
+        "delta_chirp_mass": torch.tensor([-0.004, 0.0, 0.003], dtype=torch.float64),
+        "chirp_mass_proxy": torch.tensor([1.1975, 1.1975, 1.1975], dtype=torch.float64),
+    }
+    out, log_prob_contribution = reparam.sample_and_log_prob(1, None, given)
+    assert torch.allclose(
+        out["chirp_mass"], torch.tensor([1.1935, 1.1975, 1.2005], dtype=torch.float64)
+    )
+    # A pure shift at fixed proxy: measure-preserving.
+    assert torch.equal(log_prob_contribution, torch.zeros(3))
+
+    back = reparam.inverse(
+        out, None, given={"chirp_mass_proxy": given["chirp_mass_proxy"]}
+    )
+    assert torch.allclose(back["delta_chirp_mass"], given["delta_chirp_mass"])
+    with pytest.raises(ValueError, match="proxy"):
+        reparam.inverse(out, None)
+
+
+def test_chain_replug_rebuilds_offset():
+    # In a chain the reverse fold supplies the proxy to the inverse, so log_prob
+    # at the emitted samples reproduces the stored chain density.
+    table = SampleTableFactor(
+        {
+            "delta_chirp_mass": torch.tensor([-0.002, 0.001]),
+            "chirp_mass_proxy": torch.tensor([1.4868, 1.4868]),
+        },
+        log_prob=torch.tensor([0.7, -0.3]),
+    )
+    chain = ChainComposer([table, ProxyOffsetReparam("chirp_mass")])
+    samples, log_prob = chain.sample_and_log_prob(2, context=None)
+    assert "chirp_mass" in samples and "delta_chirp_mass" not in samples
+    assert torch.allclose(log_prob, torch.tensor([0.7, -0.3]))
+
+    # Full-chain re-plug traverses the reparametrization inverse first (the fold
+    # supplies the proxy as given) and only then fails at the table root, which
+    # is not a density -- the specific error proves the inverse succeeded.
+    with pytest.raises(NotImplementedError, match="not a density"):
+        chain.log_prob(samples, context=None)
diff --git a/tests/gw/test_spin_convention_reparam.py b/tests/gw/test_spin_convention_reparam.py
index 9278b7854..c1b50c773 100644
--- a/tests/gw/test_spin_convention_reparam.py
+++ b/tests/gw/test_spin_convention_reparam.py
@@ -54,12 +54,23 @@ def test_contract():
     assert set(reparam.parameters) < set(reparam.conditioning)
 
 
-def test_inverse_is_guarded():
-    # The chain re-plug protocol hands inverse only the parameters block, which
-    # is insufficient for this bijection; the DataFrame API is the working
-    # physical -> network direction.
-    with pytest.raises(NotImplementedError, match="conditioning"):
-        SpinConventionReparam().inverse({}, None)
+def test_inverse_requires_given_and_closes_roundtrip():
+    # The physical -> network direction needs the invariant conditioning, which
+    # the reverse fold supplies as `given`; without it the call fails loudly.
+    class _Context:
+        model_metadata = _MODEL_METADATA_SC0
+
+    reparam = SpinConventionReparam()
+    with pytest.raises(ValueError, match="conditioning"):
+        reparam.inverse({}, None)
+
+    samples = _samples()
+    given = {k: torch.as_tensor(samples[k].to_numpy()) for k in reparam.conditioning}
+    physical, _ = reparam.sample_and_log_prob(1, _Context(), given)
+    inverse_given = {**given, **physical}
+    back = reparam.inverse(physical, _Context(), inverse_given)
+    for k in reparam.parameters:
+        assert np.allclose(back[k].numpy(), samples[k].to_numpy(), atol=1e-10)
 
 
 def test_none_convention_is_identity():

From 7d73d9ef637f347dd47d5d6111559f0383a0e9b6 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 15:46:37 +0100
Subject: [PATCH 44/65] Add the fixed-proxy BNS chain builder

GWSamplerContext.from_model_metadata accepts fixed_context_parameters:
for a chirp-GNPE model, the raw strain is heterodyned once with the
fixed proxy (HeterodynePhase at the head of the data preparation, before
decimation, on the base domain); only proxy entries parameterize the
data preparation, so a pinned sky position is conditioning-only. Without
pins the network-input view fails loudly while the likelihood and prior
views keep working (from-file results reconstruct their context without
pins).

GWComposedSampler.from_model becomes the general single-network builder:
a model with context_parameters (the DINGO-BNS chirp-mass prior
conditioning with a fixed sky position) requires pins for exactly those
parameters, and the chain becomes [DeltaFactor(pins), conditioned
FlowFactor, one ProxyOffsetReparam per inferred offset with a pinned
proxy]. Plain-NPE assembly is unchanged.

ProxyOffsetReparam moves to dingo.core.factors: it is domain-agnostic (a
name-keyed shift), meeting the core criterion alongside DeltaFactor.

Verified on the DINGO-BNS GW170817 model (window-fixed copy): the chain
samples in a single pass with finite log_prob, chirp_mass minus the
proxy fills the +/-0.005 training kernel on noise-only data, the
re-plug through the offset inverse reproduces the stored density at
float32 roundoff, and to_result exports the tidal prior intact.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py                         |  52 ++++++
 dingo/gw/inference/factors.py                 | 166 +++++++++++-------
 .../{gw => core}/test_proxy_offset_reparam.py |   3 +-
 tests/gw/test_bns_chain_builder.py            | 133 ++++++++++++++
 4 files changed, 287 insertions(+), 67 deletions(-)
 rename tests/{gw => core}/test_proxy_offset_reparam.py (95%)
 create mode 100644 tests/gw/test_bns_chain_builder.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index d5a673e7d..d87625696 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -512,6 +512,58 @@ def describe(self) -> dict:
         return _describe_default(self)
 
 
+class ProxyOffsetReparam(Reparametrization):
+    """
+    Reconstruct a physical parameter from a network's offset output and its proxy:
+    `X = delta_X + X_proxy`.
+
+    A proxy-conditioned network (e.g. the chirp-mass prior conditioning of
+    DINGO-BNS) infers the offset `delta_X = X - X_proxy` rather than `X` itself.
+    This step rebuilds `X`, consuming the offset column while keeping the proxy in
+    the chain (it is recorded with the samples, like the GNPE time proxies). A pure
+    shift at fixed proxy, so `log_det = 0`; `inverse` recovers the offset from the
+    proxy the reverse fold supplies.
+    """
+
+    def __init__(self, parameter_name: str):
+        """
+        Parameters
+        ----------
+        parameter_name : str
+            The physical parameter name `X`; the step reads `delta_X` and
+            `X_proxy` and produces `X`.
+        """
+        self.parameter_name = parameter_name
+        self.delta_name = f"delta_{parameter_name}"
+        self.proxy_name = f"{parameter_name}_proxy"
+        self.parameters = [parameter_name]
+        self.conditioning = [self.delta_name, self.proxy_name]
+
+    @property
+    def consumes(self) -> list[str]:
+        # The offset is replaced by the physical parameter; the proxy stays in
+        # the chain (recorded with the samples).
+        return [self.delta_name]
+
+    def forward(self, given, context):
+        return {self.parameter_name: given[self.delta_name] + given[self.proxy_name]}
+
+    def inverse(self, params, context, given=None):
+        if given is None or self.proxy_name not in given:
+            raise ValueError(
+                f"Inverting {self.parameter_name} = {self.delta_name} + "
+                f"{self.proxy_name} requires the proxy in `given`."
+            )
+        return {self.delta_name: params[self.parameter_name] - given[self.proxy_name]}
+
+    def describe(self) -> dict:
+        return {
+            "step": type(self).__name__,
+            "parameters": list(self.parameters),
+            "conditioning": list(self.conditioning),
+        }
+
+
 class TargetCorrection(ABC):
     """
     A `Step` that emits an importance-sampling target correction as a side-channel column
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index cf9a9cea1..d7c97f0e0 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -37,6 +37,7 @@
     Factor,
     FlowFactor,
     GibbsBlock,
+    ProxyOffsetReparam,
     Reparametrization,
     Standardization,
     TargetCorrection,
@@ -62,6 +63,7 @@
     GNPECoalescenceTimes,
     MaskDataForFrequencyRangeUpdate,
     PostCorrectGeocentTime,
+    HeterodynePhase,
     RepackageStrainsAndASDS,
     SelectStandardizeRepackageParameters,
     TimeShiftStrain,
@@ -92,6 +94,20 @@ def normalize(value):
     return {"minimum_frequency": minimum, "maximum_frequency": maximum}
 
 
+class _RequireFixedProxies:
+    """Data-prep placeholder for a chirp-GNPE model built without fixed proxies:
+    the network-input view is undefined without the heterodyning proxy, so
+    `prepared_data()` fails loudly. The likelihood and prior views are unaffected
+    (a from-file result reconstructs its context without pins)."""
+
+    def __call__(self, sample):
+        raise ValueError(
+            "This model heterodynes its network input with a fixed chirp-mass "
+            "proxy; construct the context with fixed_context_parameters "
+            "(e.g. {'chirp_mass_proxy': ...}) to prepare network input."
+        )
+
+
 class GWSamplerContext:
     """
     Per-event shared GW state: the data `d` and everything derived from it. Referenced
@@ -188,6 +204,7 @@ def from_model_metadata(
         event_data: dict,
         event_metadata: Optional[dict] = None,
         device: Union[torch.device, str] = "cpu",
+        fixed_context_parameters: Optional[dict] = None,
     ) -> "GWSamplerContext":
         """Build the context from a model-metadata dict -- no model required. The
         domain, one-time data-prep chain, prior, and likelihood are all defined by
@@ -205,6 +222,11 @@ def from_model_metadata(
             Per-event metadata.
         device : torch.device or str, default "cpu"
             Device for `prepared_data()` and chain-created tensors.
+        fixed_context_parameters : dict, optional
+            Pinned network-conditioning values (e.g.
+            `{"chirp_mass_proxy": 1.1975, "ra": ..., "dec": ...}`). Proxy entries
+            that parameterize the data preparation (the chirp-mass heterodyne)
+            are consumed here; the chain conditions the network on all of them.
 
         Returns
         -------
@@ -218,6 +240,28 @@ def from_model_metadata(
         detectors = data_settings["detectors"]
 
         transforms = []
+        # Chirp-mass GNPE (BNS): heterodyne the raw strain once with the fixed
+        # proxy -- before decimation (they do not commute) and on the base
+        # domain. Iterated chirp GNPE (heterodyning inside a Gibbs loop) is not
+        # implemented: it would require carrying the undecimated strain per
+        # sample.
+        gnpe_chirp = data_settings.get("gnpe_chirp")
+        if gnpe_chirp is not None:
+            fixed = {
+                k[: -len("_proxy")]: v
+                for k, v in (fixed_context_parameters or {}).items()
+                if k.endswith("_proxy")
+            }
+            if "chirp_mass" in fixed:
+                transforms.append(
+                    HeterodynePhase(
+                        domain=getattr(domain, "base_domain", domain),
+                        order=gnpe_chirp.get("order", 0),
+                        fixed_parameters=fixed,
+                    )
+                )
+            else:
+                transforms.append(_RequireFixedProxies())
         # Decimate from the base domain when using a multibanded frequency domain.
         if isinstance(domain, MultibandedFrequencyDomain):
             transforms.append(
@@ -259,6 +303,7 @@ def from_model(
         model: BasePosteriorModel,
         event_data: dict,
         event_metadata: Optional[dict] = None,
+        fixed_context_parameters: Optional[dict] = None,
     ) -> "GWSamplerContext":
         """Build the context from a model: its own metadata and its device. Data
         preparation is network-bound, so the settings come from `model.metadata`
@@ -288,7 +333,11 @@ def from_model(
                 "GWSamplerContext.from_model_metadata(_base_model_metadata(model), ...)."
             )
         return cls.from_model_metadata(
-            model.metadata, event_data, event_metadata, device=model.device
+            model.metadata,
+            event_data,
+            event_metadata,
+            device=model.device,
+            fixed_context_parameters=fixed_context_parameters,
         )
 
     def derive(
@@ -1202,58 +1251,6 @@ def inverse(self, params, context, given=None):
         }
 
 
-class ProxyOffsetReparam(Reparametrization):
-    """
-    Reconstruct a physical parameter from a network's offset output and its proxy:
-    `X = delta_X + X_proxy`.
-
-    A proxy-conditioned network (e.g. the chirp-mass prior conditioning of
-    DINGO-BNS) infers the offset `delta_X = X - X_proxy` rather than `X` itself.
-    This step rebuilds `X`, consuming the offset column while keeping the proxy in
-    the chain (it is recorded with the samples, like the GNPE time proxies). A pure
-    shift at fixed proxy, so `log_det = 0`; `inverse` recovers the offset from the
-    proxy the reverse fold supplies.
-    """
-
-    def __init__(self, parameter_name: str):
-        """
-        Parameters
-        ----------
-        parameter_name : str
-            The physical parameter name `X`; the step reads `delta_X` and
-            `X_proxy` and produces `X`.
-        """
-        self.parameter_name = parameter_name
-        self.delta_name = f"delta_{parameter_name}"
-        self.proxy_name = f"{parameter_name}_proxy"
-        self.parameters = [parameter_name]
-        self.conditioning = [self.delta_name, self.proxy_name]
-
-    @property
-    def consumes(self) -> list[str]:
-        # The offset is replaced by the physical parameter; the proxy stays in
-        # the chain (recorded with the samples).
-        return [self.delta_name]
-
-    def forward(self, given, context):
-        return {self.parameter_name: given[self.delta_name] + given[self.proxy_name]}
-
-    def inverse(self, params, context, given=None):
-        if given is None or self.proxy_name not in given:
-            raise ValueError(
-                f"Inverting {self.parameter_name} = {self.delta_name} + "
-                f"{self.proxy_name} requires the proxy in `given`."
-            )
-        return {self.delta_name: params[self.parameter_name] - given[self.proxy_name]}
-
-    def describe(self) -> dict:
-        return {
-            "step": type(self).__name__,
-            "parameters": list(self.parameters),
-            "conditioning": list(self.conditioning),
-        }
-
-
 class GNPEKernelCorrection(TargetCorrection):
     """
     The single-step GNPE kernel correction as a target-side chain step.
@@ -1285,6 +1282,19 @@ def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
     return {"ra": "ra@t_ref"} if "ra" in inference_parameters else {}
 
 
+def _proxy_offset_steps(
+    inference_parameters: list[str], context_parameters: list[str]
+) -> list:
+    """One offset reconstruction (`X = delta_X + X_proxy`) per `delta_X` the
+    network infers whose proxy it conditions on."""
+    return [
+        ProxyOffsetReparam(p[len("delta_") :])
+        for p in inference_parameters
+        if p.startswith("delta_")
+        and p[len("delta_") :] + "_proxy" in context_parameters
+    ]
+
+
 def _ra_reparam_steps(inference_parameters: list[str]) -> list:
     """The `RAReparam` step, appended to a chain only when the model infers `ra`."""
     return [RAReparam()] if "ra" in inference_parameters else []
@@ -1412,34 +1422,60 @@ def from_model(
         model: BasePosteriorModel,
         event_data: dict,
         event_metadata: Optional[dict] = None,
+        fixed_context_parameters: Optional[dict] = None,
     ) -> "GWComposedSampler":
-        """Build a plain-NPE GW sampler from a model and event data: the flow exposes
-        `ra` as `ra@t_ref`, followed by an `RAReparam` to the event frame.
+        """Build a single-network GW sampler from a model and event data.
+
+        For a plain NPE model the chain is the flow (exposing `ra` as `ra@t_ref`)
+        followed by an `RAReparam` to the event frame. A model with
+        `context_parameters` (e.g. the DINGO-BNS chirp-mass prior conditioning
+        with a fixed sky position) requires `fixed_context_parameters` pinning
+        all of them: the chain is then rooted in a `DeltaFactor` of the pins,
+        the flow conditions on them, and each inferred offset `delta_X` with a
+        pinned proxy is reconstructed by a `ProxyOffsetReparam`
+        (`X = delta_X + X_proxy`). Proxies that parameterize the data
+        preparation (the chirp-mass heterodyne) are applied once inside the
+        context.
 
         Parameters
         ----------
         model : BasePosteriorModel
-            The NPE model.
+            The model.
         event_data : dict
             The raw event data (strain + ASDs).
         event_metadata : dict, optional
             Per-event metadata.
+        fixed_context_parameters : dict, optional
+            Pinned values for the model's `context_parameters`, e.g.
+            `{"chirp_mass_proxy": 1.1975, "ra": 3.446, "dec": -0.408}`.
 
         Returns
         -------
         GWComposedSampler
         """
-        context = GWSamplerContext.from_model(model, event_data, event_metadata)
+        context = GWSamplerContext.from_model(
+            model,
+            event_data,
+            event_metadata,
+            fixed_context_parameters=fixed_context_parameters,
+        )
         metadata = _base_model_metadata(model)
-        inference_parameters = metadata["train_settings"]["data"][
-            "inference_parameters"
-        ]
+        data_settings = metadata["train_settings"]["data"]
+        inference_parameters = data_settings["inference_parameters"]
         factor = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
-        steps = (
-            [factor]
-            + _ra_reparam_steps(inference_parameters)
-            + _delta_prior_steps(context.prior, inference_parameters)
-        )
+        steps = [factor]
+        context_parameters = data_settings.get("context_parameters") or []
+        if context_parameters:
+            if set(fixed_context_parameters or {}) != set(context_parameters):
+                raise ValueError(
+                    f"The model conditions on {context_parameters}; provide "
+                    f"fixed_context_parameters with exactly these keys (for "
+                    f"iterated proxies, use the GNPE builders instead)."
+                )
+            steps = [DeltaFactor(fixed_context_parameters)] + steps
+        steps += _proxy_offset_steps(inference_parameters, context_parameters)
+        steps += _ra_reparam_steps(inference_parameters)
+        steps += _delta_prior_steps(context.prior, inference_parameters)
         return cls(
             composer=ChainComposer(steps),
             context=context,
diff --git a/tests/gw/test_proxy_offset_reparam.py b/tests/core/test_proxy_offset_reparam.py
similarity index 95%
rename from tests/gw/test_proxy_offset_reparam.py
rename to tests/core/test_proxy_offset_reparam.py
index 66f23bc75..aec256066 100644
--- a/tests/gw/test_proxy_offset_reparam.py
+++ b/tests/core/test_proxy_offset_reparam.py
@@ -5,8 +5,7 @@
 import pytest
 import torch
 
-from dingo.core.factors import ChainComposer, SampleTableFactor
-from dingo.gw.inference.factors import ProxyOffsetReparam
+from dingo.core.factors import ChainComposer, ProxyOffsetReparam, SampleTableFactor
 
 
 def test_contract():
diff --git a/tests/gw/test_bns_chain_builder.py b/tests/gw/test_bns_chain_builder.py
new file mode 100644
index 000000000..d4093f733
--- /dev/null
+++ b/tests/gw/test_bns_chain_builder.py
@@ -0,0 +1,133 @@
+"""Unit tests for the fixed-proxy (DINGO-BNS) chain construction: the heterodyne
+insertion into the context's data preparation, the pin requirement, and the
+chain assembly `[DeltaFactor(pins), conditioned FlowFactor, ProxyOffsetReparam]`.
+Mock metadata only; the real-model smoke lives in the local harnesses."""
+
+import numpy as np
+import pytest
+
+from dingo.core.factors import DeltaFactor, FlowFactor, ProxyOffsetReparam
+from dingo.gw.inference.factors import (
+    GWComposedSampler,
+    GWSamplerContext,
+    _proxy_offset_steps,
+)
+from dingo.gw.transforms import HeterodynePhase
+
+_PINS = {"chirp_mass_proxy": 1.1975, "ra": 3.44616, "dec": -0.408084}
+
+_BNS_METADATA = {
+    "dataset_settings": {
+        "domain": {
+            "type": "UniformFrequencyDomain",
+            "f_min": 20.0,
+            "f_max": 256.0,
+            "delta_f": 0.25,
+        },
+        "waveform_generator": {"approximant": "IMRPhenomPv2_NRTidal", "f_ref": 100.0},
+        "intrinsic_prior": {
+            "chirp_mass": (
+                "bilby.gw.prior.UniformInComponentsChirpMass(minimum=1.0, maximum=2.0)"
+            ),
+            "mass_ratio": (
+                "bilby.gw.prior.UniformInComponentsMassRatio(minimum=0.5, maximum=1.0)"
+            ),
+        },
+    },
+    "train_settings": {
+        "data": {
+            "detectors": ["H1"],
+            "ref_time": 1187008882.42,
+            "extrinsic_prior": {"ra": "default", "dec": "default"},
+            "inference_parameters": ["delta_chirp_mass", "mass_ratio"],
+            "context_parameters": ["ra", "dec", "chirp_mass_proxy"],
+            "gnpe_chirp": {
+                "kernel": {
+                    "chirp_mass": (
+                        "bilby.core.prior.Uniform(minimum=-0.005, maximum=0.005)"
+                    )
+                },
+                "order": 0,
+            },
+            "standardization": {
+                "mean": {
+                    "delta_chirp_mass": 0.0,
+                    "mass_ratio": 0.75,
+                    "ra": np.pi,
+                    "dec": 0.0,
+                    "chirp_mass_proxy": 1.5,
+                },
+                "std": {
+                    "delta_chirp_mass": 0.003,
+                    "mass_ratio": 0.15,
+                    "ra": 1.8,
+                    "dec": 0.6,
+                    "chirp_mass_proxy": 0.3,
+                },
+            },
+        }
+    },
+}
+
+
+class _StubBNSModel:
+    metadata = _BNS_METADATA
+    device = "cpu"
+
+
+def _event_data(n_bins=1025):
+    strain = np.ones(n_bins, dtype=complex)
+    return {"waveform": {"H1": strain}, "asds": {"H1": np.ones(n_bins)}}
+
+
+def test_heterodyne_inserted_before_decimation_with_fixed_proxy():
+    ctx = GWSamplerContext.from_model_metadata(
+        _BNS_METADATA, _event_data(), fixed_context_parameters=_PINS
+    )
+    first = ctx._data_prep.transforms[0]
+    assert isinstance(first, HeterodynePhase)
+    # Only proxy entries parameterize the data preparation; the pinned sky
+    # position is conditioning only.
+    assert first.fixed_parameters == {"chirp_mass": _PINS["chirp_mass_proxy"]}
+    assert first.order == 0 and first.inverse is False
+
+
+def test_missing_proxy_makes_prepared_data_fail_loudly():
+    # Without pins the context still constructs (the likelihood and prior views
+    # must work for from-file results), but the network-input view is undefined.
+    ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
+    assert ctx.prior is not None
+    with pytest.raises(ValueError, match="chirp-mass"):
+        ctx.prepared_data()
+
+
+def test_proxy_offset_step_selection():
+    steps = _proxy_offset_steps(
+        ["delta_chirp_mass", "mass_ratio"], ["ra", "dec", "chirp_mass_proxy"]
+    )
+    assert len(steps) == 1 and steps[0].parameters == ["chirp_mass"]
+    # No offset step without a matching pinned proxy.
+    assert _proxy_offset_steps(["delta_chirp_mass"], ["ra"]) == []
+
+
+def test_from_model_assembles_pinned_chain():
+    sampler = GWComposedSampler.from_model(
+        _StubBNSModel(), _event_data(), fixed_context_parameters=_PINS
+    )
+    kinds = [type(s).__name__ for s in sampler.composer.steps]
+    assert kinds == ["DeltaFactor", "FlowFactor", "ProxyOffsetReparam"]
+    delta, flow, offset = sampler.composer.steps
+    assert set(delta.parameters) == set(_PINS)
+    assert flow.conditioning == ["ra", "dec", "chirp_mass_proxy"]
+    assert offset.consumes == ["delta_chirp_mass"]
+
+
+def test_from_model_requires_all_pins():
+    with pytest.raises(ValueError, match="conditions on"):
+        GWComposedSampler.from_model(_StubBNSModel(), _event_data())
+    with pytest.raises(ValueError, match="conditions on"):
+        GWComposedSampler.from_model(
+            _StubBNSModel(),
+            _event_data(),
+            fixed_context_parameters={"chirp_mass_proxy": 1.2},
+        )

From a63a6839a0bdeedb4b3a10901681d5aa9b891ccc Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 17:22:11 +0100
Subject: [PATCH 45/65] Draw the pinned data preparation from the chain

The chirp-mass heterodyne now runs in parameters mode: prepared_data
receives the chain's conditioning from the conditioned factors, verifies
the required columns are constant, injects the values into the transform
chain under their physical names (extracted as Python floats, preserving
the float64 heterodyne path), and keys its cache on them -- a later call
with different values fails rather than silently serving stale data. The
pin values therefore have a single owner, the chain's root DeltaFactor,
consumed by both the network conditioning and the data preparation, so
the two cannot disagree. The context records only the conditioning names
its preparation is a function of.

The composer gains the point-mass rule: DeltaFactor declares
point_mass = True, and the requested sample count lands on the first
stage that carries multiplicity -- a root prefix of pins emits a single
row each, and the flow draws N in one conditioned call (one embedding),
with the prefix rows expanded to match. Chains with stochastic roots
(NPE, GNPE, sample tables) are unchanged -- verified bit-exact by the
pipe regressions -- and the BNS chain produces the same draws to
float32 kernel-order noise (measured ~1e-5).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py              |  79 +++++++++++-----
 dingo/gw/inference/factors.py      | 140 ++++++++++++++++++-----------
 tests/core/test_factors.py         |  33 +++++++
 tests/gw/test_bns_chain_builder.py |  57 +++++++++---
 4 files changed, 222 insertions(+), 87 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index d87625696..e70d68cda 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -78,9 +78,12 @@ class SamplerContext(Protocol):
     event_metadata: Optional[dict]
     device: Union[torch.device, str]
 
-    def prepared_data(self) -> torch.Tensor:
+    def prepared_data(self, conditioning=None) -> torch.Tensor:
         """The data representation the factors condition on
-        (whiten/decimate/repackage/...), computed once and cached."""
+        (whiten/decimate/repackage/...), computed once and cached. When the
+        preparation is a function of chain conditioning (e.g. a heterodyning
+        proxy), the calling factor passes its conditioning and the cache is
+        keyed on the consumed values."""
         ...
 
     def likelihood(self):
@@ -133,6 +136,15 @@ def chunk_and_concat(
     return samples, log_prob
 
 
+def _interleave_rows(samples, total, n):
+    """Repeat each carried chain row `n` times (row-major), keeping the summed
+    log-prob aligned."""
+    samples = {k: v.repeat_interleave(n, 0) for k, v in samples.items()}
+    if torch.is_tensor(total):
+        total = total.repeat_interleave(n, 0)
+    return samples, total
+
+
 def _describe_default(step) -> dict:
     """Default provenance descriptor for a chain step: class name, the parameter
     block it emits, and what it conditions on. Literal-only (round-trips through
@@ -306,7 +318,9 @@ def sample_and_log_prob(self, num_samples, context, given=None):
             z = z.squeeze(0)
             log_prob = log_prob.squeeze(0)
         else:
-            data = context.prepared_data()
+            # The conditioning may parameterize the data preparation itself (e.g.
+            # the chirp-mass heterodyne), so the context receives it.
+            data = context.prepared_data(conditioning=given)
             # Standardize the (physical) conditioning the network was trained on, giving
             # B = N context rows; expand the shared data to match. The network draws
             # num_samples per row -> (N, num_samples, dim).
@@ -339,14 +353,15 @@ def log_prob(self, theta_i, context, given=None):
         net_context: tuple[torch.Tensor, ...]
         if self.unconditional:
             net_context = ()
-        else:
+        elif not self.context_parameters:
             data = context.prepared_data()
             data = data.expand(num_samples, *data.shape)
-            if not self.context_parameters:
-                net_context = (data,)
-            else:
-                ctx = self.standardization.standardize(given, self.context_parameters)
-                net_context = (data, ctx)
+            net_context = (data,)
+        else:
+            data = context.prepared_data(conditioning=given)
+            data = data.expand(num_samples, *data.shape)
+            ctx = self.standardization.standardize(given, self.context_parameters)
+            net_context = (data, ctx)
         self.model.network.eval()
         with torch.no_grad():
             log_prob = self.model.log_prob(z, *net_context)
@@ -363,8 +378,15 @@ class DeltaFactor(Factor):
     Used as the chain root for prior-conditioning or known proxies (where later factors
     condition on the pinned values), and as a non-root filler for delta-prior parameters
     a model does not infer (one constant per current row).
+
+    As a point mass, n draws are one atom repeated: `point_mass = True` tells the
+    composer that a root prefix of such steps carries no multiplicity, so the
+    requested sample count lands on the first stage that does (e.g. the flow), which
+    then draws it in a single conditioned call.
     """
 
+    point_mass = True
+
     def __init__(self, values: dict[str, float]):
         self.values = values
         self.parameters = list(values)
@@ -711,30 +733,39 @@ def _run_chain_once(
         self, base: int, context: SamplerContext
     ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]:
         """One pass of the whole chain for `base` root samples. Returns the samples and
-        the summed proposal log-prob, or `None` if any step is density-free."""
+        the summed proposal log-prob, or `None` if any step is density-free.
+
+        The base count lands on the first stage that carries multiplicity: a root
+        prefix of point-mass steps (fixed pins) emits a single row each -- n draws
+        from a point mass are one atom repeated -- and the first non-point-mass
+        stage draws `base` (times its fan-out) in one conditioned call, with the
+        prefix's carried rows expanded to match."""
         samples: dict[str, torch.Tensor] = {}
         total: torch.Tensor | float = 0.0
         has_density = True
+        base_pending: Optional[int] = base
         for i, stage in enumerate(self.stages):
             step = stage.step
-            if i == 0:
-                n = base  # root: draw the base count
+            if base_pending is not None:
+                if getattr(step, "point_mass", False):
+                    n = 1  # a point mass carries no multiplicity
+                else:
+                    n = base_pending * (stage.fan_out if step.conditioning else 1)
+                    base_pending = None
             elif step.conditioning:
                 n = stage.fan_out  # fan_out samples per conditioning row
             else:
-                # Unconditioned non-root step (e.g. a fixed/delta factor): draw one value
-                # per current row -- it fills the batch rather than fanning out.
+                # Unconditioned non-root step (e.g. a fixed/delta filler): draw one
+                # value per current row -- it fills the batch rather than fanning out.
                 n = len(next(iter(samples.values())))
             given = {k: samples[k] for k in step.conditioning}
             block, lp = step.sample_and_log_prob(n, context, given)
-            # A conditioned fan-out (fan_out > 1) expands the batch: repeat each carried row
-            # to align with the step's fan_out sub-rows (the block is flattened row-major).
-            # 1:1 stages and unconditioned fillers leave the batch untouched.
-            if i > 0 and step.conditioning and stage.fan_out > 1:
-                fan = stage.fan_out
-                samples = {k: v.repeat_interleave(fan, 0) for k, v in samples.items()}
-                if torch.is_tensor(total):
-                    total = total.repeat_interleave(fan, 0)
+            # A conditioned multi-draw (fan-out, or the base count landing past a
+            # point-mass prefix) expands the batch: repeat each carried row to align
+            # with the step's sub-rows (the block is flattened row-major). 1:1 stages
+            # and unconditioned fillers leave the batch untouched.
+            if step.conditioning and samples and n > 1:
+                samples, total = _interleave_rows(samples, total, n)
             samples.update(block)
             # Drop any intermediates the step consumed: a reparametrization replaces its
             # inputs (in-place bijection); a target correction drops the columns it read.
@@ -745,6 +776,10 @@ def _run_chain_once(
                 has_density = False
             elif has_density:
                 total = total + lp
+        # Degenerate all-point-mass chain: no stage carried multiplicity, so expand
+        # the single emitted row to the requested count.
+        if base_pending is not None and base_pending > 1 and samples:
+            samples, total = _interleave_rows(samples, total, base_pending)
         return samples, (total if has_density else None)
 
     def log_prob(
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index d7c97f0e0..40d460f34 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -94,20 +94,6 @@ def normalize(value):
     return {"minimum_frequency": minimum, "maximum_frequency": maximum}
 
 
-class _RequireFixedProxies:
-    """Data-prep placeholder for a chirp-GNPE model built without fixed proxies:
-    the network-input view is undefined without the heterodyning proxy, so
-    `prepared_data()` fails loudly. The likelihood and prior views are unaffected
-    (a from-file result reconstructs its context without pins)."""
-
-    def __call__(self, sample):
-        raise ValueError(
-            "This model heterodynes its network input with a fixed chirp-mass "
-            "proxy; construct the context with fixed_context_parameters "
-            "(e.g. {'chirp_mass_proxy': ...}) to prepare network input."
-        )
-
-
 class GWSamplerContext:
     """
     Per-event shared GW state: the data `d` and everything derived from it. Referenced
@@ -144,6 +130,7 @@ def __init__(
         device: Union[torch.device, str] = "cpu",
         use_base_domain: bool = False,
         wfg_delta_f: Optional[float] = None,
+        data_prep_conditioning: Optional[list[str]] = None,
     ):
         """
         Parameters
@@ -173,6 +160,12 @@ def __init__(
         wfg_delta_f : float, optional
             Override for the waveform-generator domain `delta_f` (an updated
             `T = 1/delta_f` cannot be handled by domain projection). Set via `derive`.
+        data_prep_conditioning : list[str], optional
+            Names of the chain-conditioning parameters the data preparation is a
+            function of (e.g. `["chirp_mass_proxy"]` for a heterodyning model).
+            `prepared_data` requires their values, injects them into the
+            transform chain, and keys its cache on them; the values themselves
+            have a single owner -- the chain.
         """
         self.domain = domain
         self._data_prep = data_prep
@@ -182,6 +175,8 @@ def __init__(
         self.device = device
         self.use_base_domain = use_base_domain
         self.wfg_delta_f = wfg_delta_f
+        self.data_prep_conditioning = list(data_prep_conditioning or [])
+        self._prepared_key: Optional[dict] = None
         self._prepared: Optional[torch.Tensor] = None
         self._prior: Optional[PriorDict] = None
         self._likelihood: Optional[StationaryGaussianGWLikelihood] = None
@@ -204,7 +199,6 @@ def from_model_metadata(
         event_data: dict,
         event_metadata: Optional[dict] = None,
         device: Union[torch.device, str] = "cpu",
-        fixed_context_parameters: Optional[dict] = None,
     ) -> "GWSamplerContext":
         """Build the context from a model-metadata dict -- no model required. The
         domain, one-time data-prep chain, prior, and likelihood are all defined by
@@ -222,11 +216,6 @@ def from_model_metadata(
             Per-event metadata.
         device : torch.device or str, default "cpu"
             Device for `prepared_data()` and chain-created tensors.
-        fixed_context_parameters : dict, optional
-            Pinned network-conditioning values (e.g.
-            `{"chirp_mass_proxy": 1.1975, "ra": ..., "dec": ...}`). Proxy entries
-            that parameterize the data preparation (the chirp-mass heterodyne)
-            are consumed here; the chain conditions the network on all of them.
 
         Returns
         -------
@@ -240,28 +229,24 @@ def from_model_metadata(
         detectors = data_settings["detectors"]
 
         transforms = []
-        # Chirp-mass GNPE (BNS): heterodyne the raw strain once with the fixed
-        # proxy -- before decimation (they do not commute) and on the base
-        # domain. Iterated chirp GNPE (heterodyning inside a Gibbs loop) is not
+        # Chirp-mass GNPE (BNS): heterodyne the raw strain -- before decimation
+        # (they do not commute) and on the base domain. The transform draws the
+        # chirp mass from the sample's "parameters", which `prepared_data`
+        # injects from the chain's conditioning: the proxy value has a single
+        # owner (the chain's DeltaFactor), and the preparation is a function of
+        # it. Iterated chirp GNPE (heterodyning inside a Gibbs loop) is not
         # implemented: it would require carrying the undecimated strain per
         # sample.
         gnpe_chirp = data_settings.get("gnpe_chirp")
+        data_prep_conditioning = []
         if gnpe_chirp is not None:
-            fixed = {
-                k[: -len("_proxy")]: v
-                for k, v in (fixed_context_parameters or {}).items()
-                if k.endswith("_proxy")
-            }
-            if "chirp_mass" in fixed:
-                transforms.append(
-                    HeterodynePhase(
-                        domain=getattr(domain, "base_domain", domain),
-                        order=gnpe_chirp.get("order", 0),
-                        fixed_parameters=fixed,
-                    )
+            data_prep_conditioning = [k + "_proxy" for k in gnpe_chirp["kernel"]]
+            transforms.append(
+                HeterodynePhase(
+                    domain=getattr(domain, "base_domain", domain),
+                    order=gnpe_chirp.get("order", 0),
                 )
-            else:
-                transforms.append(_RequireFixedProxies())
+            )
         # Decimate from the base domain when using a multibanded frequency domain.
         if isinstance(domain, MultibandedFrequencyDomain):
             transforms.append(
@@ -295,6 +280,7 @@ def from_model_metadata(
             event_metadata=event_metadata,
             model_metadata=metadata,
             device=device,
+            data_prep_conditioning=data_prep_conditioning,
         )
 
     @classmethod
@@ -303,7 +289,6 @@ def from_model(
         model: BasePosteriorModel,
         event_data: dict,
         event_metadata: Optional[dict] = None,
-        fixed_context_parameters: Optional[dict] = None,
     ) -> "GWSamplerContext":
         """Build the context from a model: its own metadata and its device. Data
         preparation is network-bound, so the settings come from `model.metadata`
@@ -337,7 +322,6 @@ def from_model(
             event_data,
             event_metadata,
             device=model.device,
-            fixed_context_parameters=fixed_context_parameters,
         )
 
     def derive(
@@ -415,15 +399,67 @@ def derive(
                 self.use_base_domain if use_base_domain is None else use_base_domain
             ),
             wfg_delta_f=wfg_delta_f,
+            data_prep_conditioning=self.data_prep_conditioning,
         )
 
-    def prepared_data(self) -> torch.Tensor:
+    def prepared_data(self, conditioning=None) -> torch.Tensor:
         """One-time data preprocessing, computed once and cached. An event
         frequency-range update is validated against the training crop license
-        before the first preparation."""
-        if self._prepared is None:
-            self._validate_frequency_range()
-            self._prepared = self._data_prep(self.event_data)
+        before the first preparation.
+
+        When the preparation is a function of chain conditioning
+        (`data_prep_conditioning`, e.g. the chirp-mass heterodyne), the caller
+        supplies the conditioning columns; the required values must be constant
+        across rows (a context serves one representation), are injected into the
+        transform chain under their physical names, and key the cache -- a later
+        call with different values fails rather than silently serving stale
+        data.
+
+        Parameters
+        ----------
+        conditioning : dict[str, torch.Tensor], optional
+            The chain conditioning available to the calling factor. Ignored
+            when the preparation needs none of it.
+        """
+        if not self.data_prep_conditioning:
+            if self._prepared is None:
+                self._validate_frequency_range()
+                self._prepared = self._data_prep(self.event_data)
+            return self._prepared
+
+        conditioning = conditioning or {}
+        key = {}
+        for name in self.data_prep_conditioning:
+            if name not in conditioning:
+                raise ValueError(
+                    f"This model's data preparation is a function of the chain "
+                    f"conditioning `{name}`, which the caller does not provide "
+                    f"(pass it to prepared_data, e.g. from the chain's pins)."
+                )
+            column = torch.as_tensor(conditioning[name]).reshape(-1)
+            if not torch.all(column == column[0]):
+                raise ValueError(
+                    f"Data preparation requires a constant `{name}`, but the "
+                    f"conditioning varies across rows. A context serves one "
+                    f"representation; per-value batching (e.g. a chirp-mass "
+                    f"sweep) needs one context per value."
+                )
+            # Extract as a Python float: the transform chain computes the
+            # heterodyne phase in float64, and a float32 chain column must not
+            # degrade it.
+            key[name[: -len("_proxy")]] = float(column[0])
+
+        if self._prepared is not None:
+            if key != self._prepared_key:
+                raise ValueError(
+                    f"Data already prepared with {self._prepared_key}, but the "
+                    f"chain now conditions on {key}. A context serves one "
+                    f"representation; build a new context for a different value."
+                )
+            return self._prepared
+        self._validate_frequency_range()
+        self._prepared = self._data_prep({**self.event_data, "parameters": key})
+        self._prepared_key = key
         return self._prepared
 
     def _validate_frequency_range(self):
@@ -992,7 +1028,7 @@ def sample_and_log_prob(self, num_samples, context, given=None):
         proxies = {p: given[p] for p in self.proxy_parameters}
         n_rows = next(iter(proxies.values())).shape[0]
         x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = context.prepared_data().clone()
+        d = context.prepared_data(conditioning=given).clone()
         x["data"] = d.expand(n_rows, *d.shape)
         x = self.transform_pre(x)
         self.model.network.eval()
@@ -1026,7 +1062,7 @@ def log_prob(self, theta_i, context, given=None):
         proxies = {p: given[p] for p in self.proxy_parameters}
         n_rows = next(iter(proxies.values())).shape[0]
         x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = context.prepared_data().clone()
+        d = context.prepared_data(conditioning=given).clone()
         x["data"] = d.expand(n_rows, *d.shape)
         x = self.transform_pre(x)
         theta_net = {
@@ -1434,8 +1470,9 @@ def from_model(
         the flow conditions on them, and each inferred offset `delta_X` with a
         pinned proxy is reconstructed by a `ProxyOffsetReparam`
         (`X = delta_X + X_proxy`). Proxies that parameterize the data
-        preparation (the chirp-mass heterodyne) are applied once inside the
-        context.
+        preparation (the chirp-mass heterodyne) reach it through the chain:
+        the conditioned flow passes them to `prepared_data`, which consumes
+        and caches on them -- the pins have a single owner, the chain's root.
 
         Parameters
         ----------
@@ -1453,12 +1490,7 @@ def from_model(
         -------
         GWComposedSampler
         """
-        context = GWSamplerContext.from_model(
-            model,
-            event_data,
-            event_metadata,
-            fixed_context_parameters=fixed_context_parameters,
-        )
+        context = GWSamplerContext.from_model(model, event_data, event_metadata)
         metadata = _base_model_metadata(model)
         data_settings = metadata["train_settings"]["data"]
         inference_parameters = data_settings["inference_parameters"]
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 2c0dfedec..0442f2167 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -658,3 +658,36 @@ def test_unconditional_flow_factor():
     comp = ChainComposer([factor, _ConstFactor("y", conditioning=["H1_time_proxy"])])
     out = comp.sample(3, context=None)
     assert out["y"].shape == (3,)
+
+
+def test_point_mass_prefix_passes_base_to_first_stochastic_stage():
+    # A root prefix of point-mass steps (fixed pins) carries no multiplicity:
+    # it emits a single row, and the base count lands on the first stage that
+    # does carry it, as one conditioned call with num_samples = base. The
+    # prefix's carried rows are then expanded to match.
+    calls = []
+
+    class _RecordingFactor(Factor):
+        parameters = ["x"]
+        conditioning = ["a"]
+
+        def sample_and_log_prob(self, num_samples, context, given=None):
+            n_rows = len(next(iter(given.values())))
+            calls.append((num_samples, n_rows))
+            n = n_rows * num_samples
+            return {"x": torch.arange(n, dtype=torch.float32)}, torch.zeros(n)
+
+        def log_prob(self, theta_i, context, given=None):
+            return torch.zeros_like(next(iter(theta_i.values())))
+
+    chain = ChainComposer([DeltaFactor({"a": 2.0}), _RecordingFactor()])
+    samples, log_prob = chain.sample_and_log_prob(5, context=None)
+    assert calls == [(5, 1)]
+    assert samples["a"].shape == (5,) and torch.all(samples["a"] == 2.0)
+    assert samples["x"].shape == (5,) and log_prob.shape == (5,)
+
+
+def test_all_point_mass_chain_expands_to_base():
+    chain = ChainComposer([DeltaFactor({"a": 2.0, "b": -1.0})])
+    samples, log_prob = chain.sample_and_log_prob(4, context=None)
+    assert samples["a"].shape == (4,) and log_prob.shape == (4,)
diff --git a/tests/gw/test_bns_chain_builder.py b/tests/gw/test_bns_chain_builder.py
index d4093f733..0702b9add 100644
--- a/tests/gw/test_bns_chain_builder.py
+++ b/tests/gw/test_bns_chain_builder.py
@@ -80,27 +80,62 @@ def _event_data(n_bins=1025):
     return {"waveform": {"H1": strain}, "asds": {"H1": np.ones(n_bins)}}
 
 
-def test_heterodyne_inserted_before_decimation_with_fixed_proxy():
-    ctx = GWSamplerContext.from_model_metadata(
-        _BNS_METADATA, _event_data(), fixed_context_parameters=_PINS
-    )
+def test_heterodyne_inserted_before_decimation_draws_from_conditioning():
+    ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
     first = ctx._data_prep.transforms[0]
     assert isinstance(first, HeterodynePhase)
-    # Only proxy entries parameterize the data preparation; the pinned sky
-    # position is conditioning only.
-    assert first.fixed_parameters == {"chirp_mass": _PINS["chirp_mass_proxy"]}
+    # Parameters mode: the transform draws the chirp mass from the sample dict;
+    # no value is baked into the context. The context records only the names its
+    # preparation is a function of.
+    assert first.fixed_parameters is None
     assert first.order == 0 and first.inverse is False
+    assert ctx.data_prep_conditioning == ["chirp_mass_proxy"]
 
 
-def test_missing_proxy_makes_prepared_data_fail_loudly():
-    # Without pins the context still constructs (the likelihood and prior views
-    # must work for from-file results), but the network-input view is undefined.
+def test_prepared_data_requires_conditioning():
+    # Without conditioning the context still constructs (the likelihood and
+    # prior views must work for from-file results), but the network-input view
+    # is a function of the chain conditioning and fails loudly without it.
     ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
     assert ctx.prior is not None
-    with pytest.raises(ValueError, match="chirp-mass"):
+    with pytest.raises(ValueError, match="chirp_mass_proxy"):
         ctx.prepared_data()
 
 
+def _conditioning(value=1.1975, n=4):
+    import torch
+
+    return {
+        "chirp_mass_proxy": torch.full((n,), value),
+        "ra": torch.full((n,), _PINS["ra"]),
+        "dec": torch.full((n,), _PINS["dec"]),
+    }
+
+
+def test_prepared_data_consumes_conditioning_and_keys_the_cache():
+    import torch
+
+    ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
+    prepared = ctx.prepared_data(conditioning=_conditioning())
+    assert torch.is_tensor(prepared)
+    # Same values: the cached tensor is served.
+    assert ctx.prepared_data(conditioning=_conditioning()) is prepared
+    # A different pin is a different representation: fail rather than serve
+    # stale data.
+    with pytest.raises(ValueError, match="one\\s+representation"):
+        ctx.prepared_data(conditioning=_conditioning(value=1.30))
+
+
+def test_prepared_data_rejects_varying_conditioning():
+    import torch
+
+    ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
+    conditioning = _conditioning()
+    conditioning["chirp_mass_proxy"] = torch.tensor([1.19, 1.20, 1.21, 1.22])
+    with pytest.raises(ValueError, match="constant"):
+        ctx.prepared_data(conditioning=conditioning)
+
+
 def test_proxy_offset_step_selection():
     steps = _proxy_offset_steps(
         ["delta_chirp_mass", "mass_ratio"], ["ra", "dec", "chirp_mass_proxy"]

From ccbfdfdeaa04a93d35a780c02ab2df9261d1c81d Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 7 Jul 2026 21:47:23 +0100
Subject: [PATCH 46/65] Wire fixed context parameters into dingo_pipe

Add the --fixed-context-parameters option (the demo-INI dict form) and
pass it through SamplingInput to the composed single-network builder, so
context-conditioned models (single-step chirp-mass GNPE, e.g. binary
neutron stars) run through the pipe. Combining it with model-init raises.
Verified end-to-end on real GW170817 data (demo INI adapted to the exact
base-domain likelihood; efficiency 11.4%, pins and full chain provenance
in the result) and bit-exact on the no-pins NPE pipe regression.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/pipe/parser.py                        | 11 +++++++
 dingo/pipe/sampling.py                      | 15 ++++++++-
 tests/pipe/test_fixed_context_parameters.py | 34 +++++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 tests/pipe/test_fixed_context_parameters.py

diff --git a/dingo/pipe/parser.py b/dingo/pipe/parser.py
index 86691bf17..59269162f 100644
--- a/dingo/pipe/parser.py
+++ b/dingo/pipe/parser.py
@@ -1549,6 +1549,17 @@ def create_parser(top_level=True, usage=None):
         help="For GNPE models, whether to recover the log probability by training an "
         "unconditional flow for the GNPE proxies.",
     )
+    sampler_parser.add(
+        "--fixed-context-parameters",
+        type=nonestr,
+        default=None,
+        help="Dictionary of fixed values for the model's context parameters, e.g. "
+        "{chirp_mass_proxy: 1.19786, ra: 3.44616, dec: -0.408084}. Required for "
+        "single-network models conditioned on context parameters (single-step "
+        "GNPE, e.g. binary neutron stars): proxy entries parameterize the data "
+        "preparation (chirp-mass heterodyning), and all pinned values are "
+        "recorded with the samples.",
+    )
     sampler_parser.add(
         "--device",
         type=str,
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index 2df334f94..ddde8cdf7 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -57,6 +57,11 @@ def __init__(self, args, unknown_args):
         self.recover_log_prob = args.recover_log_prob
         self.device = args.device
         self.num_gnpe_iterations = args.num_gnpe_iterations
+        self.fixed_context_parameters = (
+            convert_string_to_dict(args.fixed_context_parameters)
+            if args.fixed_context_parameters is not None
+            else None
+        )
         self.num_samples = args.num_samples
         self.batch_size = args.batch_size
         self.density_recovery_settings = args.density_recovery_settings
@@ -128,6 +133,11 @@ def _load_sampler(self):
         )
 
         if self.model_init is not None:
+            if self.fixed_context_parameters is not None:
+                raise ValueError(
+                    "fixed-context-parameters pins the context of a single-network "
+                    "model; it cannot be combined with model-init (iterative GNPE)."
+                )
             self.gnpe = True
             self._main_model = model
             init_model = build_model_from_kwargs(
@@ -145,7 +155,10 @@ def _load_sampler(self):
         else:
             self.gnpe = False
             self.dingo_sampler = GWComposedSampler.from_model(
-                model, self.context, event_metadata=self.event_metadata
+                model,
+                self.context,
+                event_metadata=self.event_metadata,
+                fixed_context_parameters=self.fixed_context_parameters,
             )
         self.dingo_sampler.provenance_extra["models"] = self._model_paths()
 
diff --git a/tests/pipe/test_fixed_context_parameters.py b/tests/pipe/test_fixed_context_parameters.py
new file mode 100644
index 000000000..a378c9616
--- /dev/null
+++ b/tests/pipe/test_fixed_context_parameters.py
@@ -0,0 +1,34 @@
+"""Tests for the fixed-context-parameters option in the dingo_pipe parser."""
+
+from bilby_pipe.utils import convert_string_to_dict, parse_args
+
+from dingo.pipe.parser import create_parser
+
+DEMO_FORM = "{chirp_mass_proxy: 1.19786, ra: 3.44616, dec: -0.408084}"
+DEMO_DICT = {"chirp_mass_proxy": 1.19786, "ra": 3.44616, "dec": -0.408084}
+
+
+def _parse(tmp_path, ini_body="model = model.pt\n", extra=()):
+    ini = tmp_path / "test.ini"
+    ini.write_text(ini_body)
+    args, _ = parse_args([str(ini), *extra], create_parser(top_level=False))
+    return args
+
+
+def test_default_is_none(tmp_path):
+    assert _parse(tmp_path).fixed_context_parameters is None
+
+
+def test_cli_dict_form_parses(tmp_path):
+    args = _parse(tmp_path, extra=["--fixed-context-parameters", DEMO_FORM])
+    assert convert_string_to_dict(args.fixed_context_parameters) == DEMO_DICT
+
+
+def test_ini_dict_form_parses(tmp_path):
+    # The published binary-neutron-star-demo INI form, read from the INI file
+    # and converted the way SamplingInput does.
+    args = _parse(
+        tmp_path,
+        ini_body=f"model = model.pt\nfixed-context-parameters = {DEMO_FORM}\n",
+    )
+    assert convert_string_to_dict(args.fixed_context_parameters) == DEMO_DICT

From f17f74b4b93f8bbc5b5da03b8da32a83932f03e6 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 12 Jul 2026 19:58:04 +0100
Subject: [PATCH 47/65] Add the chirp-mass scan with row-aligned data
 preparation

Sweep the chirp-mass proxy over the training prior on the ordinary chain
machinery -- a SampleTableFactor grid root and a conditioned FlowFactor
drawing per grid row -- and take the maximum-likelihood draw as the
trigger value (Dax et al., Nature 639, 49 (2025)). The scan runs the
grid in blocks and selects with a phase-marginalized likelihood that now
optionally reports the matched-filter SNR (return_aux_snr).

To support this, prepared_data gains a uniform contract: with
conditioning, the result is row-aligned -- one data row per conditioning
row -- with a pinned value prepared once behind a value-keyed memo and
varying values prepared in a single batch-native pass. Columns the
preparation does not consume condition the network alone.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py              |  34 ++---
 dingo/gw/inference/factors.py      | 115 +++++++++-------
 dingo/gw/inference/scan.py         | 203 +++++++++++++++++++++++++++++
 dingo/gw/likelihood.py             |  36 ++++-
 tests/gw/test_bns_chain_builder.py |  90 +++++++++++--
 5 files changed, 401 insertions(+), 77 deletions(-)
 create mode 100644 dingo/gw/inference/scan.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index e70d68cda..056e795df 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -80,10 +80,13 @@ class SamplerContext(Protocol):
 
     def prepared_data(self, conditioning=None) -> torch.Tensor:
         """The data representation the factors condition on
-        (whiten/decimate/repackage/...), computed once and cached. When the
-        preparation is a function of chain conditioning (e.g. a heterodyning
-        proxy), the calling factor passes its conditioning and the cache is
-        keyed on the consumed values."""
+        (whiten/decimate/repackage/...). Without `conditioning`: the single
+        shared representation, computed once and cached. With `conditioning`
+        (the chain columns available to a conditioned factor): row-aligned,
+        one data row per conditioning row. The context consumes only the
+        columns its preparation depends on (e.g. a heterodyning proxy);
+        unconsumed columns condition the network alone, and the shared
+        representation is viewed across their rows."""
         ...
 
     def likelihood(self):
@@ -136,6 +139,11 @@ def chunk_and_concat(
     return samples, log_prob
 
 
+def _n_rows(block: dict) -> int:
+    """Row count of a column block (all columns share one length)."""
+    return len(next(iter(block.values())))
+
+
 def _interleave_rows(samples, total, n):
     """Repeat each carried chain row `n` times (row-major), keeping the summed
     log-prob aligned."""
@@ -318,19 +326,16 @@ def sample_and_log_prob(self, num_samples, context, given=None):
             z = z.squeeze(0)
             log_prob = log_prob.squeeze(0)
         else:
-            # The conditioning may parameterize the data preparation itself (e.g.
-            # the chirp-mass heterodyne), so the context receives it.
-            data = context.prepared_data(conditioning=given)
-            # Standardize the (physical) conditioning the network was trained on, giving
-            # B = N context rows; expand the shared data to match. The network draws
-            # num_samples per row -> (N, num_samples, dim).
+            # Standardize the (physical) conditioning the network was trained on,
+            # giving B = N context rows; the context returns data row-aligned to
+            # them. The network draws num_samples per row -> (N, num_samples, dim).
             # TODO: the embedding runs once per row, so row-identical data with
             # row-varying conditioning (the intrinsic/extrinsic split) recomputes N
             # identical data embeddings. Fixing this needs an embed/fuse split on the
             # model (FlowWrapper) so the cached embedding can be reused across rows.
             ctx = self.standardization.standardize(given, self.context_parameters)
             n_rows = ctx.shape[0]
-            data = data.expand(n_rows, *data.shape)
+            data = context.prepared_data(conditioning=given)
             with torch.no_grad():
                 z, log_prob = self.model.sample_and_log_prob(
                     data, ctx, num_samples=num_samples
@@ -348,7 +353,7 @@ def log_prob(self, theta_i, context, given=None):
         theta_net = {
             net: theta_i[self.aliases.get(net, net)] for net in self._net_parameters
         }
-        num_samples = next(iter(theta_net.values())).shape[0]
+        num_samples = _n_rows(theta_net)
         z = self.standardization.standardize(theta_net, self._net_parameters)
         net_context: tuple[torch.Tensor, ...]
         if self.unconditional:
@@ -359,7 +364,6 @@ def log_prob(self, theta_i, context, given=None):
             net_context = (data,)
         else:
             data = context.prepared_data(conditioning=given)
-            data = data.expand(num_samples, *data.shape)
             ctx = self.standardization.standardize(given, self.context_parameters)
             net_context = (data, ctx)
         self.model.network.eval()
@@ -452,7 +456,7 @@ def __init__(self, table: dict, log_prob=None):
     def sample_and_log_prob(self, num_samples, context, given=None):
         """Emit the table; `num_samples` must equal the table length (a fixed table
         cannot be chunked, so run the chain with `batch_size=None`)."""
-        n = len(next(iter(self.table.values())))
+        n = _n_rows(self.table)
         if num_samples != n:
             raise ValueError(
                 f"A sample table is fixed: num_samples must equal the table length "
@@ -757,7 +761,7 @@ def _run_chain_once(
             else:
                 # Unconditioned non-root step (e.g. a fixed/delta filler): draw one
                 # value per current row -- it fills the batch rather than fanning out.
-                n = len(next(iter(samples.values())))
+                n = _n_rows(samples)
             given = {k: samples[k] for k in step.conditioning}
             block, lp = step.sample_and_log_prob(n, context, given)
             # A conditioned multi-draw (fan-out, or the base count landing past a
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
index 40d460f34..78ee4dcd7 100644
--- a/dingo/gw/inference/factors.py
+++ b/dingo/gw/inference/factors.py
@@ -42,6 +42,7 @@
     Standardization,
     TargetCorrection,
     _base_model_metadata,
+    _n_rows,
 )
 from dingo.core.multiprocessing import apply_func_with_multiprocessing
 from dingo.core.posterior_models import BasePosteriorModel
@@ -403,32 +404,60 @@ def derive(
         )
 
     def prepared_data(self, conditioning=None) -> torch.Tensor:
-        """One-time data preprocessing, computed once and cached. An event
-        frequency-range update is validated against the training crop license
-        before the first preparation.
-
-        When the preparation is a function of chain conditioning
-        (`data_prep_conditioning`, e.g. the chirp-mass heterodyne), the caller
-        supplies the conditioning columns; the required values must be constant
-        across rows (a context serves one representation), are injected into the
-        transform chain under their physical names, and key the cache -- a later
-        call with different values fails rather than silently serving stale
-        data.
+        """The network-input data representation of this event.
+
+        Without `conditioning`: the single shared representation, computed once
+        and cached. With `conditioning` (the chain columns available to a
+        conditioned factor): row-aligned, one data row per conditioning row.
+        Only the columns named in `data_prep_conditioning` are consumed by the
+        preparation (e.g. the chirp-mass heterodyne proxy, injected under its
+        physical name); the remaining columns condition the network only, and a
+        context that consumes nothing serves the shared representation viewed
+        across the rows. A constant consumed value (a pinned proxy) is prepared
+        once and viewed across the rows; varying values (a sweep) run through
+        the batch-native transform chain in one vectorized pass, uncached -- a
+        caller sweeping more rows than memory allows blocks its own request.
+
+        An event frequency-range update is validated against the training crop
+        license before any preparation.
 
         Parameters
         ----------
         conditioning : dict[str, torch.Tensor], optional
-            The chain conditioning available to the calling factor. Ignored
-            when the preparation needs none of it.
+            The chain conditioning available to the calling factor, one value
+            per row. May contain columns irrelevant to the preparation.
         """
         if not self.data_prep_conditioning:
             if self._prepared is None:
                 self._validate_frequency_range()
                 self._prepared = self._data_prep(self.event_data)
-            return self._prepared
+            if conditioning is None:
+                return self._prepared
+            return self._prepared.expand(_n_rows(conditioning), *self._prepared.shape)
 
+        columns = self._conditioning_columns(conditioning)
+        self._validate_frequency_range()
+        n_rows = _n_rows(columns)
+        if all(torch.all(c == c[0]) for c in columns.values()):
+            # N rows of one pinned value: prepare once, view it across the rows.
+            key = {name: float(c[0]) for name, c in columns.items()}
+            if key != self._prepared_key:
+                self._prepared = self._data_prep({**self.event_data, "parameters": key})
+                self._prepared_key = key
+            return self._prepared.expand(n_rows, *self._prepared.shape)
+        parameters = {name: column.numpy() for name, column in columns.items()}
+        return self._data_prep(
+            {**self._broadcast_event(n_rows), "parameters": parameters}
+        )
+
+    def _conditioning_columns(self, conditioning) -> dict[str, torch.Tensor]:
+        """Collect the conditioning columns the preparation consumes, keyed by
+        their physical names (the `_proxy` suffix names the chain column; the
+        transform chain reads the physical parameter), as float64 (the
+        heterodyne phase is computed in float64, and a float32 chain column
+        must not degrade it)."""
         conditioning = conditioning or {}
-        key = {}
+        columns = {}
         for name in self.data_prep_conditioning:
             if name not in conditioning:
                 raise ValueError(
@@ -436,31 +465,23 @@ def prepared_data(self, conditioning=None) -> torch.Tensor:
                     f"conditioning `{name}`, which the caller does not provide "
                     f"(pass it to prepared_data, e.g. from the chain's pins)."
                 )
-            column = torch.as_tensor(conditioning[name]).reshape(-1)
-            if not torch.all(column == column[0]):
-                raise ValueError(
-                    f"Data preparation requires a constant `{name}`, but the "
-                    f"conditioning varies across rows. A context serves one "
-                    f"representation; per-value batching (e.g. a chirp-mass "
-                    f"sweep) needs one context per value."
-                )
-            # Extract as a Python float: the transform chain computes the
-            # heterodyne phase in float64, and a float32 chain column must not
-            # degrade it.
-            key[name[: -len("_proxy")]] = float(column[0])
-
-        if self._prepared is not None:
-            if key != self._prepared_key:
-                raise ValueError(
-                    f"Data already prepared with {self._prepared_key}, but the "
-                    f"chain now conditions on {key}. A context serves one "
-                    f"representation; build a new context for a different value."
-                )
-            return self._prepared
-        self._validate_frequency_range()
-        self._prepared = self._data_prep({**self.event_data, "parameters": key})
-        self._prepared_key = key
-        return self._prepared
+            columns[name[: -len("_proxy")]] = torch.as_tensor(
+                conditioning[name], dtype=torch.float64
+            ).reshape(-1)
+        return columns
+
+    def _broadcast_event(self, n_rows: int) -> dict:
+        """The event arrays broadcast (as read-only views) across `n_rows`
+        rows, forming the batched sample dict for a single transform-chain
+        pass."""
+        return {
+            part: (
+                {k: np.broadcast_to(v, (n_rows, *np.shape(v))) for k, v in data.items()}
+                if isinstance(data, dict)
+                else data
+            )
+            for part, data in self.event_data.items()
+        }
 
     def _validate_frequency_range(self):
         """Validate an event frequency-range update: hard bounds against the (base)
@@ -1026,9 +1047,9 @@ def sample_and_log_prob(self, num_samples, context, given=None):
         if num_samples != 1:
             raise ValueError("GNPE is 1:1; draw one sample per proxy (fan_out=1).")
         proxies = {p: given[p] for p in self.proxy_parameters}
-        n_rows = next(iter(proxies.values())).shape[0]
+        n_rows = _n_rows(proxies)
         x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = context.prepared_data(conditioning=given).clone()
+        d = context.prepared_data().clone()
         x["data"] = d.expand(n_rows, *d.shape)
         x = self.transform_pre(x)
         self.model.network.eval()
@@ -1056,13 +1077,13 @@ def sample_and_log_prob(self, num_samples, context, given=None):
     def log_prob(self, theta_i, context, given=None):
         """Evaluate the network density `log q(theta | theta_hat, d)` in physical space
         at given `theta_i` (exposed / aliased names), one row per proxy row in `given`.
-        Runs the same proxies-present data preparation as sampling (time-shift by the
-        proxies, standardize the conditioning), then scores the standardized parameters
-        under the network."""
+        Applies the same per-row view of the data as sampling (the shared
+        representation time-shifted by the proxies, the conditioning standardized),
+        then scores the standardized parameters under the network."""
         proxies = {p: given[p] for p in self.proxy_parameters}
-        n_rows = next(iter(proxies.values())).shape[0]
+        n_rows = _n_rows(proxies)
         x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = context.prepared_data(conditioning=given).clone()
+        d = context.prepared_data().clone()
         x["data"] = d.expand(n_rows, *d.shape)
         x = self.transform_pre(x)
         theta_net = {
diff --git a/dingo/gw/inference/scan.py b/dingo/gw/inference/scan.py
new file mode 100644
index 000000000..2f2377b99
--- /dev/null
+++ b/dingo/gw/inference/scan.py
@@ -0,0 +1,203 @@
+"""Chirp-mass scan for prior-conditioned (DINGO-BNS) networks.
+
+A chirp-mass-conditioned network infers parameters relative to a proxy value
+that must be supplied per event. When no external trigger value is available,
+this scan determines it from the data: sweep the proxy over the training
+chirp-mass prior, draw a few samples per grid point in a single network pass,
+and take the chirp mass of the maximum-likelihood draw (Dax et al., Nature 639,
+49 (2025), Methods).
+
+The sweep runs on the ordinary chain machinery: a `SampleTableFactor` roots the
+chain in the pin table (one row per grid point), the conditioned `FlowFactor`
+draws `num_samples` per row against per-row prepared data (the row-varying
+heterodyne), and the usual reconstruction steps follow. The winner is selected
+with a phase-marginalized likelihood on the (base-domain) event data.
+"""
+
+import math
+
+import numpy as np
+import pandas as pd
+from bilby.gw.prior import BBHPriorDict
+
+from dingo.core.factors import ChainComposer, FlowFactor, SampleTableFactor, Stage
+from dingo.gw.inference.factors import (
+    GWSamplerContext,
+    _base_model_metadata,
+    _delta_prior_steps,
+    _proxy_offset_steps,
+    _ra_aliases,
+    _ra_reparam_steps,
+)
+
+
+def chirp_mass_scan_grid(model_metadata: dict, overlap_factor: int = 2) -> np.ndarray:
+    """Proxy grid spanning the training chirp-mass prior at `kernel_width /
+    overlap_factor` spacing, with the bounds inset by the kernel edges so that
+    every chirp mass in the prior is within the kernel of some grid point.
+
+    Parameters
+    ----------
+    model_metadata : dict
+        The model metadata (with `dataset_settings` and `train_settings`).
+    overlap_factor : int, default 2
+        Grid points per kernel width.
+
+    Returns
+    -------
+    np.ndarray
+        The proxy grid (float64).
+    """
+    prior = BBHPriorDict(dict(model_metadata["dataset_settings"]["intrinsic_prior"]))[
+        "chirp_mass"
+    ]
+    kernel = BBHPriorDict(
+        dict(model_metadata["train_settings"]["data"]["gnpe_chirp"]["kernel"])
+    )["chirp_mass"]
+    num_points = math.ceil(
+        (prior.maximum - prior.minimum)
+        / (kernel.maximum - kernel.minimum)
+        * overlap_factor
+    )
+    return np.linspace(
+        prior.minimum - kernel.minimum, prior.maximum - kernel.maximum, num_points
+    )
+
+
+def chirp_mass_scan(
+    model,
+    event_data: dict,
+    event_metadata: dict = None,
+    fixed_context_parameters: dict = None,
+    *,
+    num_samples: int = 10,
+    overlap_factor: int = 2,
+    block_size: int = 32,
+    num_processes: int = 1,
+) -> dict:
+    """Scan for the trigger chirp mass of an event.
+
+    Draws `num_samples` posterior samples at each grid value of the chirp-mass
+    proxy (one batched network pass over the whole grid), evaluates a
+    phase-marginalized likelihood for every within-prior draw, and returns the
+    maximum-likelihood draw's chirp mass as the trigger value.
+
+    Parameters
+    ----------
+    model : BasePosteriorModel
+        A chirp-mass-conditioned model (trained with `gnpe_chirp`).
+    event_data : dict
+        The raw event data (strain + ASDs).
+    event_metadata : dict, optional
+        Per-event metadata.
+    fixed_context_parameters : dict, optional
+        Fixed values for the model's remaining (non-proxy) context parameters,
+        e.g. the sky position for a sky-conditioned network. The scan fills
+        `chirp_mass_proxy` itself.
+    num_samples : int, default 10
+        Draws per grid point.
+    overlap_factor : int, default 2
+        Grid points per kernel width.
+    block_size : int, default 32
+        Grid points per sweep block; bounds the transient memory of the
+        row-wise data preparation and the network batch.
+    num_processes : int, default 1
+        Parallel processes for the likelihood evaluations.
+
+    Returns
+    -------
+    dict
+        `chirp_mass_trigger` (the winner), `snr` and `max_log_likelihood` (its
+        trigger quality), `grid`, `samples` (all within-prior draws with their
+        `log_likelihood` and `snr`), and `settings` (the scan configuration).
+    """
+    metadata = _base_model_metadata(model)
+    data_settings = metadata["train_settings"]["data"]
+    inference_parameters = data_settings["inference_parameters"]
+    context_parameters = data_settings.get("context_parameters") or []
+    if (
+        "chirp_mass_proxy" not in context_parameters
+        or "gnpe_chirp" not in data_settings
+    ):
+        raise ValueError(
+            "The chirp-mass scan requires a model conditioned on chirp_mass_proxy "
+            "(trained with gnpe_chirp)."
+        )
+    pins = dict(fixed_context_parameters or {})
+    pins.pop("chirp_mass_proxy", None)
+    expected = set(context_parameters) - {"chirp_mass_proxy"}
+    if set(pins) != expected:
+        raise ValueError(
+            f"The scan fills chirp_mass_proxy; provide fixed_context_parameters "
+            f"for the remaining context parameters {sorted(expected)}, "
+            f"got {sorted(pins)}."
+        )
+
+    grid = chirp_mass_scan_grid(metadata, overlap_factor)
+
+    context = GWSamplerContext.from_model(model, event_data, event_metadata)
+    flow = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
+    tail_steps = (
+        _proxy_offset_steps(inference_parameters, context_parameters)
+        + _ra_reparam_steps(inference_parameters)
+        + _delta_prior_steps(context.prior, inference_parameters)
+    )
+    # Sweep the grid in blocks: a fixed table roots the chain (its length is
+    # the base count, so composer-level chunking does not apply), and the block
+    # size bounds both the transient base-domain memory of the row-wise data
+    # preparation and the network batch. The chain is density-free (no stored
+    # table log-prob) -- the scan selects on the likelihood, not on a proposal
+    # density.
+    frames = []
+    for block in np.array_split(grid, math.ceil(len(grid) / block_size)):
+        # Chain columns carry the network dtype (float32); the data preparation
+        # extracts pin values at this precision, exactly as for fixed-proxy
+        # pins.
+        table = {"chirp_mass_proxy": block.astype(np.float32)}
+        table.update(
+            {k: np.full(len(block), v, dtype=np.float32) for k, v in pins.items()}
+        )
+        composer = ChainComposer(
+            [SampleTableFactor(table), Stage(flow, fan_out=num_samples)] + tail_steps
+        )
+        out, _ = composer.sample_and_log_prob(len(block), context)
+        frames.append(pd.DataFrame({k: v.numpy() for k, v in out.items()}))
+    theta = pd.concat(frames, ignore_index=True)
+
+    # Restrict to draws within the prior (unphysical combinations cannot be
+    # evaluated); the constraint factors are irrelevant for an argmax.
+    prior_keys = [k for k in context.prior if k in theta.columns]
+    log_prior = context.prior.ln_prob(theta[prior_keys], axis=0)
+    theta = theta.iloc[np.flatnonzero(np.isfinite(log_prior))]
+
+    # Exact likelihood on the event data -- on the base domain for multibanded
+    # models -- marginalized over the (unsampled) phase.
+    likelihood_context = (
+        context.derive(use_base_domain=True)
+        if hasattr(context.domain, "base_domain")
+        else context
+    )
+    likelihood = likelihood_context.likelihood(
+        phase_marginalization_kwargs={"approximation_22_mode": True}
+    )
+    likelihood.return_aux_snr = True
+    theta_likelihood = theta[[c for c in theta.columns if c != "chirp_mass_proxy"]]
+    log_likelihood, snr = likelihood.log_likelihood_multi(
+        theta_likelihood, num_processes=num_processes
+    ).T
+
+    winner = int(np.argmax(log_likelihood))
+    samples = theta.assign(log_likelihood=log_likelihood, snr=snr)
+    return {
+        "chirp_mass_trigger": float(samples["chirp_mass"].iloc[winner]),
+        "snr": float(snr[winner]),
+        "max_log_likelihood": float(log_likelihood[winner]),
+        "grid": grid,
+        "samples": samples,
+        "settings": {
+            "num_samples": num_samples,
+            "overlap_factor": overlap_factor,
+            "block_size": block_size,
+            "grid_size": len(grid),
+        },
+    }
diff --git a/dingo/gw/likelihood.py b/dingo/gw/likelihood.py
index ea9927322..57d7ec54f 100644
--- a/dingo/gw/likelihood.py
+++ b/dingo/gw/likelihood.py
@@ -178,6 +178,12 @@ def __init__(
         # Initialize calibration marginalization using the setter from GWSignal.
         self.calibration_marginalization_kwargs = calibration_marginalization_kwargs
 
+        # When set, log_likelihood returns (log_likelihood, snr) pairs, where snr
+        # is the matched-filter signal-to-noise ratio of the waveform against the
+        # data (phase-maximized under phase marginalization). Used by the
+        # chirp-mass scan to report trigger quality.
+        self.return_aux_snr = False
+
     def initialize_time_marginalization(self, t_lower, t_upper, n_fft=1):
         """
         Initialize time marginalization. Time marginalization can be performed via FFT,
@@ -314,7 +320,13 @@ def _log_likelihood(self, theta):
                 for d_ifo, mu_ifo in zip(d.values(), mu.values())
             ],
         )
-        return self.log_Zn + kappa2 - 1 / 2.0 * rho2opt
+        log_likelihood = self.log_Zn + kappa2 - 1 / 2.0 * rho2opt
+
+        if self.return_aux_snr:
+            snr = kappa2 / rho2opt**0.5
+            return log_likelihood, snr
+
+        return log_likelihood
 
     def log_likelihood_phase_grid(
         self, theta: dict, phases: Optional[np.ndarray] = None
@@ -335,6 +347,8 @@ def log_likelihood_phase_grid(
     def _log_likelihood_phase_grid_manual(
         self, theta: dict, phases: Optional[np.ndarray] = None
     ) -> np.ndarray:
+        if self.return_aux_snr:
+            raise NotImplementedError
         if self.phase_marginalization:
             raise ValueError(
                 "Can't compute likelihood on a phase grid for "
@@ -356,6 +370,8 @@ def _log_likelihood_phase_grid_manual(
 
     def _log_likelihood_phase_grid_mode_decomposed(self, theta, phases=None):
         # TODO: Implement for time marginalization
+        if self.return_aux_snr:
+            raise NotImplementedError
         if self.phase_marginalization:
             raise ValueError(
                 "Can't compute likelihood on a phase grid for "
@@ -512,7 +528,13 @@ def _log_likelihood_phase_marginalized(self, theta):
                     for d_ifo, mu_ifo in zip(d.values(), mu.values())
                 ]
             )
-            return self.log_Zn + ln_i0(np.abs(kappa2C)) - 1 / 2.0 * rho2opt
+            log_likelihood = self.log_Zn + ln_i0(np.abs(kappa2C)) - 1 / 2.0 * rho2opt
+
+            if self.return_aux_snr:
+                snr = ln_i0(np.abs(kappa2C)) / rho2opt**0.5
+                return log_likelihood, snr
+
+            return log_likelihood
 
         else:
             log_likelihoods_phase_grid = self.log_likelihood_phase_grid(theta)
@@ -596,7 +618,13 @@ def _log_likelihood_time_marginalized(self, theta):
         alpha = np.max(exponent)
         kappa2 = alpha + np.log(np.sum(np.exp(exponent - alpha)))
 
-        return self.log_Zn + kappa2 - 1 / 2.0 * rho2opt
+        log_likelihood = self.log_Zn + kappa2 - 1 / 2.0 * rho2opt
+
+        if self.return_aux_snr:
+            snr = kappa2 / rho2opt**0.5
+            return log_likelihood, snr
+
+        return log_likelihood
 
     def _log_likelihood_calibration_marginalized(self, theta):
         """
@@ -611,6 +639,8 @@ def _log_likelihood_calibration_marginalized(self, theta):
         -------
         log_likelihood: float
         """
+        if self.return_aux_snr:
+            raise NotImplementedError
 
         # Step 1: Compute whitened GW strain mu(theta) for parameters theta.
         mu = self.signal(theta)["waveform"]
diff --git a/tests/gw/test_bns_chain_builder.py b/tests/gw/test_bns_chain_builder.py
index 0702b9add..3fe78824a 100644
--- a/tests/gw/test_bns_chain_builder.py
+++ b/tests/gw/test_bns_chain_builder.py
@@ -112,28 +112,60 @@ def _conditioning(value=1.1975, n=4):
     }
 
 
-def test_prepared_data_consumes_conditioning_and_keys_the_cache():
+def test_prepared_data_pinned_is_row_aligned_and_memoized():
     import torch
 
     ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
     prepared = ctx.prepared_data(conditioning=_conditioning())
-    assert torch.is_tensor(prepared)
-    # Same values: the cached tensor is served.
-    assert ctx.prepared_data(conditioning=_conditioning()) is prepared
-    # A different pin is a different representation: fail rather than serve
-    # stale data.
-    with pytest.raises(ValueError, match="one\\s+representation"):
-        ctx.prepared_data(conditioning=_conditioning(value=1.30))
+    # Row-aligned: one data row per conditioning row, the single pinned
+    # preparation viewed across them.
+    assert prepared.shape[0] == 4
+    cached = ctx._prepared
+    assert torch.equal(ctx.prepared_data(conditioning=_conditioning()), prepared)
+    assert ctx._prepared is cached
+    # A different pin recomputes correctly (the memo is keyed on the value).
+    other = ctx.prepared_data(conditioning=_conditioning(value=1.30))
+    assert ctx._prepared is not cached
+    assert not torch.equal(other, prepared)
+
+
+def test_prepared_data_varying_matches_single_preparations():
+    import torch
+
+    values = [1.19, 1.20, 1.21, 1.22]
+    ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
+    conditioning = _conditioning()
+    conditioning["chirp_mass_proxy"] = torch.tensor(values, dtype=torch.float64)
+    batch = ctx.prepared_data(conditioning=conditioning)
+    assert batch.shape[0] == len(values)
+    # Row i equals the pinned preparation at value i (compared at equal pin
+    # precision: the pinned path extracts the float64 value).
+    for i, value in enumerate(values):
+        single_conditioning = _conditioning()
+        single_conditioning["chirp_mass_proxy"] = torch.full(
+            (4,), value, dtype=torch.float64
+        )
+        single = GWSamplerContext.from_model_metadata(
+            _BNS_METADATA, _event_data()
+        ).prepared_data(conditioning=single_conditioning)
+        assert torch.equal(batch[i], single[0])
+    # A varying pass is not memoized, and the pinned path still works after.
+    assert ctx._prepared is None
+    assert ctx.prepared_data(conditioning=_conditioning()).shape[0] == 4
 
 
-def test_prepared_data_rejects_varying_conditioning():
+def test_prepared_data_varying_rows_follow_values():
     import torch
 
     ctx = GWSamplerContext.from_model_metadata(_BNS_METADATA, _event_data())
     conditioning = _conditioning()
-    conditioning["chirp_mass_proxy"] = torch.tensor([1.19, 1.20, 1.21, 1.22])
-    with pytest.raises(ValueError, match="constant"):
-        ctx.prepared_data(conditioning=conditioning)
+    conditioning["chirp_mass_proxy"] = torch.tensor(
+        [1.19, 1.19, 1.21, 1.21], dtype=torch.float64
+    )
+    batch = ctx.prepared_data(conditioning=conditioning)
+    assert torch.equal(batch[0], batch[1])
+    assert torch.equal(batch[2], batch[3])
+    assert not torch.equal(batch[0], batch[2])
 
 
 def test_proxy_offset_step_selection():
@@ -166,3 +198,37 @@ def test_from_model_requires_all_pins():
             _event_data(),
             fixed_context_parameters={"chirp_mass_proxy": 1.2},
         )
+
+
+def test_chirp_mass_scan_grid_spans_prior_at_kernel_spacing():
+    from dingo.gw.inference.scan import chirp_mass_scan_grid
+
+    grid = chirp_mass_scan_grid(_BNS_METADATA, overlap_factor=2)
+    # Prior [1.0, 2.0], kernel width 0.01, overlap 2 -> 200 points inset by the
+    # kernel edges, spaced at most half a kernel width apart.
+    assert len(grid) == 200
+    assert grid[0] == pytest.approx(1.005)
+    assert grid[-1] == pytest.approx(1.995)
+    assert np.diff(grid).max() <= 0.005 + 1e-12
+
+
+def test_chirp_mass_scan_validates_model_and_pins():
+    import copy
+
+    from dingo.gw.inference.scan import chirp_mass_scan
+
+    # The non-proxy context parameters must be pinned by the caller.
+    with pytest.raises(ValueError, match="remaining context parameters"):
+        chirp_mass_scan(_StubBNSModel(), _event_data(), fixed_context_parameters={})
+
+    # A model without chirp-mass conditioning cannot be scanned.
+    metadata_no_chirp = copy.deepcopy(_BNS_METADATA)
+    del metadata_no_chirp["train_settings"]["data"]["gnpe_chirp"]
+    metadata_no_chirp["train_settings"]["data"]["context_parameters"] = ["ra", "dec"]
+
+    class _StubNonChirpModel:
+        metadata = metadata_no_chirp
+        device = "cpu"
+
+    with pytest.raises(ValueError, match="chirp_mass_proxy"):
+        chirp_mass_scan(_StubNonChirpModel(), _event_data())

From ffbc07c24867c8093c21d88e0d82c5d8682f62f8 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 12 Jul 2026 21:37:16 +0100
Subject: [PATCH 48/65] Wire the chirp-mass scan into dingo_pipe

Add the --chirp-mass-scan option ('true' for model-derived defaults, a
dictionary for overrides; mutually exclusive with a pinned
chirp_mass_proxy and with model-init). The sampling stage runs the scan
once before building the sampler, fills the winning trigger into the
fixed context parameters, and records the scan settings and winner
alongside the chain provenance. Verified end-to-end on GW170817 with
only the sky position pinned: the scanned trigger matches the posterior
median, and the evidence agrees with the externally-pinned run within
Monte Carlo error.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/pipe/parser.py                        | 12 +++++
 dingo/pipe/sampling.py                      | 49 +++++++++++++++++++++
 tests/pipe/test_fixed_context_parameters.py | 21 +++++++++
 3 files changed, 82 insertions(+)

diff --git a/dingo/pipe/parser.py b/dingo/pipe/parser.py
index 59269162f..4fa698b24 100644
--- a/dingo/pipe/parser.py
+++ b/dingo/pipe/parser.py
@@ -1560,6 +1560,18 @@ def create_parser(top_level=True, usage=None):
         "preparation (chirp-mass heterodyning), and all pinned values are "
         "recorded with the samples.",
     )
+    sampler_parser.add(
+        "--chirp-mass-scan",
+        type=nonestr,
+        default=None,
+        help="Determine the trigger chirp mass from the data before sampling, "
+        "instead of pinning chirp_mass_proxy in fixed-context-parameters "
+        "(mutually exclusive with such a pin). 'true' scans with defaults "
+        "derived from the model (grid from the training prior and kernel, 10 "
+        "draws per point, maximum-likelihood winner); a dictionary overrides "
+        "individual settings, e.g. {num_samples: 10, overlap_factor: 2, "
+        "block_size: 32}.",
+    )
     sampler_parser.add(
         "--device",
         type=str,
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index ddde8cdf7..b652d6acf 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -17,6 +17,7 @@
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
 from dingo.gw.data.event_dataset import EventDataset
 from dingo.gw.inference.factors import GNPEKernelFactor, GWComposedSampler
+from dingo.gw.inference.scan import chirp_mass_scan
 from dingo.pipe.default_settings import DENSITY_RECOVERY_SETTINGS
 from dingo.pipe.parser import create_parser
 
@@ -62,6 +63,12 @@ def __init__(self, args, unknown_args):
             if args.fixed_context_parameters is not None
             else None
         )
+        if args.chirp_mass_scan is None:
+            self.chirp_mass_scan_settings = None
+        elif args.chirp_mass_scan.strip().lower() == "true":
+            self.chirp_mass_scan_settings = {}
+        else:
+            self.chirp_mass_scan_settings = convert_string_to_dict(args.chirp_mass_scan)
         self.num_samples = args.num_samples
         self.batch_size = args.batch_size
         self.density_recovery_settings = args.density_recovery_settings
@@ -132,6 +139,46 @@ def _load_sampler(self):
             filename=self.model, device=self.device, load_training_info=False
         )
 
+        scan_provenance = None
+        if self.chirp_mass_scan_settings is not None:
+            if self.model_init is not None:
+                raise ValueError(
+                    "chirp-mass-scan applies to single-network models; it cannot "
+                    "be combined with model-init (iterative GNPE)."
+                )
+            if (
+                self.fixed_context_parameters is not None
+                and "chirp_mass_proxy" in self.fixed_context_parameters
+            ):
+                raise ValueError(
+                    "chirp-mass-scan and a pinned chirp_mass_proxy in "
+                    "fixed-context-parameters are mutually exclusive; provide one."
+                )
+            scan_kwargs = dict(self.chirp_mass_scan_settings)
+            scan_kwargs.setdefault("num_processes", self.request_cpus)
+            scan = chirp_mass_scan(
+                model,
+                self.context,
+                event_metadata=self.event_metadata,
+                fixed_context_parameters=self.fixed_context_parameters,
+                **scan_kwargs,
+            )
+            logger.info(
+                f"Chirp-mass scan: trigger {scan['chirp_mass_trigger']:.5f}, "
+                f"snr {scan['snr']:.2f}, "
+                f"max log likelihood {scan['max_log_likelihood']:.2f}."
+            )
+            self.fixed_context_parameters = {
+                **(self.fixed_context_parameters or {}),
+                "chirp_mass_proxy": scan["chirp_mass_trigger"],
+            }
+            scan_provenance = {
+                **scan["settings"],
+                "chirp_mass_trigger": scan["chirp_mass_trigger"],
+                "snr": scan["snr"],
+                "max_log_likelihood": scan["max_log_likelihood"],
+            }
+
         if self.model_init is not None:
             if self.fixed_context_parameters is not None:
                 raise ValueError(
@@ -161,6 +208,8 @@ def _load_sampler(self):
                 fixed_context_parameters=self.fixed_context_parameters,
             )
         self.dingo_sampler.provenance_extra["models"] = self._model_paths()
+        if scan_provenance is not None:
+            self.dingo_sampler.provenance_extra["chirp_mass_scan"] = scan_provenance
 
     @property
     def density_recovery_settings(self):
diff --git a/tests/pipe/test_fixed_context_parameters.py b/tests/pipe/test_fixed_context_parameters.py
index a378c9616..07d5273e5 100644
--- a/tests/pipe/test_fixed_context_parameters.py
+++ b/tests/pipe/test_fixed_context_parameters.py
@@ -32,3 +32,24 @@ def test_ini_dict_form_parses(tmp_path):
         ini_body=f"model = model.pt\nfixed-context-parameters = {DEMO_FORM}\n",
     )
     assert convert_string_to_dict(args.fixed_context_parameters) == DEMO_DICT
+
+
+def test_chirp_mass_scan_default_is_none(tmp_path):
+    assert _parse(tmp_path).chirp_mass_scan is None
+
+
+def test_chirp_mass_scan_true_form_parses(tmp_path):
+    args = _parse(tmp_path, ini_body="model = model.pt\nchirp-mass-scan = true\n")
+    assert args.chirp_mass_scan.strip().lower() == "true"
+
+
+def test_chirp_mass_scan_dict_form_parses(tmp_path):
+    args = _parse(
+        tmp_path,
+        ini_body="model = model.pt\n"
+        "chirp-mass-scan = {num_samples: 10, overlap_factor: 2}\n",
+    )
+    assert convert_string_to_dict(args.chirp_mass_scan) == {
+        "num_samples": 10,
+        "overlap_factor": 2,
+    }

From 767249059c5e2300fd7b39d51fbda8a9bad27365 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Sun, 12 Jul 2026 22:10:22 +0100
Subject: [PATCH 49/65] Split the GW inference module by responsibility

dingo/gw/inference/factors.py (1673 lines) becomes three modules with
one concern each, moved verbatim:

  context.py  GWSamplerContext and the frequency-range update helper
  steps.py    the GNPE factors, SyntheticPhaseFactor, and the
              coordinate reparametrizations
  sampler.py  GWComposedSampler and its chain builders

All imports are retargeted repo-wide (pipe, result, importance
sampling, scan, tests, docs); no compatibility shim is kept, since the
module was introduced on this branch and has no external consumers.
Pure relocation -- no code changes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py                         |    4 +-
 .../importance_sampling/importance_weights.py |    2 +-
 dingo/gw/inference/context.py                 |  636 +++++++
 dingo/gw/inference/factors.py                 | 1673 -----------------
 dingo/gw/inference/sampler.py                 |  366 ++++
 dingo/gw/inference/scan.py                    |   11 +-
 dingo/gw/inference/steps.py                   |  697 +++++++
 dingo/gw/result.py                            |    8 +-
 dingo/pipe/sampling.py                        |    3 +-
 docs/source/dingo_pipe.md                     |    2 +-
 docs/source/example_injection.md              |    2 +-
 docs/source/inference.md                      |    2 +-
 tests/gw/test_bns_chain_builder.py            |    7 +-
 tests/gw/test_gw_sampler_context.py           |    8 +-
 tests/gw/test_ra_reparam.py                   |    2 +-
 tests/gw/test_spin_convention_reparam.py      |    2 +-
 tests/gw/test_synthetic_phase.py              |    3 +-
 17 files changed, 1729 insertions(+), 1699 deletions(-)
 create mode 100644 dingo/gw/inference/context.py
 delete mode 100644 dingo/gw/inference/factors.py
 create mode 100644 dingo/gw/inference/sampler.py
 create mode 100644 dingo/gw/inference/steps.py

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 056e795df..b14c6603a 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -68,7 +68,7 @@ class SamplerContext(Protocol):
     Per-event shared state referenced by every factor: the data `d` and quantities
     derived from it (the prepared-data representation, the likelihood, event metadata).
     Concrete implementations are domain-specific; see
-    `dingo.gw.inference.factors.GWSamplerContext`.
+    `dingo.gw.inference.context.GWSamplerContext`.
 
     `device` is the torch device the chain runs on: steps that create fresh tensors
     (rather than transforming existing ones) create them here, so their outputs can
@@ -861,7 +861,7 @@ class GibbsBlock:
     and returns `None` for the log-prob -- the cyclic dependency has no tractable marginal
     (recoverable by fitting an unconditional density to the samples and taking one
     `ChainComposer` step). Dingo uses this only for multi-iteration GNPE (the GNPE factors
-    in `dingo.gw.inference.factors`), but the loop is generic.
+    in `dingo.gw.inference.steps`), but the loop is generic.
 
     Batching is handled by the enclosing `ChainComposer`: it chunks the walkers and runs
     the whole loop per chunk (`chunk_and_concat`).
diff --git a/dingo/gw/importance_sampling/importance_weights.py b/dingo/gw/importance_sampling/importance_weights.py
index eecd35d92..7db6ce59b 100644
--- a/dingo/gw/importance_sampling/importance_weights.py
+++ b/dingo/gw/importance_sampling/importance_weights.py
@@ -13,7 +13,7 @@
 from dingo.core.posterior_models import NormalizingFlowPosteriorModel
 from dingo.gw.result import Result
 from dingo.core.factors import ChainComposer, FlowFactor
-from dingo.gw.inference.factors import GWComposedSampler
+from dingo.gw.inference.sampler import GWComposedSampler
 from dingo.gw.importance_sampling.diagnostics import plot_diagnostics
 
 
diff --git a/dingo/gw/inference/context.py b/dingo/gw/inference/context.py
new file mode 100644
index 000000000..15296c13d
--- /dev/null
+++ b/dingo/gw/inference/context.py
@@ -0,0 +1,636 @@
+"""Per-event sampler context for gravitational-wave inference: the event
+data and its derived views -- the network-input representation, the prior,
+and the likelihood."""
+
+from __future__ import annotations
+import logging
+from typing import Optional, Union
+import numpy as np
+import torch
+import yaml
+from bilby.core.prior import PriorDict, Uniform
+from torchvision.transforms import Compose
+from dingo.core.factors import _n_rows
+from dingo.core.posterior_models import BasePosteriorModel
+from dingo.core.transforms import GetItem
+from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
+from dingo.gw.frequency_updates import (
+    _validate_maximum_frequency,
+    _validate_minimum_frequency,
+)
+from dingo.gw.gwutils import get_extrinsic_prior_dict
+from dingo.gw.likelihood import StationaryGaussianGWLikelihood
+from dingo.gw.prior import build_prior_with_defaults
+from dingo.gw.transforms import (
+    DecimateWaveformsAndASDS,
+    MaskDataForFrequencyRangeUpdate,
+    HeterodynePhase,
+    RepackageStrainsAndASDS,
+    ToTorch,
+    WhitenAndScaleStrain,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _frequency_range_update(domain, event_metadata) -> Optional[dict]:
+    """The event's requested frequency range when it differs from the data-domain
+    bounds, else `None`. Values may be floats or per-detector dicts; defaults are
+    the domain bounds. A request merely *wider* than the domain also triggers --
+    data generation writes base-domain bounds, which can exceed a multibanded
+    domain's quantized band edge (e.g. 1099.0 vs 1098.875) -- and the resulting
+    mask is then an identity."""
+    if event_metadata is None:
+        return None
+    minimum = event_metadata.get("minimum_frequency", domain.f_min)
+    maximum = event_metadata.get("maximum_frequency", domain.f_max)
+
+    def normalize(value):
+        return set(value.values()) if isinstance(value, dict) else {value}
+
+    if normalize(minimum) == {domain.f_min} and normalize(maximum) == {domain.f_max}:
+        return None
+    return {"minimum_frequency": minimum, "maximum_frequency": maximum}
+
+
+class GWSamplerContext:
+    """
+    Per-event shared GW state: the data `d` and everything derived from it. Referenced
+    by every factor in a chain, and serialized as the transport state between pipe
+    stages.
+
+    This implements the `dingo.core.factors.SamplerContext` protocol. It owns the
+    one-time data preprocessing (`prepared_data`) and builds the exact likelihood
+    (`likelihood`) used by likelihood-based factors (synthetic phase) and importance
+    sampling.
+
+    Event metadata lives here (not on individual factors): it is a property of the data,
+    and it drives frequency cropping, the t_ref/RA correction, and the likelihood.
+
+    The data representation is part of a context's identity: importance-sampling
+    settings updates produce a *derived* context (`derive`) rather than arguments to
+    `likelihood()`. Contexts are treated as immutable -- to change the representation,
+    derive a new one.
+
+    Note that the representation vocabulary here (frequency domains, multibanded
+    decimation, the base-domain likelihood view, frequency-range masking) is specific
+    to this domain family. A new domain family should get its own context class
+    implementing the same interface (`prepared_data` / `likelihood` / `prior` /
+    `derive`) rather than extending this one.
+    """
+
+    def __init__(
+        self,
+        domain,
+        data_prep: Compose,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+        model_metadata: Optional[dict] = None,
+        device: Union[torch.device, str] = "cpu",
+        use_base_domain: bool = False,
+        wfg_delta_f: Optional[float] = None,
+        data_prep_conditioning: Optional[list[str]] = None,
+    ):
+        """
+        Parameters
+        ----------
+        domain : Domain
+            The data frequency domain.
+        data_prep : Compose
+            The one-time data-preprocessing transform chain (whiten / decimate /
+            repackage).
+        event_data : dict
+            The raw event data `d` (strain + ASDs per detector), i.e. `EventDataset.data`.
+            Consumed lazily by `prepared_data()` and reused for the likelihood.
+        event_metadata : dict, optional
+            Per-event metadata; drives frequency cropping, the RA correction, and the
+            likelihood reference time.
+        model_metadata : dict, optional
+            The metadata of the model defining this analysis (dataset + train
+            settings); the source for the prior, the likelihood, the detector
+            names, and the reference time.
+        device : torch.device or str, default "cpu"
+            The torch device the chain runs on (the model device); steps that create
+            fresh tensors (e.g. `DeltaFactor`) create them here.
+        use_base_domain : bool, default False
+            For a multibanded domain, evaluate the likelihood on the base
+            (undecimated) frequency domain rather than the decimated one. Set via
+            `derive`.
+        wfg_delta_f : float, optional
+            Override for the waveform-generator domain `delta_f` (an updated
+            `T = 1/delta_f` cannot be handled by domain projection). Set via `derive`.
+        data_prep_conditioning : list[str], optional
+            Names of the chain-conditioning parameters the data preparation is a
+            function of (e.g. `["chirp_mass_proxy"]` for a heterodyning model).
+            `prepared_data` requires their values, injects them into the
+            transform chain, and keys its cache on them; the values themselves
+            have a single owner -- the chain.
+        """
+        self.domain = domain
+        self._data_prep = data_prep
+        self.model_metadata = model_metadata
+        self.event_data = event_data
+        self.event_metadata = event_metadata
+        self.device = device
+        self.use_base_domain = use_base_domain
+        self.wfg_delta_f = wfg_delta_f
+        self.data_prep_conditioning = list(data_prep_conditioning or [])
+        self._prepared_key: Optional[dict] = None
+        self._prepared: Optional[torch.Tensor] = None
+        self._prior: Optional[PriorDict] = None
+        self._likelihood: Optional[StationaryGaussianGWLikelihood] = None
+        self._likelihood_settings: Optional[dict] = None
+
+    @property
+    def detectors(self) -> list[str]:
+        """Detector names, read from the model metadata."""
+        return self.model_metadata["train_settings"]["data"]["detectors"]
+
+    @property
+    def t_ref(self) -> float:
+        """Training reference GPS time, read from the model metadata."""
+        return self.model_metadata["train_settings"]["data"]["ref_time"]
+
+    @classmethod
+    def from_model_metadata(
+        cls,
+        metadata: dict,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+        device: Union[torch.device, str] = "cpu",
+    ) -> "GWSamplerContext":
+        """Build the context from a model-metadata dict -- no model required. The
+        domain, one-time data-prep chain, prior, and likelihood are all defined by
+        the metadata (e.g. a saved `Result.settings`); `device` only sets where
+        `prepared_data()` and chain-created tensors live.
+
+        Parameters
+        ----------
+        metadata : dict
+            Conditional-model metadata (`dataset_settings` + `train_settings`),
+            e.g. the `settings` of a saved `Result`.
+        event_data : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+        device : torch.device or str, default "cpu"
+            Device for `prepared_data()` and chain-created tensors.
+
+        Returns
+        -------
+        GWSamplerContext
+        """
+        data_settings = metadata["train_settings"]["data"]
+
+        domain = build_domain(metadata["dataset_settings"]["domain"])
+        if "domain_update" in data_settings:
+            domain.update(data_settings["domain_update"])
+        detectors = data_settings["detectors"]
+
+        transforms = []
+        # Chirp-mass GNPE (BNS): heterodyne the raw strain -- before decimation
+        # (they do not commute) and on the base domain. The transform draws the
+        # chirp mass from the sample's "parameters", which `prepared_data`
+        # injects from the chain's conditioning: the proxy value has a single
+        # owner (the chain's DeltaFactor), and the preparation is a function of
+        # it. Iterated chirp GNPE (heterodyning inside a Gibbs loop) is not
+        # implemented: it would require carrying the undecimated strain per
+        # sample.
+        gnpe_chirp = data_settings.get("gnpe_chirp")
+        data_prep_conditioning = []
+        if gnpe_chirp is not None:
+            data_prep_conditioning = [k + "_proxy" for k in gnpe_chirp["kernel"]]
+            transforms.append(
+                HeterodynePhase(
+                    domain=getattr(domain, "base_domain", domain),
+                    order=gnpe_chirp.get("order", 0),
+                )
+            )
+        # Decimate from the base domain when using a multibanded frequency domain.
+        if isinstance(domain, MultibandedFrequencyDomain):
+            transforms.append(
+                DecimateWaveformsAndASDS(domain, decimation_mode="whitened")
+            )
+        # Whiten and scale (the network expects standardized data).
+        transforms.append(WhitenAndScaleStrain(domain.noise_std))
+        # Event frequency-range update: mask the whitened strain/ASDs outside the
+        # requested range. Must precede repackaging (ranges may be per-detector).
+        # The request is validated against the training crop license the first time
+        # prepared_data() runs -- deliberately not here: the license governs the
+        # network-input view only, and contexts are also reconstructed for
+        # likelihood-only use (e.g. from saved importance-sampling results), where
+        # a range that is illegal as network input is legal for ASD masking.
+        range_update = _frequency_range_update(domain, event_metadata)
+        if range_update is not None:
+            transforms.append(
+                MaskDataForFrequencyRangeUpdate(domain=domain, **range_update)
+            )
+        # Repackage strains/ASDs into an array, move to torch, extract the waveform.
+        transforms += [
+            RepackageStrainsAndASDS(ifos=detectors, first_index=domain.min_idx),
+            ToTorch(device=device),
+            GetItem("waveform"),
+        ]
+
+        return cls(
+            domain=domain,
+            data_prep=Compose(transforms),
+            event_data=event_data,
+            event_metadata=event_metadata,
+            model_metadata=metadata,
+            device=device,
+            data_prep_conditioning=data_prep_conditioning,
+        )
+
+    @classmethod
+    def from_model(
+        cls,
+        model: BasePosteriorModel,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+    ) -> "GWSamplerContext":
+        """Build the context from a model: its own metadata and its device. Data
+        preparation is network-bound, so the settings come from `model.metadata`
+        (for a conditional model this equals the base analysis metadata). An
+        unconditional model prepares no data, so no context can be built from one;
+        for the prior/likelihood views alone, use
+        `from_model_metadata(_base_model_metadata(model), ...)`.
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The (conditional) model whose metadata defines the domain and
+            preprocessing.
+        event_data : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+
+        Returns
+        -------
+        GWSamplerContext
+        """
+        if model.metadata["train_settings"]["data"].get("unconditional", False):
+            raise ValueError(
+                "An unconditional model has no data preparation, so a context "
+                "cannot be built from it. For the prior/likelihood views, use "
+                "GWSamplerContext.from_model_metadata(_base_model_metadata(model), ...)."
+            )
+        return cls.from_model_metadata(
+            model.metadata,
+            event_data,
+            event_metadata,
+            device=model.device,
+        )
+
+    def derive(
+        self,
+        updates: Optional[dict] = None,
+        use_base_domain: Optional[bool] = None,
+    ) -> "GWSamplerContext":
+        """Derive a context for the same event under a different data representation.
+
+        The derived context shares the event payload (`event_data`, `event_metadata`),
+        the analysis metadata, and the reference time by construction -- parameter
+        meaning is preserved, so importance weights between the two representations
+        remain well-defined. Only the representation changes: an updated duration `T`
+        rebuilds the data domain at `delta_f = 1/T` (and enters waveform generation as
+        `wfg_delta_f`), and `use_base_domain` switches a multibanded likelihood to the
+        undecimated domain. Frequency-range updates act through `event_metadata` (ASD
+        masking in the likelihood) and need no derived state. Caches start fresh.
+
+        Parameters
+        ----------
+        updates : dict, optional
+            Importance-sampling settings updates, as recorded in
+            `Result.importance_sampling_metadata["updates"]`. May contain `T`,
+            `minimum_frequency`, `maximum_frequency` (and non-domain keys, which are
+            ignored here), but no other quantities that define a new domain.
+        use_base_domain : bool, optional
+            For a multibanded domain, evaluate the likelihood on the base
+            (undecimated) frequency domain. `None` keeps this context's setting.
+
+        Returns
+        -------
+        GWSamplerContext
+        """
+        domain = self.domain
+        wfg_delta_f = self.wfg_delta_f
+        if updates and any(
+            k in updates for k in ("minimum_frequency", "maximum_frequency", "T")
+        ):
+            # TODO: Make compatible with MultibandedFrequencyDomain.
+            if isinstance(domain, MultibandedFrequencyDomain):
+                raise NotImplementedError()
+
+            updates = updates.copy()
+            # A duration update is generation-level on both sides: the pipe
+            # regenerates the event data at the new T, and the waveform must be
+            # generated natively at delta_f = 1/T (a resolution change is not a
+            # projection of existing samples). wfg_delta_f persists on the derived
+            # context, reaching every downstream likelihood build and further
+            # derivations.
+            if "T" in updates:
+                updates["delta_f"] = 1.0 / updates["T"]
+                wfg_delta_f = updates["delta_f"]
+                print(
+                    f"Updating waveform generation delta_f from "
+                    f'{self.model_metadata["dataset_settings"]["domain"]["delta_f"]} '
+                    f"to {wfg_delta_f}."
+                )
+
+            domain_dict = domain.domain_dict  # Existing settings
+            domain_dict.update(
+                (k, updates[k]) for k in set(domain_dict).intersection(updates)
+            )
+            print("Rebuilding domain as follows:")
+            print(yaml.dump(domain_dict, default_flow_style=False, sort_keys=False))
+            domain = build_domain(domain_dict)
+
+        return type(self)(
+            domain=domain,
+            data_prep=self._data_prep,
+            event_data=self.event_data,
+            event_metadata=self.event_metadata,
+            model_metadata=self.model_metadata,
+            device=self.device,
+            use_base_domain=(
+                self.use_base_domain if use_base_domain is None else use_base_domain
+            ),
+            wfg_delta_f=wfg_delta_f,
+            data_prep_conditioning=self.data_prep_conditioning,
+        )
+
+    def prepared_data(self, conditioning=None) -> torch.Tensor:
+        """The network-input data representation of this event.
+
+        Without `conditioning`: the single shared representation, computed once
+        and cached. With `conditioning` (the chain columns available to a
+        conditioned factor): row-aligned, one data row per conditioning row.
+        Only the columns named in `data_prep_conditioning` are consumed by the
+        preparation (e.g. the chirp-mass heterodyne proxy, injected under its
+        physical name); the remaining columns condition the network only, and a
+        context that consumes nothing serves the shared representation viewed
+        across the rows. A constant consumed value (a pinned proxy) is prepared
+        once and viewed across the rows; varying values (a sweep) run through
+        the batch-native transform chain in one vectorized pass, uncached -- a
+        caller sweeping more rows than memory allows blocks its own request.
+
+        An event frequency-range update is validated against the training crop
+        license before any preparation.
+
+        Parameters
+        ----------
+        conditioning : dict[str, torch.Tensor], optional
+            The chain conditioning available to the calling factor, one value
+            per row. May contain columns irrelevant to the preparation.
+        """
+        if not self.data_prep_conditioning:
+            if self._prepared is None:
+                self._validate_frequency_range()
+                self._prepared = self._data_prep(self.event_data)
+            if conditioning is None:
+                return self._prepared
+            return self._prepared.expand(_n_rows(conditioning), *self._prepared.shape)
+
+        columns = self._conditioning_columns(conditioning)
+        self._validate_frequency_range()
+        n_rows = _n_rows(columns)
+        if all(torch.all(c == c[0]) for c in columns.values()):
+            # N rows of one pinned value: prepare once, view it across the rows.
+            key = {name: float(c[0]) for name, c in columns.items()}
+            if key != self._prepared_key:
+                self._prepared = self._data_prep({**self.event_data, "parameters": key})
+                self._prepared_key = key
+            return self._prepared.expand(n_rows, *self._prepared.shape)
+        parameters = {name: column.numpy() for name, column in columns.items()}
+        return self._data_prep(
+            {**self._broadcast_event(n_rows), "parameters": parameters}
+        )
+
+    def _conditioning_columns(self, conditioning) -> dict[str, torch.Tensor]:
+        """Collect the conditioning columns the preparation consumes, keyed by
+        their physical names (the `_proxy` suffix names the chain column; the
+        transform chain reads the physical parameter), as float64 (the
+        heterodyne phase is computed in float64, and a float32 chain column
+        must not degrade it)."""
+        conditioning = conditioning or {}
+        columns = {}
+        for name in self.data_prep_conditioning:
+            if name not in conditioning:
+                raise ValueError(
+                    f"This model's data preparation is a function of the chain "
+                    f"conditioning `{name}`, which the caller does not provide "
+                    f"(pass it to prepared_data, e.g. from the chain's pins)."
+                )
+            columns[name[: -len("_proxy")]] = torch.as_tensor(
+                conditioning[name], dtype=torch.float64
+            ).reshape(-1)
+        return columns
+
+    def _broadcast_event(self, n_rows: int) -> dict:
+        """The event arrays broadcast (as read-only views) across `n_rows`
+        rows, forming the batched sample dict for a single transform-chain
+        pass."""
+        return {
+            part: (
+                {k: np.broadcast_to(v, (n_rows, *np.shape(v))) for k, v in data.items()}
+                if isinstance(data, dict)
+                else data
+            )
+            for part, data in self.event_data.items()
+        }
+
+    def _validate_frequency_range(self):
+        """Validate an event frequency-range update: hard bounds against the (base)
+        domain, and narrowing only when the network was trained with random strain
+        cropping covering the requested range. Applies to the network-input view
+        only -- the likelihood view applies the range independently via ASD
+        masking."""
+        update = _frequency_range_update(self.domain, self.event_metadata)
+        if update is None:
+            return
+        domain = getattr(self.domain, "base_domain", self.domain)
+        crop_settings = self.model_metadata["train_settings"]["data"].get(
+            "random_strain_cropping"
+        )
+        _validate_minimum_frequency(
+            update["minimum_frequency"], self.detectors, domain, crop_settings
+        )
+        _validate_maximum_frequency(
+            update["maximum_frequency"], self.detectors, domain, crop_settings
+        )
+
+    @property
+    def prior(self) -> PriorDict:
+        """The static prior over all parameters, built once from the model metadata
+        (intrinsic + extrinsic priors with Dingo defaults).
+
+        This is the event-independent prior fixed at training time. Importance-sampling
+        prior-bound updates and the time / phase split-off for marginalized networks are
+        applied downstream (they depend on the evolving analysis state), not here.
+        """
+        if self._prior is None:
+            data_settings = self.model_metadata["train_settings"]["data"]
+            intrinsic_prior = self.model_metadata["dataset_settings"]["intrinsic_prior"]
+            extrinsic_prior = get_extrinsic_prior_dict(data_settings["extrinsic_prior"])
+            self._prior = build_prior_with_defaults(
+                {**intrinsic_prior, **extrinsic_prior}
+            )
+        return self._prior
+
+    def likelihood(
+        self,
+        time_marginalization_kwargs: Optional[dict] = None,
+        phase_marginalization_kwargs: Optional[dict] = None,
+        calibration_marginalization_kwargs: Optional[dict] = None,
+    ) -> StationaryGaussianGWLikelihood:
+        """
+        Build the exact GW likelihood on this event's data, in physical parameter space.
+
+        The network's standardized, decimated view of the data is `prepared_data()`; the
+        likelihood instead takes the raw event data (`event_data`) and builds its own
+        representation -- decimating to the multibanded domain unless `use_base_domain`,
+        and masking the ASDs to the event's frequency range. Its reference time is the
+        event time (the training-frame right-ascension correction is already applied to the
+        samples), or the training reference time when no event time is set.
+
+        The data representation (the domain, `use_base_domain`, `wfg_delta_f`, the
+        event's frequency range) is context state -- importance-sampling settings
+        updates enter by deriving a new context (`derive`), not as arguments here.
+        Marginalization is per-request and enters as arguments.
+
+        The most recently built likelihood is cached with its arguments: a repeated
+        call with the same arguments returns the shared instance (the synthetic-phase
+        factor requests one per chain chunk), and a call with different arguments
+        builds a replacement. In the exact synthetic-phase mode the consumer assigns
+        a `phase_grid` attribute on the shared instance.
+
+        Parameters
+        ----------
+        time_marginalization_kwargs : dict, optional
+            Analytically marginalize over `geocent_time`. `t_lower` / `t_upper` are
+            filled from the network's (uniform) time prior when not already provided
+            (a caller with an updated prior passes its own bounds). Requires a
+            time-marginalized network.
+        phase_marginalization_kwargs : dict, optional
+            Analytically marginalize over `phase`. Requires a uniform [0, 2 pi) phase prior.
+        calibration_marginalization_kwargs : dict, optional
+            Marginalize over detector calibration uncertainty.
+
+        Returns
+        -------
+        StationaryGaussianGWLikelihood
+        """
+        # Capture the call's arguments for the cache comparison; this must stay the
+        # first statement so locals() holds exactly the arguments. The comparison
+        # happens before the marginalization bounds are filled in below (which is
+        # deterministic given the same arguments). The representation fields join the
+        # key as a guard against off-contract in-place mutation.
+        settings = dict(locals())
+        settings.pop("self")
+        settings["use_base_domain"] = self.use_base_domain
+        settings["wfg_delta_f"] = self.wfg_delta_f
+        if settings == self._likelihood_settings:
+            return self._likelihood
+
+        if self.event_data is None:
+            raise ValueError(
+                "Building the likelihood requires event data (strain + ASDs), "
+                "which this context does not carry."
+            )
+
+        # Marginalizing over a parameter needs the (uniform) prior the network
+        # marginalized over; use it to parameterize the requested marginalization.
+        # Bounds already provided by the caller win: the importance-sampling layer
+        # validates against its evolved prior (prior updates, time/phase
+        # split-offs), which this sample-free context cannot see, and passes
+        # explicit bounds. The fill below covers standalone (chain) use from the
+        # network's static prior.
+        if time_marginalization_kwargs is not None and not (
+            "t_lower" in time_marginalization_kwargs
+            and "t_upper" in time_marginalization_kwargs
+        ):
+            time_prior = self._marginalized_prior("geocent_time")
+            if time_prior is None:
+                raise NotImplementedError(
+                    "Time marginalization is not compatible with a non-marginalized "
+                    "network."
+                )
+            if type(time_prior) != Uniform:
+                raise NotImplementedError(
+                    "Only uniform time prior is supported for time marginalization."
+                )
+            time_marginalization_kwargs = {
+                **time_marginalization_kwargs,
+                "t_lower": time_prior.minimum,
+                "t_upper": time_prior.maximum,
+            }
+        if phase_marginalization_kwargs is not None:
+            phase_prior = self._marginalized_prior("phase")
+            if not (
+                isinstance(phase_prior, Uniform)
+                and (phase_prior._minimum, phase_prior._maximum) == (0, 2 * np.pi)
+            ):
+                raise ValueError(
+                    f"Phase prior should be uniform [0, 2pi) for phase marginalization, "
+                    f"but is {phase_prior}."
+                )
+
+        dataset_settings = self.model_metadata["dataset_settings"]
+        # WaveformGenerator domain -- deliberately the dataset domain WITHOUT any
+        # domain_update, so waveform generation mirrors how the training set was
+        # generated (possibly wider than the network's data domain); the likelihood
+        # projects the generated waveform onto the data domain. An updated
+        # T = 1/delta_f enters here: unlike range masking or decimation, a
+        # resolution change needs samples that do not exist on the old grid, so it
+        # cannot be handled by domain projection and must apply at generation.
+        wfg_domain_dict = dataset_settings["domain"]
+        if self.wfg_delta_f is not None:
+            wfg_domain_dict = {**wfg_domain_dict, "delta_f": self.wfg_delta_f}
+        wfg_domain = build_domain(wfg_domain_dict)
+
+        # Likelihood reference time: the event time (the training-frame RA correction has
+        # already been applied to the samples), falling back to the training reference.
+        if self.event_metadata is not None and "time_event" in self.event_metadata:
+            t_ref = self.event_metadata["time_event"]
+        else:
+            t_ref = self.t_ref
+
+        frequency_update = dict(
+            minimum_frequency=self._frequency("minimum_frequency", self.domain.f_min),
+            maximum_frequency=self._frequency("maximum_frequency", self.domain.f_max),
+        )
+
+        likelihood = StationaryGaussianGWLikelihood(
+            wfg_kwargs=dataset_settings["waveform_generator"],
+            wfg_domain=wfg_domain,
+            data_domain=self.domain,
+            event_data=self.event_data,
+            t_ref=t_ref,
+            time_marginalization_kwargs=time_marginalization_kwargs,
+            phase_marginalization_kwargs=phase_marginalization_kwargs,
+            calibration_marginalization_kwargs=calibration_marginalization_kwargs,
+            use_base_domain=self.use_base_domain,
+            frequency_update=frequency_update,
+        )
+        self._likelihood = likelihood
+        self._likelihood_settings = settings
+        return likelihood
+
+    def _frequency(self, key: str, default: float):
+        """The event's frequency-range override for `key` (min/max), else `default`."""
+        if self.event_metadata is None:
+            return default
+        return self.event_metadata.get(key, default)
+
+    def _marginalized_prior(self, name: str):
+        """The prior over `name` if the network marginalized it (i.e. `name` is not an
+        inference parameter), else `None` -- used to parameterize likelihood
+        marginalization over time / phase."""
+        if (
+            name
+            in self.model_metadata["train_settings"]["data"]["inference_parameters"]
+        ):
+            return None
+        return self.prior.get(name)
diff --git a/dingo/gw/inference/factors.py b/dingo/gw/inference/factors.py
deleted file mode 100644
index 78ee4dcd7..000000000
--- a/dingo/gw/inference/factors.py
+++ /dev/null
@@ -1,1673 +0,0 @@
-"""
-Gravitational-wave factors, steps, and samplers for the factorized sampler.
-
-The domain-agnostic pieces (`Factor`, `FlowFactor`, `Reparametrization`,
-`TargetCorrection`, and the composers) live in `dingo.core.factors`. This module adds
-the GW-specific ones: `GWSamplerContext` (the data-preprocessing conditioning map), the
-GNPE factors and steps (`GNPEKernelFactor`, `GNPEFlowFactor`, `GNPEKernelCorrection`),
-the `RAReparam` sky-frame reparametrization, and `GWComposedSampler`, which assembles
-the chain for plain NPE, multi-iteration GNPE, or single-step GNPE.
-"""
-
-from __future__ import annotations
-
-import copy
-import logging
-import time
-from pathlib import Path
-from typing import Optional, Union
-
-import numpy as np
-import pandas as pd
-import torch
-import yaml
-from astropy.time import Time
-from bilby.core.prior import DeltaFunction, PriorDict, Uniform
-from bilby.gw.detector import InterferometerList
-from torchvision.transforms import Compose
-
-from dingo.core.density import (
-    interpolated_log_prob_multi,
-    interpolated_sample_and_log_prob_multi,
-)
-from dingo.core.factors import (
-    ChainComposer,
-    ComposedSampler,
-    DeltaFactor,
-    Factor,
-    FlowFactor,
-    GibbsBlock,
-    ProxyOffsetReparam,
-    Reparametrization,
-    Standardization,
-    TargetCorrection,
-    _base_model_metadata,
-    _n_rows,
-)
-from dingo.core.multiprocessing import apply_func_with_multiprocessing
-from dingo.core.posterior_models import BasePosteriorModel
-from dingo.core.transforms import GetItem, RenameKey
-from dingo.gw.conversion import change_spin_conversion_phase
-from dingo.gw.domains import build_domain, MultibandedFrequencyDomain
-from dingo.gw.frequency_updates import (
-    _validate_maximum_frequency,
-    _validate_minimum_frequency,
-)
-from dingo.gw.gwutils import get_extrinsic_prior_dict
-from dingo.gw.likelihood import StationaryGaussianGWLikelihood
-from dingo.gw.prior import build_prior_with_defaults
-from dingo.gw.transforms import (
-    CopyToExtrinsicParameters,
-    DecimateWaveformsAndASDS,
-    GetDetectorTimes,
-    GNPEBase,
-    GNPECoalescenceTimes,
-    MaskDataForFrequencyRangeUpdate,
-    PostCorrectGeocentTime,
-    HeterodynePhase,
-    RepackageStrainsAndASDS,
-    SelectStandardizeRepackageParameters,
-    TimeShiftStrain,
-    ToTorch,
-    WhitenAndScaleStrain,
-)
-
-logger = logging.getLogger(__name__)
-
-
-def _frequency_range_update(domain, event_metadata) -> Optional[dict]:
-    """The event's requested frequency range when it differs from the data-domain
-    bounds, else `None`. Values may be floats or per-detector dicts; defaults are
-    the domain bounds. A request merely *wider* than the domain also triggers --
-    data generation writes base-domain bounds, which can exceed a multibanded
-    domain's quantized band edge (e.g. 1099.0 vs 1098.875) -- and the resulting
-    mask is then an identity."""
-    if event_metadata is None:
-        return None
-    minimum = event_metadata.get("minimum_frequency", domain.f_min)
-    maximum = event_metadata.get("maximum_frequency", domain.f_max)
-
-    def normalize(value):
-        return set(value.values()) if isinstance(value, dict) else {value}
-
-    if normalize(minimum) == {domain.f_min} and normalize(maximum) == {domain.f_max}:
-        return None
-    return {"minimum_frequency": minimum, "maximum_frequency": maximum}
-
-
-class GWSamplerContext:
-    """
-    Per-event shared GW state: the data `d` and everything derived from it. Referenced
-    by every factor in a chain, and serialized as the transport state between pipe
-    stages.
-
-    This implements the `dingo.core.factors.SamplerContext` protocol. It owns the
-    one-time data preprocessing (`prepared_data`) and builds the exact likelihood
-    (`likelihood`) used by likelihood-based factors (synthetic phase) and importance
-    sampling.
-
-    Event metadata lives here (not on individual factors): it is a property of the data,
-    and it drives frequency cropping, the t_ref/RA correction, and the likelihood.
-
-    The data representation is part of a context's identity: importance-sampling
-    settings updates produce a *derived* context (`derive`) rather than arguments to
-    `likelihood()`. Contexts are treated as immutable -- to change the representation,
-    derive a new one.
-
-    Note that the representation vocabulary here (frequency domains, multibanded
-    decimation, the base-domain likelihood view, frequency-range masking) is specific
-    to this domain family. A new domain family should get its own context class
-    implementing the same interface (`prepared_data` / `likelihood` / `prior` /
-    `derive`) rather than extending this one.
-    """
-
-    def __init__(
-        self,
-        domain,
-        data_prep: Compose,
-        event_data: dict,
-        event_metadata: Optional[dict] = None,
-        model_metadata: Optional[dict] = None,
-        device: Union[torch.device, str] = "cpu",
-        use_base_domain: bool = False,
-        wfg_delta_f: Optional[float] = None,
-        data_prep_conditioning: Optional[list[str]] = None,
-    ):
-        """
-        Parameters
-        ----------
-        domain : Domain
-            The data frequency domain.
-        data_prep : Compose
-            The one-time data-preprocessing transform chain (whiten / decimate /
-            repackage).
-        event_data : dict
-            The raw event data `d` (strain + ASDs per detector), i.e. `EventDataset.data`.
-            Consumed lazily by `prepared_data()` and reused for the likelihood.
-        event_metadata : dict, optional
-            Per-event metadata; drives frequency cropping, the RA correction, and the
-            likelihood reference time.
-        model_metadata : dict, optional
-            The metadata of the model defining this analysis (dataset + train
-            settings); the source for the prior, the likelihood, the detector
-            names, and the reference time.
-        device : torch.device or str, default "cpu"
-            The torch device the chain runs on (the model device); steps that create
-            fresh tensors (e.g. `DeltaFactor`) create them here.
-        use_base_domain : bool, default False
-            For a multibanded domain, evaluate the likelihood on the base
-            (undecimated) frequency domain rather than the decimated one. Set via
-            `derive`.
-        wfg_delta_f : float, optional
-            Override for the waveform-generator domain `delta_f` (an updated
-            `T = 1/delta_f` cannot be handled by domain projection). Set via `derive`.
-        data_prep_conditioning : list[str], optional
-            Names of the chain-conditioning parameters the data preparation is a
-            function of (e.g. `["chirp_mass_proxy"]` for a heterodyning model).
-            `prepared_data` requires their values, injects them into the
-            transform chain, and keys its cache on them; the values themselves
-            have a single owner -- the chain.
-        """
-        self.domain = domain
-        self._data_prep = data_prep
-        self.model_metadata = model_metadata
-        self.event_data = event_data
-        self.event_metadata = event_metadata
-        self.device = device
-        self.use_base_domain = use_base_domain
-        self.wfg_delta_f = wfg_delta_f
-        self.data_prep_conditioning = list(data_prep_conditioning or [])
-        self._prepared_key: Optional[dict] = None
-        self._prepared: Optional[torch.Tensor] = None
-        self._prior: Optional[PriorDict] = None
-        self._likelihood: Optional[StationaryGaussianGWLikelihood] = None
-        self._likelihood_settings: Optional[dict] = None
-
-    @property
-    def detectors(self) -> list[str]:
-        """Detector names, read from the model metadata."""
-        return self.model_metadata["train_settings"]["data"]["detectors"]
-
-    @property
-    def t_ref(self) -> float:
-        """Training reference GPS time, read from the model metadata."""
-        return self.model_metadata["train_settings"]["data"]["ref_time"]
-
-    @classmethod
-    def from_model_metadata(
-        cls,
-        metadata: dict,
-        event_data: dict,
-        event_metadata: Optional[dict] = None,
-        device: Union[torch.device, str] = "cpu",
-    ) -> "GWSamplerContext":
-        """Build the context from a model-metadata dict -- no model required. The
-        domain, one-time data-prep chain, prior, and likelihood are all defined by
-        the metadata (e.g. a saved `Result.settings`); `device` only sets where
-        `prepared_data()` and chain-created tensors live.
-
-        Parameters
-        ----------
-        metadata : dict
-            Conditional-model metadata (`dataset_settings` + `train_settings`),
-            e.g. the `settings` of a saved `Result`.
-        event_data : dict
-            The raw event data (strain + ASDs).
-        event_metadata : dict, optional
-            Per-event metadata.
-        device : torch.device or str, default "cpu"
-            Device for `prepared_data()` and chain-created tensors.
-
-        Returns
-        -------
-        GWSamplerContext
-        """
-        data_settings = metadata["train_settings"]["data"]
-
-        domain = build_domain(metadata["dataset_settings"]["domain"])
-        if "domain_update" in data_settings:
-            domain.update(data_settings["domain_update"])
-        detectors = data_settings["detectors"]
-
-        transforms = []
-        # Chirp-mass GNPE (BNS): heterodyne the raw strain -- before decimation
-        # (they do not commute) and on the base domain. The transform draws the
-        # chirp mass from the sample's "parameters", which `prepared_data`
-        # injects from the chain's conditioning: the proxy value has a single
-        # owner (the chain's DeltaFactor), and the preparation is a function of
-        # it. Iterated chirp GNPE (heterodyning inside a Gibbs loop) is not
-        # implemented: it would require carrying the undecimated strain per
-        # sample.
-        gnpe_chirp = data_settings.get("gnpe_chirp")
-        data_prep_conditioning = []
-        if gnpe_chirp is not None:
-            data_prep_conditioning = [k + "_proxy" for k in gnpe_chirp["kernel"]]
-            transforms.append(
-                HeterodynePhase(
-                    domain=getattr(domain, "base_domain", domain),
-                    order=gnpe_chirp.get("order", 0),
-                )
-            )
-        # Decimate from the base domain when using a multibanded frequency domain.
-        if isinstance(domain, MultibandedFrequencyDomain):
-            transforms.append(
-                DecimateWaveformsAndASDS(domain, decimation_mode="whitened")
-            )
-        # Whiten and scale (the network expects standardized data).
-        transforms.append(WhitenAndScaleStrain(domain.noise_std))
-        # Event frequency-range update: mask the whitened strain/ASDs outside the
-        # requested range. Must precede repackaging (ranges may be per-detector).
-        # The request is validated against the training crop license the first time
-        # prepared_data() runs -- deliberately not here: the license governs the
-        # network-input view only, and contexts are also reconstructed for
-        # likelihood-only use (e.g. from saved importance-sampling results), where
-        # a range that is illegal as network input is legal for ASD masking.
-        range_update = _frequency_range_update(domain, event_metadata)
-        if range_update is not None:
-            transforms.append(
-                MaskDataForFrequencyRangeUpdate(domain=domain, **range_update)
-            )
-        # Repackage strains/ASDs into an array, move to torch, extract the waveform.
-        transforms += [
-            RepackageStrainsAndASDS(ifos=detectors, first_index=domain.min_idx),
-            ToTorch(device=device),
-            GetItem("waveform"),
-        ]
-
-        return cls(
-            domain=domain,
-            data_prep=Compose(transforms),
-            event_data=event_data,
-            event_metadata=event_metadata,
-            model_metadata=metadata,
-            device=device,
-            data_prep_conditioning=data_prep_conditioning,
-        )
-
-    @classmethod
-    def from_model(
-        cls,
-        model: BasePosteriorModel,
-        event_data: dict,
-        event_metadata: Optional[dict] = None,
-    ) -> "GWSamplerContext":
-        """Build the context from a model: its own metadata and its device. Data
-        preparation is network-bound, so the settings come from `model.metadata`
-        (for a conditional model this equals the base analysis metadata). An
-        unconditional model prepares no data, so no context can be built from one;
-        for the prior/likelihood views alone, use
-        `from_model_metadata(_base_model_metadata(model), ...)`.
-
-        Parameters
-        ----------
-        model : BasePosteriorModel
-            The (conditional) model whose metadata defines the domain and
-            preprocessing.
-        event_data : dict
-            The raw event data (strain + ASDs).
-        event_metadata : dict, optional
-            Per-event metadata.
-
-        Returns
-        -------
-        GWSamplerContext
-        """
-        if model.metadata["train_settings"]["data"].get("unconditional", False):
-            raise ValueError(
-                "An unconditional model has no data preparation, so a context "
-                "cannot be built from it. For the prior/likelihood views, use "
-                "GWSamplerContext.from_model_metadata(_base_model_metadata(model), ...)."
-            )
-        return cls.from_model_metadata(
-            model.metadata,
-            event_data,
-            event_metadata,
-            device=model.device,
-        )
-
-    def derive(
-        self,
-        updates: Optional[dict] = None,
-        use_base_domain: Optional[bool] = None,
-    ) -> "GWSamplerContext":
-        """Derive a context for the same event under a different data representation.
-
-        The derived context shares the event payload (`event_data`, `event_metadata`),
-        the analysis metadata, and the reference time by construction -- parameter
-        meaning is preserved, so importance weights between the two representations
-        remain well-defined. Only the representation changes: an updated duration `T`
-        rebuilds the data domain at `delta_f = 1/T` (and enters waveform generation as
-        `wfg_delta_f`), and `use_base_domain` switches a multibanded likelihood to the
-        undecimated domain. Frequency-range updates act through `event_metadata` (ASD
-        masking in the likelihood) and need no derived state. Caches start fresh.
-
-        Parameters
-        ----------
-        updates : dict, optional
-            Importance-sampling settings updates, as recorded in
-            `Result.importance_sampling_metadata["updates"]`. May contain `T`,
-            `minimum_frequency`, `maximum_frequency` (and non-domain keys, which are
-            ignored here), but no other quantities that define a new domain.
-        use_base_domain : bool, optional
-            For a multibanded domain, evaluate the likelihood on the base
-            (undecimated) frequency domain. `None` keeps this context's setting.
-
-        Returns
-        -------
-        GWSamplerContext
-        """
-        domain = self.domain
-        wfg_delta_f = self.wfg_delta_f
-        if updates and any(
-            k in updates for k in ("minimum_frequency", "maximum_frequency", "T")
-        ):
-            # TODO: Make compatible with MultibandedFrequencyDomain.
-            if isinstance(domain, MultibandedFrequencyDomain):
-                raise NotImplementedError()
-
-            updates = updates.copy()
-            # A duration update is generation-level on both sides: the pipe
-            # regenerates the event data at the new T, and the waveform must be
-            # generated natively at delta_f = 1/T (a resolution change is not a
-            # projection of existing samples). wfg_delta_f persists on the derived
-            # context, reaching every downstream likelihood build and further
-            # derivations.
-            if "T" in updates:
-                updates["delta_f"] = 1.0 / updates["T"]
-                wfg_delta_f = updates["delta_f"]
-                print(
-                    f"Updating waveform generation delta_f from "
-                    f'{self.model_metadata["dataset_settings"]["domain"]["delta_f"]} '
-                    f"to {wfg_delta_f}."
-                )
-
-            domain_dict = domain.domain_dict  # Existing settings
-            domain_dict.update(
-                (k, updates[k]) for k in set(domain_dict).intersection(updates)
-            )
-            print("Rebuilding domain as follows:")
-            print(yaml.dump(domain_dict, default_flow_style=False, sort_keys=False))
-            domain = build_domain(domain_dict)
-
-        return type(self)(
-            domain=domain,
-            data_prep=self._data_prep,
-            event_data=self.event_data,
-            event_metadata=self.event_metadata,
-            model_metadata=self.model_metadata,
-            device=self.device,
-            use_base_domain=(
-                self.use_base_domain if use_base_domain is None else use_base_domain
-            ),
-            wfg_delta_f=wfg_delta_f,
-            data_prep_conditioning=self.data_prep_conditioning,
-        )
-
-    def prepared_data(self, conditioning=None) -> torch.Tensor:
-        """The network-input data representation of this event.
-
-        Without `conditioning`: the single shared representation, computed once
-        and cached. With `conditioning` (the chain columns available to a
-        conditioned factor): row-aligned, one data row per conditioning row.
-        Only the columns named in `data_prep_conditioning` are consumed by the
-        preparation (e.g. the chirp-mass heterodyne proxy, injected under its
-        physical name); the remaining columns condition the network only, and a
-        context that consumes nothing serves the shared representation viewed
-        across the rows. A constant consumed value (a pinned proxy) is prepared
-        once and viewed across the rows; varying values (a sweep) run through
-        the batch-native transform chain in one vectorized pass, uncached -- a
-        caller sweeping more rows than memory allows blocks its own request.
-
-        An event frequency-range update is validated against the training crop
-        license before any preparation.
-
-        Parameters
-        ----------
-        conditioning : dict[str, torch.Tensor], optional
-            The chain conditioning available to the calling factor, one value
-            per row. May contain columns irrelevant to the preparation.
-        """
-        if not self.data_prep_conditioning:
-            if self._prepared is None:
-                self._validate_frequency_range()
-                self._prepared = self._data_prep(self.event_data)
-            if conditioning is None:
-                return self._prepared
-            return self._prepared.expand(_n_rows(conditioning), *self._prepared.shape)
-
-        columns = self._conditioning_columns(conditioning)
-        self._validate_frequency_range()
-        n_rows = _n_rows(columns)
-        if all(torch.all(c == c[0]) for c in columns.values()):
-            # N rows of one pinned value: prepare once, view it across the rows.
-            key = {name: float(c[0]) for name, c in columns.items()}
-            if key != self._prepared_key:
-                self._prepared = self._data_prep({**self.event_data, "parameters": key})
-                self._prepared_key = key
-            return self._prepared.expand(n_rows, *self._prepared.shape)
-        parameters = {name: column.numpy() for name, column in columns.items()}
-        return self._data_prep(
-            {**self._broadcast_event(n_rows), "parameters": parameters}
-        )
-
-    def _conditioning_columns(self, conditioning) -> dict[str, torch.Tensor]:
-        """Collect the conditioning columns the preparation consumes, keyed by
-        their physical names (the `_proxy` suffix names the chain column; the
-        transform chain reads the physical parameter), as float64 (the
-        heterodyne phase is computed in float64, and a float32 chain column
-        must not degrade it)."""
-        conditioning = conditioning or {}
-        columns = {}
-        for name in self.data_prep_conditioning:
-            if name not in conditioning:
-                raise ValueError(
-                    f"This model's data preparation is a function of the chain "
-                    f"conditioning `{name}`, which the caller does not provide "
-                    f"(pass it to prepared_data, e.g. from the chain's pins)."
-                )
-            columns[name[: -len("_proxy")]] = torch.as_tensor(
-                conditioning[name], dtype=torch.float64
-            ).reshape(-1)
-        return columns
-
-    def _broadcast_event(self, n_rows: int) -> dict:
-        """The event arrays broadcast (as read-only views) across `n_rows`
-        rows, forming the batched sample dict for a single transform-chain
-        pass."""
-        return {
-            part: (
-                {k: np.broadcast_to(v, (n_rows, *np.shape(v))) for k, v in data.items()}
-                if isinstance(data, dict)
-                else data
-            )
-            for part, data in self.event_data.items()
-        }
-
-    def _validate_frequency_range(self):
-        """Validate an event frequency-range update: hard bounds against the (base)
-        domain, and narrowing only when the network was trained with random strain
-        cropping covering the requested range. Applies to the network-input view
-        only -- the likelihood view applies the range independently via ASD
-        masking."""
-        update = _frequency_range_update(self.domain, self.event_metadata)
-        if update is None:
-            return
-        domain = getattr(self.domain, "base_domain", self.domain)
-        crop_settings = self.model_metadata["train_settings"]["data"].get(
-            "random_strain_cropping"
-        )
-        _validate_minimum_frequency(
-            update["minimum_frequency"], self.detectors, domain, crop_settings
-        )
-        _validate_maximum_frequency(
-            update["maximum_frequency"], self.detectors, domain, crop_settings
-        )
-
-    @property
-    def prior(self) -> PriorDict:
-        """The static prior over all parameters, built once from the model metadata
-        (intrinsic + extrinsic priors with Dingo defaults).
-
-        This is the event-independent prior fixed at training time. Importance-sampling
-        prior-bound updates and the time / phase split-off for marginalized networks are
-        applied downstream (they depend on the evolving analysis state), not here.
-        """
-        if self._prior is None:
-            data_settings = self.model_metadata["train_settings"]["data"]
-            intrinsic_prior = self.model_metadata["dataset_settings"]["intrinsic_prior"]
-            extrinsic_prior = get_extrinsic_prior_dict(data_settings["extrinsic_prior"])
-            self._prior = build_prior_with_defaults(
-                {**intrinsic_prior, **extrinsic_prior}
-            )
-        return self._prior
-
-    def likelihood(
-        self,
-        time_marginalization_kwargs: Optional[dict] = None,
-        phase_marginalization_kwargs: Optional[dict] = None,
-        calibration_marginalization_kwargs: Optional[dict] = None,
-    ) -> StationaryGaussianGWLikelihood:
-        """
-        Build the exact GW likelihood on this event's data, in physical parameter space.
-
-        The network's standardized, decimated view of the data is `prepared_data()`; the
-        likelihood instead takes the raw event data (`event_data`) and builds its own
-        representation -- decimating to the multibanded domain unless `use_base_domain`,
-        and masking the ASDs to the event's frequency range. Its reference time is the
-        event time (the training-frame right-ascension correction is already applied to the
-        samples), or the training reference time when no event time is set.
-
-        The data representation (the domain, `use_base_domain`, `wfg_delta_f`, the
-        event's frequency range) is context state -- importance-sampling settings
-        updates enter by deriving a new context (`derive`), not as arguments here.
-        Marginalization is per-request and enters as arguments.
-
-        The most recently built likelihood is cached with its arguments: a repeated
-        call with the same arguments returns the shared instance (the synthetic-phase
-        factor requests one per chain chunk), and a call with different arguments
-        builds a replacement. In the exact synthetic-phase mode the consumer assigns
-        a `phase_grid` attribute on the shared instance.
-
-        Parameters
-        ----------
-        time_marginalization_kwargs : dict, optional
-            Analytically marginalize over `geocent_time`. `t_lower` / `t_upper` are
-            filled from the network's (uniform) time prior when not already provided
-            (a caller with an updated prior passes its own bounds). Requires a
-            time-marginalized network.
-        phase_marginalization_kwargs : dict, optional
-            Analytically marginalize over `phase`. Requires a uniform [0, 2 pi) phase prior.
-        calibration_marginalization_kwargs : dict, optional
-            Marginalize over detector calibration uncertainty.
-
-        Returns
-        -------
-        StationaryGaussianGWLikelihood
-        """
-        # Capture the call's arguments for the cache comparison; this must stay the
-        # first statement so locals() holds exactly the arguments. The comparison
-        # happens before the marginalization bounds are filled in below (which is
-        # deterministic given the same arguments). The representation fields join the
-        # key as a guard against off-contract in-place mutation.
-        settings = dict(locals())
-        settings.pop("self")
-        settings["use_base_domain"] = self.use_base_domain
-        settings["wfg_delta_f"] = self.wfg_delta_f
-        if settings == self._likelihood_settings:
-            return self._likelihood
-
-        if self.event_data is None:
-            raise ValueError(
-                "Building the likelihood requires event data (strain + ASDs), "
-                "which this context does not carry."
-            )
-
-        # Marginalizing over a parameter needs the (uniform) prior the network
-        # marginalized over; use it to parameterize the requested marginalization.
-        # Bounds already provided by the caller win: the importance-sampling layer
-        # validates against its evolved prior (prior updates, time/phase
-        # split-offs), which this sample-free context cannot see, and passes
-        # explicit bounds. The fill below covers standalone (chain) use from the
-        # network's static prior.
-        if time_marginalization_kwargs is not None and not (
-            "t_lower" in time_marginalization_kwargs
-            and "t_upper" in time_marginalization_kwargs
-        ):
-            time_prior = self._marginalized_prior("geocent_time")
-            if time_prior is None:
-                raise NotImplementedError(
-                    "Time marginalization is not compatible with a non-marginalized "
-                    "network."
-                )
-            if type(time_prior) != Uniform:
-                raise NotImplementedError(
-                    "Only uniform time prior is supported for time marginalization."
-                )
-            time_marginalization_kwargs = {
-                **time_marginalization_kwargs,
-                "t_lower": time_prior.minimum,
-                "t_upper": time_prior.maximum,
-            }
-        if phase_marginalization_kwargs is not None:
-            phase_prior = self._marginalized_prior("phase")
-            if not (
-                isinstance(phase_prior, Uniform)
-                and (phase_prior._minimum, phase_prior._maximum) == (0, 2 * np.pi)
-            ):
-                raise ValueError(
-                    f"Phase prior should be uniform [0, 2pi) for phase marginalization, "
-                    f"but is {phase_prior}."
-                )
-
-        dataset_settings = self.model_metadata["dataset_settings"]
-        # WaveformGenerator domain -- deliberately the dataset domain WITHOUT any
-        # domain_update, so waveform generation mirrors how the training set was
-        # generated (possibly wider than the network's data domain); the likelihood
-        # projects the generated waveform onto the data domain. An updated
-        # T = 1/delta_f enters here: unlike range masking or decimation, a
-        # resolution change needs samples that do not exist on the old grid, so it
-        # cannot be handled by domain projection and must apply at generation.
-        wfg_domain_dict = dataset_settings["domain"]
-        if self.wfg_delta_f is not None:
-            wfg_domain_dict = {**wfg_domain_dict, "delta_f": self.wfg_delta_f}
-        wfg_domain = build_domain(wfg_domain_dict)
-
-        # Likelihood reference time: the event time (the training-frame RA correction has
-        # already been applied to the samples), falling back to the training reference.
-        if self.event_metadata is not None and "time_event" in self.event_metadata:
-            t_ref = self.event_metadata["time_event"]
-        else:
-            t_ref = self.t_ref
-
-        frequency_update = dict(
-            minimum_frequency=self._frequency("minimum_frequency", self.domain.f_min),
-            maximum_frequency=self._frequency("maximum_frequency", self.domain.f_max),
-        )
-
-        likelihood = StationaryGaussianGWLikelihood(
-            wfg_kwargs=dataset_settings["waveform_generator"],
-            wfg_domain=wfg_domain,
-            data_domain=self.domain,
-            event_data=self.event_data,
-            t_ref=t_ref,
-            time_marginalization_kwargs=time_marginalization_kwargs,
-            phase_marginalization_kwargs=phase_marginalization_kwargs,
-            calibration_marginalization_kwargs=calibration_marginalization_kwargs,
-            use_base_domain=self.use_base_domain,
-            frequency_update=frequency_update,
-        )
-        self._likelihood = likelihood
-        self._likelihood_settings = settings
-        return likelihood
-
-    def _frequency(self, key: str, default: float):
-        """The event's frequency-range override for `key` (min/max), else `default`."""
-        if self.event_metadata is None:
-            return default
-        return self.event_metadata.get(key, default)
-
-    def _marginalized_prior(self, name: str):
-        """The prior over `name` if the network marginalized it (i.e. `name` is not an
-        inference parameter), else `None` -- used to parameterize likelihood
-        marginalization over time / phase."""
-        if (
-            name
-            in self.model_metadata["train_settings"]["data"]["inference_parameters"]
-        ):
-            return None
-        return self.prior.get(name)
-
-
-def _to_numpy(v) -> np.ndarray:
-    """Detach a torch tensor (or coerce anything) to a numpy array."""
-    if torch.is_tensor(v):
-        return v.detach().cpu().numpy()
-    return np.asarray(v)
-
-
-class SyntheticPhaseFactor(Factor):
-    """
-    Reconstruct the coalescence phase for a phase-marginalized network:
-    `q(phase | theta_rest, d)` from the likelihood on a phase grid. The terminal factor of
-    the chain.
-
-    For each incoming `theta_rest` it builds the phase-full likelihood
-    (`context.likelihood()`) and evaluates `log L` on a grid over `[0, 2 pi)`, exploiting
-    that the waveform modes computed once at `phase = 0` each transform as `exp(-i m phase)`
-    -- so the whole grid follows from a single waveform evaluation. The grid is
-    exponentiated into a conditional phase distribution, a uniform floor (weight
-    `uniform_weight`) is added to keep it mass-covering, and one phase is drawn per sample
-    from the interpolated distribution. The returned proposal log-prob
-    `log q(phase | theta_rest, d)` joins the chain's proposal density; importance sampling
-    then targets the phase-full posterior (with the phase prior re-added).
-
-    Two grid modes: `approximation_22_mode=True` assumes a (2, 2)-dominated signal (the
-    whole waveform transforms as `exp(2i phase)`), giving `log L` from the complex overlap
-    `Re[(d | h(phase=0)) exp(2i phase)]`; `False` sums the modes exactly and requires the
-    waveform generator's `spin_conversion_phase = 0`. Note the entry points differ on the
-    default: this factor and `dingo_pipe`'s `PhaseRecoveryDefault` use the exact mode,
-    while `Result.sample_synthetic_phase` defaults to the (2, 2) approximation when the
-    key is omitted.
-    """
-
-    def __init__(
-        self,
-        conditioning: list[str],
-        n_grid: int = 5001,
-        approximation_22_mode: bool = False,
-        uniform_weight: float = 0.01,
-        num_processes: int = 1,
-    ):
-        """
-        Parameters
-        ----------
-        conditioning : list[str]
-            The physical parameters the likelihood needs to generate the waveform
-            (everything the chain has produced except `phase`).
-        n_grid : int, default 5001
-            Number of phase grid points on `[0, 2 pi)`.
-        approximation_22_mode : bool, default False
-            Use the (2, 2)-mode approximation instead of the exact mode sum.
-        uniform_weight : float, default 0.01
-            Weight of the uniform floor added to the phase distribution for mass coverage.
-        num_processes : int, default 1
-            Parallel processes for the per-sample likelihood evaluation and phase sampling.
-        """
-        self.parameters = ["phase"]
-        self.conditioning = list(conditioning)
-        self.n_grid = n_grid
-        self.approximation_22_mode = approximation_22_mode
-        self.uniform_weight = uniform_weight
-        self.num_processes = num_processes
-
-    def sample_and_log_prob(self, num_samples, context, given=None):
-        """Draw one phase per `theta_rest` row (`num_samples` must be 1); return the phases
-        and their proposal log-prob `log q(phase | theta_rest, d)`."""
-        if num_samples != 1:
-            raise ValueError(
-                "Synthetic phase is 1:1; draw one phase per sample (fan_out=1)."
-            )
-        reference = next(iter(given.values()))
-        device = reference.device if torch.is_tensor(reference) else None
-        n = len(reference)
-        logger.info(f"Estimating synthetic phase for {n} samples.")
-        t0 = time.time()
-        phases, phase_posterior = self._phase_profile(given, context)
-        new_phase, log_prob = interpolated_sample_and_log_prob_multi(
-            phases, phase_posterior, self.num_processes
-        )
-        logger.info(f"Done. This took {time.time() - t0:.2f} s.")
-        return (
-            {"phase": torch.as_tensor(new_phase, device=device)},
-            torch.as_tensor(log_prob, device=device),
-        )
-
-    def log_prob(self, theta_i, context, given=None):
-        """Evaluate `log q(phase | theta_rest, d)` at the given phases (re-plug / IS)."""
-        reference = next(iter(given.values()))
-        device = reference.device if torch.is_tensor(reference) else None
-        phases, phase_posterior = self._phase_profile(given, context)
-        log_prob = interpolated_log_prob_multi(
-            phases, phase_posterior, _to_numpy(theta_i["phase"]), self.num_processes
-        )
-        return torch.as_tensor(log_prob, device=device)
-
-    def describe(self) -> dict:
-        return {
-            "step": type(self).__name__,
-            "parameters": list(self.parameters),
-            "conditioning": list(self.conditioning),
-            "n_grid": self.n_grid,
-            "approximation_22_mode": self.approximation_22_mode,
-            "uniform_weight": self.uniform_weight,
-        }
-
-    def _phase_profile(self, given, context):
-        """The phase grid and the mass-covered (un-normalized) phase distribution, one row
-        per sample: evaluate `log L` on the grid, exponentiate (shifted by the per-row
-        max), and add the uniform floor."""
-        theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
-        # The context (possibly derived for importance sampling) carries the
-        # representation; the phase-full likelihood needs no arguments.
-        likelihood = context.likelihood()
-        phases = np.linspace(0, 2 * np.pi, self.n_grid)
-        if self.approximation_22_mode:
-            # Assume the waveform is (2, 2)-dominated (transforms as exp(2i phase)), so the
-            # phase-dependent log-posterior is Re[(d | h(phase=0)) exp(2i phase)].
-            theta = theta.copy()
-            theta["phase"] = 0.0
-            d_inner_h = likelihood.d_inner_h_complex_multi(theta, self.num_processes)
-            phase_log_posterior = np.outer(d_inner_h, np.exp(2j * phases)).real
-        else:
-            # Exact: each mode m contributes exp(-i m phase); needs spin_conversion_phase=0.
-            likelihood.phase_grid = phases
-            phase_log_posterior = apply_func_with_multiprocessing(
-                likelihood.log_likelihood_phase_grid,
-                theta,
-                num_processes=self.num_processes,
-            )
-        phase_posterior = np.exp(
-            phase_log_posterior - np.amax(phase_log_posterior, axis=1, keepdims=True)
-        )
-        # Uniform floor: keep q(phase) > 0 everywhere so importance sampling stays finite.
-        phase_posterior += (
-            phase_posterior.mean(axis=-1, keepdims=True) * self.uniform_weight
-        )
-        return phases, phase_posterior
-
-
-def _build_gnpe_transforms(model: BasePosteriorModel):
-    """Build the time-shift GNPE per-step transforms from a model's metadata: the
-    proxy blur, the per-row time-shift alignment applied before the network, and the
-    post-network geocent-time correction.
-
-    Returns
-    -------
-    transform_pre, transform_post : Compose
-    gnpe_parameters : list[str]
-        The GNPE input parameters (detector times).
-    inference_parameters : list[str]
-    kernel : PriorDict
-        The proxy perturbation kernel.
-    gnpe_transform : GNPECoalescenceTimes
-        The blur transform itself, shared so the kernel factor can call `sample_proxies`.
-    """
-    meta = _base_model_metadata(model)
-    data_settings = meta["train_settings"]["data"]
-    ifo_list = InterferometerList(data_settings["detectors"])
-    domain = build_domain(meta["dataset_settings"]["domain"])
-    if "domain_update" in data_settings:
-        domain.update(data_settings["domain_update"])
-
-    gnpe_time_settings = data_settings.get("gnpe_time_shifts")
-    if not gnpe_time_settings:
-        raise NotImplementedError(
-            "Only time-shift GNPE is supported here so far (no gnpe_chirp / gnpe_phase)."
-        )
-
-    gnpe_transform = GNPECoalescenceTimes(
-        ifo_list,
-        gnpe_time_settings["kernel"],
-        gnpe_time_settings["exact_equiv"],
-        inference=True,
-    )
-    transform_pre = [
-        RenameKey("data", "waveform"),
-        gnpe_transform,
-        TimeShiftStrain(ifo_list, domain),
-        SelectStandardizeRepackageParameters(
-            {"context_parameters": data_settings["context_parameters"]},
-            data_settings["standardization"],
-            device=model.device,
-        ),
-        RenameKey("waveform", "data"),
-    ]
-    gnpe_parameters: list[str] = []
-    kernel = PriorDict()
-    for transform in transform_pre:
-        if isinstance(transform, GNPEBase):
-            gnpe_parameters += transform.input_parameter_names
-            for k, v in transform.kernel.items():
-                kernel[k] = v
-
-    inference_parameters = data_settings["inference_parameters"]
-    transform_post = [
-        SelectStandardizeRepackageParameters(
-            {"inference_parameters": inference_parameters},
-            data_settings["standardization"],
-            inverse=True,
-            as_type="dict",
-        ),
-        PostCorrectGeocentTime(),
-        CopyToExtrinsicParameters(
-            "ra", "dec", "geocent_time", "chirp_mass", "mass_ratio", "phase"
-        ),
-        GetDetectorTimes(ifo_list, data_settings["ref_time"]),
-    ]
-    return (
-        Compose(transform_pre),
-        Compose(transform_post),
-        gnpe_parameters,
-        inference_parameters,
-        kernel,
-        gnpe_transform,
-    )
-
-
-class GNPEKernelFactor(Factor):
-    """
-    The GNPE perturbation kernel `p(theta_hat | theta)` as a non-network factor.
-
-    `theta` are the detector coalescence times; the kernel adds a bounded perturbation to
-    each, giving the proxies `theta_hat` the main network conditions on. The parameter
-    block is the proxies, the conditioning is the detector times.
-    `sample_and_log_prob` blurs the times into proxies (the proxy update of a Gibbs
-    sweep); `log_prob` returns the kernel density `log p(theta_hat | theta)` at the
-    proxies and the detector times. One proxy per detector-time row.
-    """
-
-    def __init__(self, gnpe_transform: GNPEBase, gnpe_parameters: list[str]):
-        """
-        Parameters
-        ----------
-        gnpe_transform : GNPEBase
-            The blur transform supplying the kernel and `sample_proxies`.
-        gnpe_parameters : list[str]
-            The detector-time parameters perturbed into proxies.
-        """
-        self.gnpe = gnpe_transform
-        self.gnpe_parameters = gnpe_parameters
-        self.parameters = [p + "_proxy" for p in gnpe_parameters]
-        self.conditioning = list(gnpe_parameters)
-        self.kernel = gnpe_transform.kernel
-
-    @classmethod
-    def from_model(cls, model: BasePosteriorModel) -> "GNPEKernelFactor":
-        """Build from the main model's metadata (the kernel / blur transform)."""
-        _, _, gnpe_parameters, _, _, gnpe_transform = _build_gnpe_transforms(model)
-        return cls(gnpe_transform, gnpe_parameters)
-
-    def sample_and_log_prob(self, num_samples, context, given=None):
-        """Blur the conditioning detector times into proxies; `num_samples` must be 1
-        (GNPE is 1:1). Returns the proxies and their kernel log-prob."""
-        if num_samples != 1:
-            raise ValueError("GNPE proxy is 1:1; use fan_out=1.")
-        times = {k: given[k] for k in self.gnpe_parameters}
-        proxies = self.gnpe.sample_proxies(times)
-        return proxies, self.log_prob(proxies, context, given)
-
-    def log_prob(self, theta_i, context, given=None):
-        """`log p(theta_hat | theta)` from the kernel, at the proxies (`theta_i`) and
-        the detector times (`given`).
-
-        The kernel is a bilby `PriorDict` -- the same object that samples the blur --
-        so the density is evaluated in numpy (converting each side first: the times
-        and proxies may live on different devices) and returned on the detector
-        times' device."""
-        reference = next(iter(given.values()))
-        device = reference.device if torch.is_tensor(reference) else None
-        diffs = {
-            k: _to_numpy(given[k]) - _to_numpy(theta_i[f"{k}_proxy"])
-            for k in self.kernel.keys()
-        }
-        return torch.as_tensor(
-            self.kernel.ln_prob(diffs, axis=0), dtype=torch.float32, device=device
-        )
-
-
-class GNPEFlowFactor(Factor):
-    """
-    The GNPE main network `q(theta | theta_hat, d)` as a factor.
-
-    Conditions on the detector-time proxies from `GNPEKernelFactor`: it shifts each
-    detector's strain by the corresponding proxy time (standardizing the network input),
-    samples the network, and recomputes the detector times from the sampled sky position
-    and geocent time. The proxies are supplied, so no blurring happens here.
-
-    The single network factor in either GNPE mode: cycled by a `GibbsBlock` for
-    multi-iteration GNPE, or a `ChainComposer` factor for single-step GNPE. The recomputed
-    detector times are emitted as extra columns (`produces`): the next Gibbs iteration
-    blurs them into fresh proxies, and single-step GNPE evaluates the kernel correction at
-    them. One sample per proxy row.
-    """
-
-    def __init__(
-        self,
-        model: BasePosteriorModel,
-        transform_pre: Compose,
-        transform_post: Compose,
-        gnpe_parameters: list[str],
-        parameters: list[str],
-        aliases: Optional[dict[str, str]] = None,
-    ):
-        """
-        Parameters
-        ----------
-        model : BasePosteriorModel
-            The GNPE main network.
-        transform_pre : Compose
-            Per-iteration pre-network transforms (proxy bookkeeping, time shift,
-            standardization).
-        transform_post : Compose
-            Post-network transforms (de-standardize, recompute detector times).
-        gnpe_parameters : list[str]
-            The detector-time parameters.
-        parameters : list[str]
-            The network's trained inference parameters.
-        aliases : dict[str, str], optional
-            Trained-name to exposed-name map (e.g. `{"ra": "ra@t_ref"}`).
-        """
-        self.model = model
-        self.transform_pre = transform_pre
-        self.transform_post = transform_post
-        self.gnpe_parameters = gnpe_parameters
-        self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
-        self.aliases = aliases or {}
-        self._net_parameters = parameters
-        self.parameters = [self.aliases.get(p, p) for p in parameters]
-        self.conditioning = list(self.proxy_parameters)
-        # For log_prob: the sampling path de-standardizes (and corrects the log-prob)
-        # inside transform_post, but evaluating at a point needs the forward map too.
-        # The model's own standardization (network-bound, like FlowFactor's).
-        std = model.metadata["train_settings"]["data"]["standardization"]
-        self.standardization = Standardization(std["mean"], std["std"])
-
-    @property
-    def produces(self) -> list[str]:
-        """Emitted columns: the inference block plus the recomputed detector times."""
-        return self.parameters + self.gnpe_parameters
-
-    @classmethod
-    def from_model(
-        cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
-    ) -> "GNPEFlowFactor":
-        """Build the GNPE per-iteration transforms from the main model's metadata.
-
-        Parameters
-        ----------
-        model : BasePosteriorModel
-            The GNPE main model.
-        aliases : dict[str, str], optional
-            Trained-name to canonical-name map (e.g. `{"ra": "ra@t_ref"}`).
-
-        Returns
-        -------
-        GNPEFlowFactor
-        """
-        pre, post, gnpe_parameters, inference_parameters, _, _ = _build_gnpe_transforms(
-            model
-        )
-        return cls(
-            model, pre, post, gnpe_parameters, inference_parameters, aliases=aliases
-        )
-
-    def sample_and_log_prob(self, num_samples, context, given=None):
-        """Sample one parameter set per proxy row; `num_samples` must be 1 (GNPE is 1:1).
-        Returns theta plus the recomputed detector times, and the network log-prob."""
-        if num_samples != 1:
-            raise ValueError("GNPE is 1:1; draw one sample per proxy (fan_out=1).")
-        proxies = {p: given[p] for p in self.proxy_parameters}
-        n_rows = _n_rows(proxies)
-        x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = context.prepared_data().clone()
-        x["data"] = d.expand(n_rows, *d.shape)
-        x = self.transform_pre(x)
-        self.model.network.eval()
-        with torch.no_grad():
-            if "context_parameters" in x:
-                y, log_prob = self.model.sample_and_log_prob(
-                    x["data"], x["context_parameters"]
-                )
-            else:
-                y, log_prob = self.model.sample_and_log_prob(x["data"])
-        # sample_and_log_prob(num_samples=1) adds a singleton dim; the batch is the proxy
-        # rows.
-        x["parameters"] = y.squeeze(1)
-        x["log_prob"] = log_prob.squeeze(1)
-        x = self.transform_post(x)
-        params = dict(x["parameters"])
-        # Expose trained names under their canonical aliases (e.g. ra -> ra@t_ref).
-        params = {self.aliases.get(k, k): v for k, v in params.items()}
-        # Surface the recomputed detector times: the next Gibbs iteration's input, and the
-        # evaluation point for GNPEKernelFactor's importance-sampling correction.
-        for k in self.gnpe_parameters:
-            params[k] = x["extrinsic_parameters"][k]
-        return params, x["log_prob"]
-
-    def log_prob(self, theta_i, context, given=None):
-        """Evaluate the network density `log q(theta | theta_hat, d)` in physical space
-        at given `theta_i` (exposed / aliased names), one row per proxy row in `given`.
-        Applies the same per-row view of the data as sampling (the shared
-        representation time-shifted by the proxies, the conditioning standardized),
-        then scores the standardized parameters under the network."""
-        proxies = {p: given[p] for p in self.proxy_parameters}
-        n_rows = _n_rows(proxies)
-        x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
-        d = context.prepared_data().clone()
-        x["data"] = d.expand(n_rows, *d.shape)
-        x = self.transform_pre(x)
-        theta_net = {
-            net: theta_i[self.aliases.get(net, net)] for net in self._net_parameters
-        }
-        # Mirror transform_post: sampling shifts geocent_time by the preferred proxy
-        # after the network (PostCorrectGeocentTime, using the extrinsic geocent_time
-        # set up by the GNPE transform), so score the network in its own output frame
-        # by applying the inverse correction first.
-        y = {
-            "parameters": dict(theta_net),
-            "extrinsic_parameters": dict(x["extrinsic_parameters"]),
-        }
-        theta_net = PostCorrectGeocentTime(inverse=True)(y)["parameters"]
-        z = self.standardization.standardize(theta_net, self._net_parameters)
-        self.model.network.eval()
-        with torch.no_grad():
-            if "context_parameters" in x:
-                log_prob = self.model.log_prob(z, x["data"], x["context_parameters"])
-            else:
-                log_prob = self.model.log_prob(z, x["data"])
-        return log_prob + self.standardization.log_det(self._net_parameters)
-
-
-class RAReparam(Reparametrization):
-    """
-    Rotate right ascension from the network's training reference frame (`ra@t_ref`) to the
-    event frame (`ra`).
-
-    The network is trained at a fixed reference time; an event at a different GPS time needs
-    the sky rotated by the sidereal-time difference. This is a measure-preserving shift
-    modulo 2*pi (`log_det = 0`), so it contributes nothing to the density. `forward`
-    produces the event-frame `ra`, `inverse` recovers `ra@t_ref`. The sidereal
-    correction is read from the shared context (`t_ref` and the event time).
-
-    The modulo makes the map a bijection on the circle, while the flow's density lives
-    on the real line: a sample drawn outside `[0, 2 pi)` is wrapped, so `inverse`
-    recovers its principal-branch representative and a re-evaluated `log_prob` refers
-    to that branch. Only tail samples outside the bounded `ra` prior are affected.
-    """
-
-    def __init__(self):
-        self.conditioning = ["ra@t_ref"]
-        self.parameters = ["ra"]
-
-    @staticmethod
-    def _correction(context) -> float:
-        """Sidereal-time difference (event minus reference) in radians; 0 when the event
-        time is unset or equal to the reference time."""
-        event_metadata = context.event_metadata
-        t_event = None if event_metadata is None else event_metadata.get("time_event")
-        t_ref = context.t_ref
-        if t_event is None or t_event == t_ref:
-            return 0.0
-        longitude_event = Time(t_event, format="gps", scale="utc").sidereal_time(
-            "apparent", "greenwich"
-        )
-        longitude_reference = Time(t_ref, format="gps", scale="utc").sidereal_time(
-            "apparent", "greenwich"
-        )
-        return (longitude_event - longitude_reference).rad
-
-    def forward(self, given, context):
-        correction = self._correction(context)
-        if correction == 0.0:
-            return {"ra": given["ra@t_ref"]}
-        # ra is a bounded angle -> float32 is plenty. The correction is a difference of
-        # absolute GPS times, so compute it in float64, but store the wrapped angle float32.
-        ra = (given["ra@t_ref"].double() + correction) % (2 * np.pi)
-        return {"ra": ra.float()}
-
-    def inverse(self, params, context, given=None):
-        correction = self._correction(context)
-        if correction == 0.0:
-            return {"ra@t_ref": params["ra"]}
-        ra_tref = (params["ra"].double() - correction) % (2 * np.pi)
-        return {"ra@t_ref": ra_tref.float()}
-
-
-class SpinConventionReparam(Reparametrization):
-    """
-    Relabel the precessing-spin angles between Dingo's internal spin-phase
-    convention and the physical (Bilby) one.
-
-    Dingo fixes the spin-conversion phase (usually to 0) so that the Cartesian
-    spins decouple from the coalescence phase; the whole density / likelihood /
-    synthetic-phase pipeline requires that convention, so stored samples keep the
-    plain names `theta_jn` / `phi_jl` in the *network* convention throughout. The
-    physical convention (spin conversion at the sample's own phase) is what Bilby
-    and PESummary mean by the same names, so the relabel happens at the export
-    boundary -- this class is its single home. Only `theta_jn` and `phi_jl`
-    change; the conversion phase and reference frequency are read from the model
-    metadata, and a model trained without a fixed conversion phase (`None`)
-    relabels to the identity.
-
-    Unlike `RAReparam` no marked intermediate name is needed: the two conventions
-    never coexist in one table -- each world's plain names denote its own
-    convention, and this bijection is the boundary crossing.
-
-    Exporting a finished weighted sample set needs no Jacobian (proposal, prior,
-    and likelihood transform together), which is how `to_physical` is used. As a
-    chain `Step` the bijection is *not* measure-preserving in the flat
-    `(theta_jn, phi_jl)` coordinates: it rotates the line of sight rigidly about
-    the orbital angular momentum, preserving the spherical measure
-    `sin(theta_jn) dtheta dphi`, so
-    `log_det = log sin(theta_jn) - log sin(theta_jn')` (verified numerically
-    against finite differences through the LAL conversion), and `inverse`
-    rebuilds the network convention from the physical one using the invariant
-    conditioning the reverse fold supplies.
-    """
-
-    def __init__(self, num_processes: int = 1):
-        """
-        Parameters
-        ----------
-        num_processes : int, default 1
-            Parallel processes for the per-sample LAL spin conversion.
-        """
-        self.parameters = ["theta_jn", "phi_jl"]
-        self.conditioning = [
-            "theta_jn",
-            "phi_jl",
-            "phase",
-            "chirp_mass",
-            "mass_ratio",
-            "a_1",
-            "a_2",
-            "tilt_1",
-            "tilt_2",
-            "phi_12",
-        ]
-        self.num_processes = num_processes
-
-    @property
-    def consumes(self) -> list[str]:
-        # The bijection overwrites theta_jn / phi_jl in place; the remaining
-        # conditioning (phase, masses, tilts, ...) is read-only and must stay in
-        # the chain -- the default conditioning-minus-parameters would drop it.
-        return []
-
-    def log_det(self, given, context):
-        """`log|det J|` of `forward`, per row. The map preserves the spherical
-        measure, so the flat-coordinate Jacobian is
-        `sin(theta_jn) / sin(theta_jn')` -- verified numerically against finite
-        differences through the LAL conversion in
-        `test_jacobian_matches_sin_ratio` (agreement ~1e-9)."""
-        converted = self.forward(given, context)
-        return self._log_det(given["theta_jn"], converted["theta_jn"])
-
-    @staticmethod
-    def _log_det(theta_jn_in, theta_jn_out):
-        # Compute in double (the LAL conversion is double precision anyway),
-        # return in the input dtype: a reparametrization preserves the chain's
-        # dtype rather than promoting the summed log_prob.
-        log_det = torch.log(torch.sin(theta_jn_in.double())) - torch.log(
-            torch.sin(theta_jn_out.double())
-        )
-        return log_det.to(theta_jn_in.dtype)
-
-    def sample_and_log_prob(self, num_samples, context, given=None):
-        """Apply `forward`; contribute `-log|det J|`. Overridden to share the
-        single LAL conversion between the transform and its Jacobian (the base
-        implementation would convert twice)."""
-        if num_samples != 1:
-            raise ValueError("A reparametrization is 1:1; use fan_out=1.")
-        out = self.forward(given, context)
-        return out, -self._log_det(given["theta_jn"], out["theta_jn"])
-
-    @staticmethod
-    def _model_convention(model_metadata: dict) -> tuple[float, Optional[float]]:
-        """The reference frequency and spin-conversion phase the model trained with."""
-        wfg_settings = model_metadata["dataset_settings"]["waveform_generator"]
-        return wfg_settings["f_ref"], wfg_settings.get("spin_conversion_phase")
-
-    def to_physical(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFrame:
-        """Relabel samples from the model's convention to the physical (Bilby) one."""
-        f_ref, sc_phase = self._model_convention(model_metadata)
-        return change_spin_conversion_phase(
-            samples, f_ref, sc_phase, None, num_processes=self.num_processes
-        )
-
-    def to_network(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFrame:
-        """Relabel samples from the physical (Bilby) convention to the model's,
-        e.g. to ingest external posteriors for comparison."""
-        f_ref, sc_phase = self._model_convention(model_metadata)
-        return change_spin_conversion_phase(
-            samples, f_ref, None, sc_phase, num_processes=self.num_processes
-        )
-
-    def forward(self, given, context):
-        # The conversion runs in double; the outputs return in the input dtype
-        # and device (cf. RAReparam: compute in float64, store the chain dtype).
-        reference = given["theta_jn"]
-        theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
-        converted = self.to_physical(theta, context.model_metadata)
-        return {
-            k: torch.as_tensor(converted[k].to_numpy()).to(
-                dtype=reference.dtype, device=reference.device
-            )
-            for k in self.parameters
-        }
-
-    def inverse(self, params, context, given=None):
-        # The physical -> network direction also needs the invariant
-        # conditioning (phase, masses, tilts), which the reverse fold supplies
-        # as `given`.
-        if given is None:
-            raise ValueError(
-                "The spin-convention inverse needs the conditioning block "
-                "(phase, masses, tilts); pass it as `given`, or convert "
-                "DataFrames with to_network()."
-            )
-        reference = params["theta_jn"]
-        rows = {**given, **params}
-        theta = pd.DataFrame({k: _to_numpy(v) for k, v in rows.items()})
-        converted = self.to_network(theta, context.model_metadata)
-        return {
-            k: torch.as_tensor(converted[k].to_numpy()).to(
-                dtype=reference.dtype, device=reference.device
-            )
-            for k in self.parameters
-        }
-
-
-class GNPEKernelCorrection(TargetCorrection):
-    """
-    The single-step GNPE kernel correction as a target-side chain step.
-
-    Emits `delta_log_prob_target = log p(theta_hat | theta)` -- the GNPE kernel evaluated
-    at the proxies and the detector times recomputed from theta -- for importance sampling
-    on the joint proposal `q(theta, theta_hat | d)`. Contributes 0 to the proposal
-    density and consumes the intermediate detector times.
-    """
-
-    def __init__(self, kernel_factor: GNPEKernelFactor):
-        self.kernel_factor = kernel_factor
-        self.parameters = ["delta_log_prob_target"]
-        self.conditioning = list(kernel_factor.parameters) + list(
-            kernel_factor.gnpe_parameters
-        )
-        self.consumes = list(kernel_factor.gnpe_parameters)
-
-    def correction(self, given, context):
-        proxies = {p: given[p] for p in self.kernel_factor.parameters}
-        times = {k: given[k] for k in self.kernel_factor.gnpe_parameters}
-        correction = self.kernel_factor.log_prob(proxies, context, times)
-        return {"delta_log_prob_target": correction}
-
-
-def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
-    """The RA frame alias (`ra` -> `ra@t_ref`), applied only when the model infers
-    `ra`; paired with an `RAReparam` step that maps it back to the event frame."""
-    return {"ra": "ra@t_ref"} if "ra" in inference_parameters else {}
-
-
-def _proxy_offset_steps(
-    inference_parameters: list[str], context_parameters: list[str]
-) -> list:
-    """One offset reconstruction (`X = delta_X + X_proxy`) per `delta_X` the
-    network infers whose proxy it conditions on."""
-    return [
-        ProxyOffsetReparam(p[len("delta_") :])
-        for p in inference_parameters
-        if p.startswith("delta_")
-        and p[len("delta_") :] + "_proxy" in context_parameters
-    ]
-
-
-def _ra_reparam_steps(inference_parameters: list[str]) -> list:
-    """The `RAReparam` step, appended to a chain only when the model infers `ra`."""
-    return [RAReparam()] if "ra" in inference_parameters else []
-
-
-def _delta_prior_steps(prior, inference_parameters: list[str]) -> list:
-    """Delta-prior parameters the chain does not produce, as a single `DeltaFactor` step
-    (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0).
-
-    Parameters
-    ----------
-    prior : PriorDict
-        The static prior (`GWSamplerContext.prior`); its delta-function entries that are
-        not inference parameters become the pinned constants.
-    inference_parameters : list of str
-        The inferred parameter names.
-    """
-    fixed = {
-        k: p.peak
-        for k, p in prior.items()
-        if isinstance(p, DeltaFunction) and k not in inference_parameters
-    }
-    return [DeltaFactor(fixed)] if fixed else []
-
-
-def _assert_consistent_gnpe_data_prep(init_model, main_model):
-    """Assert the init and main GNPE models agree on the data-preprocessing view.
-
-    Multi-iteration GNPE shares one `GWSamplerContext` (built from the main model)
-    between the init and main factors, so both read the same `prepared_data()` and
-    reference time. That is only valid when the two models agree on everything that
-    determines those: the domain, the detectors, and the reference time. Raises
-    `ValueError` on any mismatch.
-
-    Parameters
-    ----------
-    init_model, main_model : BasePosteriorModel
-        The GNPE init and main networks.
-    """
-    init = _base_model_metadata(init_model)
-    main = _base_model_metadata(main_model)
-    fields = {
-        "domain": (
-            init["dataset_settings"]["domain"],
-            main["dataset_settings"]["domain"],
-        ),
-        "domain_update": (
-            init["train_settings"]["data"].get("domain_update"),
-            main["train_settings"]["data"].get("domain_update"),
-        ),
-        "detectors": (
-            init["train_settings"]["data"]["detectors"],
-            main["train_settings"]["data"]["detectors"],
-        ),
-        "ref_time": (
-            init["train_settings"]["data"]["ref_time"],
-            main["train_settings"]["data"]["ref_time"],
-        ),
-    }
-    mismatched = {k: (i, m) for k, (i, m) in fields.items() if i != m}
-    if mismatched:
-        details = "; ".join(
-            f"{k}: init={i!r} vs main={m!r}" for k, (i, m) in mismatched.items()
-        )
-        raise ValueError(
-            f"GNPE init and main models disagree on the data-preprocessing view "
-            f"({details}). They share one context, so they must agree on the domain, "
-            f"detectors, and reference time."
-        )
-
-
-class GWComposedSampler(ComposedSampler):
-    """
-    GW builder and exporter over the generic `ComposedSampler` runner. The `from_*`
-    constructors assemble the chain for plain NPE, multi-iteration GNPE, or single-step
-    GNPE from model metadata; `to_result` exports the samples to a gw `Result`. All
-    GW-specific processing (RA frame, fixed parameters, kernel correction) is expressed as
-    chain steps, so there is no post-processing.
-    """
-
-    def __init__(
-        self,
-        composer: ChainComposer,
-        context: GWSamplerContext,
-        metadata: dict,
-        inference_parameters: list[str],
-    ):
-        """
-        Parameters
-        ----------
-        composer : ChainComposer
-            The assembled chain of steps.
-        context : GWSamplerContext
-            Per-event shared state.
-        metadata : dict
-            Model metadata, carried through to the exported `Result`.
-        inference_parameters : list[str]
-            The inferred parameter names.
-        """
-        super().__init__(composer, context)
-        self.metadata = metadata
-        self.inference_parameters = inference_parameters
-        # Extra provenance merged into settings["sampler"] by to_result -- e.g. the
-        # pipe records model checkpoint paths and the density-recovery recipe.
-        # Literal-only values (the settings dict round-trips through str/literal_eval).
-        self.provenance_extra: dict = {}
-
-    def sampler_provenance(self) -> dict:
-        """Provenance of how the samples were made, stored as `settings["sampler"]`
-        in the exported `Result`: the executed chain in order (one descriptor per
-        step, via `Step.describe()`), plus anything in `provenance_extra`. The
-        block is purely a record -- nothing consumes it at load time -- and the
-        `version` field allows future consumers (e.g. chain reconstruction from
-        file) to evolve the format safely."""
-        return {
-            "version": 1,
-            "implementation": "composed",
-            "chain": [step.describe() for step in self.composer.steps],
-            **copy.deepcopy(self.provenance_extra),
-        }
-
-    @classmethod
-    def from_model(
-        cls,
-        model: BasePosteriorModel,
-        event_data: dict,
-        event_metadata: Optional[dict] = None,
-        fixed_context_parameters: Optional[dict] = None,
-    ) -> "GWComposedSampler":
-        """Build a single-network GW sampler from a model and event data.
-
-        For a plain NPE model the chain is the flow (exposing `ra` as `ra@t_ref`)
-        followed by an `RAReparam` to the event frame. A model with
-        `context_parameters` (e.g. the DINGO-BNS chirp-mass prior conditioning
-        with a fixed sky position) requires `fixed_context_parameters` pinning
-        all of them: the chain is then rooted in a `DeltaFactor` of the pins,
-        the flow conditions on them, and each inferred offset `delta_X` with a
-        pinned proxy is reconstructed by a `ProxyOffsetReparam`
-        (`X = delta_X + X_proxy`). Proxies that parameterize the data
-        preparation (the chirp-mass heterodyne) reach it through the chain:
-        the conditioned flow passes them to `prepared_data`, which consumes
-        and caches on them -- the pins have a single owner, the chain's root.
-
-        Parameters
-        ----------
-        model : BasePosteriorModel
-            The model.
-        event_data : dict
-            The raw event data (strain + ASDs).
-        event_metadata : dict, optional
-            Per-event metadata.
-        fixed_context_parameters : dict, optional
-            Pinned values for the model's `context_parameters`, e.g.
-            `{"chirp_mass_proxy": 1.1975, "ra": 3.446, "dec": -0.408}`.
-
-        Returns
-        -------
-        GWComposedSampler
-        """
-        context = GWSamplerContext.from_model(model, event_data, event_metadata)
-        metadata = _base_model_metadata(model)
-        data_settings = metadata["train_settings"]["data"]
-        inference_parameters = data_settings["inference_parameters"]
-        factor = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
-        steps = [factor]
-        context_parameters = data_settings.get("context_parameters") or []
-        if context_parameters:
-            if set(fixed_context_parameters or {}) != set(context_parameters):
-                raise ValueError(
-                    f"The model conditions on {context_parameters}; provide "
-                    f"fixed_context_parameters with exactly these keys (for "
-                    f"iterated proxies, use the GNPE builders instead)."
-                )
-            steps = [DeltaFactor(fixed_context_parameters)] + steps
-        steps += _proxy_offset_steps(inference_parameters, context_parameters)
-        steps += _ra_reparam_steps(inference_parameters)
-        steps += _delta_prior_steps(context.prior, inference_parameters)
-        return cls(
-            composer=ChainComposer(steps),
-            context=context,
-            metadata=metadata,
-            inference_parameters=inference_parameters,
-        )
-
-    @classmethod
-    def from_gnpe_models(
-        cls,
-        init_model: BasePosteriorModel,
-        main_model: BasePosteriorModel,
-        event_data: dict,
-        event_metadata: Optional[dict] = None,
-        num_iterations: int = 30,
-    ) -> "GWComposedSampler":
-        """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
-        init model's data preprocessing, an init `FlowFactor` to seed, and a single
-        `GibbsBlock` step -- cycling the GNPE kernel and main-network factors -- in a
-        `ChainComposer`, then an `RAReparam` to the event frame. Returns samples without
-        a log_prob (Gibbs breaks density access).
-
-        Parameters
-        ----------
-        init_model : BasePosteriorModel
-            The init network (detector times); seeds the Gibbs loop and defines the data
-            preprocessing.
-        main_model : BasePosteriorModel
-            The GNPE main network.
-        event_data : dict
-            The raw event data (strain + ASDs).
-        event_metadata : dict, optional
-            Per-event metadata.
-        num_iterations : int, default 30
-            Number of Gibbs sweeps.
-
-        Returns
-        -------
-        GWComposedSampler
-        """
-        _assert_consistent_gnpe_data_prep(init_model, main_model)
-        # Build the context from the main model: it owns the analysis (likelihood,
-        # prior, inference parameters). The init model shares the data domain and
-        # preprocessing (asserted above), so prepared_data() is identical either way.
-        context = GWSamplerContext.from_model(main_model, event_data, event_metadata)
-        metadata = _base_model_metadata(main_model)
-        inference_parameters = metadata["train_settings"]["data"][
-            "inference_parameters"
-        ]
-        init_factor = FlowFactor.from_model(init_model)
-        kernel_factor = GNPEKernelFactor.from_model(main_model)
-        flow_factor = GNPEFlowFactor.from_model(
-            main_model, aliases=_ra_aliases(inference_parameters)
-        )
-        gibbs = GibbsBlock(init_factor, [kernel_factor, flow_factor], num_iterations)
-        steps = (
-            [gibbs]
-            + _ra_reparam_steps(inference_parameters)
-            + _delta_prior_steps(context.prior, inference_parameters)
-        )
-        return cls(
-            ChainComposer(steps),
-            context,
-            metadata,
-            inference_parameters,
-        )
-
-    @classmethod
-    def from_singlestep_gnpe(
-        cls,
-        main_model: BasePosteriorModel,
-        proxy_source: Factor,
-        event_data: dict,
-        event_metadata: Optional[dict] = None,
-    ) -> "GWComposedSampler":
-        """Build a single-step (density-preserving) time-GNPE sampler: a `ChainComposer`
-        of `[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]`. The chain is
-        autoregressive, so log_prob is preserved, and `GNPEKernelCorrection` emits the
-        `delta_log_prob_target` correction that importance sampling adds to the target.
-
-        Parameters
-        ----------
-        main_model : BasePosteriorModel
-            The GNPE main network.
-        proxy_source : Factor
-            Supplies the detector-time proxies -- a `DeltaFactor` for prior conditioning
-            (BNS), or an unconditional NDE for density recovery.
-        event_data : dict
-            The raw event data (strain + ASDs).
-        event_metadata : dict, optional
-            Per-event metadata.
-
-        Returns
-        -------
-        GWComposedSampler
-        """
-        context = GWSamplerContext.from_model(main_model, event_data, event_metadata)
-        metadata = _base_model_metadata(main_model)
-        inference_parameters = metadata["train_settings"]["data"][
-            "inference_parameters"
-        ]
-        flow_factor = GNPEFlowFactor.from_model(
-            main_model, aliases=_ra_aliases(inference_parameters)
-        )
-        kernel_factor = GNPEKernelFactor.from_model(main_model)
-        steps = (
-            [proxy_source, flow_factor, GNPEKernelCorrection(kernel_factor)]
-            + _ra_reparam_steps(inference_parameters)
-            + _delta_prior_steps(context.prior, inference_parameters)
-        )
-        return cls(ChainComposer(steps), context, metadata, inference_parameters)
-
-    def to_result(self):
-        """Export to a gw `Result` (samples + raw event data + metadata), so the
-        existing post-processing pipeline -- synthetic phase, importance sampling,
-        evidence, plotting -- runs on the factorized sampler's output unchanged.
-
-        The raw event-data dict (`GWSamplerContext.event_data`) is stored as the
-        `Result` context (serialized), and the live `GWSamplerContext` is passed as
-        `sampler_context` so `Result` can pull the prior (and, later, the likelihood)
-        from it rather than rebuilding them from metadata.
-        """
-        from dingo.gw.result import Result
-
-        settings = copy.deepcopy(self.metadata)
-        settings["sampler"] = self.sampler_provenance()
-        data_dict = {
-            "samples": self.samples,
-            "context": self.context.event_data,
-            "event_metadata": self.context.event_metadata,
-            "importance_sampling_metadata": None,
-            "log_evidence": None,
-            "log_noise_evidence": None,
-            "settings": settings,
-        }
-        return Result(dictionary=data_dict, sampler_context=self.context)
-
-    def to_hdf5(self, label="result", outdir="."):
-        """Export via `to_result` and save to `<outdir>/<label>.hdf5`."""
-        result = self.to_result()
-        Path(outdir).mkdir(parents=True, exist_ok=True)
-        result.to_file(file_name=Path(outdir, label + ".hdf5"))
diff --git a/dingo/gw/inference/sampler.py b/dingo/gw/inference/sampler.py
new file mode 100644
index 000000000..bbc4791a9
--- /dev/null
+++ b/dingo/gw/inference/sampler.py
@@ -0,0 +1,366 @@
+"""The composed gravitational-wave sampler and its chain builders."""
+
+from __future__ import annotations
+import copy
+import logging
+from pathlib import Path
+from typing import Optional
+from bilby.core.prior import DeltaFunction
+from dingo.core.factors import (
+    ChainComposer,
+    ComposedSampler,
+    DeltaFactor,
+    Factor,
+    FlowFactor,
+    GibbsBlock,
+    ProxyOffsetReparam,
+    _base_model_metadata,
+)
+from dingo.core.posterior_models import BasePosteriorModel
+from dingo.gw.inference.context import GWSamplerContext
+from dingo.gw.inference.steps import (
+    GNPEFlowFactor,
+    GNPEKernelCorrection,
+    GNPEKernelFactor,
+    RAReparam,
+)
+
+
+def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
+    """The RA frame alias (`ra` -> `ra@t_ref`), applied only when the model infers
+    `ra`; paired with an `RAReparam` step that maps it back to the event frame."""
+    return {"ra": "ra@t_ref"} if "ra" in inference_parameters else {}
+
+
+def _proxy_offset_steps(
+    inference_parameters: list[str], context_parameters: list[str]
+) -> list:
+    """One offset reconstruction (`X = delta_X + X_proxy`) per `delta_X` the
+    network infers whose proxy it conditions on."""
+    return [
+        ProxyOffsetReparam(p[len("delta_") :])
+        for p in inference_parameters
+        if p.startswith("delta_")
+        and p[len("delta_") :] + "_proxy" in context_parameters
+    ]
+
+
+def _ra_reparam_steps(inference_parameters: list[str]) -> list:
+    """The `RAReparam` step, appended to a chain only when the model infers `ra`."""
+    return [RAReparam()] if "ra" in inference_parameters else []
+
+
+def _delta_prior_steps(prior, inference_parameters: list[str]) -> list:
+    """Delta-prior parameters the chain does not produce, as a single `DeltaFactor` step
+    (or none). These are pinned constants (e.g. an aligned-spin component fixed to 0).
+
+    Parameters
+    ----------
+    prior : PriorDict
+        The static prior (`GWSamplerContext.prior`); its delta-function entries that are
+        not inference parameters become the pinned constants.
+    inference_parameters : list of str
+        The inferred parameter names.
+    """
+    fixed = {
+        k: p.peak
+        for k, p in prior.items()
+        if isinstance(p, DeltaFunction) and k not in inference_parameters
+    }
+    return [DeltaFactor(fixed)] if fixed else []
+
+
+def _assert_consistent_gnpe_data_prep(init_model, main_model):
+    """Assert the init and main GNPE models agree on the data-preprocessing view.
+
+    Multi-iteration GNPE shares one `GWSamplerContext` (built from the main model)
+    between the init and main factors, so both read the same `prepared_data()` and
+    reference time. That is only valid when the two models agree on everything that
+    determines those: the domain, the detectors, and the reference time. Raises
+    `ValueError` on any mismatch.
+
+    Parameters
+    ----------
+    init_model, main_model : BasePosteriorModel
+        The GNPE init and main networks.
+    """
+    init = _base_model_metadata(init_model)
+    main = _base_model_metadata(main_model)
+    fields = {
+        "domain": (
+            init["dataset_settings"]["domain"],
+            main["dataset_settings"]["domain"],
+        ),
+        "domain_update": (
+            init["train_settings"]["data"].get("domain_update"),
+            main["train_settings"]["data"].get("domain_update"),
+        ),
+        "detectors": (
+            init["train_settings"]["data"]["detectors"],
+            main["train_settings"]["data"]["detectors"],
+        ),
+        "ref_time": (
+            init["train_settings"]["data"]["ref_time"],
+            main["train_settings"]["data"]["ref_time"],
+        ),
+    }
+    mismatched = {k: (i, m) for k, (i, m) in fields.items() if i != m}
+    if mismatched:
+        details = "; ".join(
+            f"{k}: init={i!r} vs main={m!r}" for k, (i, m) in mismatched.items()
+        )
+        raise ValueError(
+            f"GNPE init and main models disagree on the data-preprocessing view "
+            f"({details}). They share one context, so they must agree on the domain, "
+            f"detectors, and reference time."
+        )
+
+
+class GWComposedSampler(ComposedSampler):
+    """
+    GW builder and exporter over the generic `ComposedSampler` runner. The `from_*`
+    constructors assemble the chain for plain NPE, multi-iteration GNPE, or single-step
+    GNPE from model metadata; `to_result` exports the samples to a gw `Result`. All
+    GW-specific processing (RA frame, fixed parameters, kernel correction) is expressed as
+    chain steps, so there is no post-processing.
+    """
+
+    def __init__(
+        self,
+        composer: ChainComposer,
+        context: GWSamplerContext,
+        metadata: dict,
+        inference_parameters: list[str],
+    ):
+        """
+        Parameters
+        ----------
+        composer : ChainComposer
+            The assembled chain of steps.
+        context : GWSamplerContext
+            Per-event shared state.
+        metadata : dict
+            Model metadata, carried through to the exported `Result`.
+        inference_parameters : list[str]
+            The inferred parameter names.
+        """
+        super().__init__(composer, context)
+        self.metadata = metadata
+        self.inference_parameters = inference_parameters
+        # Extra provenance merged into settings["sampler"] by to_result -- e.g. the
+        # pipe records model checkpoint paths and the density-recovery recipe.
+        # Literal-only values (the settings dict round-trips through str/literal_eval).
+        self.provenance_extra: dict = {}
+
+    def sampler_provenance(self) -> dict:
+        """Provenance of how the samples were made, stored as `settings["sampler"]`
+        in the exported `Result`: the executed chain in order (one descriptor per
+        step, via `Step.describe()`), plus anything in `provenance_extra`. The
+        block is purely a record -- nothing consumes it at load time -- and the
+        `version` field allows future consumers (e.g. chain reconstruction from
+        file) to evolve the format safely."""
+        return {
+            "version": 1,
+            "implementation": "composed",
+            "chain": [step.describe() for step in self.composer.steps],
+            **copy.deepcopy(self.provenance_extra),
+        }
+
+    @classmethod
+    def from_model(
+        cls,
+        model: BasePosteriorModel,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+        fixed_context_parameters: Optional[dict] = None,
+    ) -> "GWComposedSampler":
+        """Build a single-network GW sampler from a model and event data.
+
+        For a plain NPE model the chain is the flow (exposing `ra` as `ra@t_ref`)
+        followed by an `RAReparam` to the event frame. A model with
+        `context_parameters` (e.g. the DINGO-BNS chirp-mass prior conditioning
+        with a fixed sky position) requires `fixed_context_parameters` pinning
+        all of them: the chain is then rooted in a `DeltaFactor` of the pins,
+        the flow conditions on them, and each inferred offset `delta_X` with a
+        pinned proxy is reconstructed by a `ProxyOffsetReparam`
+        (`X = delta_X + X_proxy`). Proxies that parameterize the data
+        preparation (the chirp-mass heterodyne) reach it through the chain:
+        the conditioned flow passes them to `prepared_data`, which consumes
+        and caches on them -- the pins have a single owner, the chain's root.
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The model.
+        event_data : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+        fixed_context_parameters : dict, optional
+            Pinned values for the model's `context_parameters`, e.g.
+            `{"chirp_mass_proxy": 1.1975, "ra": 3.446, "dec": -0.408}`.
+
+        Returns
+        -------
+        GWComposedSampler
+        """
+        context = GWSamplerContext.from_model(model, event_data, event_metadata)
+        metadata = _base_model_metadata(model)
+        data_settings = metadata["train_settings"]["data"]
+        inference_parameters = data_settings["inference_parameters"]
+        factor = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
+        steps = [factor]
+        context_parameters = data_settings.get("context_parameters") or []
+        if context_parameters:
+            if set(fixed_context_parameters or {}) != set(context_parameters):
+                raise ValueError(
+                    f"The model conditions on {context_parameters}; provide "
+                    f"fixed_context_parameters with exactly these keys (for "
+                    f"iterated proxies, use the GNPE builders instead)."
+                )
+            steps = [DeltaFactor(fixed_context_parameters)] + steps
+        steps += _proxy_offset_steps(inference_parameters, context_parameters)
+        steps += _ra_reparam_steps(inference_parameters)
+        steps += _delta_prior_steps(context.prior, inference_parameters)
+        return cls(
+            composer=ChainComposer(steps),
+            context=context,
+            metadata=metadata,
+            inference_parameters=inference_parameters,
+        )
+
+    @classmethod
+    def from_gnpe_models(
+        cls,
+        init_model: BasePosteriorModel,
+        main_model: BasePosteriorModel,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+        num_iterations: int = 30,
+    ) -> "GWComposedSampler":
+        """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
+        init model's data preprocessing, an init `FlowFactor` to seed, and a single
+        `GibbsBlock` step -- cycling the GNPE kernel and main-network factors -- in a
+        `ChainComposer`, then an `RAReparam` to the event frame. Returns samples without
+        a log_prob (Gibbs breaks density access).
+
+        Parameters
+        ----------
+        init_model : BasePosteriorModel
+            The init network (detector times); seeds the Gibbs loop and defines the data
+            preprocessing.
+        main_model : BasePosteriorModel
+            The GNPE main network.
+        event_data : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+        num_iterations : int, default 30
+            Number of Gibbs sweeps.
+
+        Returns
+        -------
+        GWComposedSampler
+        """
+        _assert_consistent_gnpe_data_prep(init_model, main_model)
+        # Build the context from the main model: it owns the analysis (likelihood,
+        # prior, inference parameters). The init model shares the data domain and
+        # preprocessing (asserted above), so prepared_data() is identical either way.
+        context = GWSamplerContext.from_model(main_model, event_data, event_metadata)
+        metadata = _base_model_metadata(main_model)
+        inference_parameters = metadata["train_settings"]["data"][
+            "inference_parameters"
+        ]
+        init_factor = FlowFactor.from_model(init_model)
+        kernel_factor = GNPEKernelFactor.from_model(main_model)
+        flow_factor = GNPEFlowFactor.from_model(
+            main_model, aliases=_ra_aliases(inference_parameters)
+        )
+        gibbs = GibbsBlock(init_factor, [kernel_factor, flow_factor], num_iterations)
+        steps = (
+            [gibbs]
+            + _ra_reparam_steps(inference_parameters)
+            + _delta_prior_steps(context.prior, inference_parameters)
+        )
+        return cls(
+            ChainComposer(steps),
+            context,
+            metadata,
+            inference_parameters,
+        )
+
+    @classmethod
+    def from_singlestep_gnpe(
+        cls,
+        main_model: BasePosteriorModel,
+        proxy_source: Factor,
+        event_data: dict,
+        event_metadata: Optional[dict] = None,
+    ) -> "GWComposedSampler":
+        """Build a single-step (density-preserving) time-GNPE sampler: a `ChainComposer`
+        of `[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]`. The chain is
+        autoregressive, so log_prob is preserved, and `GNPEKernelCorrection` emits the
+        `delta_log_prob_target` correction that importance sampling adds to the target.
+
+        Parameters
+        ----------
+        main_model : BasePosteriorModel
+            The GNPE main network.
+        proxy_source : Factor
+            Supplies the detector-time proxies -- a `DeltaFactor` for prior conditioning
+            (BNS), or an unconditional NDE for density recovery.
+        event_data : dict
+            The raw event data (strain + ASDs).
+        event_metadata : dict, optional
+            Per-event metadata.
+
+        Returns
+        -------
+        GWComposedSampler
+        """
+        context = GWSamplerContext.from_model(main_model, event_data, event_metadata)
+        metadata = _base_model_metadata(main_model)
+        inference_parameters = metadata["train_settings"]["data"][
+            "inference_parameters"
+        ]
+        flow_factor = GNPEFlowFactor.from_model(
+            main_model, aliases=_ra_aliases(inference_parameters)
+        )
+        kernel_factor = GNPEKernelFactor.from_model(main_model)
+        steps = (
+            [proxy_source, flow_factor, GNPEKernelCorrection(kernel_factor)]
+            + _ra_reparam_steps(inference_parameters)
+            + _delta_prior_steps(context.prior, inference_parameters)
+        )
+        return cls(ChainComposer(steps), context, metadata, inference_parameters)
+
+    def to_result(self):
+        """Export to a gw `Result` (samples + raw event data + metadata), so the
+        existing post-processing pipeline -- synthetic phase, importance sampling,
+        evidence, plotting -- runs on the factorized sampler's output unchanged.
+
+        The raw event-data dict (`GWSamplerContext.event_data`) is stored as the
+        `Result` context (serialized), and the live `GWSamplerContext` is passed as
+        `sampler_context` so `Result` can pull the prior (and, later, the likelihood)
+        from it rather than rebuilding them from metadata.
+        """
+        from dingo.gw.result import Result
+
+        settings = copy.deepcopy(self.metadata)
+        settings["sampler"] = self.sampler_provenance()
+        data_dict = {
+            "samples": self.samples,
+            "context": self.context.event_data,
+            "event_metadata": self.context.event_metadata,
+            "importance_sampling_metadata": None,
+            "log_evidence": None,
+            "log_noise_evidence": None,
+            "settings": settings,
+        }
+        return Result(dictionary=data_dict, sampler_context=self.context)
+
+    def to_hdf5(self, label="result", outdir="."):
+        """Export via `to_result` and save to `<outdir>/<label>.hdf5`."""
+        result = self.to_result()
+        Path(outdir).mkdir(parents=True, exist_ok=True)
+        result.to_file(file_name=Path(outdir, label + ".hdf5"))
diff --git a/dingo/gw/inference/scan.py b/dingo/gw/inference/scan.py
index 2f2377b99..da1449395 100644
--- a/dingo/gw/inference/scan.py
+++ b/dingo/gw/inference/scan.py
@@ -20,10 +20,15 @@
 import pandas as pd
 from bilby.gw.prior import BBHPriorDict
 
-from dingo.core.factors import ChainComposer, FlowFactor, SampleTableFactor, Stage
-from dingo.gw.inference.factors import (
-    GWSamplerContext,
+from dingo.core.factors import (
+    ChainComposer,
+    FlowFactor,
+    SampleTableFactor,
+    Stage,
     _base_model_metadata,
+)
+from dingo.gw.inference.context import GWSamplerContext
+from dingo.gw.inference.sampler import (
     _delta_prior_steps,
     _proxy_offset_steps,
     _ra_aliases,
diff --git a/dingo/gw/inference/steps.py b/dingo/gw/inference/steps.py
new file mode 100644
index 000000000..003a3e684
--- /dev/null
+++ b/dingo/gw/inference/steps.py
@@ -0,0 +1,697 @@
+"""Gravitational-wave chain steps: the GNPE factors, the synthetic-phase
+factor, and the coordinate reparametrizations."""
+
+from __future__ import annotations
+import logging
+import time
+from typing import Optional
+import numpy as np
+import pandas as pd
+import torch
+from astropy.time import Time
+from bilby.core.prior import PriorDict
+from bilby.gw.detector import InterferometerList
+from torchvision.transforms import Compose
+from dingo.core.density import (
+    interpolated_log_prob_multi,
+    interpolated_sample_and_log_prob_multi,
+)
+from dingo.core.factors import (
+    Factor,
+    Reparametrization,
+    Standardization,
+    TargetCorrection,
+    _base_model_metadata,
+    _n_rows,
+)
+from dingo.core.multiprocessing import apply_func_with_multiprocessing
+from dingo.core.posterior_models import BasePosteriorModel
+from dingo.core.transforms import RenameKey
+from dingo.gw.conversion import change_spin_conversion_phase
+from dingo.gw.domains import build_domain
+from dingo.gw.transforms import (
+    CopyToExtrinsicParameters,
+    GetDetectorTimes,
+    GNPEBase,
+    GNPECoalescenceTimes,
+    PostCorrectGeocentTime,
+    SelectStandardizeRepackageParameters,
+    TimeShiftStrain,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _to_numpy(v) -> np.ndarray:
+    """Detach a torch tensor (or coerce anything) to a numpy array."""
+    if torch.is_tensor(v):
+        return v.detach().cpu().numpy()
+    return np.asarray(v)
+
+
+class SyntheticPhaseFactor(Factor):
+    """
+    Reconstruct the coalescence phase for a phase-marginalized network:
+    `q(phase | theta_rest, d)` from the likelihood on a phase grid. The terminal factor of
+    the chain.
+
+    For each incoming `theta_rest` it builds the phase-full likelihood
+    (`context.likelihood()`) and evaluates `log L` on a grid over `[0, 2 pi)`, exploiting
+    that the waveform modes computed once at `phase = 0` each transform as `exp(-i m phase)`
+    -- so the whole grid follows from a single waveform evaluation. The grid is
+    exponentiated into a conditional phase distribution, a uniform floor (weight
+    `uniform_weight`) is added to keep it mass-covering, and one phase is drawn per sample
+    from the interpolated distribution. The returned proposal log-prob
+    `log q(phase | theta_rest, d)` joins the chain's proposal density; importance sampling
+    then targets the phase-full posterior (with the phase prior re-added).
+
+    Two grid modes: `approximation_22_mode=True` assumes a (2, 2)-dominated signal (the
+    whole waveform transforms as `exp(2i phase)`), giving `log L` from the complex overlap
+    `Re[(d | h(phase=0)) exp(2i phase)]`; `False` sums the modes exactly and requires the
+    waveform generator's `spin_conversion_phase = 0`. Note the entry points differ on the
+    default: this factor and `dingo_pipe`'s `PhaseRecoveryDefault` use the exact mode,
+    while `Result.sample_synthetic_phase` defaults to the (2, 2) approximation when the
+    key is omitted.
+    """
+
+    def __init__(
+        self,
+        conditioning: list[str],
+        n_grid: int = 5001,
+        approximation_22_mode: bool = False,
+        uniform_weight: float = 0.01,
+        num_processes: int = 1,
+    ):
+        """
+        Parameters
+        ----------
+        conditioning : list[str]
+            The physical parameters the likelihood needs to generate the waveform
+            (everything the chain has produced except `phase`).
+        n_grid : int, default 5001
+            Number of phase grid points on `[0, 2 pi)`.
+        approximation_22_mode : bool, default False
+            Use the (2, 2)-mode approximation instead of the exact mode sum.
+        uniform_weight : float, default 0.01
+            Weight of the uniform floor added to the phase distribution for mass coverage.
+        num_processes : int, default 1
+            Parallel processes for the per-sample likelihood evaluation and phase sampling.
+        """
+        self.parameters = ["phase"]
+        self.conditioning = list(conditioning)
+        self.n_grid = n_grid
+        self.approximation_22_mode = approximation_22_mode
+        self.uniform_weight = uniform_weight
+        self.num_processes = num_processes
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        """Draw one phase per `theta_rest` row (`num_samples` must be 1); return the phases
+        and their proposal log-prob `log q(phase | theta_rest, d)`."""
+        if num_samples != 1:
+            raise ValueError(
+                "Synthetic phase is 1:1; draw one phase per sample (fan_out=1)."
+            )
+        reference = next(iter(given.values()))
+        device = reference.device if torch.is_tensor(reference) else None
+        n = len(reference)
+        logger.info(f"Estimating synthetic phase for {n} samples.")
+        t0 = time.time()
+        phases, phase_posterior = self._phase_profile(given, context)
+        new_phase, log_prob = interpolated_sample_and_log_prob_multi(
+            phases, phase_posterior, self.num_processes
+        )
+        logger.info(f"Done. This took {time.time() - t0:.2f} s.")
+        return (
+            {"phase": torch.as_tensor(new_phase, device=device)},
+            torch.as_tensor(log_prob, device=device),
+        )
+
+    def log_prob(self, theta_i, context, given=None):
+        """Evaluate `log q(phase | theta_rest, d)` at the given phases (re-plug / IS)."""
+        reference = next(iter(given.values()))
+        device = reference.device if torch.is_tensor(reference) else None
+        phases, phase_posterior = self._phase_profile(given, context)
+        log_prob = interpolated_log_prob_multi(
+            phases, phase_posterior, _to_numpy(theta_i["phase"]), self.num_processes
+        )
+        return torch.as_tensor(log_prob, device=device)
+
+    def describe(self) -> dict:
+        return {
+            "step": type(self).__name__,
+            "parameters": list(self.parameters),
+            "conditioning": list(self.conditioning),
+            "n_grid": self.n_grid,
+            "approximation_22_mode": self.approximation_22_mode,
+            "uniform_weight": self.uniform_weight,
+        }
+
+    def _phase_profile(self, given, context):
+        """The phase grid and the mass-covered (un-normalized) phase distribution, one row
+        per sample: evaluate `log L` on the grid, exponentiate (shifted by the per-row
+        max), and add the uniform floor."""
+        theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
+        # The context (possibly derived for importance sampling) carries the
+        # representation; the phase-full likelihood needs no arguments.
+        likelihood = context.likelihood()
+        phases = np.linspace(0, 2 * np.pi, self.n_grid)
+        if self.approximation_22_mode:
+            # Assume the waveform is (2, 2)-dominated (transforms as exp(2i phase)), so the
+            # phase-dependent log-posterior is Re[(d | h(phase=0)) exp(2i phase)].
+            theta = theta.copy()
+            theta["phase"] = 0.0
+            d_inner_h = likelihood.d_inner_h_complex_multi(theta, self.num_processes)
+            phase_log_posterior = np.outer(d_inner_h, np.exp(2j * phases)).real
+        else:
+            # Exact: each mode m contributes exp(-i m phase); needs spin_conversion_phase=0.
+            likelihood.phase_grid = phases
+            phase_log_posterior = apply_func_with_multiprocessing(
+                likelihood.log_likelihood_phase_grid,
+                theta,
+                num_processes=self.num_processes,
+            )
+        phase_posterior = np.exp(
+            phase_log_posterior - np.amax(phase_log_posterior, axis=1, keepdims=True)
+        )
+        # Uniform floor: keep q(phase) > 0 everywhere so importance sampling stays finite.
+        phase_posterior += (
+            phase_posterior.mean(axis=-1, keepdims=True) * self.uniform_weight
+        )
+        return phases, phase_posterior
+
+
+def _build_gnpe_transforms(model: BasePosteriorModel):
+    """Build the time-shift GNPE per-step transforms from a model's metadata: the
+    proxy blur, the per-row time-shift alignment applied before the network, and the
+    post-network geocent-time correction.
+
+    Returns
+    -------
+    transform_pre, transform_post : Compose
+    gnpe_parameters : list[str]
+        The GNPE input parameters (detector times).
+    inference_parameters : list[str]
+    kernel : PriorDict
+        The proxy perturbation kernel.
+    gnpe_transform : GNPECoalescenceTimes
+        The blur transform itself, shared so the kernel factor can call `sample_proxies`.
+    """
+    meta = _base_model_metadata(model)
+    data_settings = meta["train_settings"]["data"]
+    ifo_list = InterferometerList(data_settings["detectors"])
+    domain = build_domain(meta["dataset_settings"]["domain"])
+    if "domain_update" in data_settings:
+        domain.update(data_settings["domain_update"])
+
+    gnpe_time_settings = data_settings.get("gnpe_time_shifts")
+    if not gnpe_time_settings:
+        raise NotImplementedError(
+            "Only time-shift GNPE is supported here so far (no gnpe_chirp / gnpe_phase)."
+        )
+
+    gnpe_transform = GNPECoalescenceTimes(
+        ifo_list,
+        gnpe_time_settings["kernel"],
+        gnpe_time_settings["exact_equiv"],
+        inference=True,
+    )
+    transform_pre = [
+        RenameKey("data", "waveform"),
+        gnpe_transform,
+        TimeShiftStrain(ifo_list, domain),
+        SelectStandardizeRepackageParameters(
+            {"context_parameters": data_settings["context_parameters"]},
+            data_settings["standardization"],
+            device=model.device,
+        ),
+        RenameKey("waveform", "data"),
+    ]
+    gnpe_parameters: list[str] = []
+    kernel = PriorDict()
+    for transform in transform_pre:
+        if isinstance(transform, GNPEBase):
+            gnpe_parameters += transform.input_parameter_names
+            for k, v in transform.kernel.items():
+                kernel[k] = v
+
+    inference_parameters = data_settings["inference_parameters"]
+    transform_post = [
+        SelectStandardizeRepackageParameters(
+            {"inference_parameters": inference_parameters},
+            data_settings["standardization"],
+            inverse=True,
+            as_type="dict",
+        ),
+        PostCorrectGeocentTime(),
+        CopyToExtrinsicParameters(
+            "ra", "dec", "geocent_time", "chirp_mass", "mass_ratio", "phase"
+        ),
+        GetDetectorTimes(ifo_list, data_settings["ref_time"]),
+    ]
+    return (
+        Compose(transform_pre),
+        Compose(transform_post),
+        gnpe_parameters,
+        inference_parameters,
+        kernel,
+        gnpe_transform,
+    )
+
+
+class GNPEKernelFactor(Factor):
+    """
+    The GNPE perturbation kernel `p(theta_hat | theta)` as a non-network factor.
+
+    `theta` are the detector coalescence times; the kernel adds a bounded perturbation to
+    each, giving the proxies `theta_hat` the main network conditions on. The parameter
+    block is the proxies, the conditioning is the detector times.
+    `sample_and_log_prob` blurs the times into proxies (the proxy update of a Gibbs
+    sweep); `log_prob` returns the kernel density `log p(theta_hat | theta)` at the
+    proxies and the detector times. One proxy per detector-time row.
+    """
+
+    def __init__(self, gnpe_transform: GNPEBase, gnpe_parameters: list[str]):
+        """
+        Parameters
+        ----------
+        gnpe_transform : GNPEBase
+            The blur transform supplying the kernel and `sample_proxies`.
+        gnpe_parameters : list[str]
+            The detector-time parameters perturbed into proxies.
+        """
+        self.gnpe = gnpe_transform
+        self.gnpe_parameters = gnpe_parameters
+        self.parameters = [p + "_proxy" for p in gnpe_parameters]
+        self.conditioning = list(gnpe_parameters)
+        self.kernel = gnpe_transform.kernel
+
+    @classmethod
+    def from_model(cls, model: BasePosteriorModel) -> "GNPEKernelFactor":
+        """Build from the main model's metadata (the kernel / blur transform)."""
+        _, _, gnpe_parameters, _, _, gnpe_transform = _build_gnpe_transforms(model)
+        return cls(gnpe_transform, gnpe_parameters)
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        """Blur the conditioning detector times into proxies; `num_samples` must be 1
+        (GNPE is 1:1). Returns the proxies and their kernel log-prob."""
+        if num_samples != 1:
+            raise ValueError("GNPE proxy is 1:1; use fan_out=1.")
+        times = {k: given[k] for k in self.gnpe_parameters}
+        proxies = self.gnpe.sample_proxies(times)
+        return proxies, self.log_prob(proxies, context, given)
+
+    def log_prob(self, theta_i, context, given=None):
+        """`log p(theta_hat | theta)` from the kernel, at the proxies (`theta_i`) and
+        the detector times (`given`).
+
+        The kernel is a bilby `PriorDict` -- the same object that samples the blur --
+        so the density is evaluated in numpy (converting each side first: the times
+        and proxies may live on different devices) and returned on the detector
+        times' device."""
+        reference = next(iter(given.values()))
+        device = reference.device if torch.is_tensor(reference) else None
+        diffs = {
+            k: _to_numpy(given[k]) - _to_numpy(theta_i[f"{k}_proxy"])
+            for k in self.kernel.keys()
+        }
+        return torch.as_tensor(
+            self.kernel.ln_prob(diffs, axis=0), dtype=torch.float32, device=device
+        )
+
+
+class GNPEFlowFactor(Factor):
+    """
+    The GNPE main network `q(theta | theta_hat, d)` as a factor.
+
+    Conditions on the detector-time proxies from `GNPEKernelFactor`: it shifts each
+    detector's strain by the corresponding proxy time (standardizing the network input),
+    samples the network, and recomputes the detector times from the sampled sky position
+    and geocent time. The proxies are supplied, so no blurring happens here.
+
+    The single network factor in either GNPE mode: cycled by a `GibbsBlock` for
+    multi-iteration GNPE, or a `ChainComposer` factor for single-step GNPE. The recomputed
+    detector times are emitted as extra columns (`produces`): the next Gibbs iteration
+    blurs them into fresh proxies, and single-step GNPE evaluates the kernel correction at
+    them. One sample per proxy row.
+    """
+
+    def __init__(
+        self,
+        model: BasePosteriorModel,
+        transform_pre: Compose,
+        transform_post: Compose,
+        gnpe_parameters: list[str],
+        parameters: list[str],
+        aliases: Optional[dict[str, str]] = None,
+    ):
+        """
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The GNPE main network.
+        transform_pre : Compose
+            Per-iteration pre-network transforms (proxy bookkeeping, time shift,
+            standardization).
+        transform_post : Compose
+            Post-network transforms (de-standardize, recompute detector times).
+        gnpe_parameters : list[str]
+            The detector-time parameters.
+        parameters : list[str]
+            The network's trained inference parameters.
+        aliases : dict[str, str], optional
+            Trained-name to exposed-name map (e.g. `{"ra": "ra@t_ref"}`).
+        """
+        self.model = model
+        self.transform_pre = transform_pre
+        self.transform_post = transform_post
+        self.gnpe_parameters = gnpe_parameters
+        self.proxy_parameters = [p + "_proxy" for p in gnpe_parameters]
+        self.aliases = aliases or {}
+        self._net_parameters = parameters
+        self.parameters = [self.aliases.get(p, p) for p in parameters]
+        self.conditioning = list(self.proxy_parameters)
+        # For log_prob: the sampling path de-standardizes (and corrects the log-prob)
+        # inside transform_post, but evaluating at a point needs the forward map too.
+        # The model's own standardization (network-bound, like FlowFactor's).
+        std = model.metadata["train_settings"]["data"]["standardization"]
+        self.standardization = Standardization(std["mean"], std["std"])
+
+    @property
+    def produces(self) -> list[str]:
+        """Emitted columns: the inference block plus the recomputed detector times."""
+        return self.parameters + self.gnpe_parameters
+
+    @classmethod
+    def from_model(
+        cls, model: BasePosteriorModel, aliases: Optional[dict[str, str]] = None
+    ) -> "GNPEFlowFactor":
+        """Build the GNPE per-iteration transforms from the main model's metadata.
+
+        Parameters
+        ----------
+        model : BasePosteriorModel
+            The GNPE main model.
+        aliases : dict[str, str], optional
+            Trained-name to canonical-name map (e.g. `{"ra": "ra@t_ref"}`).
+
+        Returns
+        -------
+        GNPEFlowFactor
+        """
+        pre, post, gnpe_parameters, inference_parameters, _, _ = _build_gnpe_transforms(
+            model
+        )
+        return cls(
+            model, pre, post, gnpe_parameters, inference_parameters, aliases=aliases
+        )
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        """Sample one parameter set per proxy row; `num_samples` must be 1 (GNPE is 1:1).
+        Returns theta plus the recomputed detector times, and the network log-prob."""
+        if num_samples != 1:
+            raise ValueError("GNPE is 1:1; draw one sample per proxy (fan_out=1).")
+        proxies = {p: given[p] for p in self.proxy_parameters}
+        n_rows = _n_rows(proxies)
+        x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
+        d = context.prepared_data().clone()
+        x["data"] = d.expand(n_rows, *d.shape)
+        x = self.transform_pre(x)
+        self.model.network.eval()
+        with torch.no_grad():
+            if "context_parameters" in x:
+                y, log_prob = self.model.sample_and_log_prob(
+                    x["data"], x["context_parameters"]
+                )
+            else:
+                y, log_prob = self.model.sample_and_log_prob(x["data"])
+        # sample_and_log_prob(num_samples=1) adds a singleton dim; the batch is the proxy
+        # rows.
+        x["parameters"] = y.squeeze(1)
+        x["log_prob"] = log_prob.squeeze(1)
+        x = self.transform_post(x)
+        params = dict(x["parameters"])
+        # Expose trained names under their canonical aliases (e.g. ra -> ra@t_ref).
+        params = {self.aliases.get(k, k): v for k, v in params.items()}
+        # Surface the recomputed detector times: the next Gibbs iteration's input, and the
+        # evaluation point for GNPEKernelFactor's importance-sampling correction.
+        for k in self.gnpe_parameters:
+            params[k] = x["extrinsic_parameters"][k]
+        return params, x["log_prob"]
+
+    def log_prob(self, theta_i, context, given=None):
+        """Evaluate the network density `log q(theta | theta_hat, d)` in physical space
+        at given `theta_i` (exposed / aliased names), one row per proxy row in `given`.
+        Applies the same per-row view of the data as sampling (the shared
+        representation time-shifted by the proxies, the conditioning standardized),
+        then scores the standardized parameters under the network."""
+        proxies = {p: given[p] for p in self.proxy_parameters}
+        n_rows = _n_rows(proxies)
+        x = {"extrinsic_parameters": dict(proxies), "parameters": {}}
+        d = context.prepared_data().clone()
+        x["data"] = d.expand(n_rows, *d.shape)
+        x = self.transform_pre(x)
+        theta_net = {
+            net: theta_i[self.aliases.get(net, net)] for net in self._net_parameters
+        }
+        # Mirror transform_post: sampling shifts geocent_time by the preferred proxy
+        # after the network (PostCorrectGeocentTime, using the extrinsic geocent_time
+        # set up by the GNPE transform), so score the network in its own output frame
+        # by applying the inverse correction first.
+        y = {
+            "parameters": dict(theta_net),
+            "extrinsic_parameters": dict(x["extrinsic_parameters"]),
+        }
+        theta_net = PostCorrectGeocentTime(inverse=True)(y)["parameters"]
+        z = self.standardization.standardize(theta_net, self._net_parameters)
+        self.model.network.eval()
+        with torch.no_grad():
+            if "context_parameters" in x:
+                log_prob = self.model.log_prob(z, x["data"], x["context_parameters"])
+            else:
+                log_prob = self.model.log_prob(z, x["data"])
+        return log_prob + self.standardization.log_det(self._net_parameters)
+
+
+class RAReparam(Reparametrization):
+    """
+    Rotate right ascension from the network's training reference frame (`ra@t_ref`) to the
+    event frame (`ra`).
+
+    The network is trained at a fixed reference time; an event at a different GPS time needs
+    the sky rotated by the sidereal-time difference. This is a measure-preserving shift
+    modulo 2*pi (`log_det = 0`), so it contributes nothing to the density. `forward`
+    produces the event-frame `ra`, `inverse` recovers `ra@t_ref`. The sidereal
+    correction is read from the shared context (`t_ref` and the event time).
+
+    The modulo makes the map a bijection on the circle, while the flow's density lives
+    on the real line: a sample drawn outside `[0, 2 pi)` is wrapped, so `inverse`
+    recovers its principal-branch representative and a re-evaluated `log_prob` refers
+    to that branch. Only tail samples outside the bounded `ra` prior are affected.
+    """
+
+    def __init__(self):
+        self.conditioning = ["ra@t_ref"]
+        self.parameters = ["ra"]
+
+    @staticmethod
+    def _correction(context) -> float:
+        """Sidereal-time difference (event minus reference) in radians; 0 when the event
+        time is unset or equal to the reference time."""
+        event_metadata = context.event_metadata
+        t_event = None if event_metadata is None else event_metadata.get("time_event")
+        t_ref = context.t_ref
+        if t_event is None or t_event == t_ref:
+            return 0.0
+        longitude_event = Time(t_event, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        longitude_reference = Time(t_ref, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        return (longitude_event - longitude_reference).rad
+
+    def forward(self, given, context):
+        correction = self._correction(context)
+        if correction == 0.0:
+            return {"ra": given["ra@t_ref"]}
+        # ra is a bounded angle -> float32 is plenty. The correction is a difference of
+        # absolute GPS times, so compute it in float64, but store the wrapped angle float32.
+        ra = (given["ra@t_ref"].double() + correction) % (2 * np.pi)
+        return {"ra": ra.float()}
+
+    def inverse(self, params, context, given=None):
+        correction = self._correction(context)
+        if correction == 0.0:
+            return {"ra@t_ref": params["ra"]}
+        ra_tref = (params["ra"].double() - correction) % (2 * np.pi)
+        return {"ra@t_ref": ra_tref.float()}
+
+
+class SpinConventionReparam(Reparametrization):
+    """
+    Relabel the precessing-spin angles between Dingo's internal spin-phase
+    convention and the physical (Bilby) one.
+
+    Dingo fixes the spin-conversion phase (usually to 0) so that the Cartesian
+    spins decouple from the coalescence phase; the whole density / likelihood /
+    synthetic-phase pipeline requires that convention, so stored samples keep the
+    plain names `theta_jn` / `phi_jl` in the *network* convention throughout. The
+    physical convention (spin conversion at the sample's own phase) is what Bilby
+    and PESummary mean by the same names, so the relabel happens at the export
+    boundary -- this class is its single home. Only `theta_jn` and `phi_jl`
+    change; the conversion phase and reference frequency are read from the model
+    metadata, and a model trained without a fixed conversion phase (`None`)
+    relabels to the identity.
+
+    Unlike `RAReparam` no marked intermediate name is needed: the two conventions
+    never coexist in one table -- each world's plain names denote its own
+    convention, and this bijection is the boundary crossing.
+
+    Exporting a finished weighted sample set needs no Jacobian (proposal, prior,
+    and likelihood transform together), which is how `to_physical` is used. As a
+    chain `Step` the bijection is *not* measure-preserving in the flat
+    `(theta_jn, phi_jl)` coordinates: it rotates the line of sight rigidly about
+    the orbital angular momentum, preserving the spherical measure
+    `sin(theta_jn) dtheta dphi`, so
+    `log_det = log sin(theta_jn) - log sin(theta_jn')` (verified numerically
+    against finite differences through the LAL conversion), and `inverse`
+    rebuilds the network convention from the physical one using the invariant
+    conditioning the reverse fold supplies.
+    """
+
+    def __init__(self, num_processes: int = 1):
+        """
+        Parameters
+        ----------
+        num_processes : int, default 1
+            Parallel processes for the per-sample LAL spin conversion.
+        """
+        self.parameters = ["theta_jn", "phi_jl"]
+        self.conditioning = [
+            "theta_jn",
+            "phi_jl",
+            "phase",
+            "chirp_mass",
+            "mass_ratio",
+            "a_1",
+            "a_2",
+            "tilt_1",
+            "tilt_2",
+            "phi_12",
+        ]
+        self.num_processes = num_processes
+
+    @property
+    def consumes(self) -> list[str]:
+        # The bijection overwrites theta_jn / phi_jl in place; the remaining
+        # conditioning (phase, masses, tilts, ...) is read-only and must stay in
+        # the chain -- the default conditioning-minus-parameters would drop it.
+        return []
+
+    def log_det(self, given, context):
+        """`log|det J|` of `forward`, per row. The map preserves the spherical
+        measure, so the flat-coordinate Jacobian is
+        `sin(theta_jn) / sin(theta_jn')` -- verified numerically against finite
+        differences through the LAL conversion in
+        `test_jacobian_matches_sin_ratio` (agreement ~1e-9)."""
+        converted = self.forward(given, context)
+        return self._log_det(given["theta_jn"], converted["theta_jn"])
+
+    @staticmethod
+    def _log_det(theta_jn_in, theta_jn_out):
+        # Compute in double (the LAL conversion is double precision anyway),
+        # return in the input dtype: a reparametrization preserves the chain's
+        # dtype rather than promoting the summed log_prob.
+        log_det = torch.log(torch.sin(theta_jn_in.double())) - torch.log(
+            torch.sin(theta_jn_out.double())
+        )
+        return log_det.to(theta_jn_in.dtype)
+
+    def sample_and_log_prob(self, num_samples, context, given=None):
+        """Apply `forward`; contribute `-log|det J|`. Overridden to share the
+        single LAL conversion between the transform and its Jacobian (the base
+        implementation would convert twice)."""
+        if num_samples != 1:
+            raise ValueError("A reparametrization is 1:1; use fan_out=1.")
+        out = self.forward(given, context)
+        return out, -self._log_det(given["theta_jn"], out["theta_jn"])
+
+    @staticmethod
+    def _model_convention(model_metadata: dict) -> tuple[float, Optional[float]]:
+        """The reference frequency and spin-conversion phase the model trained with."""
+        wfg_settings = model_metadata["dataset_settings"]["waveform_generator"]
+        return wfg_settings["f_ref"], wfg_settings.get("spin_conversion_phase")
+
+    def to_physical(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFrame:
+        """Relabel samples from the model's convention to the physical (Bilby) one."""
+        f_ref, sc_phase = self._model_convention(model_metadata)
+        return change_spin_conversion_phase(
+            samples, f_ref, sc_phase, None, num_processes=self.num_processes
+        )
+
+    def to_network(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFrame:
+        """Relabel samples from the physical (Bilby) convention to the model's,
+        e.g. to ingest external posteriors for comparison."""
+        f_ref, sc_phase = self._model_convention(model_metadata)
+        return change_spin_conversion_phase(
+            samples, f_ref, None, sc_phase, num_processes=self.num_processes
+        )
+
+    def forward(self, given, context):
+        # The conversion runs in double; the outputs return in the input dtype
+        # and device (cf. RAReparam: compute in float64, store the chain dtype).
+        reference = given["theta_jn"]
+        theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
+        converted = self.to_physical(theta, context.model_metadata)
+        return {
+            k: torch.as_tensor(converted[k].to_numpy()).to(
+                dtype=reference.dtype, device=reference.device
+            )
+            for k in self.parameters
+        }
+
+    def inverse(self, params, context, given=None):
+        # The physical -> network direction also needs the invariant
+        # conditioning (phase, masses, tilts), which the reverse fold supplies
+        # as `given`.
+        if given is None:
+            raise ValueError(
+                "The spin-convention inverse needs the conditioning block "
+                "(phase, masses, tilts); pass it as `given`, or convert "
+                "DataFrames with to_network()."
+            )
+        reference = params["theta_jn"]
+        rows = {**given, **params}
+        theta = pd.DataFrame({k: _to_numpy(v) for k, v in rows.items()})
+        converted = self.to_network(theta, context.model_metadata)
+        return {
+            k: torch.as_tensor(converted[k].to_numpy()).to(
+                dtype=reference.dtype, device=reference.device
+            )
+            for k in self.parameters
+        }
+
+
+class GNPEKernelCorrection(TargetCorrection):
+    """
+    The single-step GNPE kernel correction as a target-side chain step.
+
+    Emits `delta_log_prob_target = log p(theta_hat | theta)` -- the GNPE kernel evaluated
+    at the proxies and the detector times recomputed from theta -- for importance sampling
+    on the joint proposal `q(theta, theta_hat | d)`. Contributes 0 to the proposal
+    density and consumes the intermediate detector times.
+    """
+
+    def __init__(self, kernel_factor: GNPEKernelFactor):
+        self.kernel_factor = kernel_factor
+        self.parameters = ["delta_log_prob_target"]
+        self.conditioning = list(kernel_factor.parameters) + list(
+            kernel_factor.gnpe_parameters
+        )
+        self.consumes = list(kernel_factor.gnpe_parameters)
+
+    def correction(self, given, context):
+        proxies = {p: given[p] for p in self.kernel_factor.parameters}
+        times = {k: given[k] for k in self.kernel_factor.gnpe_parameters}
+        correction = self.kernel_factor.log_prob(proxies, context, times)
+        return {"delta_log_prob_target": correction}
diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index be8972336..69b8064e0 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -179,7 +179,7 @@ def _build_context(self):
         # event data itself.
         if self.settings is None:
             return None
-        from dingo.gw.inference.factors import GWSamplerContext
+        from dingo.gw.inference.context import GWSamplerContext
 
         try:
             # base_metadata resolves the unconditional ("base") indirection, so
@@ -487,7 +487,7 @@ def _sample_synthetic_phase_chain(
         updates). Out-of-prior rows get `phase = 0` and `log_prob = nan` (they
         receive zero weight in importance sampling regardless)."""
         from dingo.core.factors import ChainComposer, SampleTableFactor
-        from dingo.gw.inference.factors import SyntheticPhaseFactor
+        from dingo.gw.inference.steps import SyntheticPhaseFactor
 
         theta_within = theta.iloc[np.flatnonzero(within_prior)]
         table = SampleTableFactor(
@@ -606,7 +606,7 @@ def get_samples_bilby_phase(self, num_processes=1):
         pd.DataFrame
             Samples
         """
-        from dingo.gw.inference.factors import SpinConventionReparam
+        from dingo.gw.inference.steps import SpinConventionReparam
 
         return SpinConventionReparam(num_processes=num_processes).to_physical(
             self.samples, self.base_metadata
@@ -672,7 +672,7 @@ def get_pesummary_samples(
                 samples.loc[:, col] += self.t_ref
 
         # Redefine the spin angles to the physical (Bilby) convention.
-        from dingo.gw.inference.factors import SpinConventionReparam
+        from dingo.gw.inference.steps import SpinConventionReparam
 
         samples = SpinConventionReparam(num_processes=num_processes).to_physical(
             samples, self.base_metadata
diff --git a/dingo/pipe/sampling.py b/dingo/pipe/sampling.py
index b652d6acf..fdea11867 100644
--- a/dingo/pipe/sampling.py
+++ b/dingo/pipe/sampling.py
@@ -16,7 +16,8 @@
 from dingo.core.factors import FlowFactor
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
 from dingo.gw.data.event_dataset import EventDataset
-from dingo.gw.inference.factors import GNPEKernelFactor, GWComposedSampler
+from dingo.gw.inference.sampler import GWComposedSampler
+from dingo.gw.inference.steps import GNPEKernelFactor
 from dingo.gw.inference.scan import chirp_mass_scan
 from dingo.pipe.default_settings import DENSITY_RECOVERY_SETTINGS
 from dingo.pipe.parser import create_parser
diff --git a/docs/source/dingo_pipe.md b/docs/source/dingo_pipe.md
index fce9f4eb5..a6686ba1c 100644
--- a/docs/source/dingo_pipe.md
+++ b/docs/source/dingo_pipe.md
@@ -94,7 +94,7 @@ Injections mirror the [approach of `bilby_pipe`](https://lscsoft.docs.ligo.org/b
 
 ## Sampling
 
-The next step is sampling from the Dingo model. The model is loaded into a [GWComposedSampler](dingo.gw.inference.factors.GWComposedSampler), built from the EventDataset prepared in the previous step. (If using [GNPE](gnpe) it is necessary to specify a `model-init`.) `num-samples` samples are then generated in batches of size `batch-size`. The samples (and context) are stored in a [Result](dingo.gw.result.Result) object and saved in HDF5 format.
+The next step is sampling from the Dingo model. The model is loaded into a [GWComposedSampler](dingo.gw.inference.sampler.GWComposedSampler), built from the EventDataset prepared in the previous step. (If using [GNPE](gnpe) it is necessary to specify a `model-init`.) `num-samples` samples are then generated in batches of size `batch-size`. The samples (and context) are stored in a [Result](dingo.gw.result.Result) object and saved in HDF5 format.
 
 If using GNPE, one can optionally specify `num-gnpe-iterations` (it defaults to 30). Importantly, obtaining the log probability when using GNPE requires an [extra step of training an unconditional flow](result.md#density-recovery). This is done using the `recover-log-prob` flag, which defaults to `True`. The default density recovery settings can be overwritten by providing a `density-recovery-settings` dictionary in the `.ini` file.
 
diff --git a/docs/source/example_injection.md b/docs/source/example_injection.md
index f92fd25c5..d50a8bd2c 100644
--- a/docs/source/example_injection.md
+++ b/docs/source/example_injection.md
@@ -56,7 +56,7 @@ strain_data = injection_generator.injection(theta)
 Then one can create a injections and do inference on them.
 
 ```
-from dingo.gw.inference.factors import GWComposedSampler
+from dingo.gw.inference.sampler import GWComposedSampler
 
 sampler = GWComposedSampler.from_gnpe_models(
     init_pm, main_pm, strain_data, num_iterations=30
diff --git a/docs/source/inference.md b/docs/source/inference.md
index c1103c4a8..f98647299 100644
--- a/docs/source/inference.md
+++ b/docs/source/inference.md
@@ -11,7 +11,7 @@ chain of conditional factors (the network flow, deterministic reparametrizations
 such as the sky-frame rotation, and fillers for fixed parameters).
 
 ```{eval-rst}
-.. autoclass:: dingo.gw.inference.factors.GWComposedSampler
+.. autoclass:: dingo.gw.inference.sampler.GWComposedSampler
     :members:
     :inherited-members:
     :show-inheritance:
diff --git a/tests/gw/test_bns_chain_builder.py b/tests/gw/test_bns_chain_builder.py
index 3fe78824a..3651490c7 100644
--- a/tests/gw/test_bns_chain_builder.py
+++ b/tests/gw/test_bns_chain_builder.py
@@ -7,11 +7,8 @@
 import pytest
 
 from dingo.core.factors import DeltaFactor, FlowFactor, ProxyOffsetReparam
-from dingo.gw.inference.factors import (
-    GWComposedSampler,
-    GWSamplerContext,
-    _proxy_offset_steps,
-)
+from dingo.gw.inference.context import GWSamplerContext
+from dingo.gw.inference.sampler import GWComposedSampler, _proxy_offset_steps
 from dingo.gw.transforms import HeterodynePhase
 
 _PINS = {"chirp_mass_proxy": 1.1975, "ra": 3.44616, "dec": -0.408084}
diff --git a/tests/gw/test_gw_sampler_context.py b/tests/gw/test_gw_sampler_context.py
index 6d6572170..871d15c29 100644
--- a/tests/gw/test_gw_sampler_context.py
+++ b/tests/gw/test_gw_sampler_context.py
@@ -11,9 +11,9 @@
 import pytest
 import torch
 
-import dingo.gw.inference.factors as factors_module
+import dingo.gw.inference.context as context_module
 from dingo.gw.domains import build_domain
-from dingo.gw.inference.factors import GWSamplerContext
+from dingo.gw.inference.context import GWSamplerContext
 
 _DOMAIN_SETTINGS = {
     "type": "FrequencyDomain",
@@ -50,7 +50,7 @@ def __init__(self, **kwargs):
 @pytest.fixture
 def context(monkeypatch):
     monkeypatch.setattr(
-        factors_module, "StationaryGaussianGWLikelihood", _StubLikelihood
+        context_module, "StationaryGaussianGWLikelihood", _StubLikelihood
     )
     _StubLikelihood.constructions = 0
     return GWSamplerContext(
@@ -199,7 +199,7 @@ def test_sampler_provenance_block():
     import ast
 
     from dingo.core.factors import ChainComposer, DeltaFactor
-    from dingo.gw.inference.factors import GWComposedSampler
+    from dingo.gw.inference.sampler import GWComposedSampler
 
     sampler = GWComposedSampler(
         ChainComposer([DeltaFactor({"a": 1.0})]), None, {}, ["a"]
diff --git a/tests/gw/test_ra_reparam.py b/tests/gw/test_ra_reparam.py
index 83b4c8b70..6018d4acb 100644
--- a/tests/gw/test_ra_reparam.py
+++ b/tests/gw/test_ra_reparam.py
@@ -8,7 +8,7 @@
 import numpy as np
 import torch
 
-from dingo.gw.inference.factors import RAReparam
+from dingo.gw.inference.steps import RAReparam
 
 
 def _context(t_ref, t_event):
diff --git a/tests/gw/test_spin_convention_reparam.py b/tests/gw/test_spin_convention_reparam.py
index c1b50c773..4bc4862f8 100644
--- a/tests/gw/test_spin_convention_reparam.py
+++ b/tests/gw/test_spin_convention_reparam.py
@@ -8,7 +8,7 @@
 import pytest
 import torch
 
-from dingo.gw.inference.factors import SpinConventionReparam
+from dingo.gw.inference.steps import SpinConventionReparam
 
 _MODEL_METADATA_SC0 = {
     "dataset_settings": {
diff --git a/tests/gw/test_synthetic_phase.py b/tests/gw/test_synthetic_phase.py
index c9f92eee3..523d0a860 100644
--- a/tests/gw/test_synthetic_phase.py
+++ b/tests/gw/test_synthetic_phase.py
@@ -7,7 +7,8 @@
 import torch
 from bilby.core.utils import random as bilby_random
 
-from dingo.gw.inference.factors import GWSamplerContext, SyntheticPhaseFactor
+from dingo.gw.inference.context import GWSamplerContext
+from dingo.gw.inference.steps import SyntheticPhaseFactor
 
 
 class _MockLikelihood:

From f1f956b6f67ffb4500189840e501c51345041534 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 13 Jul 2026 08:49:59 +0100
Subject: [PATCH 50/65] Rotate a pinned right ascension into the network's
 training frame

A parameter that is frame-corrected at the output must be inversely
corrected at the input: a sky position pinned at the event time now
enters the network as ra@t_ref via RAToTrainingFrame, and the trailing
RAToEventFrame (formerly RAReparam; renamed as a directional pair)
restores the event-frame value in the samples, so likelihood and
provenance stay physical. The rotation is exactly zero when the event
time equals the training reference time, as for all current networks.

Supporting changes: FlowFactor applies its alias map to conditioning
names as well as outputs; the composer validator models consumed
columns; and the chain fold defers the base count past 1:1 steps
(reparametrizations, target corrections) to the first sampling stage,
with a mock-level regression test for that chain shape.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py                         | 44 ++++++++++++++-----
 dingo/gw/inference/sampler.py                 | 41 +++++++++++------
 dingo/gw/inference/scan.py                    | 16 +++++--
 dingo/gw/inference/steps.py                   | 27 ++++++++++--
 tests/core/test_factors.py                    | 19 ++++++++
 tests/gw/test_bns_chain_builder.py            | 37 ++++++++++++++--
 .../{test_ra_reparam.py => test_ra_frames.py} | 12 ++---
 7 files changed, 157 insertions(+), 39 deletions(-)
 rename tests/gw/{test_ra_reparam.py => test_ra_frames.py} (83%)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index b14c6603a..2a9484249 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -302,7 +302,9 @@ def from_model(
         return cls(
             model=model,
             parameters=data_settings["inference_parameters"],
-            conditioning=list(context_parameters),
+            # The chain may carry a frame-corrected alias of a trained conditioning
+            # name (e.g. `ra@t_ref` for a pinned event-frame `ra`).
+            conditioning=[(aliases or {}).get(n, n) for n in context_parameters],
             context_parameters=list(context_parameters),
             aliases=aliases,
         )
@@ -333,7 +335,9 @@ def sample_and_log_prob(self, num_samples, context, given=None):
             # row-varying conditioning (the intrinsic/extrinsic split) recomputes N
             # identical data embeddings. Fixing this needs an embed/fuse split on the
             # model (FlowWrapper) so the cached embedding can be reused across rows.
-            ctx = self.standardization.standardize(given, self.context_parameters)
+            ctx = self.standardization.standardize(
+                self._network_conditioning(given), self.context_parameters
+            )
             n_rows = ctx.shape[0]
             data = context.prepared_data(conditioning=given)
             with torch.no_grad():
@@ -364,13 +368,20 @@ def log_prob(self, theta_i, context, given=None):
             net_context = (data,)
         else:
             data = context.prepared_data(conditioning=given)
-            ctx = self.standardization.standardize(given, self.context_parameters)
+            ctx = self.standardization.standardize(
+                self._network_conditioning(given), self.context_parameters
+            )
             net_context = (data, ctx)
         self.model.network.eval()
         with torch.no_grad():
             log_prob = self.model.log_prob(z, *net_context)
         return log_prob + self.standardization.log_det(self._net_parameters)
 
+    def _network_conditioning(self, given):
+        """The conditioning values keyed by the network's trained names (the chain
+        may carry a frame-corrected alias, e.g. `ra@t_ref` for a trained `ra`)."""
+        return {n: given[self.aliases.get(n, n)] for n in self.context_parameters}
+
     def describe(self) -> dict:
         return {**_describe_default(self), "unconditional": self.unconditional}
 
@@ -474,7 +485,8 @@ def log_prob(self, theta_i, context, given=None):
 class Reparametrization(ABC):
     """
     A deterministic bijection `Step`: it transforms existing parameters (no sampling)
-    and contributes `-log|det J|` to the proposal density.
+    and contributes `-log|det J|` to the proposal density. It is 1:1 -- one output
+    row per input row -- so it carries no sample multiplicity in a chain.
 
     Unlike a `Factor` it is 1:1 and invertible -- `forward` maps the conditioning block
     to the `parameters` block, `inverse` maps back (for re-plug / importance sampling),
@@ -486,6 +498,7 @@ class Reparametrization(ABC):
 
     parameters: list[str]
     conditioning: list[str]
+    one_to_one = True
 
     @abstractmethod
     def forward(
@@ -606,6 +619,7 @@ class TargetCorrection(ABC):
 
     parameters: list[str]
     conditioning: list[str]
+    one_to_one = True
     consumes: list[str]
 
     @abstractmethod
@@ -687,7 +701,10 @@ def _validate(self):
         column -- except a `Reparametrization` replacing its own inputs, which is
         invertible, so `log_prob` can restore the overwritten state. A step's
         emitted columns default to its `parameters`, but a step may emit
-        side-channel columns too (`produces`)."""
+        side-channel columns too (`produces`). Consumed columns leave the
+        produced set, mirroring the fold, so a later step may re-emit them
+        (e.g. `RAToEventFrame` restoring a pinned `ra` that `RAToTrainingFrame`
+        consumed)."""
         produced: set[str] = set()
         for step in self.steps:
             missing = [c for c in step.conditioning if c not in produced]
@@ -708,6 +725,7 @@ def _validate(self):
                     f"replace columns (its inverse can rebuild them for log_prob)."
                 )
             produced.update(emitted)
+            produced.difference_update(getattr(step, "consumes", ()))
 
     @property
     def steps(self) -> list[Step]:
@@ -741,9 +759,10 @@ def _run_chain_once(
 
         The base count lands on the first stage that carries multiplicity: a root
         prefix of point-mass steps (fixed pins) emits a single row each -- n draws
-        from a point mass are one atom repeated -- and the first non-point-mass
-        stage draws `base` (times its fan-out) in one conditioned call, with the
-        prefix's carried rows expanded to match."""
+        from a point mass are one atom repeated -- and 1:1 steps (reparametrizations,
+        target corrections) transform those rows in place. The first sampling stage
+        draws `base` (times its fan-out) in one conditioned call, with the prefix's
+        carried rows expanded to match."""
         samples: dict[str, torch.Tensor] = {}
         total: torch.Tensor | float = 0.0
         has_density = True
@@ -751,8 +770,13 @@ def _run_chain_once(
         for i, stage in enumerate(self.stages):
             step = stage.step
             if base_pending is not None:
-                if getattr(step, "point_mass", False):
-                    n = 1  # a point mass carries no multiplicity
+                if getattr(step, "point_mass", False) or getattr(
+                    step, "one_to_one", False
+                ):
+                    # Neither a point mass (one emitted row) nor a 1:1 transform
+                    # carries multiplicity; the base count waits for the first
+                    # sampling stage.
+                    n = 1
                 else:
                     n = base_pending * (stage.fan_out if step.conditioning else 1)
                     base_pending = None
diff --git a/dingo/gw/inference/sampler.py b/dingo/gw/inference/sampler.py
index bbc4791a9..3d992bf26 100644
--- a/dingo/gw/inference/sampler.py
+++ b/dingo/gw/inference/sampler.py
@@ -19,16 +19,17 @@
 from dingo.core.posterior_models import BasePosteriorModel
 from dingo.gw.inference.context import GWSamplerContext
 from dingo.gw.inference.steps import (
+    RAToTrainingFrame,
     GNPEFlowFactor,
     GNPEKernelCorrection,
     GNPEKernelFactor,
-    RAReparam,
+    RAToEventFrame,
 )
 
 
 def _ra_aliases(inference_parameters: list[str]) -> dict[str, str]:
     """The RA frame alias (`ra` -> `ra@t_ref`), applied only when the model infers
-    `ra`; paired with an `RAReparam` step that maps it back to the event frame."""
+    `ra`; paired with an `RAToEventFrame` step that maps it back to the event frame."""
     return {"ra": "ra@t_ref"} if "ra" in inference_parameters else {}
 
 
@@ -45,9 +46,19 @@ def _proxy_offset_steps(
     ]
 
 
-def _ra_reparam_steps(inference_parameters: list[str]) -> list:
-    """The `RAReparam` step, appended to a chain only when the model infers `ra`."""
-    return [RAReparam()] if "ra" in inference_parameters else []
+def _ra_to_event_steps(inference_parameters: list[str]) -> list:
+    """The `RAToEventFrame` step, appended to a chain only when the model infers `ra`."""
+    return [RAToEventFrame()] if "ra" in inference_parameters else []
+
+
+def _ra_adjustments(context_parameters: list[str]) -> tuple[list, list]:
+    """The RA frame pair for a pinned sky position: rotate the pinned event-frame
+    `ra` into the training frame before the network, and back into the event
+    frame after it. A parameter that is frame-corrected on the output side is
+    inversely corrected on the input side."""
+    if "ra" not in context_parameters:
+        return [], []
+    return [RAToTrainingFrame()], [RAToEventFrame()]
 
 
 def _delta_prior_steps(prior, inference_parameters: list[str]) -> list:
@@ -177,7 +188,7 @@ def from_model(
         """Build a single-network GW sampler from a model and event data.
 
         For a plain NPE model the chain is the flow (exposing `ra` as `ra@t_ref`)
-        followed by an `RAReparam` to the event frame. A model with
+        followed by an `RAToEventFrame` to the event frame. A model with
         `context_parameters` (e.g. the DINGO-BNS chirp-mass prior conditioning
         with a fixed sky position) requires `fixed_context_parameters` pinning
         all of them: the chain is then rooted in a `DeltaFactor` of the pins,
@@ -208,9 +219,12 @@ def from_model(
         metadata = _base_model_metadata(model)
         data_settings = metadata["train_settings"]["data"]
         inference_parameters = data_settings["inference_parameters"]
-        factor = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
-        steps = [factor]
         context_parameters = data_settings.get("context_parameters") or []
+        factor = FlowFactor.from_model(
+            model, aliases=_ra_aliases(inference_parameters + context_parameters)
+        )
+        ra_to_training, ra_to_event = _ra_adjustments(context_parameters)
+        steps = ra_to_training + [factor]
         if context_parameters:
             if set(fixed_context_parameters or {}) != set(context_parameters):
                 raise ValueError(
@@ -220,7 +234,8 @@ def from_model(
                 )
             steps = [DeltaFactor(fixed_context_parameters)] + steps
         steps += _proxy_offset_steps(inference_parameters, context_parameters)
-        steps += _ra_reparam_steps(inference_parameters)
+        steps += ra_to_event
+        steps += _ra_to_event_steps(inference_parameters)
         steps += _delta_prior_steps(context.prior, inference_parameters)
         return cls(
             composer=ChainComposer(steps),
@@ -241,7 +256,7 @@ def from_gnpe_models(
         """Build a multi-iteration time-GNPE sampler from an init + main model pair: the
         init model's data preprocessing, an init `FlowFactor` to seed, and a single
         `GibbsBlock` step -- cycling the GNPE kernel and main-network factors -- in a
-        `ChainComposer`, then an `RAReparam` to the event frame. Returns samples without
+        `ChainComposer`, then an `RAToEventFrame` to the event frame. Returns samples without
         a log_prob (Gibbs breaks density access).
 
         Parameters
@@ -279,7 +294,7 @@ def from_gnpe_models(
         gibbs = GibbsBlock(init_factor, [kernel_factor, flow_factor], num_iterations)
         steps = (
             [gibbs]
-            + _ra_reparam_steps(inference_parameters)
+            + _ra_to_event_steps(inference_parameters)
             + _delta_prior_steps(context.prior, inference_parameters)
         )
         return cls(
@@ -298,7 +313,7 @@ def from_singlestep_gnpe(
         event_metadata: Optional[dict] = None,
     ) -> "GWComposedSampler":
         """Build a single-step (density-preserving) time-GNPE sampler: a `ChainComposer`
-        of `[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAReparam]`. The chain is
+        of `[proxy_source, GNPEFlowFactor, GNPEKernelCorrection, RAToEventFrame]`. The chain is
         autoregressive, so log_prob is preserved, and `GNPEKernelCorrection` emits the
         `delta_log_prob_target` correction that importance sampling adds to the target.
 
@@ -329,7 +344,7 @@ def from_singlestep_gnpe(
         kernel_factor = GNPEKernelFactor.from_model(main_model)
         steps = (
             [proxy_source, flow_factor, GNPEKernelCorrection(kernel_factor)]
-            + _ra_reparam_steps(inference_parameters)
+            + _ra_to_event_steps(inference_parameters)
             + _delta_prior_steps(context.prior, inference_parameters)
         )
         return cls(ChainComposer(steps), context, metadata, inference_parameters)
diff --git a/dingo/gw/inference/scan.py b/dingo/gw/inference/scan.py
index da1449395..8bc893fbb 100644
--- a/dingo/gw/inference/scan.py
+++ b/dingo/gw/inference/scan.py
@@ -30,9 +30,10 @@
 from dingo.gw.inference.context import GWSamplerContext
 from dingo.gw.inference.sampler import (
     _delta_prior_steps,
+    _ra_adjustments,
     _proxy_offset_steps,
     _ra_aliases,
-    _ra_reparam_steps,
+    _ra_to_event_steps,
 )
 
 
@@ -141,10 +142,14 @@ def chirp_mass_scan(
     grid = chirp_mass_scan_grid(metadata, overlap_factor)
 
     context = GWSamplerContext.from_model(model, event_data, event_metadata)
-    flow = FlowFactor.from_model(model, aliases=_ra_aliases(inference_parameters))
+    flow = FlowFactor.from_model(
+        model, aliases=_ra_aliases(inference_parameters + context_parameters)
+    )
+    ra_to_training, ra_to_event = _ra_adjustments(context_parameters)
     tail_steps = (
         _proxy_offset_steps(inference_parameters, context_parameters)
-        + _ra_reparam_steps(inference_parameters)
+        + ra_to_event
+        + _ra_to_event_steps(inference_parameters)
         + _delta_prior_steps(context.prior, inference_parameters)
     )
     # Sweep the grid in blocks: a fixed table roots the chain (its length is
@@ -163,7 +168,10 @@ def chirp_mass_scan(
             {k: np.full(len(block), v, dtype=np.float32) for k, v in pins.items()}
         )
         composer = ChainComposer(
-            [SampleTableFactor(table), Stage(flow, fan_out=num_samples)] + tail_steps
+            [SampleTableFactor(table)]
+            + ra_to_training
+            + [Stage(flow, fan_out=num_samples)]
+            + tail_steps
         )
         out, _ = composer.sample_and_log_prob(len(block), context)
         frames.append(pd.DataFrame({k: v.numpy() for k, v in out.items()}))
diff --git a/dingo/gw/inference/steps.py b/dingo/gw/inference/steps.py
index 003a3e684..42372eee9 100644
--- a/dingo/gw/inference/steps.py
+++ b/dingo/gw/inference/steps.py
@@ -472,7 +472,7 @@ def log_prob(self, theta_i, context, given=None):
         return log_prob + self.standardization.log_det(self._net_parameters)
 
 
-class RAReparam(Reparametrization):
+class RAToEventFrame(Reparametrization):
     """
     Rotate right ascension from the network's training reference frame (`ra@t_ref`) to the
     event frame (`ra`).
@@ -527,6 +527,27 @@ def inverse(self, params, context, given=None):
         return {"ra@t_ref": ra_tref.float()}
 
 
+class RAToTrainingFrame(RAToEventFrame):
+    """
+    Rotate a pinned event-frame right ascension (`ra`) into the network's training
+    reference frame (`ra@t_ref`) before it conditions the network: the input-side
+    mirror of `RAToEventFrame`. A parameter that is frame-corrected on the output side
+    must be inversely corrected on the input side, so a sky position pinned at the
+    event time is presented to the network in the frame it was trained in; a
+    trailing `RAToEventFrame` restores the event-frame value in the samples.
+    """
+
+    def __init__(self):
+        self.conditioning = ["ra"]
+        self.parameters = ["ra@t_ref"]
+
+    def forward(self, given, context):
+        return super().inverse({"ra": given["ra"]}, context)
+
+    def inverse(self, params, context, given=None):
+        return super().forward({"ra@t_ref": params["ra@t_ref"]}, context)
+
+
 class SpinConventionReparam(Reparametrization):
     """
     Relabel the precessing-spin angles between Dingo's internal spin-phase
@@ -543,7 +564,7 @@ class SpinConventionReparam(Reparametrization):
     metadata, and a model trained without a fixed conversion phase (`None`)
     relabels to the identity.
 
-    Unlike `RAReparam` no marked intermediate name is needed: the two conventions
+    Unlike `RAToEventFrame` no marked intermediate name is needed: the two conventions
     never coexist in one table -- each world's plain names denote its own
     convention, and this bijection is the boundary crossing.
 
@@ -639,7 +660,7 @@ def to_network(self, samples: pd.DataFrame, model_metadata: dict) -> pd.DataFram
 
     def forward(self, given, context):
         # The conversion runs in double; the outputs return in the input dtype
-        # and device (cf. RAReparam: compute in float64, store the chain dtype).
+        # and device (cf. RAToEventFrame: compute in float64, store the chain dtype).
         reference = given["theta_jn"]
         theta = pd.DataFrame({k: _to_numpy(v) for k, v in given.items()})
         converted = self.to_physical(theta, context.model_metadata)
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index 0442f2167..d58fa5664 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -291,6 +291,25 @@ def log_det(self, given, context):
         return torch.full((n,), self._ld)
 
 
+def test_point_mass_prefix_then_reparam_defers_base_count():
+    # Regression: a 1:1 step directly after a point-mass root must not swallow
+    # the base count -- it transforms the single pinned row, and the first
+    # sampling stage draws the base (the pinned-sky BNS chain
+    # [DeltaFactor, RAToTrainingFrame, flow]).
+    comp = ChainComposer(
+        [
+            DeltaFactor({"u": 2.0}),
+            _MockReparam(shift=10.0, log_det_val=0.5),
+            _ConstFactor("x", conditioning=["v"]),
+        ]
+    )
+    out, log_prob = comp.sample_and_log_prob(5, context=None)
+    assert len(out["x"]) == 5
+    assert torch.equal(out["v"], torch.full((5,), 12.0))
+    assert "u" not in out
+    assert log_prob.shape == (5,)
+
+
 def test_reparam_step_consumes_input_and_contributes_neg_logdet():
     # A reparam is an in-place bijection: it replaces its input column and contributes
     # -log_det to the chain density.
diff --git a/tests/gw/test_bns_chain_builder.py b/tests/gw/test_bns_chain_builder.py
index 3651490c7..6fdbfe79c 100644
--- a/tests/gw/test_bns_chain_builder.py
+++ b/tests/gw/test_bns_chain_builder.py
@@ -179,13 +179,44 @@ def test_from_model_assembles_pinned_chain():
         _StubBNSModel(), _event_data(), fixed_context_parameters=_PINS
     )
     kinds = [type(s).__name__ for s in sampler.composer.steps]
-    assert kinds == ["DeltaFactor", "FlowFactor", "ProxyOffsetReparam"]
-    delta, flow, offset = sampler.composer.steps
+    # The pinned event-frame ra is rotated into the training frame for the
+    # network and back into the event frame afterwards.
+    assert kinds == [
+        "DeltaFactor",
+        "RAToTrainingFrame",
+        "FlowFactor",
+        "ProxyOffsetReparam",
+        "RAToEventFrame",
+    ]
+    delta, ra_in, flow, offset, ra_out = sampler.composer.steps
     assert set(delta.parameters) == set(_PINS)
-    assert flow.conditioning == ["ra", "dec", "chirp_mass_proxy"]
+    assert flow.conditioning == ["ra@t_ref", "dec", "chirp_mass_proxy"]
     assert offset.consumes == ["delta_chirp_mass"]
 
 
+def test_pinned_ra_rotates_to_training_frame_and_back():
+    import torch
+
+    from dingo.gw.inference.steps import RAToEventFrame, RAToTrainingFrame
+
+    class _Ctx:
+        event_metadata = {"time_event": 1187008882.42 + 6 * 3600}
+        t_ref = 1187008882.42
+        device = "cpu"
+
+    ctx = _Ctx()
+    ra = torch.tensor([3.44616], dtype=torch.float32)
+    to_training = RAToTrainingFrame().forward({"ra": ra}, ctx)
+    # Six hours of Earth rotation is a genuinely different frame.
+    assert not torch.allclose(to_training["ra@t_ref"], ra)
+    back = RAToEventFrame().forward(to_training, ctx)
+    assert torch.allclose(back["ra"], ra, atol=1e-6)
+    # Zero shift (bitwise) when the event time equals the reference time.
+    ctx.event_metadata = {"time_event": _Ctx.t_ref}
+    same = RAToTrainingFrame().forward({"ra": ra}, ctx)
+    assert torch.equal(same["ra@t_ref"], ra)
+
+
 def test_from_model_requires_all_pins():
     with pytest.raises(ValueError, match="conditions on"):
         GWComposedSampler.from_model(_StubBNSModel(), _event_data())
diff --git a/tests/gw/test_ra_reparam.py b/tests/gw/test_ra_frames.py
similarity index 83%
rename from tests/gw/test_ra_reparam.py
rename to tests/gw/test_ra_frames.py
index 6018d4acb..06b746559 100644
--- a/tests/gw/test_ra_reparam.py
+++ b/tests/gw/test_ra_frames.py
@@ -1,4 +1,4 @@
-"""RAReparam: the right-ascension frame reparametrization (network training frame
+"""RAToEventFrame: the right-ascension frame reparametrization (network training frame
 ``ra@t_ref`` <-> event frame ``ra``). Covers the forward/inverse round trip -- the inverse
 direction the model-based parity harnesses do not exercise -- and the no-op case (event
 time equal to, or absent from, the reference time)."""
@@ -8,17 +8,17 @@
 import numpy as np
 import torch
 
-from dingo.gw.inference.steps import RAReparam
+from dingo.gw.inference.steps import RAToEventFrame
 
 
 def _context(t_ref, t_event):
-    """A minimal stand-in for the SamplerContext fields RAReparam reads."""
+    """A minimal stand-in for the SamplerContext fields RAToEventFrame reads."""
     event_metadata = None if t_event is None else {"time_event": t_event}
     return types.SimpleNamespace(t_ref=t_ref, event_metadata=event_metadata)
 
 
 def test_ra_reparam_round_trip():
-    rp = RAReparam()
+    rp = RAToEventFrame()
     ctx = _context(t_ref=1126259462.4, t_event=1264316116.4)  # differ -> real rotation
     ra_tref = torch.rand(1000, dtype=torch.float32) * (2 * np.pi)
 
@@ -31,7 +31,7 @@ def test_ra_reparam_round_trip():
 
 
 def test_ra_reparam_is_noop_without_a_distinct_event_time():
-    rp = RAReparam()
+    rp = RAToEventFrame()
     ra_tref = torch.rand(100, dtype=torch.float32) * (2 * np.pi)
     for ctx in (_context(1126259462.4, 1126259462.4), _context(1126259462.4, None)):
         out = rp.forward({"ra@t_ref": ra_tref}, ctx)["ra"]
@@ -39,6 +39,6 @@ def test_ra_reparam_is_noop_without_a_distinct_event_time():
 
 
 def test_ra_reparam_declared_names():
-    rp = RAReparam()
+    rp = RAToEventFrame()
     assert rp.conditioning == ["ra@t_ref"]
     assert rp.parameters == ["ra"]

From a033c8725fa08b9ad1b03e3f442ad8a68a762332 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 13 Jul 2026 17:10:59 +0100
Subject: [PATCH 51/65] Add the sampling-chains concepts page

New advanced-guide page documenting the factorized sampler: the
product-of-conditionals model, step types, the sampler context,
multiplicity rules, the reverse density fold, provenance, and chain
building, with mermaid figures for the chain, the context, and the
combined system. Also fixes two unterminated inline-code spans in the
ChainComposer docstring that produced docutils warnings under autodoc.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py          |   5 +-
 docs/source/index.md           |   1 +
 docs/source/sampling_chains.md | 440 +++++++++++++++++++++++++++++++++
 3 files changed, 444 insertions(+), 2 deletions(-)
 create mode 100644 docs/source/sampling_chains.md

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index 2a9484249..cff2b3e3f 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -678,7 +678,7 @@ class Stage:
 
 class ChainComposer:
     """
-    Autoregressive composer over an ordered list of `Stage`s.
+    Autoregressive composer over an ordered list of `Stage` entries.
 
     Folds the steps in declared order -- a topological order of the conditioning DAG --
     expanding each by its fan-out and summing the proposal log-probs. A step is a
@@ -688,7 +688,8 @@ class ChainComposer:
     conditioning, synthetic phase, intrinsic/extrinsic splits, and -- via `GibbsBlock` --
     multi-iteration GNPE.
 
-    Accepts bare steps (wrapped as `Stage(step, fan_out=1)`) or explicit `Stage`s.
+    Accepts bare steps (wrapped as `Stage(step, fan_out=1)`) or explicit `Stage`
+    entries.
     """
 
     def __init__(self, stages: list[Union["Stage", Step]]):
diff --git a/docs/source/index.md b/docs/source/index.md
index 9a254e776..3f2975e31 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -60,6 +60,7 @@ As training a network from scratch can be expensive, we intend to also distribut
    network_architecture
    training
    inference
+   sampling_chains
    gnpe
    result
    dingo_pipe
diff --git a/docs/source/sampling_chains.md b/docs/source/sampling_chains.md
new file mode 100644
index 000000000..cc7a3c855
--- /dev/null
+++ b/docs/source/sampling_chains.md
@@ -0,0 +1,440 @@
+# Sampling chains
+
+Sampling the posterior involves more than a pass through a neural network: fixed
+parameters are filled in from the prior, coordinates are rotated between reference
+frames, GNPE introduces proxy variables, and a marginalized phase may be reconstructed
+from the likelihood. Dingo's *factorized sampler* expresses all of this as an explicit
+chain of steps acting on a growing table of named parameter columns. Each step is a
+conditional distribution over the block it produces, and the chain samples the
+posterior as an ordered product,
+
+$$
+q(\theta_1, \ldots, \theta_n | d) = \prod_i q_i(\theta_i | \theta_{<i}, d).
+$$
+
+A *factor* is a stochastic conditional. A deterministic step enters as a Dirac delta:
+$\delta(\theta_i - f(\theta_{<i}))$ for a reparametrization, and $\delta(\theta_i - c)$
+for a pinned value. Sampling records the log density of the final table, which
+importance sampling later divides by, and each step contributes one term,
+
+$$
+\log q(\theta | d) = \sum_i \Delta_i,
+$$
+
+where $\Delta_i = \log q_i$ for a factor, $-\log\lvert\det J_i\rvert$ for a delta
+whose inputs the step consumes (the change of variables), and $0$ for a delta whose
+value remains in the table (a pinned value is conditioned on, not integrated over).
+These cases correspond to the [step types](#steps) below. Two example chains:
+
+* **Plain NPE** (`FlowFactor → RAToEventFrame`): the network, followed by a rotation
+  of the right ascension from the training reference frame to the event frame.
+* **DINGO-BNS prior conditioning** (`DeltaFactor → RAToTrainingFrame → FlowFactor →
+  ProxyOffsetReparam → RAToEventFrame`): pinned proxy and sky values feed a
+  conditioned network, whose offset output is then reconstructed into the physical
+  chirp mass.
+
+At construction the composer checks the chain for consistency: every conditioning
+column must be produced by an earlier step. It then runs the steps in order,
+accumulating the table and the sum:
+
+```{mermaid}
+flowchart TB
+    subgraph comp ["ChainComposer"]
+        direction TB
+        s1["step 1<br/>q#8321;(#952;#8321; | d)"]
+        s2["step 2<br/>q#8322;(#952;#8322; | #952;#8321;, d)"]
+        dots["#8230;"]
+        sn["step n<br/>q#8345;(#952;#8345; | #952;#8321;, #8230;, #952;#8345;#8331;#8321;, d)"]
+        s1 -- "#916;#8321;" --> s2
+        s2 -- "#916;#8321; + #916;#8322;" --> dots
+        dots --> sn
+    end
+    out(["samples + log_prob"])
+
+    sn -- "#916;#8321; + #8230; + #916;#8345; = log q" --> out
+
+    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
+    classDef ghost fill:none,stroke:none
+    class s1,s2,sn step
+    class dots ghost
+```
+
+The generic machinery lives in `dingo.core.factors` (steps, stages, the composer, the
+runner); the gravitational-wave steps, the per-event context, and the chain builders
+live in `dingo.gw.inference`. The [](inference.md) builders assemble the standard
+chains from model metadata, but a chain is ordinary Python and can equally be
+assembled by hand (see [](#building-and-running-a-chain)).
+
+## Steps
+
+A chain entry is a *step*: an object with `parameters` (the columns it emits),
+`conditioning` (the earlier columns it reads), and a `sample_and_log_prob` method
+(the `Step` protocol). Steps never receive event data directly; the data enters
+through the shared [sampler context](#the-sampler-context). There are three types of
+step, plus one density-free sampling block:
+
+| Step                 | Emits                                | Log-prob term $\Delta_i$                 | Examples                                                       |
+|----------------------|--------------------------------------|------------------------------------------|----------------------------------------------------------------|
+| `Factor`             | a sampled parameter block            | its conditional log density               | `FlowFactor`, `DeltaFactor`, `SampleTableFactor`               |
+| `Reparametrization`  | a deterministic transform of columns | $-\log \lvert \det J \rvert$ (usually 0)  | `RAToEventFrame`, `ProxyOffsetReparam`, `SpinConventionReparam` |
+| `TargetCorrection`   | a target-side annotation column      | 0                                         | `GNPEKernelCorrection`                                         |
+| `GibbsBlock`         | the blocks of an internal Gibbs loop | none (the chain becomes density-free)     | multi-iteration [GNPE](gnpe.md)                                |
+
+### Factors
+
+A `Factor` is a conditional distribution $q_i(\theta_i | \theta_{<i}, d)$ over one
+parameter block: it draws `num_samples` samples per conditioning row and returns its
+own log density. Its interface is in physical parameter space; a network's
+standardization is applied internally.
+
+`FlowFactor`
+: Wraps a posterior model (NPE flow, FMPE, ...), in three conditioning shapes: an
+  *unconditional* model (a density-recovery NDE) takes no input at all; a
+  *data-conditional* model draws from the shared `prepared_data()`; and a model with
+  `context_parameters` (GNPE proxies, prior-conditioning pins) additionally conditions
+  on chain columns. A factor may expose a trained parameter name under an alias
+  (`ra → ra@t_ref`), so that a downstream step can convert reference frames by name.
+
+`DeltaFactor`
+: A point mass $\delta(\theta_i - c)$ pinning parameters to fixed values, contributing
+  zero log probability: the pinned block is conditioned on, not integrated over. Used
+  as the chain root for prior conditioning, where later factors condition on the pins,
+  and as a filler for delta-prior parameters the network does not infer.
+
+`SampleTableFactor`
+: A chain root emitting a fixed table of existing samples, together with their stored
+  log probability. This is how a chain continues from earlier samples: synthetic phase
+  runs as a chain rooted in the proposal samples, and the chirp-mass scan roots its
+  sweep in a table with one row per grid point.
+
+`SyntheticPhaseFactor`, `GNPEKernelFactor`, `GNPEFlowFactor`
+: The gravitational-wave factors (`dingo.gw.inference.steps`): the likelihood-based
+  phase reconstruction for phase-marginalized networks (see
+  [synthetic phase](result.md#synthetic-phase)), the GNPE blur kernel
+  $p(\hat\theta | \theta)$, and the proxy-conditioned GNPE main network.
+
+### Reparametrizations
+
+A `Reparametrization` is a deterministic bijection: `forward` maps its conditioning
+columns to new ones, replacing the inputs it `consumes`; `inverse` rebuilds them; and
+the proposal density picks up $-\log\lvert\det J\rvert$, which is zero for the
+measure-preserving maps used here. It is 1:1 (one output row per input row), so it
+carries no sample multiplicity.
+
+* `RAToEventFrame` rotates the right ascension from the network's training reference
+  frame (`ra@t_ref`) to the event frame (`ra`) by the sidereal-time difference, which
+  is exactly zero when the event time equals the training reference time.
+  `RAToTrainingFrame` is its input-side mirror: a pinned event-frame sky position is
+  presented to the network in the frame it was trained in, and a trailing
+  `RAToEventFrame` restores the event-frame value in the samples.
+* `ProxyOffsetReparam` reconstructs a physical parameter from a network's offset
+  output and its proxy, $X = \delta_X + X_\mathrm{proxy}$, consuming the offset column
+  while keeping the proxy in the chain.
+* `SpinConventionReparam` relabels the precessing-spin angles between Dingo's internal
+  spin-phase convention and the physical (Bilby) one.
+
+### Target corrections
+
+A `TargetCorrection` contributes nothing to the proposal density. It emits a
+side-channel column, `delta_log_prob_target`, which importance sampling adds to the
+*target* log density. This covers cases where the target is not simply
+$\pi(\theta)\,\mathcal{L}(\theta)$. Because the emitted column is an annotation
+rather than a parameter block, the step stands outside the product of conditionals
+and contributes $\Delta_i = 0$. The one instance is `GNPEKernelCorrection`: in
+single-step GNPE the proposal is the joint $q(\hat\theta)\,q(\theta | d, \hat\theta)$
+over parameters and proxies, so the matching target acquires the kernel term
+$p(\hat\theta | \theta)$. The step recomputes the detector times from $\theta$,
+evaluates the kernel there, and records the result with the samples. Unlike a
+reparametrization it has no inverse, so it may consume only side-channel
+intermediates, never sampled parameters.
+
+### Density-free blocks
+
+A `GibbsBlock` runs blocked Gibbs sampling internally, seeding from an init factor and
+then sweeping its factor list for `num_iterations` iterations. It yields no log
+probability, since the cyclic dependency has no tractable marginal. A chain containing
+one is density-free (its samples carry no `log_prob`), and the density must be
+[recovered](result.md#density-recovery) afterwards. Dingo uses this only for
+multi-iteration [GNPE](gnpe.md).
+
+## The sampler context
+
+Every step reads the event through a shared, per-event `GWSamplerContext`: the data
+$d$, the per-event metadata, and everything derived from them. It provides three
+views.
+
+`prepared_data(conditioning=None)`
+: The network-input representation of the event: heterodyning, decimation, whitening,
+  frequency masking, and repackaging, as the model's metadata prescribes. Called
+  without conditioning, it is computed once and cached. Called with conditioning (the
+  chain columns available to a conditioned factor), it is *row-aligned*, with one data
+  row per conditioning row. The context consumes only the columns its preparation is a
+  function of (for a heterodyning model, `chirp_mass_proxy`); the remaining columns
+  condition the network alone. A constant consumed value (a pinned proxy) is prepared
+  once and viewed across the rows; varying values (a sweep) run through the
+  batch-native transform chain in one vectorized pass.
+
+`prior`
+: The static prior fixed at training time (intrinsic and extrinsic, with Dingo
+  defaults), built once from the model metadata. Importance-sampling prior updates and
+  the time/phase split-off for marginalized networks are applied downstream, not here.
+
+`likelihood(...)`
+: The exact likelihood on the raw event data, used by likelihood-based factors
+  (synthetic phase) and importance sampling. It builds its own data representation
+  rather than reusing the network-input view, and its reference time is the event
+  time. Marginalization settings (time, phase, calibration) are passed per request.
+
+As data flow, including the conditioning columns that a chain feeds back into the
+data preparation:
+
+```{mermaid}
+flowchart TB
+    d[("event data d")]
+    em["event metadata"]
+    md["model metadata"]
+    subgraph ctx ["GWSamplerContext"]
+        direction LR
+        pd["prepared_data"]
+        pr["prior"]
+        lk["likelihood"]
+    end
+    chain["chain steps"]
+    isamp["importance sampling"]
+
+    d --> ctx
+    em --> ctx
+    md --> ctx
+    pd -- "data rows (row-aligned)" --> chain
+    chain -. "conditioning columns (optional)" .-> pd
+    lk -.-> isamp
+    pr -.-> isamp
+
+    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
+    class d,em,md,pd,pr,lk ctxstyle
+```
+
+Alongside the strain data, the context carries the per-event **event metadata**. This
+includes the event time, which sets the likelihood reference time and the sidereal
+right-ascension correction (`RAToEventFrame`), together with any per-event analysis
+settings. Settings can be changed at inference time in three ways, none of which
+touches the trained model. A *frequency-range update* (a per-detector minimum or
+maximum frequency in the event metadata) narrows the network-input view; it is
+validated against the training-time random strain cropping, which must cover the
+request, and the likelihood applies the same range independently through ASD masking.
+A *representation update* (an updated duration, or a base-domain likelihood for
+importance sampling) produces a derived context, described below. A *prior update* is
+applied at the importance-sampling stage and never modifies the context.
+
+A context is treated as **immutable**: the data representation is part of its
+identity. Importance-sampling settings updates, such as an updated duration or a
+likelihood evaluated on the undecimated base domain, produce a *derived* context via
+`derive()`. The derived context shares the event payload and analysis metadata, so
+parameter meaning is preserved and importance weights between the two representations
+remain well defined.
+
+A context is built with `GWSamplerContext.from_model(model, event_data,
+event_metadata)`, or from a metadata dictionary alone with `from_model_metadata`; the
+latter is how a saved `Result` reconstructs the prior, domain, and likelihood from its
+stored settings. The chain's torch device is `context.device`. Steps that create fresh
+tensors, such as the pins of a `DeltaFactor`, create them on this device so that their
+outputs can join a chain running on a GPU.
+
+Putting the pieces together:
+
+```{mermaid}
+flowchart TB
+    d[("event data d")] --> ctx["GWSamplerContext"]
+    ctx -- "data views" --> comp["ChainComposer"]
+    comp -. "conditioning parameters" .-> ctx
+    comp --> out(["samples + log_prob"])
+    out --> isamp["importance sampling"]
+    ctx -.-> isamp
+
+    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
+    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
+    class ctx ctxstyle
+    class comp step
+```
+
+The chain draws its data views through the context and feeds conditioning parameters
+back into the preparation. The likelihood and prior views serve the downstream
+importance sampling in the same way, so a derived context re-targets the whole
+analysis consistently.
+
+```{note}
+The representation vocabulary here (frequency domains, multibanded decimation,
+base-domain likelihoods) is specific to this domain family. A new domain family
+should be given its own context class implementing the same interface
+(`prepared_data` / `prior` / `likelihood` / `derive`, the
+`dingo.core.factors.SamplerContext` protocol) rather than extending this one.
+```
+
+## Stages, fan-out, and multiplicity
+
+A `ChainComposer` holds an ordered list of `Stage(step, fan_out)` entries (bare steps
+are wrapped with `fan_out=1`). The declared order must be a topological order of the
+conditioning DAG, which is validated at construction: every conditioning column must
+be produced by an earlier step, and no step may overwrite an existing column. The
+exception is a reparametrization replacing its own inputs: the map is invertible, so
+the overwritten state can be restored when evaluating `log_prob`. A consumed column
+leaves the produced set, so a later step may re-emit it (as when `RAToEventFrame`
+restores the pinned `ra` that `RAToTrainingFrame` consumed).
+
+Sampling follows three multiplicity rules:
+
+1. **The requested count lands on the first stage that carries multiplicity.** A root
+   prefix of point masses (pins) emits a single row each, since $n$ draws from a
+   point mass are one atom repeated, and 1:1 steps transform rows in place. The first
+   sampling stage then draws the full count in a single conditioned call. In the
+   DINGO-BNS chain above, the `DeltaFactor` and the frame rotation produce one row,
+   and the flow draws all `num_samples` conditioned on it.
+2. **Fan-out multiplies rows.** A non-root stage with `fan_out=k` draws $k$ samples
+   per conditioning row, and the carried columns are repeated row-major to stay
+   aligned; the chain returns `num_samples × expansion` rows, where the expansion is
+   the product of the non-root fan-outs. The chirp-mass scan uses this, with a
+   grid-table root followed by `Stage(flow, fan_out=num_samples)` drawing per grid
+   row.
+3. **Unconditioned non-root steps fill rather than fan.** A delta filler emits one
+   constant per current row.
+
+`batch_size` chunks the requested count, capping peak memory at one chunk. A chain
+rooted in a `SampleTableFactor` is fixed-size and runs with `batch_size=None`; a
+caller sweeping a table, such as the chirp-mass scan, manages its own blocking of the
+grid instead.
+
+## The proposal density: forward and reverse
+
+Sampling folds the chain forward, summing each factor's log density and each
+reparametrization's Jacobian term into the proposal log probability, reported in
+physical parameter space.
+
+`ChainComposer.log_prob(samples, context)` evaluates the same density at given
+samples, as needed to re-evaluate saved samples or to importance sample. The steps are
+folded in exact *reverse* chain order, so the columns are restored to the state each
+step saw during sampling: a reparametrization rebuilds its consumed inputs via
+`inverse` (for example, `ra@t_ref` from the event-frame `ra`), a factor adds its
+`log_prob` at the restored conditioning, and a target correction contributes nothing.
+The reverse fold is what makes single-step GNPE and prior conditioning
+density-preserving end to end.
+
+The delta view of deterministic steps unifies this bookkeeping. A deterministic step
+is a point-mass kernel, and its delta is resolved in one of two ways. If the step
+*consumes* its inputs, the delta is integrated out against them and leaves the
+Jacobian term. This requires the map to be invertible, which is why `inverse` is part
+of the `Reparametrization` contract, and why a `TargetCorrection` (which has no
+inverse) may consume only side-channel intermediates. If the value instead *remains*
+in the table, the target conditions on the same value and the delta cancels in the
+importance weight; the analysis is then conditional on the pinned values. A
+`GibbsBlock` is also a kernel, but one whose output density is intractable, which is
+what "density-free" means.
+
+## Provenance
+
+A `Result` exported from a composed sampler records how its samples were made under
+`settings["sampler"]`: the executed chain in order, one descriptor per step, plus any
+entries added by the caller. `dingo_pipe` records the model checkpoint paths
+(`models`), the density-recovery recipe, and the chirp-mass-scan record.
+
+```python
+{"version": 1,
+ "implementation": "composed",
+ "chain": [
+     {"step": "DeltaFactor",
+      "parameters": ["chirp_mass_proxy", "ra", "dec"], "conditioning": [],
+      "values": {"chirp_mass_proxy": 1.1976, "ra": 3.446, "dec": -0.408}},
+     {"step": "RAToTrainingFrame", ...},
+     {"step": "GNPEFlowFactor", ...},
+ ],
+ "models": {"model": "model.pt"}}
+```
+
+The block is a record rather than a recipe: nothing consumes it at load time. It is
+also **literal-only**, meaning every value round-trips through
+`str`/`ast.literal_eval` in the saved settings. The `version` field allows the format
+to evolve safely.
+
+## Building and running a chain
+
+The standard chains are assembled from model metadata by the `GWComposedSampler`
+builders (see [](inference.md) for usage):
+
+`from_model(model, event_data, event_metadata, fixed_context_parameters=None)`
+: A single-network chain: plain NPE, or prior conditioning for a model with
+  `context_parameters`, with the pinned values as the chain root.
+
+`from_gnpe_models(init_model, main_model, event_data, event_metadata, num_iterations=30)`
+: Multi-iteration time GNPE: a `GibbsBlock` cycling the kernel and main-network
+  factors. The chain is density-free.
+
+`from_singlestep_gnpe(main_model, proxy_source, event_data, event_metadata)`
+: Single-step, density-preserving GNPE. The `proxy_source` supplies the proxies (a
+  `DeltaFactor` for prior conditioning, or an unconditional NDE for
+  [density recovery](result.md#density-recovery)), followed by the main network and
+  the kernel correction.
+
+A chain is ordinary Python, and the builders have no privileged machinery:
+
+```python
+from dingo.core.factors import ChainComposer, FlowFactor
+from dingo.gw.inference.context import GWSamplerContext
+from dingo.gw.inference.steps import RAToEventFrame
+
+context = GWSamplerContext.from_model(model, event_data, event_metadata)
+flow = FlowFactor.from_model(model, aliases={"ra": "ra@t_ref"})
+composer = ChainComposer([flow, RAToEventFrame()])
+
+samples = composer.sample(10_000, context, batch_size=5_000)
+log_prob = composer.log_prob(samples, context)  # the reverse fold, re-evaluated
+```
+
+This is the chain that `from_model` assembles for a plain NPE model (plus a
+`DeltaFactor` filler for any delta-prior parameters). `GWComposedSampler` adds the
+DataFrame runner (`run_sampler`) and the `Result` export (`to_result` / `to_hdf5`).
+
+### Writing a new step
+
+1. **Pick the type**: a step that samples new parameters is a `Factor`; one that
+   transforms existing ones deterministically and invertibly is a
+   `Reparametrization`; one that annotates the importance-sampling target is a
+   `TargetCorrection`.
+2. **Declare the interface**: `parameters` (the columns emitted) and `conditioning`
+   (the columns read), plus `consumes` / `produces` for side-channel columns.
+3. **Implement the contract.** A factor implements `sample_and_log_prob` (returning
+   `num_samples` draws per conditioning row, flattened in row-major order) and
+   `log_prob`, both in physical parameter space. A reparametrization implements
+   `forward` and `inverse` (and `log_det` when the map is not measure-preserving);
+   the inverse must rebuild exactly the consumed columns, since the reverse fold
+   depends on it. A target correction implements `correction`.
+4. **Read data only through the context**, so the step remains valid under a derived
+   context.
+5. **Override `describe()`** if the step has salient configuration, keeping the
+   descriptor literal-only.
+
+## API
+
+```{eval-rst}
+.. autoclass:: dingo.core.factors.Factor
+    :members:
+    :show-inheritance:
+
+.. autoclass:: dingo.core.factors.Reparametrization
+    :members:
+    :show-inheritance:
+
+.. autoclass:: dingo.core.factors.TargetCorrection
+    :members:
+    :show-inheritance:
+
+.. autoclass:: dingo.core.factors.Stage
+    :members:
+    :show-inheritance:
+
+.. autoclass:: dingo.core.factors.ChainComposer
+    :members:
+    :show-inheritance:
+
+.. autoclass:: dingo.gw.inference.context.GWSamplerContext
+    :members:
+    :show-inheritance:
+```

From 1e0b20d9521855f93c00d1120c10e9de9b2dcbc5 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 13 Jul 2026 17:44:01 +0100
Subject: [PATCH 52/65] Present reparametrizations as changes of variables

Refine the density presentation on the sampling-chains page: the
product of conditionals is owned by the factors, which may be
stochastic or point masses (conditioned on, not integrated over), while
a reparametrization is a bijective change of variables that removes its
inputs from the table and contributes a Jacobian term rather than a new
factor. The Dirac-delta reading of deterministic steps is dropped
except for point masses, where it is exact. Also states the composer's
consistency check in the intro and introduces each step type by what it
does rather than what it is not.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/source/sampling_chains.md | 59 ++++++++++++++++++----------------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/docs/source/sampling_chains.md b/docs/source/sampling_chains.md
index cc7a3c855..b537436ed 100644
--- a/docs/source/sampling_chains.md
+++ b/docs/source/sampling_chains.md
@@ -4,34 +4,34 @@ Sampling the posterior involves more than a pass through a neural network: fixed
 parameters are filled in from the prior, coordinates are rotated between reference
 frames, GNPE introduces proxy variables, and a marginalized phase may be reconstructed
 from the likelihood. Dingo's *factorized sampler* expresses all of this as an explicit
-chain of steps acting on a growing table of named parameter columns. Each step is a
-conditional distribution over the block it produces, and the chain samples the
-posterior as an ordered product,
+chain of steps acting on a growing table of named parameter columns. The *factors*
+among the steps sample the posterior as an ordered product of conditionals,
 
 $$
 q(\theta_1, \ldots, \theta_n | d) = \prod_i q_i(\theta_i | \theta_{<i}, d).
 $$
 
-A *factor* is a stochastic conditional. A deterministic step enters as a Dirac delta:
-$\delta(\theta_i - f(\theta_{<i}))$ for a reparametrization, and $\delta(\theta_i - c)$
-for a pinned value. Sampling records the log density of the final table, which
-importance sampling later divides by, and each step contributes one term,
+A factor may be stochastic (a network) or a point mass pinning parameters to fixed
+values, which are conditioned on rather than integrated over. A *reparametrization*
+is a bijective change of variables: it replaces sampled columns with transformed
+ones, so its inputs disappear from the table and the density picks up a Jacobian
+term rather than a new factor. Sampling records the log density of the final table,
+which importance sampling later divides by, and each step contributes one term,
 
 $$
 \log q(\theta | d) = \sum_i \Delta_i,
 $$
 
-where $\Delta_i = \log q_i$ for a factor, $-\log\lvert\det J_i\rvert$ for a delta
-whose inputs the step consumes (the change of variables), and $0$ for a delta whose
-value remains in the table (a pinned value is conditioned on, not integrated over).
-These cases correspond to the [step types](#steps) below. Two example chains:
+where $\Delta_i = \log q_i$ for a factor (identically zero for a point mass) and
+$-\log\lvert\det J_i\rvert$ for a reparametrization. These cases correspond to the
+[step types](#steps) below. Two example chains:
 
 * **Plain NPE** (`FlowFactor → RAToEventFrame`): the network, followed by a rotation
   of the right ascension from the training reference frame to the event frame.
-* **DINGO-BNS prior conditioning** (`DeltaFactor → RAToTrainingFrame → FlowFactor →
-  ProxyOffsetReparam → RAToEventFrame`): pinned proxy and sky values feed a
-  conditioned network, whose offset output is then reconstructed into the physical
-  chirp mass.
+* **DINGO-BNS prior conditioning** (`DeltaFactor → RAToTrainingFrame →
+  FlowFactor → ProxyOffsetReparam → RAToEventFrame`): pinned proxy and sky values
+  feed a conditioned network, whose offset output is then reconstructed into the
+  physical chirp mass.
 
 At construction the composer checks the chain for consistency: every conditioning
 column must be produced by an earlier step. It then runs the steps in order,
@@ -59,6 +59,11 @@ flowchart TB
     class dots ghost
 ```
 
+In the figure every step is written as a conditional $q_i$, which covers all
+factors, point masses included. A reparametrization instead replaces columns in
+place, adding its Jacobian term to the sum but no factor to the product (see
+[the proposal density](#the-proposal-density-forward-and-reverse)).
+
 The generic machinery lives in `dingo.core.factors` (steps, stages, the composer, the
 runner); the gravitational-wave steps, the per-event context, and the chain builders
 live in `dingo.gw.inference`. The [](inference.md) builders assemble the standard
@@ -135,9 +140,9 @@ carries no sample multiplicity.
 
 ### Target corrections
 
-A `TargetCorrection` contributes nothing to the proposal density. It emits a
-side-channel column, `delta_log_prob_target`, which importance sampling adds to the
-*target* log density. This covers cases where the target is not simply
+A `TargetCorrection` emits a side-channel column, `delta_log_prob_target`, which
+importance sampling adds to the *target* log density; it contributes nothing to the
+proposal. This covers cases where the target is not simply
 $\pi(\theta)\,\mathcal{L}(\theta)$. Because the emitted column is an annotation
 rather than a parameter block, the step stands outside the product of conditionals
 and contributes $\Delta_i = 0$. The one instance is `GNPEKernelCorrection`: in
@@ -318,16 +323,14 @@ step saw during sampling: a reparametrization rebuilds its consumed inputs via
 The reverse fold is what makes single-step GNPE and prior conditioning
 density-preserving end to end.
 
-The delta view of deterministic steps unifies this bookkeeping. A deterministic step
-is a point-mass kernel, and its delta is resolved in one of two ways. If the step
-*consumes* its inputs, the delta is integrated out against them and leaves the
-Jacobian term. This requires the map to be invertible, which is why `inverse` is part
-of the `Reparametrization` contract, and why a `TargetCorrection` (which has no
-inverse) may consume only side-channel intermediates. If the value instead *remains*
-in the table, the target conditions on the same value and the delta cancels in the
-importance weight; the analysis is then conditional on the pinned values. A
-`GibbsBlock` is also a kernel, but one whose output density is intractable, which is
-what "density-free" means.
+Two special cases complete the bookkeeping. A point-mass factor contributes zero to
+the proposal, and the target conditions on the same pinned value, so the analysis is
+conditional on the pins rather than integrating over them. A reparametrization
+removes the columns it consumes, which is lawful only for an invertible map: the
+change of variables supplies the density on the new columns, and the reverse fold
+restores the old ones through `inverse`. This is also why a `TargetCorrection`,
+which has no inverse, may consume only side-channel intermediates. A `GibbsBlock`
+yields samples whose density is intractable, which is what "density-free" means.
 
 ## Provenance
 

From 8f6940d5dda9f4fbec3f4d7f9d6f230d023a6957 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 13 Jul 2026 17:44:43 +0100
Subject: [PATCH 53/65] Add the binary-neutron-star inference page

New advanced-guide page for DINGO-BNS inference: phase heterodyning as
single-step chirp-mass GNPE, prior conditioning with pinned context
parameters, the fixed-proxy chain, the chirp-mass scan, and the
dingo_pipe options (fixed-context-parameters, chirp-mass-scan) with a
GW170817 configuration example. Adds the Nature 2025 reference and
links the page from the sampling-chains examples.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/source/bns.md             | 178 +++++++++++++++++++++++++++++++++
 docs/source/index.md           |   1 +
 docs/source/refs.bib           |  14 +++
 docs/source/sampling_chains.md |  12 +--
 4 files changed, 199 insertions(+), 6 deletions(-)
 create mode 100644 docs/source/bns.md

diff --git a/docs/source/bns.md b/docs/source/bns.md
new file mode 100644
index 000000000..710441a97
--- /dev/null
+++ b/docs/source/bns.md
@@ -0,0 +1,178 @@
+# Binary neutron stars
+
+Dingo analyzes binary neutron star (BNS) events with the DINGO-BNS method of
+{footcite:p}`Dax:2024mcn`. BNS signals pose two problems for plain NPE. Their long
+inspirals require a fine frequency resolution, so the data are far larger than for
+binary black holes, and their chirp mass is measured so precisely that a network
+covering the full training prior would spend nearly all of its capacity on parameter
+values excluded by any individual event. DINGO-BNS addresses both with a single
+device: a chirp-mass proxy that simplifies the data (phase heterodyning) and narrows
+the effective prior (prior conditioning). The proxy is fixed per event, so sampling is
+single-step GNPE: one pass through the network, with the density preserved and
+importance sampling available directly.
+
+```{note}
+This page describes inference with a trained chirp-mass-conditioned network.
+Training such a network is not yet supported; the training configuration will be
+documented when it is, and no pre-trained BNS network is distributed yet. Also not
+yet available: the accelerated heterodyned and decimated likelihood of
+{footcite:p}`Dax:2024mcn` (importance sampling uses the exact likelihood), the
+synthetic-phase `compute_likelihood` fast path, and per-event time scans for
+pre-merger networks.
+```
+
+## Phase heterodyning
+
+At BNS frequency resolutions the strain oscillates rapidly in frequency, which makes
+poor input for a network. Multiplying the data by $\exp(i \phi(f; \tilde{\mathcal{M}}))$
+with the leading-order chirp phase
+
+$$
+\phi(f; \tilde{\mathcal{M}}) = \frac{3}{128}
+\left(\frac{\pi G \tilde{\mathcal{M}} f}{c^3}\right)^{-5/3}
+$$
+
+removes the dominant phase evolution for a reference chirp mass
+$\tilde{\mathcal{M}}$ close to the true value. The residual oscillations are slow,
+and the multibanded frequency domain can then decimate the data far more
+aggressively.
+
+The reference value is the *chirp-mass proxy*. During training it is drawn by
+blurring the true chirp mass with a narrow kernel, and the network conditions on it;
+this is GNPE with the chirp mass as the proxy parameter (see [GNPE](gnpe.md)). At
+inference the proxy is fixed per event, so a single iteration suffices and the
+density is preserved. A chirp-mass-conditioned model records the kernel and the
+phase order in its metadata under `gnpe_chirp`, and the
+[sampler context](sampling_chains.md#the-sampler-context) reads this to prepare data
+as a function of `chirp_mass_proxy`. Heterodyning is applied to the raw strain before
+decimation (the two operations do not commute).
+
+## Prior conditioning
+
+The proxy plays a second role. Because the network conditions on
+$\tilde{\mathcal{M}}$, it effectively learns a family of posteriors under narrow
+chirp-mass priors centered on the proxy, $q(\theta | d, \tilde{\mathcal{M}})$.
+Setting the proxy at inference selects the member of the family appropriate to the
+event, so one network amortizes over events while retaining the resolution of an
+event-specific narrow prior. The network infers the offset
+`delta_chirp_mass` $= \mathcal{M} - \tilde{\mathcal{M}}$ rather than the chirp mass
+itself; the chain reconstructs the physical value with a `ProxyOffsetReparam` step.
+
+The chain:
+
+```{mermaid}
+flowchart TB
+    pins["DeltaFactor<br/>chirp_mass_proxy"]
+    flow["FlowFactor<br/>draws delta_chirp_mass, #8230;"]
+    off["ProxyOffsetReparam<br/>chirp_mass = delta_chirp_mass + chirp_mass_proxy"]
+    out(["samples + log_prob"])
+    ctx["GWSamplerContext<br/>heterodyne #8594; decimate #8594; whiten"]
+
+    pins --> flow --> off --> out
+    pins -. "chirp_mass_proxy" .-> ctx
+    ctx -. "prepared data" .-> flow
+
+    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
+    classDef reparam fill:#e2f0e6,stroke:#27ae60,color:#1a1a1a
+    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
+    class pins,flow step
+    class off reparam
+    class ctx ctxstyle
+```
+
+The pinned values have a single owner, the chain root, and are recorded with the
+samples. The heterodyne receives the proxy through the row-aligned `prepared_data`
+contract of the [sampler context](sampling_chains.md#the-sampler-context). Since the
+chain contains no Gibbs block, the samples carry their log probability and
+importance sampling proceeds without a density-recovery step.
+
+A network may condition on further context parameters, and any of them can be pinned
+in the same way. The reference configuration of {footcite:p}`Dax:2024mcn` also fixes
+the sky position: a pinned right ascension is given in the event frame and rotated
+into the network's training frame before conditioning, which inserts an
+`RAToTrainingFrame` step before the network and a trailing `RAToEventFrame` that
+restores the event-frame value in the samples. The rotation is exactly zero when the
+network's reference time equals the trigger time.
+
+## The chirp-mass scan
+
+When no external estimate of the chirp mass is available, the trigger value can be
+determined from the data (see the Methods of {footcite:p}`Dax:2024mcn`). The scan
+sweeps the proxy over the training chirp-mass prior on a grid whose spacing is set by
+the kernel width, draws a few posterior samples at each grid point in a single
+batched network pass, evaluates a phase-marginalized likelihood for every draw within
+the prior, and takes the chirp mass of the maximum-likelihood draw as the trigger
+value.
+
+The sweep runs on the ordinary chain machinery: a fixed table with one row per grid
+point roots the chain, the data preparation heterodynes each row at its own proxy
+value, and the network draws per row. The scan result (trigger value, signal-to-noise
+ratio, maximum log likelihood, and the scan settings) is recorded in the sampler
+provenance. For a GW170817-like event the scan costs about a minute of CPU time.
+
+## Running through dingo_pipe
+
+Two [dingo_pipe](dingo_pipe.md) sampler options control BNS inference:
+
+fixed-context-parameters
+: Dictionary pinning the model's context parameters, e.g.
+  `{chirp_mass_proxy: 1.19786, ra: 3.44616, dec: -0.408084}`. A single-network model
+  with context parameters requires all of them pinned, unless the chirp-mass proxy is
+  supplied by the scan. Cannot be combined with `model-init` (iterative GNPE).
+
+chirp-mass-scan
+: Set to `true` to determine the trigger chirp mass from the data with defaults
+  derived from the model (grid from the training prior and kernel, 10 draws per grid
+  point). A dictionary overrides individual settings, e.g.
+  `{num_samples: 10, overlap_factor: 2, block_size: 32}`; `num_processes` defaults to
+  `request-cpus`. Mutually exclusive with a pinned `chirp_mass_proxy`; the remaining
+  context parameters are still supplied via `fixed-context-parameters`.
+
+```{code-block} ini
+---
+caption: Sampler and data sections of a GW170817 configuration.
+---
+################################################################################
+##  Sampler arguments
+################################################################################
+
+model = /path/to/bns_model.pt
+device = 'cuda'
+num-samples = 50000
+batch-size = 50000
+fixed-context-parameters = {chirp_mass_proxy: 1.19786, ra: 3.44616, dec: -0.408084}
+# Alternatively, determine the chirp mass from the data:
+# chirp-mass-scan = true
+# fixed-context-parameters = {ra: 3.44616, dec: -0.408084}
+
+importance-sample = true
+importance-sampling-settings = {synthetic_phase: {approximation_22_mode: true, n_grid: 5001, uniform_weight: 0.01}}
+
+################################################################################
+## Data generation arguments
+################################################################################
+
+trigger-time = 1187008882.4
+label = GW170817
+outdir = outdir_GW170817
+channel-dict = {H1:GWOSC, L1:GWOSC, V1:GWOSC}
+psd-length = 128
+```
+
+Importance sampling follows the standard [workflow](result.md). For a multibanded
+model the likelihood is evaluated on the undecimated base domain by default
+(`use_base_domain`, set automatically and adjustable in
+`importance-sampling-settings`). Phase-marginalized networks reconstruct the phase
+synthetically before reweighting; setting `approximation_22_mode: true` treats the
+signal as dominated by the $(2, 2)$ mode, which is appropriate for BNS and
+substantially faster than the mode-summed default.
+
+As an indication of expected performance, analyses of GW170817 on public data with a
+development network reach sample efficiencies of roughly 10% and a log Bayes factor
+relative to noise of about +500. A scan run recovers the trigger chirp mass and
+matches the evidence of the pinned run within Monte Carlo uncertainty, as expected
+from the prior-conditioning construction.
+
+```{eval-rst}
+.. footbibliography::
+```
diff --git a/docs/source/index.md b/docs/source/index.md
index 3f2975e31..faf741b7b 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -64,6 +64,7 @@ As training a network from scratch can be expensive, we intend to also distribut
    gnpe
    result
    dingo_pipe
+   bns
    asimov
    
 .. toctree::
diff --git a/docs/source/refs.bib b/docs/source/refs.bib
index 06eeb2b10..de6ced378 100644
--- a/docs/source/refs.bib
+++ b/docs/source/refs.bib
@@ -74,3 +74,17 @@ @article{Wildberger:2022agw
     year = "2022",
     journal = {},
 }
+
+@article{Dax:2024mcn,
+    author = {Dax, Maximilian and Green, Stephen R. and Gair, Jonathan and Gupte, Nihar and P\"urrer, Michael and Raymond, Vivien and Wildberger, Jonas and Macke, Jakob H. and Buonanno, Alessandra and Sch\"olkopf, Bernhard},
+    title = "{Real-time inference for binary neutron star mergers using machine learning}",
+    eprint = "2407.09602",
+    archivePrefix = "arXiv",
+    primaryClass = "gr-qc",
+    doi = "10.1038/s41586-025-08593-z",
+    journal = "Nature",
+    volume = "639",
+    number = "8053",
+    pages = "49--53",
+    year = "2025"
+}
diff --git a/docs/source/sampling_chains.md b/docs/source/sampling_chains.md
index b537436ed..4ebc62357 100644
--- a/docs/source/sampling_chains.md
+++ b/docs/source/sampling_chains.md
@@ -28,7 +28,7 @@ $-\log\lvert\det J_i\rvert$ for a reparametrization. These cases correspond to t
 
 * **Plain NPE** (`FlowFactor → RAToEventFrame`): the network, followed by a rotation
   of the right ascension from the training reference frame to the event frame.
-* **DINGO-BNS prior conditioning** (`DeltaFactor → RAToTrainingFrame →
+* **[DINGO-BNS](bns.md) prior conditioning** (`DeltaFactor → RAToTrainingFrame →
   FlowFactor → ProxyOffsetReparam → RAToEventFrame`): pinned proxy and sky values
   feed a conditioned network, whose offset output is then reconstructed into the
   physical chirp mass.
@@ -297,9 +297,9 @@ Sampling follows three multiplicity rules:
 2. **Fan-out multiplies rows.** A non-root stage with `fan_out=k` draws $k$ samples
    per conditioning row, and the carried columns are repeated row-major to stay
    aligned; the chain returns `num_samples × expansion` rows, where the expansion is
-   the product of the non-root fan-outs. The chirp-mass scan uses this, with a
-   grid-table root followed by `Stage(flow, fan_out=num_samples)` drawing per grid
-   row.
+   the product of the non-root fan-outs. The
+   [chirp-mass scan](bns.md#the-chirp-mass-scan) uses this, with a grid-table root
+   followed by `Stage(flow, fan_out=num_samples)` drawing per grid row.
 3. **Unconditioned non-root steps fill rather than fan.** A delta filler emits one
    constant per current row.
 
@@ -363,8 +363,8 @@ The standard chains are assembled from model metadata by the `GWComposedSampler`
 builders (see [](inference.md) for usage):
 
 `from_model(model, event_data, event_metadata, fixed_context_parameters=None)`
-: A single-network chain: plain NPE, or prior conditioning for a model with
-  `context_parameters`, with the pinned values as the chain root.
+: A single-network chain: plain NPE, or [prior conditioning](bns.md) for a model
+  with `context_parameters`, with the pinned values as the chain root.
 
 `from_gnpe_models(init_model, main_model, event_data, event_metadata, num_iterations=30)`
 : Multi-iteration time GNPE: a `GibbsBlock` cycling the kernel and main-network

From 8e0fe45636069637f2af7e4b7e25cfc58963826b Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 13 Jul 2026 17:57:13 +0100
Subject: [PATCH 54/65] Rewrite the inference and injection-example docs

inference.md becomes a proper workflow page: the two inference routes,
the event_data contract, the sampler builders, running and output, and
the refreshed Injection section. The injection example is updated to
the current API and was executed end to end against a trained GNPE
model pair; its ASD line now restricts the dataset to the model's
detectors. The GNPE inference section replaces the removed
dingo_analyze_event script with the dingo_pipe and Python routes.
build_model_from_kwargs is re-exported from dingo.core.posterior_models
so the documented import works.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/posterior_models/__init__.py |   1 +
 docs/source/example_injection.md        |  65 +++++++--------
 docs/source/gnpe.md                     |  31 ++++---
 docs/source/inference.md                | 104 ++++++++++++++++++------
 4 files changed, 130 insertions(+), 71 deletions(-)

diff --git a/dingo/core/posterior_models/__init__.py b/dingo/core/posterior_models/__init__.py
index 148346f98..c682cede8 100644
--- a/dingo/core/posterior_models/__init__.py
+++ b/dingo/core/posterior_models/__init__.py
@@ -3,3 +3,4 @@
 from dingo.core.posterior_models.cflow_base import ContinuousFlowPosteriorModel
 from dingo.core.posterior_models.flow_matching import FlowMatchingPosteriorModel
 from dingo.core.posterior_models.score_matching import ScoreDiffusionPosteriorModel
+from dingo.core.posterior_models.build_model import build_model_from_kwargs
diff --git a/docs/source/example_injection.md b/docs/source/example_injection.md
index d50a8bd2c..2ea843c46 100644
--- a/docs/source/example_injection.md
+++ b/docs/source/example_injection.md
@@ -1,61 +1,58 @@
-# Inference on an injection 
+# Inference on an injection
 
-A simple example is creating an injection consistent with what the
-network was trained on, and then running Dingo on it. First one can instantiate
-the {py:class}`dingo.gw.injection.Injection` using the metadata from the
-{py:class}`dingo.core.models.posterior_model.PosteriorModel` (the trained network). An ASD dataset also needs to be specified,
-one can take the fiducial asd dataset the network was trained on. 
+A simple end-to-end test of a trained model is to create an injection consistent
+with the training data and run inference on it. The
+{py:class}`~dingo.gw.injection.Injection` is instantiated from the metadata of the
+trained network (see [inference](inference.md)). An ASD dataset must also be
+specified; here we take the fiducial dataset the network was trained on.
 
-```
-from dingo.core.models import PosteriorModel
+```python
+from dingo.core.posterior_models import build_model_from_kwargs
 import dingo.gw.injection as injection
 from dingo.gw.noise.asd_dataset import ASDDataset
 
-main_pm = PosteriorModel(
-    device="cuda",
-    model_filename="/path/to/main_network", 
-    load_training_info=False
+main_pm = build_model_from_kwargs(
+    filename="/path/to/main_network.pt", device="cuda", load_training_info=False
 )
-
-init_pm = PosteriorModel(
-    device='cuda',
-    model_filename="/path/to/init_network",
-    load_training_info=False
+init_pm = build_model_from_kwargs(
+    filename="/path/to/init_network.pt", device="cuda", load_training_info=False
 )
 
 injection_generator = injection.Injection.from_posterior_model_metadata(main_pm.metadata)
 asd_fname = main_pm.metadata["train_settings"]["training"]["stage_0"]["asd_dataset_path"]
-asd_dataset = ASDDataset(file_name=asd_fname)
-injection_generator.asd = {k:v[0] for k,v in asd_dataset.asds.items()}
+detectors = main_pm.metadata["train_settings"]["data"]["detectors"]
+injection_generator.asd = ASDDataset(file_name=asd_fname, ifos=detectors)
 
 intrinsic_parameters = {
     "chirp_mass": 35,
     "mass_ratio": 0.5,
-    "a_1": .3,
-    "a_2": .5,
-    "tilt_1": 0.,
-    "tilt_2": 0.,
-    "phi_jl": 0.,
-    "phi_12": 0.
+    "a_1": 0.3,
+    "a_2": 0.5,
+    "tilt_1": 0.0,
+    "tilt_2": 0.0,
+    "phi_jl": 0.0,
+    "phi_12": 0.0,
 }
 
 extrinsic_parameters = {
-    'phase': 0.,
-    'theta_jn': 2.3,
-    'geocent_time': 0.,
-    'luminosity_distance': 400.,
-    'ra': 0.,
-    'dec': 0.,
-    'psi': 0.,
+    "phase": 0.0,
+    "theta_jn": 2.3,
+    "geocent_time": 0.0,
+    "luminosity_distance": 400.0,
+    "ra": 0.0,
+    "dec": 0.0,
+    "psi": 0.0,
 }
 
 theta = {**intrinsic_parameters, **extrinsic_parameters}
 strain_data = injection_generator.injection(theta)
 ```
 
-Then one can create a injections and do inference on them.
+This example uses a [GNPE](gnpe.md) model pair, so the sampler is built from the
+initialization and main networks; for a plain-NPE model, use
+`GWComposedSampler.from_model(model, strain_data)` instead.
 
-```
+```python
 from dingo.gw.inference.sampler import GWComposedSampler
 
 sampler = GWComposedSampler.from_gnpe_models(
diff --git a/docs/source/gnpe.md b/docs/source/gnpe.md
index e82ae51b7..f0e458012 100644
--- a/docs/source/gnpe.md
+++ b/docs/source/gnpe.md
@@ -91,17 +91,28 @@ Complete example config files for both networks are provided in the /examples fo
 
 ### Inference
 
-The inference script must be pointed to both trained networks in order to sample using GNPE.
-```bash
-dingo_analyze_event
-  --model model
-  --model_init model_init
-  --gps_time_event gps_time_event
-  --num_samples num_samples
-  --num_gnpe_iterations num_gnpe_iterations
-  --batch_size batch_size
+GNPE inference requires both trained networks. With [dingo_pipe](dingo_pipe.md), the
+initialization network is passed alongside the main model in the `.ini` file:
+
+```ini
+model = /path/to/main_model.pt
+model-init = /path/to/init_model.pt
+num-gnpe-iterations = 30
+```
+
+The number of Gibbs iterations defaults to 30. Because Gibbs sampling provides no log
+probability, dingo_pipe recovers the density afterwards by training an unconditional
+flow for the proxies (the `recover-log-prob` flag, on by default; see
+[density recovery](result.md#density-recovery)).
+
+In Python, the sampler is built from the model pair directly (see
+[inference](inference.md)):
+
+```python
+sampler = GWComposedSampler.from_gnpe_models(
+    init_model, main_model, event_data, num_iterations=30
+)
 ```
-The number of Gibbs iterations is also specified here (typically 30 is appropriate). This script will save the final samples from each Gibbs chain.
 
 
 
diff --git a/docs/source/inference.md b/docs/source/inference.md
index f98647299..985063c39 100644
--- a/docs/source/inference.md
+++ b/docs/source/inference.md
@@ -1,43 +1,87 @@
 # Inference
 
-With a trained network, inference can be performed on injections or real data. For 
-injections, see the [discussion in the examples](example_injection.md). For real data, we 
-recommend to use [dingo_pipe](dingo_pipe.md).
+With a trained network, inference means drawing posterior samples for a specific
+stretch of data. There are two routes. For real events we recommend
+[dingo_pipe](dingo_pipe.md), which downloads and prepares the data, runs the sampler,
+and importance samples in one workflow. This page describes the underlying Python
+interface, which is useful for injections, custom data, and interactive work.
 
-## The `GWComposedSampler` class
+## Event data
 
-Inference uses the `GWComposedSampler` class, which represents the posterior as a
-chain of conditional factors (the network flow, deterministic reparametrizations
-such as the sky-frame rotation, and fillers for fixed parameters).
-
-```{eval-rst}
-.. autoclass:: dingo.gw.inference.sampler.GWComposedSampler
-    :members:
-    :inherited-members:
-    :show-inheritance:
-```
-
-A plain-NPE sampler is built with `GWComposedSampler.from_model(model, event_data,
-event_metadata)`, where `model` is a `PosteriorModel` and `event_data` is the data
-to be analyzed---a dictionary with the following keys:
+A sampler analyzes an `event_data` dictionary with the following entries:
 
 waveform
-: (unwhitened) strain data in each detector
+: (unwhitened) strain data for each detector
 
 asds
-: noise ASDs estimated in each detector at the time of the event
+: noise amplitude spectral densities estimated for each detector at the time of the
+  event
 
 parameters (optional)
-: for injections, the true parameters of the signal (for saving; ignored for sampling)
+: for injections, the true parameters of the signal (stored with the results; ignored
+  for sampling)
+
+The data must be consistent with the model: the same detectors, and frequency content
+covering the model's data domain. Data prepared by dingo_pipe is stored as an
+{py:class}`~dingo.gw.data.event_dataset.EventDataset`, whose `data` attribute has
+exactly this form, and the [Injection](#injections) class below produces it directly.
 
-The `run_sampler()` method then draws the requested samples from the posterior conditioned on the data. All processing---de-standardization inside the network factor, and the correction for the rotation of the Earth between the network reference time and the event time---is expressed as chain steps, and the samples are stored as a DataFrame in `GWComposedSampler.samples`. The DataFrame contains columns for each inference parameter, as well as the log probability of the sample under the posterior model.
+## Building and running a sampler
+
+Inference uses the `GWComposedSampler` class, which represents the posterior as a
+[chain of steps](sampling_chains.md). It is built from a trained model and the event
+data:
+
+```python
+from dingo.core.posterior_models import build_model_from_kwargs
+from dingo.gw.inference.sampler import GWComposedSampler
+
+model = build_model_from_kwargs(
+    filename="/path/to/model.pt", device="cuda", load_training_info=False
+)
+sampler = GWComposedSampler.from_model(model, event_data, event_metadata)
+sampler.run_sampler(num_samples=50_000, batch_size=10_000)
+```
 
-The `GWComposedSampler.metadata` attribute contains all settings that went into producing the samples. The `to_result()` method returns a [Result](dingo.gw.result.Result) containing the samples, settings (including structured sampler provenance under `settings["sampler"]`), and data; `to_hdf5()` saves it directly.
+The constructors cover the standard analysis types:
 
+`from_model`
+: A single-network model: plain NPE, or a prior-conditioned model with
+  `fixed_context_parameters` (see [binary neutron stars](bns.md)).
+
+`from_gnpe_models`
+: Multi-iteration [GNPE](gnpe.md), from an initialization and a main model.
+
+`from_singlestep_gnpe`
+: Single-step GNPE with an explicit proxy source, used for
+  [density recovery](result.md#density-recovery).
+
+`run_sampler()` draws the requested number of samples in batches of `batch_size` and
+stores them as a DataFrame in `sampler.samples`, with one column per inference
+parameter plus `log_prob`, the log density of each sample under the model (absent for
+density-free GNPE chains). All processing, including the de-standardization of
+network outputs and the rotation of the right ascension from the training reference
+frame to the event frame, is expressed as chain steps.
+
+The `metadata` attribute contains all settings that went into producing the samples.
+`to_result()` exports a [Result](result.md) containing the samples, the settings
+(including the structured sampler provenance under `settings["sampler"]`), and the
+event data; `to_hdf5()` saves it directly. Importance sampling and plotting then
+proceed on the `Result`.
+
+```{eval-rst}
+.. autoclass:: dingo.gw.inference.sampler.GWComposedSampler
+    :members:
+    :inherited-members:
+    :show-inheritance:
+```
 
 ## Injections
 
-Injections (i.e., simulated data) are produced using the `Injection` class. It includes options for fixed or random parameters (drawn from a prior), and it returns injections in a format that can be passed directly as the `event_data` of a `GWComposedSampler`.
+Injections (simulated signals in stationary Gaussian noise) are produced with the
+`Injection` class. It supports fixed parameters (`injection(theta)`) and random
+parameters drawn from a prior (`random_injection()`), and it returns data in the
+`event_data` format above, ready to pass to a sampler.
 
 ```{eval-rst}
 .. autoclass:: dingo.gw.injection.Injection
@@ -46,9 +90,15 @@ Injections (i.e., simulated data) are produced using the `Injection` class. It i
 ```
 
 ```{hint}
-The convenience class method `from_posterior_model_metadata()` instantiates an `Injection` with all of the settings that went into the posterior model. To this class pass the PosteriorModel.metadata dictionary. It should produce injections that perfectly match the characteristics of the training data (waveform approximant, data conditioning, noise characteristics, etc.). This can be very useful for testing a trained model.
+The class method `from_posterior_model_metadata()` instantiates an `Injection` with
+the settings of a trained model (waveform approximant, data conditioning, detectors,
+priors), so that its injections match the characteristics of the training data. This
+is very useful for testing a model; see the [worked example](example_injection.md).
 ```
 
 ```{important}
-Repeated calls to `Injection.injection()`, even with the same parameters, will produce injections with different noise realizations (which therefore lead to different posteriors). For repeated analyses of the *exact same* injection (e.g., with different models or codes) it is necessary to either save the injection for re-use or fix a random seed.
-```
\ No newline at end of file
+Repeated calls to `Injection.injection()`, even with the same parameters, produce
+different noise realizations and therefore different posteriors. For repeated
+analyses of the exact same injection (e.g., with different models or codes), either
+save the injection for re-use or fix a random seed.
+```

From eafb9c234352b1a68ece5988f9d5ffb439608d3d Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Mon, 13 Jul 2026 17:59:33 +0100
Subject: [PATCH 55/65] Document single-step GNPE and the BNS pipe options

gnpe.md gains a single-step section covering the two density-preserving
uses (density-recovery re-sampling and fixed BNS proxies). dingo_pipe.md
points single-network context-parameter models at
fixed-context-parameters / chirp-mass-scan on the BNS page and notes
the base-domain likelihood default for multibanded models.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/source/dingo_pipe.md |  4 +++-
 docs/source/gnpe.md       | 11 +++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/docs/source/dingo_pipe.md b/docs/source/dingo_pipe.md
index a6686ba1c..78d822433 100644
--- a/docs/source/dingo_pipe.md
+++ b/docs/source/dingo_pipe.md
@@ -98,6 +98,8 @@ The next step is sampling from the Dingo model. The model is loaded into a [GWCo
 
 If using GNPE, one can optionally specify `num-gnpe-iterations` (it defaults to 30). Importantly, obtaining the log probability when using GNPE requires an [extra step of training an unconditional flow](result.md#density-recovery). This is done using the `recover-log-prob` flag, which defaults to `True`. The default density recovery settings can be overwritten by providing a `density-recovery-settings` dictionary in the `.ini` file.
 
+Single-network models that condition on context parameters (e.g., chirp-mass-conditioned [BNS](bns.md) networks) take their pinned values from `fixed-context-parameters`, or determine the trigger chirp mass from the data with `chirp-mass-scan`. Both options are described on the [binary neutron stars](bns.md) page.
+
 Since sampling uses GPU hardware, there is an additional key `sampling-requirements` for HTCondor requirements during the sampling stage. This is intended for specifying GPU requirements such as memory or CUDA version.
 
 ## Importance sampling
@@ -110,7 +112,7 @@ If `prior-dict-updates` is specified in the `.ini` file, then this will be used
 If extending the prior support during importance sampling, be sure that the posterior does not rail up against the prior boundary being extended.
 ```
 
-By default, dingo_pipe assumes that it is necessary to sample the phase synthetically, so it will do so before importance sampling. This can be turned off by passing an empty dictionary to `importance-sampling-settings`. Note that importance sampling itself can be switched off by setting the `importance-sample` flag to False (it defaults to True). 
+By default, dingo_pipe assumes that it is necessary to sample the phase synthetically, so it will do so before importance sampling. This can be turned off by passing an empty dictionary to `importance-sampling-settings`. Note that importance sampling itself can be switched off by setting the `importance-sample` flag to False (it defaults to True). For models using a multibanded frequency domain, the likelihood is evaluated on the undecimated base domain by default (`use_base_domain` in `importance-sampling-settings`).
 
 Importance sampling (including synthetic phase sampling) is an expensive step, so dingo_pipe allows for parallelization: this step is split over `n-parallel` jobs, each of which uses `request-cpus-importance-sampling` processes. In the backend, this makes use of the Result [split()](dingo.core.result.Result.split) and [merge()](dingo.core.result.Result.merge) methods.
 
diff --git a/docs/source/gnpe.md b/docs/source/gnpe.md
index f0e458012..471a333b0 100644
--- a/docs/source/gnpe.md
+++ b/docs/source/gnpe.md
@@ -127,6 +127,17 @@ which wraps the cyclic GNPE loop in a `GibbsBlock` chain step,
 ```
 The block seeds the chain with samples from the initialization network, perturbs them with the kernel to obtain proxy samples, and then performs `num_iterations` Gibbs sweeps (kernel and main-network factors in turn) to obtain the final samples. Because Gibbs sampling breaks density access, the chain yields no log probability; `dingo_pipe` recovers it by training an unconditional flow for the proxies and re-sampling with a single density-preserving step (see [density recovery](result.md#density-recovery)).
 
+## Single-step GNPE
+
+When reliable proxy values are available before sampling, a single GNPE iteration
+suffices, and the chain becomes autoregressive rather than iterative: the sample
+density is preserved, and no recovery step is needed. Dingo uses this in two places.
+In [density recovery](result.md#density-recovery), an unconditional flow trained on
+the Gibbs proxies supplies the proxy values, and a single pass through the main
+network (`GWComposedSampler.from_singlestep_gnpe`) yields new samples together with
+their density. For [binary neutron stars](bns.md), the chirp-mass proxy is fixed per
+event, so sampling is single-step from the start.
+
 ```{eval-rst}
 .. footbibliography::
 ```
\ No newline at end of file

From 6b1b94eb4488d3c25c7573603f5c7f2991418e2d Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 08:54:41 +0100
Subject: [PATCH 56/65] Triage the #387 tests against the composed sampler

The merge from hackathon-1 brought in #387's tests for the removed
legacy samplers. Of their 32 tests, 17 are adapted to the new API
(composed-sampler counting/batching, FlowFactor log_prob round-trips
including the data-conditional branch, the frequency validators, the
sidereal shift magnitude, delta-prior filling, and Result export round
trips), 7 are dropped as duplicates of existing coverage, and 8 are
dropped as obsolete along with the removed Sampler machinery. All other
#387 test files pass unchanged and are kept.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 tests/core/test_composed_sampler.py           | 216 +++++++++++
 tests/core/test_samplers.py                   | 261 --------------
 .../gw/inference/test_gw_composed_sampler.py  | 162 +++++++++
 tests/gw/inference/test_gw_samplers.py        | 341 ------------------
 tests/gw/test_frequency_updates.py            | 159 ++++++++
 tests/gw/test_ra_frames.py                    |  30 ++
 6 files changed, 567 insertions(+), 602 deletions(-)
 create mode 100644 tests/core/test_composed_sampler.py
 delete mode 100644 tests/core/test_samplers.py
 create mode 100644 tests/gw/inference/test_gw_composed_sampler.py
 delete mode 100644 tests/gw/inference/test_gw_samplers.py
 create mode 100644 tests/gw/test_frequency_updates.py

diff --git a/tests/core/test_composed_sampler.py b/tests/core/test_composed_sampler.py
new file mode 100644
index 000000000..5f5c16f04
--- /dev/null
+++ b/tests/core/test_composed_sampler.py
@@ -0,0 +1,216 @@
+"""Unit tests for the generic ``ComposedSampler`` runner and for ``FlowFactor`` driving
+a *real* normalizing flow.
+
+``test_factors`` exercises ``ChainComposer`` and the batching plumbing with network-free
+mock factors; this file covers the two pieces that require an actual network:
+
+* ``ComposedSampler.run_sampler`` -- the DataFrame runner over a chain (sample count,
+  column layout, and batch-size invariance).
+* ``FlowFactor`` around a genuine ``NormalizingFlowPosteriorModel`` -- the real
+  standardize/de-standardize round trip, including the change-of-variables correction
+  (``-sum(log std)``), which the deterministic fake model in ``test_factors`` cannot
+  reach. Covered in both conditioning shapes: the *unconditional* flow (context-free,
+  ``context=None``) and the *data-conditional* flow, whose data enters through a stub
+  ``SamplerContext`` serving a fixed random tensor as ``prepared_data()`` -- no trained
+  model or GW data machinery is needed for either.
+
+(Adapts the runner/log_prob coverage from the pre-sampler-revamp ``test_samplers.py``,
+which imported the now-deleted ``dingo.core.samplers.Sampler``.)
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from dingo.core.factors import ChainComposer, ComposedSampler, FlowFactor
+from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel
+
+
+INFERENCE_PARAMETERS = ["a", "b", "c"]
+
+# Non-trivial standardization so that the log_prob change-of-variables correction
+# (-sum(log(std))) is actually exercised by the round-trip test below.
+STANDARDIZATION = {
+    "mean": {"a": 1.0, "b": -2.0, "c": 0.5},
+    "std": {"a": 2.0, "b": 0.5, "c": 3.0},
+}
+
+BASE_TRANSFORM_KWARGS = {
+    "hidden_dim": 16,
+    "num_transform_blocks": 1,
+    "activation": "elu",
+    "dropout_probability": 0.0,
+    "batch_norm": False,
+    "num_bins": 4,
+    "base_transform_type": "rq-coupling",
+}
+
+
+# Network-input shape served by the fake context below; the embedding network is
+# sized to consume exactly this.
+DATA_SHAPE = (2, 3, 20)
+
+EMBEDDING_KWARGS = {
+    "input_dims": DATA_SHAPE,
+    "svd": {"size": 10},
+    "V_rb_list": None,
+    "output_dim": 8,
+    "hidden_dims": [32, 16, 8],
+    "activation": "elu",
+    "dropout": 0.0,
+    "batch_norm": False,
+    "added_context": False,
+}
+
+
+def _build_model(unconditional, context_dim=None, embedding_kwargs=None):
+    """A small (real but untrained) normalizing-flow posterior model."""
+    model_kwargs = {
+        "posterior_model_type": "normalizing_flow",
+        "posterior_kwargs": {
+            "input_dim": len(INFERENCE_PARAMETERS),
+            "context_dim": context_dim,
+            "num_flow_steps": 2,
+            "base_transform_kwargs": BASE_TRANSFORM_KWARGS,
+        },
+    }
+    if embedding_kwargs is not None:
+        model_kwargs["embedding_kwargs"] = embedding_kwargs
+    metadata = {
+        "train_settings": {
+            "model": model_kwargs,
+            "data": {
+                "inference_parameters": INFERENCE_PARAMETERS,
+                "standardization": STANDARDIZATION,
+                "unconditional": unconditional,
+            },
+        },
+    }
+    if unconditional:
+        # Unconditional models keep the precursor ("base") metadata separately; the
+        # FlowFactor must read the model's *own* standardization, not this.
+        metadata["base"] = {}
+    return NormalizingFlowPosteriorModel(metadata=metadata, device="cpu")
+
+
+def _build_unconditional_model():
+    return _build_model(unconditional=True)
+
+
+def _build_conditional_model():
+    """A data-conditional flow with an embedding network sized to ``DATA_SHAPE``."""
+    return _build_model(
+        unconditional=False,
+        context_dim=EMBEDDING_KWARGS["output_dim"],
+        embedding_kwargs=EMBEDDING_KWARGS,
+    )
+
+
+class _FakeDataContext:
+    """Minimal ``SamplerContext`` stub for the data-conditional ``FlowFactor`` branch:
+    ``prepared_data()`` serves a stored, *unbatched* network-input tensor (the factor
+    adds the batch dimension itself: ``unsqueeze(0)`` at sample time, ``expand`` in
+    ``log_prob``)."""
+
+    def __init__(self, data):
+        self._data = data
+
+    def prepared_data(self, conditioning=None):
+        assert conditioning is None  # plain NPE: no chain columns enter the data prep
+        return self._data
+
+
+@pytest.fixture()
+def sampler():
+    """A ``ComposedSampler`` over a single unconditional ``FlowFactor`` (context-free)."""
+    factor = FlowFactor.from_model(_build_unconditional_model())
+    return ComposedSampler(ChainComposer([factor]), context=None)
+
+
+def test_run_sampler_returns_correct_number_of_samples(sampler):
+    sampler.run_sampler(num_samples=17)
+    assert len(sampler.samples) == 17
+
+
+def test_samples_columns_are_inference_parameters_plus_log_prob(sampler):
+    sampler.run_sampler(num_samples=10)
+    assert list(sampler.samples.columns) == INFERENCE_PARAMETERS + ["log_prob"]
+
+
+@pytest.mark.parametrize("batch_size", [None, 3, 5, 7, 10, 25])
+def test_run_sampler_count_invariant_to_batch_size(sampler, batch_size):
+    """The number of samples must not depend on batch_size.
+
+    Covers the chunking in ComposedSampler.run_sampler / chunk_and_concat: exact
+    divisors (5, 10), divisors with a remainder (3, 7), a batch larger than the request
+    (25 -> one short batch), and the default single-batch case (None).
+    """
+    num_samples = 10
+    sampler.run_sampler(num_samples=num_samples, batch_size=batch_size)
+    assert len(sampler.samples) == num_samples
+
+
+def test_flow_factor_log_prob_round_trip_matches_sampling_log_prob():
+    """log_prob recomputed at the sampled points reproduces the sampling log_prob.
+
+    This exercises the full standardize -> de-standardize round trip through a real
+    flow, including the change-of-variables correction (-sum(log(std))). The fake model
+    used in ``test_factors`` is deterministic and cannot reach this path.
+    """
+    factor = FlowFactor.from_model(_build_unconditional_model())
+    torch.manual_seed(0)
+    samples, log_prob = factor.sample_and_log_prob(20, context=None)
+    recomputed = factor.log_prob(samples, context=None)
+    # The network runs in float32, and sampling vs. log_prob use the flow's forward
+    # vs. inverse transforms, which round-trip only to the float32 floor (observed
+    # ~1e-6); atol=1e-5 leaves ~10x margin.
+    np.testing.assert_allclose(
+        recomputed.numpy(), log_prob.numpy(), atol=1e-5
+    )
+
+
+def test_conditional_samples_depend_on_context():
+    """For a data-conditional model the samples depend on the context data.
+
+    With the torch RNG fixed, the same context reproduces samples exactly, while a
+    different context changes them -- isolating ``prepared_data()`` as the source of
+    variation. This is the ``FlowFactor`` branch that draws from the shared data
+    context (an untrained flow + embedding network suffices).
+    """
+    factor = FlowFactor.from_model(_build_conditional_model())
+    context_a = _FakeDataContext(torch.randn(*DATA_SHAPE))
+    context_b = _FakeDataContext(torch.randn(*DATA_SHAPE))
+
+    torch.manual_seed(0)
+    samples_a, _ = factor.sample_and_log_prob(8, context_a)
+    torch.manual_seed(0)
+    samples_a_repeat, _ = factor.sample_and_log_prob(8, context_a)
+    torch.manual_seed(0)
+    samples_b, _ = factor.sample_and_log_prob(8, context_b)
+
+    assert all(samples_a[p].shape == (8,) for p in INFERENCE_PARAMETERS)
+    # Same seed + same context -> identical samples.
+    assert all(
+        torch.equal(samples_a[p], samples_a_repeat[p]) for p in INFERENCE_PARAMETERS
+    )
+    # Same seed + different context -> different samples.
+    assert not all(
+        torch.equal(samples_a[p], samples_b[p]) for p in INFERENCE_PARAMETERS
+    )
+
+
+def test_conditional_log_prob_round_trip_matches_sampling_log_prob():
+    """log_prob recomputed at the sampled points reproduces the sampling log_prob,
+    through the data-conditional branch.
+
+    Sampling conditions on ``prepared_data().unsqueeze(0)`` (one shared context row,
+    squeezed back out); ``log_prob`` instead expands the data across the sample rows
+    (``data.expand(num_samples, ...)``) -- the only CI coverage of that expand path.
+    """
+    factor = FlowFactor.from_model(_build_conditional_model())
+    context = _FakeDataContext(torch.randn(*DATA_SHAPE))
+    torch.manual_seed(0)
+    samples, log_prob = factor.sample_and_log_prob(20, context)
+    recomputed = factor.log_prob(samples, context)
+    # float32 round-trip floor, as in the unconditional round-trip test above.
+    np.testing.assert_allclose(recomputed.numpy(), log_prob.numpy(), atol=1e-5)
diff --git a/tests/core/test_samplers.py b/tests/core/test_samplers.py
deleted file mode 100644
index b86918826..000000000
--- a/tests/core/test_samplers.py
+++ /dev/null
@@ -1,261 +0,0 @@
-import numpy as np
-import pandas as pd
-import pytest
-import torch
-
-from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel
-from dingo.core.result import Result
-from dingo.core.samplers import Sampler
-
-
-INFERENCE_PARAMETERS = ["a", "b", "c"]
-
-# Non-trivial standardization so that the log_prob change-of-variables correction
-# (-sum(log(std))) is actually exercised by the round-trip test below.
-STANDARDIZATION = {
-    "mean": {"a": 1.0, "b": -2.0, "c": 0.5},
-    "std": {"a": 2.0, "b": 0.5, "c": 3.0},
-}
-
-BASE_TRANSFORM_KWARGS = {
-    "hidden_dim": 16,
-    "num_transform_blocks": 1,
-    "activation": "elu",
-    "dropout_probability": 0.0,
-    "batch_norm": False,
-    "num_bins": 4,
-    "base_transform_type": "rq-coupling",
-}
-
-
-def _build_model(unconditional, context_dim=None, embedding_kwargs=None):
-    """Build a small normalizing-flow posterior model for use with a Sampler."""
-    posterior_kwargs = {
-        "input_dim": len(INFERENCE_PARAMETERS),
-        "context_dim": context_dim,
-        "num_flow_steps": 2,
-        "base_transform_kwargs": BASE_TRANSFORM_KWARGS,
-    }
-    model_kwargs = {
-        "posterior_model_type": "normalizing_flow",
-        "posterior_kwargs": posterior_kwargs,
-    }
-    if embedding_kwargs is not None:
-        model_kwargs["embedding_kwargs"] = embedding_kwargs
-
-    metadata = {
-        "train_settings": {
-            "model": model_kwargs,
-            "data": {
-                "inference_parameters": INFERENCE_PARAMETERS,
-                "standardization": STANDARDIZATION,
-                "unconditional": unconditional,
-            },
-        }
-    }
-    if unconditional:
-        # For unconditional models the Sampler keeps the precursor ("base") metadata
-        # separately; it just needs to exist.
-        metadata["base"] = {}
-
-    return NormalizingFlowPosteriorModel(metadata=metadata, device="cpu")
-
-
-@pytest.fixture()
-def unconditional_sampler():
-    """Sampler wrapping an unconditional flow.
-
-    This takes the context-free (x = []) path through run_sampler / log_prob, which
-    isolates the count, batching, and log_prob logic with minimal setup.
-    """
-    return Sampler(model=_build_model(unconditional=True))
-
-
-@pytest.fixture()
-def conditional_sampler():
-    """Sampler wrapping a conditional flow, without an embedding network.
-
-    Sufficient for exercising the context-required error paths, which are checked
-    before the network is ever called.
-    """
-    return Sampler(model=_build_model(unconditional=False, context_dim=5))
-
-
-@pytest.fixture()
-def conditional_sampler_with_context():
-    """Sampler wrapping a conditional flow *with* an embedding network.
-
-    The full conditional sampling path runs here. The base Sampler's default
-    transform_pre is the identity (a subclass would replace it); we substitute a
-    minimal callable that extracts the context tensor so that data reaches the
-    embedding network.
-    """
-    embedding_kwargs = {
-        "input_dims": (2, 3, 20),
-        "svd": {"size": 10},
-        "V_rb_list": None,
-        "output_dim": 8,
-        "hidden_dims": [32, 16, 8],
-        "activation": "elu",
-        "dropout": 0.0,
-        "batch_norm": False,
-        "added_context": False,
-    }
-    sampler = Sampler(
-        model=_build_model(
-            unconditional=False, context_dim=8, embedding_kwargs=embedding_kwargs
-        )
-    )
-    sampler.transform_pre = lambda context: context["data"]
-    return sampler
-
-
-def test_run_sampler_returns_correct_number_of_samples(unconditional_sampler):
-    unconditional_sampler.run_sampler(num_samples=17)
-    assert len(unconditional_sampler.samples) == 17
-
-
-def test_samples_columns_are_inference_parameters_plus_log_prob(unconditional_sampler):
-    unconditional_sampler.run_sampler(num_samples=10)
-    assert list(unconditional_sampler.samples.columns) == INFERENCE_PARAMETERS + [
-        "log_prob"
-    ]
-
-
-@pytest.mark.parametrize("batch_size", [None, 3, 5, 7, 10, 25])
-def test_run_sampler_count_invariant_to_batch_size(unconditional_sampler, batch_size):
-    """The number of samples must not depend on batch_size.
-
-    Covers the divmod branch in run_sampler: exact divisors (5, 10), divisors with a
-    remainder (3, 7), a batch larger than the request (25 -> one short batch), and the
-    default single-batch case (None).
-    """
-    num_samples = 10
-    unconditional_sampler.run_sampler(num_samples=num_samples, batch_size=batch_size)
-    assert len(unconditional_sampler.samples) == num_samples
-
-
-def test_conditional_run_sampler_without_context_raises(conditional_sampler):
-    with pytest.raises(ValueError, match="Context must be set"):
-        conditional_sampler.run_sampler(num_samples=5)
-
-
-def test_conditional_log_prob_without_context_raises(conditional_sampler):
-    samples = pd.DataFrame({p: [0.0] for p in INFERENCE_PARAMETERS})
-    with pytest.raises(ValueError, match="Context must be set"):
-        conditional_sampler.log_prob(samples)
-
-
-def test_unconditional_run_sampler_needs_no_context(unconditional_sampler):
-    # Unconditional models carry no context and require none in order to sample.
-    assert unconditional_sampler.context is None
-    unconditional_sampler.run_sampler(num_samples=8)
-    assert len(unconditional_sampler.samples) == 8
-
-
-def test_log_prob_round_trip_matches_sampling_log_prob(unconditional_sampler):
-    """log_prob recomputed at the sampled points reproduces the sampling log_prob.
-
-    This exercises the full standardize -> de-standardize round trip, including the
-    change-of-variables correction (-sum(log(std))). The stored log_prob column must
-    be dropped first, since log_prob() adds any log_prob already present (see
-    test_log_prob_adds_existing_log_prob_column).
-    """
-    unconditional_sampler.run_sampler(num_samples=20)
-    samples = unconditional_sampler.samples
-    recomputed = unconditional_sampler.log_prob(samples.drop(columns="log_prob"))
-    # The network runs in float32, and sampling vs. log_prob use the flow's forward
-    # vs. inverse transforms, which round-trip only to the float32 floor (observed
-    # ~1e-6); atol=1e-5 leaves ~10x margin.
-    np.testing.assert_allclose(recomputed, samples["log_prob"].to_numpy(), atol=1e-5)
-
-
-def test_log_prob_adds_existing_log_prob_column(unconditional_sampler):
-    """A log_prob column in the input is added to the recomputed network log_prob.
-
-    This is relied upon by GNPE / density-recovery, where the carried log_prob is a
-    proposal contribution. The base Sampler does not drop it.
-    """
-    unconditional_sampler.run_sampler(num_samples=6)
-    samples = unconditional_sampler.samples
-    stored = samples["log_prob"].to_numpy()
-    recomputed = unconditional_sampler.log_prob(samples.drop(columns="log_prob"))
-    combined = unconditional_sampler.log_prob(samples)
-    # float32 round-trip floor, as in test_log_prob_round_trip_matches_sampling_log_prob.
-    np.testing.assert_allclose(combined, recomputed + stored, atol=1e-5)
-
-
-def test_log_prob_accepts_dataframe_and_dict_inputs(unconditional_sampler):
-    unconditional_sampler.run_sampler(num_samples=6)
-    samples = unconditional_sampler.samples.drop(columns="log_prob")
-
-    # DataFrame -> one log_prob per row.
-    lp_dataframe = unconditional_sampler.log_prob(samples)
-    assert lp_dataframe.shape == (6,)
-
-    # dict of arrays -> one log_prob per row.
-    lp_dict = unconditional_sampler.log_prob(
-        {p: samples[p].to_numpy() for p in INFERENCE_PARAMETERS}
-    )
-    assert lp_dict.shape == (6,)
-
-    # dict of scalars -> a single log_prob.
-    lp_scalar = unconditional_sampler.log_prob(
-        {p: samples[p].iloc[0] for p in INFERENCE_PARAMETERS}
-    )
-    assert lp_scalar.shape == (1,)
-
-
-def test_to_result_round_trips_samples_and_settings(unconditional_sampler):
-    unconditional_sampler.run_sampler(num_samples=10)
-    result = unconditional_sampler.to_result()
-    assert isinstance(result, Result)
-    pd.testing.assert_frame_equal(result.samples, unconditional_sampler.samples)
-    assert result.settings == unconditional_sampler.metadata
-
-
-def test_to_hdf5_round_trips_samples(unconditional_sampler, tmp_path):
-    unconditional_sampler.run_sampler(num_samples=10)
-    unconditional_sampler.to_hdf5(label="result", outdir=str(tmp_path))
-
-    reloaded = Result(file_name=str(tmp_path / "result.hdf5"))
-    np.testing.assert_allclose(
-        reloaded.samples[INFERENCE_PARAMETERS].to_numpy(),
-        unconditional_sampler.samples[INFERENCE_PARAMETERS].to_numpy(),
-    )
-
-
-def test_context_setter_moves_parameters_to_event_metadata(conditional_sampler):
-    conditional_sampler.context = {"waveform": [1, 2, 3], "parameters": {"mass": 30.0}}
-    assert conditional_sampler.event_metadata["injection_parameters"] == {"mass": 30.0}
-    # The injected truth parameters are removed from the context itself.
-    assert "parameters" not in conditional_sampler.context
-
-
-def test_conditional_samples_depend_on_context(conditional_sampler_with_context):
-    """For a conditional model the samples depend on the context.
-
-    With the torch RNG fixed, the same context reproduces samples exactly, while a
-    different context changes them -- isolating the context as the source of variation.
-    """
-    sampler = conditional_sampler_with_context
-    context_a = {"data": torch.randn(2, 3, 20)}
-    context_b = {"data": torch.randn(2, 3, 20)}
-
-    torch.manual_seed(0)
-    sampler.context = context_a
-    sampler.run_sampler(num_samples=8)
-    samples_a = sampler.samples.to_numpy().copy()
-
-    torch.manual_seed(0)
-    sampler.context = context_a
-    sampler.run_sampler(num_samples=8)
-    samples_a_repeat = sampler.samples.to_numpy().copy()
-
-    torch.manual_seed(0)
-    sampler.context = context_b
-    sampler.run_sampler(num_samples=8)
-    samples_b = sampler.samples.to_numpy().copy()
-
-    assert np.array_equal(samples_a, samples_a_repeat)
-    assert not np.array_equal(samples_a, samples_b)
diff --git a/tests/gw/inference/test_gw_composed_sampler.py b/tests/gw/inference/test_gw_composed_sampler.py
new file mode 100644
index 000000000..555aa26fc
--- /dev/null
+++ b/tests/gw/inference/test_gw_composed_sampler.py
@@ -0,0 +1,162 @@
+"""Unit tests for ``GWComposedSampler`` assembly and export.
+
+``test_bns_chain_builder`` covers the pinned (DINGO-BNS) ``from_model`` chain; this file
+covers the two GW-builder concerns the branch's tests otherwise leave untested:
+
+* the *plain-NPE* ``from_model`` chain, including the ``DeltaFactor`` that fills the
+  network-independent delta-prior parameters (an aligned spin fixed to 0), and
+* ``to_result`` / ``to_hdf5`` -- exporting the sampled DataFrame, the raw event data, and
+  the chain provenance to a gw ``Result`` (and back through HDF5).
+
+Both use metadata-only stubs and a network-free ``DeltaFactor`` chain, so no waveform
+generation or trained network is needed.
+
+(Adapts the ``run_sampler`` export coverage and the fixed-prior post-processing check
+from the pre-sampler-revamp ``test_samplers.py`` / ``test_gw_samplers.py``, which
+imported the now-deleted ``dingo.core.samplers`` and ``dingo.gw.inference.gw_samplers``.)
+"""
+
+import numpy as np
+import pytest
+
+from dingo.core.factors import ChainComposer, DeltaFactor
+from dingo.gw.inference.context import GWSamplerContext
+from dingo.gw.inference.sampler import GWComposedSampler
+from dingo.gw.result import Result
+
+
+DETECTORS = ["H1", "L1"]
+REF_TIME = 1126259462.4
+INFERENCE_PARAMETERS = ["chirp_mass", "mass_ratio", "ra", "dec"]
+
+DOMAIN_SETTINGS = {
+    "type": "UniformFrequencyDomain",
+    "f_min": 20.0,
+    "f_max": 256.0,  # small (T=2s) so any data view is cheap
+    "delta_f": 0.5,
+}
+
+STANDARDIZATION = {
+    "mean": {p: 0.0 for p in INFERENCE_PARAMETERS},
+    "std": {p: 1.0 for p in INFERENCE_PARAMETERS},
+}
+
+INTRINSIC_PRIOR = {
+    "mass_1": "bilby.core.prior.Constraint(minimum=10, maximum=80)",
+    "mass_2": "bilby.core.prior.Constraint(minimum=10, maximum=80)",
+    "chirp_mass": "bilby.gw.prior.UniformInComponentsChirpMass(minimum=25, maximum=31)",
+    "mass_ratio": "bilby.gw.prior.UniformInComponentsMassRatio("
+    "minimum=0.125, maximum=1)",
+    "phase": "default",
+    # Aligned spins pinned to 0: delta-prior parameters the network does not infer,
+    # so the chain fills them with a DeltaFactor.
+    "a_1": 0.0,
+    "a_2": 0.0,
+}
+EXTRINSIC_PRIOR = {
+    "dec": "default",
+    "ra": "default",
+    "geocent_time": "default",
+    "luminosity_distance": "default",
+    "psi": "default",
+}
+
+
+def _metadata():
+    return {
+        "dataset_settings": {
+            "domain": DOMAIN_SETTINGS,
+            "waveform_generator": {"approximant": "IMRPhenomD", "f_ref": 20.0},
+            "intrinsic_prior": INTRINSIC_PRIOR,
+        },
+        "train_settings": {
+            "data": {
+                "detectors": DETECTORS,
+                "ref_time": REF_TIME,
+                "extrinsic_prior": EXTRINSIC_PRIOR,
+                "inference_parameters": INFERENCE_PARAMETERS,
+                "standardization": STANDARDIZATION,
+            }
+        },
+    }
+
+
+class _StubModel:
+    metadata = _metadata()
+    device = "cpu"
+
+
+def _event_data():
+    n_bins = int(DOMAIN_SETTINGS["f_max"] / DOMAIN_SETTINGS["delta_f"]) + 1
+    strain = np.ones(n_bins, dtype=complex)
+    asd = np.ones(n_bins)
+    return {
+        "waveform": {d: strain.copy() for d in DETECTORS},
+        "asds": {d: asd.copy() for d in DETECTORS},
+    }
+
+
+# ---------------------------------------------------------------------------
+# from_model: plain-NPE chain assembly + delta-prior filling
+# ---------------------------------------------------------------------------
+
+
+def test_from_model_plain_npe_chain_fills_delta_prior_parameters():
+    sampler = GWComposedSampler.from_model(_StubModel(), _event_data())
+    kinds = [type(s).__name__ for s in sampler.composer.steps]
+    # The flow (exposing ra as ra@t_ref), the RA rotation to the event frame, and a
+    # trailing DeltaFactor for the network-independent delta-prior parameters.
+    assert kinds == ["FlowFactor", "RAToEventFrame", "DeltaFactor"]
+
+    delta = sampler.composer.steps[-1]
+    # a_1 and a_2 are DeltaFunctions (0.0) in the intrinsic prior and are not inference
+    # parameters, so they are pinned as fixed constants; chirp_mass etc. are not.
+    assert set(delta.parameters) == {"a_1", "a_2"}
+    assert delta.values == {"a_1": 0.0, "a_2": 0.0}
+
+
+# ---------------------------------------------------------------------------
+# to_result / to_hdf5: export round trips
+# ---------------------------------------------------------------------------
+
+
+def _network_free_sampler():
+    """A GWComposedSampler whose chain pins every inference parameter, so run_sampler
+    produces a DataFrame with no network involved."""
+    context = GWSamplerContext.from_model_metadata(_metadata(), _event_data())
+    pins = {"chirp_mass": 30.0, "mass_ratio": 0.8, "ra": 1.0, "dec": 0.1}
+    composer = ChainComposer([DeltaFactor(pins)])
+    return GWComposedSampler(composer, context, _metadata(), INFERENCE_PARAMETERS)
+
+
+def test_to_result_round_trips_samples_and_provenance():
+    sampler = _network_free_sampler()
+    sampler.run_sampler(num_samples=10)
+
+    result = sampler.to_result()
+    assert isinstance(result, Result)
+    # The sampled DataFrame is exported unchanged.
+    assert len(result.samples) == 10
+    assert list(result.samples.columns) == list(sampler.samples.columns)
+    np.testing.assert_allclose(
+        result.samples[INFERENCE_PARAMETERS].to_numpy(),
+        sampler.samples[INFERENCE_PARAMETERS].to_numpy(),
+    )
+    # Model metadata is carried through, and the chain provenance is recorded.
+    assert result.settings["train_settings"] == sampler.metadata["train_settings"]
+    provenance = result.settings["sampler"]
+    assert provenance["implementation"] == "composed"
+    assert [step["step"] for step in provenance["chain"]] == ["DeltaFactor"]
+
+
+def test_to_hdf5_round_trips_samples(tmp_path):
+    sampler = _network_free_sampler()
+    sampler.run_sampler(num_samples=10)
+    sampler.to_hdf5(label="result", outdir=str(tmp_path))
+
+    reloaded = Result(file_name=str(tmp_path / "result.hdf5"))
+    assert len(reloaded.samples) == 10
+    np.testing.assert_allclose(
+        reloaded.samples[INFERENCE_PARAMETERS].to_numpy(),
+        sampler.samples[INFERENCE_PARAMETERS].to_numpy(),
+    )
diff --git a/tests/gw/inference/test_gw_samplers.py b/tests/gw/inference/test_gw_samplers.py
deleted file mode 100644
index dbeb0bc89..000000000
--- a/tests/gw/inference/test_gw_samplers.py
+++ /dev/null
@@ -1,341 +0,0 @@
-import numpy as np
-import pytest
-from astropy.time import Time
-from astropy.utils import iers
-
-from dingo.core.posterior_models.normalizing_flow import NormalizingFlowPosteriorModel
-from dingo.gw.domains import UniformFrequencyDomain
-from dingo.gw.inference.gw_samplers import (
-    GWSampler,
-    check_frequency_updates,
-    _validate_maximum_frequency,
-    _validate_minimum_frequency,
-)
-
-# Avoid network access (and the associated timeout) when astropy computes sidereal
-# time in _correct_reference_time; the bundled IERS data is sufficient for tests.
-iers.conf.auto_download = False
-
-
-DETECTORS = ["H1", "L1"]
-INFERENCE_PARAMETERS = ["chirp_mass", "mass_ratio", "ra", "dec"]
-
-DOMAIN_SETTINGS = {
-    "type": "UniformFrequencyDomain",
-    "f_min": 20.0,
-    "f_max": 1024.0,
-    "delta_f": 0.25,
-}
-
-
-# ---------------------------------------------------------------------------
-# Frequency-range validators (pure functions, only need a domain).
-# ---------------------------------------------------------------------------
-
-
-@pytest.fixture()
-def domain():
-    return UniformFrequencyDomain(f_min=20.0, f_max=1024.0, delta_f=0.25)
-
-
-@pytest.mark.parametrize(
-    "validate, valid_change, beyond_bound",
-    [
-        (_validate_minimum_frequency, 40.0, 10.0),  # raise f_min; below hard f_min
-        (_validate_maximum_frequency, 512.0, 2048.0),  # lower f_max; above hard f_max
-    ],
-)
-def test_frequency_validator_no_op_when_unchanged(
-    domain, validate, valid_change, beyond_bound
-):
-    # Value equal to the domain bound is a no-op and is allowed even without cropping.
-    bound = domain.f_min if validate is _validate_minimum_frequency else domain.f_max
-    assert validate(bound, DETECTORS, domain, None) is None
-
-
-@pytest.mark.parametrize(
-    "validate, valid_change, beyond_bound",
-    [
-        (_validate_minimum_frequency, 40.0, 10.0),
-        (_validate_maximum_frequency, 512.0, 2048.0),
-    ],
-)
-def test_frequency_validator_expands_float_to_all_detectors(
-    domain, validate, valid_change, beyond_bound
-):
-    # A float applies to every detector; a valid change passes with cropping on.
-    # The cap/floor must be given explicitly, else it defaults to the domain bound.
-    crop = {"cropping_probability": 0.5, "f_min_upper": 100.0, "f_max_lower": 400.0}
-    assert validate(valid_change, DETECTORS, domain, crop) is None
-
-
-@pytest.mark.parametrize(
-    "validate, valid_change, beyond_bound",
-    [
-        (_validate_minimum_frequency, 40.0, 10.0),
-        (_validate_maximum_frequency, 512.0, 2048.0),
-    ],
-)
-def test_frequency_validator_rejects_value_beyond_hard_bound(
-    domain, validate, valid_change, beyond_bound
-):
-    crop = {"cropping_probability": 0.5}
-    with pytest.raises(ValueError, match="domain.f_"):
-        validate(beyond_bound, DETECTORS, domain, crop)
-
-
-@pytest.mark.parametrize(
-    "validate, valid_change",
-    [(_validate_minimum_frequency, 40.0), (_validate_maximum_frequency, 512.0)],
-)
-def test_frequency_validator_rejects_detector_key_mismatch(
-    domain, validate, valid_change
-):
-    crop = {"cropping_probability": 0.5}
-    with pytest.raises(ValueError, match="exactly detectors"):
-        validate({"H1": valid_change}, DETECTORS, domain, crop)
-
-
-@pytest.mark.parametrize(
-    "validate, valid_change",
-    [(_validate_minimum_frequency, 40.0), (_validate_maximum_frequency, 512.0)],
-)
-def test_frequency_validator_rejects_change_when_cropping_disabled(
-    domain, validate, valid_change
-):
-    # No crop settings at all.
-    with pytest.raises(ValueError, match="[Cc]ropping"):
-        validate(valid_change, DETECTORS, domain, None)
-    # Crop settings present but probability zero.
-    with pytest.raises(ValueError, match="[Cc]ropping"):
-        validate(valid_change, DETECTORS, domain, {"cropping_probability": 0.0})
-
-
-def test_validate_minimum_frequency_rejects_value_above_cap(domain):
-    crop = {"cropping_probability": 0.5, "f_min_upper": 60.0}
-    assert _validate_minimum_frequency(50.0, DETECTORS, domain, crop) is None
-    with pytest.raises(ValueError, match="upper bound"):
-        _validate_minimum_frequency(80.0, DETECTORS, domain, crop)
-
-
-def test_validate_maximum_frequency_rejects_value_below_floor(domain):
-    crop = {"cropping_probability": 0.5, "f_max_lower": 400.0}
-    assert _validate_maximum_frequency(500.0, DETECTORS, domain, crop) is None
-    with pytest.raises(ValueError, match="lower bound"):
-        _validate_maximum_frequency(300.0, DETECTORS, domain, crop)
-
-
-def test_validate_minimum_frequency_rejects_differing_values_when_not_independent(
-    domain,
-):
-    crop = {
-        "cropping_probability": 0.5,
-        "independent_detectors": False,
-        "f_min_upper": 100.0,
-    }
-    with pytest.raises(ValueError, match="[Ii]ndependent"):
-        _validate_minimum_frequency({"H1": 40.0, "L1": 50.0}, DETECTORS, domain, crop)
-
-
-def test_check_frequency_updates_accepts_valid_and_rejects_invalid():
-    model_metadata = {
-        "train_settings": {
-            "data": {
-                "detectors": DETECTORS,
-                "random_strain_cropping": {
-                    "cropping_probability": 0.5,
-                    "f_min_upper": 100.0,
-                    "f_max_lower": 400.0,
-                },
-            }
-        },
-        "dataset_settings": {"domain": DOMAIN_SETTINGS},
-    }
-    # Valid frequency updates pass without raising.
-    assert check_frequency_updates(model_metadata, f_min=40.0, f_max=512.0) is None
-    # Beyond the hard bound raises.
-    with pytest.raises(ValueError, match="domain.f_min"):
-        check_frequency_updates(model_metadata, f_min=10.0)
-
-
-# ---------------------------------------------------------------------------
-# GWSamplerMixin methods (lightweight GWSampler; network not exercised).
-# ---------------------------------------------------------------------------
-
-
-def _build_gw_sampler(unconditional=False, domain_update=None):
-    """Build a GWSampler around a tiny flow plus minimal but valid GW metadata.
-
-    The network is never run by the methods under test here; it only needs to exist.
-    """
-    standardization = {
-        "mean": {p: 0.0 for p in INFERENCE_PARAMETERS},
-        "std": {p: 1.0 for p in INFERENCE_PARAMETERS},
-    }
-    posterior_kwargs = {
-        "input_dim": len(INFERENCE_PARAMETERS),
-        "context_dim": None,
-        "num_flow_steps": 2,
-        "base_transform_kwargs": {
-            "hidden_dim": 8,
-            "num_transform_blocks": 1,
-            "activation": "elu",
-            "dropout_probability": 0.0,
-            "batch_norm": False,
-            "num_bins": 4,
-            "base_transform_type": "rq-coupling",
-        },
-    }
-    data_settings = {
-        "unconditional": unconditional,
-        "inference_parameters": INFERENCE_PARAMETERS,
-        "standardization": standardization,
-        "detectors": DETECTORS,
-        "ref_time": 1126259462.4,
-        "extrinsic_prior": {
-            "dec": "default",
-            "ra": "default",
-            "geocent_time": "default",
-            "luminosity_distance": "default",
-            "psi": "default",
-        },
-    }
-    if domain_update is not None:
-        data_settings["domain_update"] = domain_update
-
-    metadata = {
-        "train_settings": {
-            "model": {
-                "posterior_model_type": "normalizing_flow",
-                "posterior_kwargs": posterior_kwargs,
-            },
-            "data": data_settings,
-        },
-        "dataset_settings": {
-            "domain": DOMAIN_SETTINGS,
-            "intrinsic_prior": {
-                "mass_1": "bilby.core.prior.Constraint(minimum=10, maximum=80)",
-                "mass_2": "bilby.core.prior.Constraint(minimum=10, maximum=80)",
-                "chirp_mass": "bilby.gw.prior.UniformInComponentsChirpMass("
-                "minimum=25, maximum=31)",
-                "mass_ratio": "bilby.gw.prior.UniformInComponentsMassRatio("
-                "minimum=0.125, maximum=1)",
-                "phase": "default",
-                "a_1": 0.0,
-                "a_2": 0.0,
-            },
-        },
-    }
-    if unconditional:
-        metadata["base"] = metadata
-    model = NormalizingFlowPosteriorModel(metadata=metadata, device="cpu")
-    return GWSampler(model=model)
-
-
-@pytest.fixture()
-def gw_sampler():
-    return _build_gw_sampler()
-
-
-def test_build_domain_from_metadata(gw_sampler):
-    assert isinstance(gw_sampler.domain, UniformFrequencyDomain)
-    assert gw_sampler.domain.f_min == DOMAIN_SETTINGS["f_min"]
-    assert gw_sampler.domain.f_max == DOMAIN_SETTINGS["f_max"]
-
-
-def test_build_domain_applies_domain_update():
-    sampler = _build_gw_sampler(domain_update={"f_min": 30.0})
-    assert sampler.domain.f_min == 30.0
-
-
-def test_correct_reference_time_round_trip(gw_sampler):
-    gw_sampler._event_metadata = {"time_event": gw_sampler.t_ref + 3600.0}
-    samples = {"ra": np.array([0.5, 1.5, 2.5]), "dec": np.array([0.1, 0.2, 0.3])}
-    original_ra = samples["ra"].copy()
-
-    gw_sampler._correct_reference_time(samples, inverse=False)
-    assert not np.allclose(samples["ra"], original_ra)
-    assert np.all((samples["ra"] >= 0) & (samples["ra"] < 2 * np.pi))
-
-    gw_sampler._correct_reference_time(samples, inverse=True)
-    np.testing.assert_allclose(samples["ra"], original_ra)
-
-
-def test_correct_reference_time_matches_sidereal_shift(gw_sampler):
-    """The RA shift must equal the difference in apparent sidereal time."""
-    t_event = gw_sampler.t_ref + 3600.0
-    gw_sampler._event_metadata = {"time_event": t_event}
-    samples = {"ra": np.array([0.5, 1.5, 2.5])}
-    original_ra = samples["ra"].copy()
-
-    ra_correction = (
-        Time(t_event, format="gps", scale="utc").sidereal_time("apparent", "greenwich")
-        - Time(gw_sampler.t_ref, format="gps", scale="utc").sidereal_time(
-            "apparent", "greenwich"
-        )
-    ).rad
-
-    gw_sampler._correct_reference_time(samples, inverse=False)
-    np.testing.assert_allclose(
-        samples["ra"], (original_ra + ra_correction) % (2 * np.pi)
-    )
-
-
-def test_correct_reference_time_noop_when_time_matches_reference(gw_sampler):
-    gw_sampler._event_metadata = {"time_event": gw_sampler.t_ref}
-    samples = {"ra": np.array([0.5, 1.5])}
-    original_ra = samples["ra"].copy()
-    gw_sampler._correct_reference_time(samples, inverse=False)
-    np.testing.assert_allclose(samples["ra"], original_ra)
-
-
-def test_correct_reference_time_noop_without_ra(gw_sampler):
-    gw_sampler._event_metadata = {"time_event": gw_sampler.t_ref + 3600.0}
-    samples = {"dec": np.array([0.1, 0.2])}
-    # No "ra" key: nothing to correct, and no error.
-    gw_sampler._correct_reference_time(samples, inverse=False)
-    assert "ra" not in samples
-
-
-def test_post_process_forward_adds_fixed_prior_parameters(gw_sampler):
-    samples = {p: np.zeros(5) for p in INFERENCE_PARAMETERS}
-    gw_sampler._event_metadata = None
-    gw_sampler._post_process(samples, inverse=False)
-    # a_1 and a_2 are DeltaFunctions (0.0) in the intrinsic prior; they get added.
-    for fixed in ("a_1", "a_2"):
-        assert fixed in samples
-        np.testing.assert_array_equal(samples[fixed], np.zeros(5))
-
-
-def test_post_process_inverse_drops_non_inference_parameters(gw_sampler):
-    samples = {
-        "chirp_mass": np.array([28.0]),
-        "ra": np.array([1.0]),
-        "log_prob": np.array([0.5]),
-        "extra": np.array([9.0]),
-    }
-    gw_sampler._event_metadata = None
-    gw_sampler._post_process(samples, inverse=True)
-    assert set(samples) <= set(INFERENCE_PARAMETERS)
-    assert "log_prob" not in samples
-    assert "extra" not in samples
-
-
-def test_frequency_updates_flag(gw_sampler):
-    # By default the requested range equals the domain, so no updates are flagged and
-    # the min/max frequencies report the domain bounds.
-    assert gw_sampler.minimum_frequency == gw_sampler.domain.f_min
-    assert gw_sampler.maximum_frequency == gw_sampler.domain.f_max
-    assert gw_sampler.frequency_updates is False
-
-    # A requested minimum frequency that differs from the domain flags an update.
-    # (Set the private attribute directly to bypass the validating setter, which
-    # would also rebuild the transforms.)
-    gw_sampler._minimum_frequency = 40.0
-    assert gw_sampler.frequency_updates is True
-
-
-def test_event_metadata_injects_frequency_bounds(gw_sampler):
-    metadata = gw_sampler.event_metadata
-    assert metadata["minimum_frequency"] == gw_sampler.domain.f_min
-    assert metadata["maximum_frequency"] == gw_sampler.domain.f_max
diff --git a/tests/gw/test_frequency_updates.py b/tests/gw/test_frequency_updates.py
new file mode 100644
index 000000000..339a33939
--- /dev/null
+++ b/tests/gw/test_frequency_updates.py
@@ -0,0 +1,159 @@
+"""Unit tests for the frequency-range validators (``dingo.gw.frequency_updates``):
+the pure functions that gate event- and importance-sampling-time minimum/maximum
+frequency updates against the network's training domain and its random-strain
+cropping licence.
+
+These are called at INI-parse time (``dingo_pipe``) and, via ``GWSamplerContext``,
+whenever event metadata carries ``minimum_frequency`` / ``maximum_frequency``. The
+context exercises them indirectly through ``prepared_data`` (see
+``test_gw_sampler_context``); this file covers the functions directly, including the
+per-detector cap/floor, detector-key, independent-detector, and cropping-licence
+branches the context-level tests do not reach.
+
+(Adapted from the validator block of the pre-sampler-revamp ``test_gw_samplers.py``,
+which imported them from the now-deleted ``dingo.gw.inference.gw_samplers``.)
+"""
+
+import pytest
+
+from dingo.gw.domains import UniformFrequencyDomain
+from dingo.gw.frequency_updates import (
+    check_frequency_updates,
+    _validate_maximum_frequency,
+    _validate_minimum_frequency,
+)
+
+
+DETECTORS = ["H1", "L1"]
+
+DOMAIN_SETTINGS = {
+    "type": "UniformFrequencyDomain",
+    "f_min": 20.0,
+    "f_max": 1024.0,
+    "delta_f": 0.25,
+}
+
+
+@pytest.fixture()
+def domain():
+    return UniformFrequencyDomain(f_min=20.0, f_max=1024.0, delta_f=0.25)
+
+
+@pytest.mark.parametrize(
+    "validate, valid_change, beyond_bound",
+    [
+        (_validate_minimum_frequency, 40.0, 10.0),  # raise f_min; below hard f_min
+        (_validate_maximum_frequency, 512.0, 2048.0),  # lower f_max; above hard f_max
+    ],
+)
+def test_frequency_validator_no_op_when_unchanged(
+    domain, validate, valid_change, beyond_bound
+):
+    # Value equal to the domain bound is a no-op and is allowed even without cropping.
+    bound = domain.f_min if validate is _validate_minimum_frequency else domain.f_max
+    assert validate(bound, DETECTORS, domain, None) is None
+
+
+@pytest.mark.parametrize(
+    "validate, valid_change, beyond_bound",
+    [
+        (_validate_minimum_frequency, 40.0, 10.0),
+        (_validate_maximum_frequency, 512.0, 2048.0),
+    ],
+)
+def test_frequency_validator_expands_float_to_all_detectors(
+    domain, validate, valid_change, beyond_bound
+):
+    # A float applies to every detector; a valid change passes with cropping on.
+    # The cap/floor must be given explicitly, else it defaults to the domain bound.
+    crop = {"cropping_probability": 0.5, "f_min_upper": 100.0, "f_max_lower": 400.0}
+    assert validate(valid_change, DETECTORS, domain, crop) is None
+
+
+@pytest.mark.parametrize(
+    "validate, valid_change, beyond_bound",
+    [
+        (_validate_minimum_frequency, 40.0, 10.0),
+        (_validate_maximum_frequency, 512.0, 2048.0),
+    ],
+)
+def test_frequency_validator_rejects_value_beyond_hard_bound(
+    domain, validate, valid_change, beyond_bound
+):
+    crop = {"cropping_probability": 0.5}
+    with pytest.raises(ValueError, match="domain.f_"):
+        validate(beyond_bound, DETECTORS, domain, crop)
+
+
+@pytest.mark.parametrize(
+    "validate, valid_change",
+    [(_validate_minimum_frequency, 40.0), (_validate_maximum_frequency, 512.0)],
+)
+def test_frequency_validator_rejects_detector_key_mismatch(
+    domain, validate, valid_change
+):
+    crop = {"cropping_probability": 0.5}
+    with pytest.raises(ValueError, match="exactly detectors"):
+        validate({"H1": valid_change}, DETECTORS, domain, crop)
+
+
+@pytest.mark.parametrize(
+    "validate, valid_change",
+    [(_validate_minimum_frequency, 40.0), (_validate_maximum_frequency, 512.0)],
+)
+def test_frequency_validator_rejects_change_when_cropping_disabled(
+    domain, validate, valid_change
+):
+    # No crop settings at all.
+    with pytest.raises(ValueError, match="[Cc]ropping"):
+        validate(valid_change, DETECTORS, domain, None)
+    # Crop settings present but probability zero.
+    with pytest.raises(ValueError, match="[Cc]ropping"):
+        validate(valid_change, DETECTORS, domain, {"cropping_probability": 0.0})
+
+
+def test_validate_minimum_frequency_rejects_value_above_cap(domain):
+    crop = {"cropping_probability": 0.5, "f_min_upper": 60.0}
+    assert _validate_minimum_frequency(50.0, DETECTORS, domain, crop) is None
+    with pytest.raises(ValueError, match="upper bound"):
+        _validate_minimum_frequency(80.0, DETECTORS, domain, crop)
+
+
+def test_validate_maximum_frequency_rejects_value_below_floor(domain):
+    crop = {"cropping_probability": 0.5, "f_max_lower": 400.0}
+    assert _validate_maximum_frequency(500.0, DETECTORS, domain, crop) is None
+    with pytest.raises(ValueError, match="lower bound"):
+        _validate_maximum_frequency(300.0, DETECTORS, domain, crop)
+
+
+def test_validate_minimum_frequency_rejects_differing_values_when_not_independent(
+    domain,
+):
+    crop = {
+        "cropping_probability": 0.5,
+        "independent_detectors": False,
+        "f_min_upper": 100.0,
+    }
+    with pytest.raises(ValueError, match="[Ii]ndependent"):
+        _validate_minimum_frequency({"H1": 40.0, "L1": 50.0}, DETECTORS, domain, crop)
+
+
+def test_check_frequency_updates_accepts_valid_and_rejects_invalid():
+    model_metadata = {
+        "train_settings": {
+            "data": {
+                "detectors": DETECTORS,
+                "random_strain_cropping": {
+                    "cropping_probability": 0.5,
+                    "f_min_upper": 100.0,
+                    "f_max_lower": 400.0,
+                },
+            }
+        },
+        "dataset_settings": {"domain": DOMAIN_SETTINGS},
+    }
+    # Valid frequency updates pass without raising.
+    assert check_frequency_updates(model_metadata, f_min=40.0, f_max=512.0) is None
+    # Beyond the hard bound raises.
+    with pytest.raises(ValueError, match="domain.f_min"):
+        check_frequency_updates(model_metadata, f_min=10.0)
diff --git a/tests/gw/test_ra_frames.py b/tests/gw/test_ra_frames.py
index 06b746559..dc29c13e9 100644
--- a/tests/gw/test_ra_frames.py
+++ b/tests/gw/test_ra_frames.py
@@ -7,9 +7,15 @@
 
 import numpy as np
 import torch
+from astropy.time import Time
+from astropy.utils import iers
 
 from dingo.gw.inference.steps import RAToEventFrame
 
+# Avoid network access (and the associated timeout) when astropy computes sidereal
+# time; the bundled IERS data is sufficient for tests.
+iers.conf.auto_download = False
+
 
 def _context(t_ref, t_event):
     """A minimal stand-in for the SamplerContext fields RAToEventFrame reads."""
@@ -30,6 +36,30 @@ def test_ra_reparam_round_trip():
     assert torch.allclose(back, ra_tref, atol=1e-6)  # forward then inverse recovers it
 
 
+def test_ra_reparam_matches_sidereal_shift():
+    """The forward rotation must shift `ra` by the apparent-sidereal-time difference.
+
+    The round-trip test above holds for any bijection; this pins the magnitude and
+    direction of the rotation to the physical sidereal shift (event minus reference).
+    """
+    t_ref, t_event = 1126259462.4, 1126259462.4 + 3600.0  # 1 h -> real rotation
+    rp = RAToEventFrame()
+    ctx = _context(t_ref=t_ref, t_event=t_event)
+    ra_tref = torch.rand(1000, dtype=torch.float32) * (2 * np.pi)
+
+    ra_correction = (
+        Time(t_event, format="gps", scale="utc").sidereal_time(
+            "apparent", "greenwich"
+        )
+        - Time(t_ref, format="gps", scale="utc").sidereal_time("apparent", "greenwich")
+    ).rad
+
+    ra = rp.forward({"ra@t_ref": ra_tref}, ctx)["ra"]
+    expected = ((ra_tref.double() + ra_correction) % (2 * np.pi)).float()
+    assert torch.equal(ra, expected)
+    assert torch.all((ra >= 0) & (ra < 2 * np.pi))
+
+
 def test_ra_reparam_is_noop_without_a_distinct_event_time():
     rp = RAToEventFrame()
     ra_tref = torch.rand(100, dtype=torch.float32) * (2 * np.pi)

From 7d3399044a0b03a26b7e421969394388ddf7afe0 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 08:54:41 +0100
Subject: [PATCH 57/65] Ignore CLAUDE.local.md

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .gitignore | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.gitignore b/.gitignore
index 3a5bb4048..6de552d62 100644
--- a/.gitignore
+++ b/.gitignore
@@ -148,6 +148,9 @@ dmypy.json
 # .DS_Store files for Mac OS
 .DS_Store
 
+# Claude Code personal notes
+CLAUDE.local.md
+
 # directories with temporary files
 **/tmp_files/*
 

From f78d956b984f75cc6c2e220f1f787eab59ce0ad7 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 13:54:17 +0100
Subject: [PATCH 58/65] Fix device handling for table-rooted chains on CUDA

GPU testing of the chirp-mass scan exposed two device leaks: the
SampleTableFactor emitted its fixed table on the CPU, feeding CPU
conditioning into a CUDA network, and the consumed conditioning columns
reached the numpy data-preparation world as CUDA tensors. The table now
joins the chain on context.device (the same policy as DeltaFactor), the
conditioning columns are normalized to the host alongside their float64
dtype, and the scan moves chain outputs to the CPU before numpy. Adds a
CUDA-gated regression test. Validated on an A100: the GW170817 scan
recovers the trigger (1.19758, snr 32.9) and matches the pinned run's
evidence within Monte Carlo error.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/core/factors.py         | 11 ++++++++++-
 dingo/gw/inference/context.py | 15 +++++++++------
 dingo/gw/inference/scan.py    |  2 +-
 tests/core/test_factors.py    | 19 +++++++++++++++++++
 4 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/dingo/core/factors.py b/dingo/core/factors.py
index cff2b3e3f..88e517c37 100644
--- a/dingo/core/factors.py
+++ b/dingo/core/factors.py
@@ -473,7 +473,16 @@ def sample_and_log_prob(self, num_samples, context, given=None):
                 f"A sample table is fixed: num_samples must equal the table length "
                 f"({n}), got {num_samples}. Run the chain with batch_size=None."
             )
-        return dict(self.table), self.table_log_prob
+        # The table's fresh tensors join the chain on its device (the same policy
+        # as DeltaFactor), so a table-rooted chain can condition a CUDA network.
+        device = getattr(context, "device", None)
+        table = {k: v.to(device) for k, v in self.table.items()}
+        log_prob = (
+            self.table_log_prob.to(device)
+            if self.table_log_prob is not None
+            else None
+        )
+        return table, log_prob
 
     def log_prob(self, theta_i, context, given=None):
         raise NotImplementedError(
diff --git a/dingo/gw/inference/context.py b/dingo/gw/inference/context.py
index 15296c13d..fc4d65702 100644
--- a/dingo/gw/inference/context.py
+++ b/dingo/gw/inference/context.py
@@ -411,9 +411,10 @@ def prepared_data(self, conditioning=None) -> torch.Tensor:
     def _conditioning_columns(self, conditioning) -> dict[str, torch.Tensor]:
         """Collect the conditioning columns the preparation consumes, keyed by
         their physical names (the `_proxy` suffix names the chain column; the
-        transform chain reads the physical parameter), as float64 (the
-        heterodyne phase is computed in float64, and a float32 chain column
-        must not degrade it)."""
+        transform chain reads the physical parameter), as float64 on the host
+        (the heterodyne phase is computed in float64, and a float32 chain
+        column must not degrade it; the preparation is a numpy transform
+        chain, so columns from a CUDA chain come back to the CPU here)."""
         conditioning = conditioning or {}
         columns = {}
         for name in self.data_prep_conditioning:
@@ -423,9 +424,11 @@ def _conditioning_columns(self, conditioning) -> dict[str, torch.Tensor]:
                     f"conditioning `{name}`, which the caller does not provide "
                     f"(pass it to prepared_data, e.g. from the chain's pins)."
                 )
-            columns[name[: -len("_proxy")]] = torch.as_tensor(
-                conditioning[name], dtype=torch.float64
-            ).reshape(-1)
+            columns[name[: -len("_proxy")]] = (
+                torch.as_tensor(conditioning[name], dtype=torch.float64)
+                .cpu()
+                .reshape(-1)
+            )
         return columns
 
     def _broadcast_event(self, n_rows: int) -> dict:
diff --git a/dingo/gw/inference/scan.py b/dingo/gw/inference/scan.py
index 8bc893fbb..80eeff6b1 100644
--- a/dingo/gw/inference/scan.py
+++ b/dingo/gw/inference/scan.py
@@ -174,7 +174,7 @@ def chirp_mass_scan(
             + tail_steps
         )
         out, _ = composer.sample_and_log_prob(len(block), context)
-        frames.append(pd.DataFrame({k: v.numpy() for k, v in out.items()}))
+        frames.append(pd.DataFrame({k: v.cpu().numpy() for k, v in out.items()}))
     theta = pd.concat(frames, ignore_index=True)
 
     # Restrict to draws within the prior (unphysical combinations cannot be
diff --git a/tests/core/test_factors.py b/tests/core/test_factors.py
index d58fa5664..a695e839c 100644
--- a/tests/core/test_factors.py
+++ b/tests/core/test_factors.py
@@ -710,3 +710,22 @@ def test_all_point_mass_chain_expands_to_base():
     chain = ChainComposer([DeltaFactor({"a": 2.0, "b": -1.0})])
     samples, log_prob = chain.sample_and_log_prob(4, context=None)
     assert samples["a"].shape == (4,) and log_prob.shape == (4,)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_sample_table_factor_emits_on_context_device():
+    """A table-rooted chain must follow the chain's device policy: the fixed
+    table's fresh tensors join the chain on context.device, so they can
+    condition a network living on the GPU (regression: the chirp-mass scan on
+    CUDA mixed CPU conditioning into a CUDA embedding)."""
+    from types import SimpleNamespace
+
+    from dingo.core.factors import SampleTableFactor
+
+    factor = SampleTableFactor(
+        {"x": torch.arange(4.0)}, log_prob=torch.zeros(4)
+    )
+    context = SimpleNamespace(device="cuda")
+    samples, log_prob = factor.sample_and_log_prob(4, context)
+    assert samples["x"].device.type == "cuda"
+    assert log_prob.device.type == "cuda"

From b113a4e1825d64fb18f12e1d5fb142c0178b9eae Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 13:54:18 +0100
Subject: [PATCH 59/65] Add unit tests for importance-sampling prior updates

Covers Result.update_prior directly: the string-form metadata round
trip, reweighting before and after importance sampling (hand-computable
PowerLaw-vs-Uniform tilt on identical support), and a from-file reload
rebuilding the evolved prior rather than the training prior.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 tests/gw/test_gw_result.py | 143 +++++++++++++++++++++++++++++++++++++
 1 file changed, 143 insertions(+)

diff --git a/tests/gw/test_gw_result.py b/tests/gw/test_gw_result.py
index b0aec44af..161f89d9f 100644
--- a/tests/gw/test_gw_result.py
+++ b/tests/gw/test_gw_result.py
@@ -1,6 +1,8 @@
 import numpy as np
 import pandas as pd
 import pytest
+from bilby.core.prior import Constraint, PowerLaw, PriorDict, Uniform
+from bilby.core.utils import random as bilby_random
 
 from dingo.gw.domains import UniformFrequencyDomain
 from dingo.gw.gwutils import get_extrinsic_prior_dict
@@ -219,6 +221,147 @@ def test_sample_synthetic_phase_requires_uniform_phase_prior():
         result.sample_synthetic_phase({"n_grid": 16})
 
 
+# ---------------------------------------------------------------------------
+# update_prior (importance-sampling prior reweighting)
+# ---------------------------------------------------------------------------
+#
+# The context owns the static training prior; a prior update is applied downstream
+# at the importance-sampling layer (update_prior, and _build_prior on reload). These
+# tests exercise that path on synthetic samples -- pure prior arithmetic, no waveform
+# generation. The update swaps the luminosity-distance Uniform(100, 1000) for a
+# PowerLaw on the same support, so every sample stays inside the prior and the
+# expected reweighting is exactly exp(delta log-prior), computable by hand.
+
+# String-instantiable form ({key: prior_str}), as update_prior requires for saving.
+_LD_PRIOR_UPDATE = {
+    "luminosity_distance": "bilby.core.prior.PowerLaw("
+    "alpha=2.0, minimum=100.0, maximum=1000.0, name='luminosity_distance')"
+}
+
+
+def _seeded_result(n=30, seed=42):
+    np.random.seed(seed)
+    bilby_random.seed(seed)
+    return make_gw_result(n=n)
+
+
+def _ld_log_prior_delta(result):
+    """Per-sample change in total ln-prior under _LD_PRIOR_UPDATE: only the
+    luminosity-distance factor changes, so the delta is its ln-prob difference."""
+    ld = result.samples["luminosity_distance"].to_numpy()
+    new = PowerLaw(alpha=2.0, minimum=100.0, maximum=1000.0)
+    old = Uniform(minimum=100.0, maximum=1000.0)
+    return new.ln_prob(ld) - old.ln_prob(ld)
+
+
+def test_update_prior_updates_prior_and_records_string_form():
+    result = _seeded_result()
+    training_ld = result.prior["luminosity_distance"]
+    assert isinstance(training_ld, Uniform)
+
+    # Pass a copy: bilby's PriorDict instantiates the {key: prior_str} dict *in
+    # place* (strings become Prior objects), mutating the caller's dict.
+    result.update_prior(dict(_LD_PRIOR_UPDATE))
+
+    # The live prior is evolved.
+    assert isinstance(result.prior["luminosity_distance"], PowerLaw)
+    # The update is recorded in its original string-instantiable form, and
+    # PriorDict(prior_update) reconstructs exactly the evolved prior (the round
+    # trip _build_prior performs when a saved Result is reloaded).
+    stored = result.importance_sampling_metadata["prior_update"]
+    assert stored == _LD_PRIOR_UPDATE
+    reconstructed = PriorDict(stored.copy())
+    assert reconstructed["luminosity_distance"] == result.prior["luminosity_distance"]
+    assert reconstructed["luminosity_distance"] != training_ld
+
+
+def test_update_prior_reweights_samples_before_importance_sampling():
+    # Before importance sampling (log_evidence is None), update_prior stores
+    # weights proportional to exp(log_prior_new - log_prior_old).
+    result = _seeded_result()
+    assert result.log_evidence is None
+    param_keys = [
+        k for k, v in result.prior.items() if not isinstance(v, Constraint)
+    ]
+    # ln-prior of the samples under the training prior, evaluated exactly as
+    # update_prior does internally (same PriorDict instance, so the stochastic
+    # constraint-normalization factor is cached and cancels in differences).
+    log_prior_before = result.prior.ln_prob(result.samples[param_keys], axis=0)
+
+    result.update_prior(dict(_LD_PRIOR_UPDATE))
+
+    delta = _ld_log_prior_delta(result)
+    # The log_prior column reflects the evolved prior: it shifts by exactly the
+    # luminosity-distance ln-prob difference.
+    np.testing.assert_allclose(
+        result.samples["log_prior"].to_numpy() - log_prior_before, delta, rtol=1e-8
+    )
+    # Weights are the normalized (mean 1) exponential of that shift -- and the
+    # update genuinely reweights (the PowerLaw is not flat over the samples).
+    expected = np.exp(delta - delta.max())
+    expected /= expected.mean()
+    weights = result.samples["weights"].to_numpy()
+    np.testing.assert_allclose(weights, expected, rtol=1e-8)
+    assert weights.mean() == pytest.approx(1.0)
+    assert weights.std() > 0.1
+
+
+def test_update_prior_reweights_importance_sampled_result():
+    # After importance sampling (log_evidence set), update_prior recomputes the
+    # weights and evidence. The importance-sampled state is fabricated from the
+    # sample columns (_calculate_evidence only reads log_prob / log_likelihood /
+    # log_prior), so no likelihood evaluation is needed.
+    result = _seeded_result()
+    param_keys = [
+        k for k, v in result.prior.items() if not isinstance(v, Constraint)
+    ]
+    result.samples["log_prior"] = result.prior.ln_prob(
+        result.samples[param_keys], axis=0
+    )
+    result.samples["log_likelihood"] = np.random.default_rng(0).normal(
+        0.0, 1.0, len(result.samples)
+    )
+    result._calculate_evidence()
+    assert result.log_evidence is not None
+    weights_before = result.samples["weights"].to_numpy().copy()
+    log_evidence_before = result.log_evidence
+
+    result.update_prior(dict(_LD_PRIOR_UPDATE))
+
+    # New weights are the old ones tilted by exp(delta log-prior), renormalized;
+    # the evidence is recomputed accordingly.
+    expected = weights_before * np.exp(_ld_log_prior_delta(result))
+    expected /= expected.mean()
+    np.testing.assert_allclose(
+        result.samples["weights"].to_numpy(), expected, rtol=1e-8
+    )
+    assert result.log_evidence != log_evidence_before
+    assert np.isfinite(result.log_evidence)
+
+
+def test_update_prior_round_trips_through_file(tmp_path):
+    # A reloaded Result rebuilds the *evolved* prior, not the training prior:
+    # _build_prior re-applies importance_sampling_metadata["prior_update"] on top
+    # of the context's static prior.
+    result = _seeded_result()
+    training_ld = result.prior["luminosity_distance"]
+    result.update_prior(dict(_LD_PRIOR_UPDATE))
+    file_name = str(tmp_path / "result.hdf5")
+    result.to_file(file_name=file_name)
+
+    reloaded = Result(file_name=file_name)
+    assert reloaded.importance_sampling_metadata["prior_update"] == _LD_PRIOR_UPDATE
+    assert isinstance(reloaded.prior["luminosity_distance"], PowerLaw)
+    assert reloaded.prior["luminosity_distance"] == result.prior["luminosity_distance"]
+    assert reloaded.prior["luminosity_distance"] != training_ld
+    # The reweighted samples survive the round trip alongside the evolved prior.
+    np.testing.assert_allclose(
+        reloaded.samples["weights"].to_numpy(),
+        result.samples["weights"].to_numpy(),
+        rtol=1e-12,
+    )
+
+
 # ---------------------------------------------------------------------------
 # sample_calibration_parameters
 # ---------------------------------------------------------------------------

From 76fdf1de27053232554ed4309f1a6ff677550ead Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 13:56:33 +0100
Subject: [PATCH 60/65] Stop update_prior mutating the caller's dict

bilby's PriorDict instantiates its input dict in place, replacing the
string values with Prior objects. The stored metadata was already
protected by a copy; the caller's dict was not, a latent behavior
dating back to the introduction of prior updates. Instantiate from a
copy and assert the caller's view in the tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/result.py         | 4 +++-
 tests/gw/test_gw_result.py | 8 +++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/dingo/gw/result.py b/dingo/gw/result.py
index 4931ae261..dcc2b4a09 100644
--- a/dingo/gw/result.py
+++ b/dingo/gw/result.py
@@ -255,7 +255,9 @@ def update_prior(self, prior_update):
             the Result and later instantiated.
         """
         self.importance_sampling_metadata["prior_update"] = prior_update.copy()
-        prior_update = PriorDict(prior_update)
+        # PriorDict instantiates in place, so work on a copy: the caller's dict
+        # keeps its string form.
+        prior_update = PriorDict(prior_update.copy())
 
         param_keys = [k for k, v in self.prior.items() if not isinstance(v, Constraint)]
         theta = self.samples[param_keys]
diff --git a/tests/gw/test_gw_result.py b/tests/gw/test_gw_result.py
index 161f89d9f..872ab4a13 100644
--- a/tests/gw/test_gw_result.py
+++ b/tests/gw/test_gw_result.py
@@ -259,9 +259,11 @@ def test_update_prior_updates_prior_and_records_string_form():
     training_ld = result.prior["luminosity_distance"]
     assert isinstance(training_ld, Uniform)
 
-    # Pass a copy: bilby's PriorDict instantiates the {key: prior_str} dict *in
-    # place* (strings become Prior objects), mutating the caller's dict.
-    result.update_prior(dict(_LD_PRIOR_UPDATE))
+    prior_update = dict(_LD_PRIOR_UPDATE)
+    result.update_prior(prior_update)
+    # The caller's dict keeps its string form (update_prior copies before
+    # bilby's in-place instantiation).
+    assert prior_update == _LD_PRIOR_UPDATE
 
     # The live prior is evolved.
     assert isinstance(result.prior["luminosity_distance"], PowerLaw)

From 2db75f32cc4322615060bdbe2a5bc4631a1c3e15 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 14:16:18 +0100
Subject: [PATCH 61/65] Take the DINGO-BNS reference from INSPIRE

Replaces the hand-written entry with the INSPIRE export (same texkey
and fields, verified against the published paper; adds the LIGO report
number).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/source/refs.bib | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/source/refs.bib b/docs/source/refs.bib
index de6ced378..03b728893 100644
--- a/docs/source/refs.bib
+++ b/docs/source/refs.bib
@@ -76,11 +76,12 @@ @article{Wildberger:2022agw
 }
 
 @article{Dax:2024mcn,
-    author = {Dax, Maximilian and Green, Stephen R. and Gair, Jonathan and Gupte, Nihar and P\"urrer, Michael and Raymond, Vivien and Wildberger, Jonas and Macke, Jakob H. and Buonanno, Alessandra and Sch\"olkopf, Bernhard},
+    author = {Dax, Maximilian and Green, Stephen R. and Gair, Jonathan and Gupte, Nihar and P{\"u}rrer, Michael and Raymond, Vivien and Wildberger, Jonas and Macke, Jakob H. and Buonanno, Alessandra and Sch{\"o}lkopf, Bernhard},
     title = "{Real-time inference for binary neutron star mergers using machine learning}",
     eprint = "2407.09602",
     archivePrefix = "arXiv",
     primaryClass = "gr-qc",
+    reportNumber = "LIGO-P2400294",
     doi = "10.1038/s41586-025-08593-z",
     journal = "Nature",
     volume = "639",

From f20edf52dff8c8c7d0dc9a99256a8e71c29c2e8e Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 16:42:58 +0100
Subject: [PATCH 62/65] Temporarily omit the docs figures to probe the
 ultrareview size cap

Removes the four mermaid blocks and their captions (100 diff lines),
bringing the PR to 9,988 changed lines against a suspected 10,000-line
limit. To be reverted after the review runs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/source/bns.md             | 22 ----------
 docs/source/sampling_chains.md | 80 +---------------------------------
 2 files changed, 1 insertion(+), 101 deletions(-)

diff --git a/docs/source/bns.md b/docs/source/bns.md
index 710441a97..22b96fad0 100644
--- a/docs/source/bns.md
+++ b/docs/source/bns.md
@@ -58,28 +58,6 @@ event-specific narrow prior. The network infers the offset
 `delta_chirp_mass` $= \mathcal{M} - \tilde{\mathcal{M}}$ rather than the chirp mass
 itself; the chain reconstructs the physical value with a `ProxyOffsetReparam` step.
 
-The chain:
-
-```{mermaid}
-flowchart TB
-    pins["DeltaFactor<br/>chirp_mass_proxy"]
-    flow["FlowFactor<br/>draws delta_chirp_mass, #8230;"]
-    off["ProxyOffsetReparam<br/>chirp_mass = delta_chirp_mass + chirp_mass_proxy"]
-    out(["samples + log_prob"])
-    ctx["GWSamplerContext<br/>heterodyne #8594; decimate #8594; whiten"]
-
-    pins --> flow --> off --> out
-    pins -. "chirp_mass_proxy" .-> ctx
-    ctx -. "prepared data" .-> flow
-
-    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
-    classDef reparam fill:#e2f0e6,stroke:#27ae60,color:#1a1a1a
-    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
-    class pins,flow step
-    class off reparam
-    class ctx ctxstyle
-```
-
 The pinned values have a single owner, the chain root, and are recorded with the
 samples. The heterodyne receives the proxy through the row-aligned `prepared_data`
 contract of the [sampler context](sampling_chains.md#the-sampler-context). Since the
diff --git a/docs/source/sampling_chains.md b/docs/source/sampling_chains.md
index 4ebc62357..1a3b793cd 100644
--- a/docs/source/sampling_chains.md
+++ b/docs/source/sampling_chains.md
@@ -35,34 +35,7 @@ $-\log\lvert\det J_i\rvert$ for a reparametrization. These cases correspond to t
 
 At construction the composer checks the chain for consistency: every conditioning
 column must be produced by an earlier step. It then runs the steps in order,
-accumulating the table and the sum:
-
-```{mermaid}
-flowchart TB
-    subgraph comp ["ChainComposer"]
-        direction TB
-        s1["step 1<br/>q#8321;(#952;#8321; | d)"]
-        s2["step 2<br/>q#8322;(#952;#8322; | #952;#8321;, d)"]
-        dots["#8230;"]
-        sn["step n<br/>q#8345;(#952;#8345; | #952;#8321;, #8230;, #952;#8345;#8331;#8321;, d)"]
-        s1 -- "#916;#8321;" --> s2
-        s2 -- "#916;#8321; + #916;#8322;" --> dots
-        dots --> sn
-    end
-    out(["samples + log_prob"])
-
-    sn -- "#916;#8321; + #8230; + #916;#8345; = log q" --> out
-
-    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
-    classDef ghost fill:none,stroke:none
-    class s1,s2,sn step
-    class dots ghost
-```
-
-In the figure every step is written as a conditional $q_i$, which covers all
-factors, point masses included. A reparametrization instead replaces columns in
-place, adding its Jacobian term to the sum but no factor to the product (see
-[the proposal density](#the-proposal-density-forward-and-reverse)).
+accumulating the table and the sum.
 
 The generic machinery lives in `dingo.core.factors` (steps, stages, the composer, the
 runner); the gravitational-wave steps, the per-event context, and the chain builders
@@ -190,35 +163,6 @@ views.
   rather than reusing the network-input view, and its reference time is the event
   time. Marginalization settings (time, phase, calibration) are passed per request.
 
-As data flow, including the conditioning columns that a chain feeds back into the
-data preparation:
-
-```{mermaid}
-flowchart TB
-    d[("event data d")]
-    em["event metadata"]
-    md["model metadata"]
-    subgraph ctx ["GWSamplerContext"]
-        direction LR
-        pd["prepared_data"]
-        pr["prior"]
-        lk["likelihood"]
-    end
-    chain["chain steps"]
-    isamp["importance sampling"]
-
-    d --> ctx
-    em --> ctx
-    md --> ctx
-    pd -- "data rows (row-aligned)" --> chain
-    chain -. "conditioning columns (optional)" .-> pd
-    lk -.-> isamp
-    pr -.-> isamp
-
-    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
-    class d,em,md,pd,pr,lk ctxstyle
-```
-
 Alongside the strain data, the context carries the per-event **event metadata**. This
 includes the event time, which sets the likelihood reference time and the sidereal
 right-ascension correction (`RAToEventFrame`), together with any per-event analysis
@@ -245,28 +189,6 @@ stored settings. The chain's torch device is `context.device`. Steps that create
 tensors, such as the pins of a `DeltaFactor`, create them on this device so that their
 outputs can join a chain running on a GPU.
 
-Putting the pieces together:
-
-```{mermaid}
-flowchart TB
-    d[("event data d")] --> ctx["GWSamplerContext"]
-    ctx -- "data views" --> comp["ChainComposer"]
-    comp -. "conditioning parameters" .-> ctx
-    comp --> out(["samples + log_prob"])
-    out --> isamp["importance sampling"]
-    ctx -.-> isamp
-
-    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
-    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
-    class ctx ctxstyle
-    class comp step
-```
-
-The chain draws its data views through the context and feeds conditioning parameters
-back into the preparation. The likelihood and prior views serve the downstream
-importance sampling in the same way, so a derived context re-targets the whole
-analysis consistently.
-
 ```{note}
 The representation vocabulary here (frequency domains, multibanded decimation,
 base-domain likelihoods) is specific to this domain family. A new domain family

From 0591e4d022e030dff9e3f2c1d8f9cf1aaa12089e Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 17:06:50 +0100
Subject: [PATCH 63/65] Revert "Temporarily omit the docs figures to probe the
 ultrareview size cap"

This reverts commit f20edf52dff8c8c7d0dc9a99256a8e71c29c2e8e.
---
 docs/source/bns.md             | 22 ++++++++++
 docs/source/sampling_chains.md | 80 +++++++++++++++++++++++++++++++++-
 2 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/docs/source/bns.md b/docs/source/bns.md
index 22b96fad0..710441a97 100644
--- a/docs/source/bns.md
+++ b/docs/source/bns.md
@@ -58,6 +58,28 @@ event-specific narrow prior. The network infers the offset
 `delta_chirp_mass` $= \mathcal{M} - \tilde{\mathcal{M}}$ rather than the chirp mass
 itself; the chain reconstructs the physical value with a `ProxyOffsetReparam` step.
 
+The chain:
+
+```{mermaid}
+flowchart TB
+    pins["DeltaFactor<br/>chirp_mass_proxy"]
+    flow["FlowFactor<br/>draws delta_chirp_mass, #8230;"]
+    off["ProxyOffsetReparam<br/>chirp_mass = delta_chirp_mass + chirp_mass_proxy"]
+    out(["samples + log_prob"])
+    ctx["GWSamplerContext<br/>heterodyne #8594; decimate #8594; whiten"]
+
+    pins --> flow --> off --> out
+    pins -. "chirp_mass_proxy" .-> ctx
+    ctx -. "prepared data" .-> flow
+
+    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
+    classDef reparam fill:#e2f0e6,stroke:#27ae60,color:#1a1a1a
+    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
+    class pins,flow step
+    class off reparam
+    class ctx ctxstyle
+```
+
 The pinned values have a single owner, the chain root, and are recorded with the
 samples. The heterodyne receives the proxy through the row-aligned `prepared_data`
 contract of the [sampler context](sampling_chains.md#the-sampler-context). Since the
diff --git a/docs/source/sampling_chains.md b/docs/source/sampling_chains.md
index 1a3b793cd..4ebc62357 100644
--- a/docs/source/sampling_chains.md
+++ b/docs/source/sampling_chains.md
@@ -35,7 +35,34 @@ $-\log\lvert\det J_i\rvert$ for a reparametrization. These cases correspond to t
 
 At construction the composer checks the chain for consistency: every conditioning
 column must be produced by an earlier step. It then runs the steps in order,
-accumulating the table and the sum.
+accumulating the table and the sum:
+
+```{mermaid}
+flowchart TB
+    subgraph comp ["ChainComposer"]
+        direction TB
+        s1["step 1<br/>q#8321;(#952;#8321; | d)"]
+        s2["step 2<br/>q#8322;(#952;#8322; | #952;#8321;, d)"]
+        dots["#8230;"]
+        sn["step n<br/>q#8345;(#952;#8345; | #952;#8321;, #8230;, #952;#8345;#8331;#8321;, d)"]
+        s1 -- "#916;#8321;" --> s2
+        s2 -- "#916;#8321; + #916;#8322;" --> dots
+        dots --> sn
+    end
+    out(["samples + log_prob"])
+
+    sn -- "#916;#8321; + #8230; + #916;#8345; = log q" --> out
+
+    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
+    classDef ghost fill:none,stroke:none
+    class s1,s2,sn step
+    class dots ghost
+```
+
+In the figure every step is written as a conditional $q_i$, which covers all
+factors, point masses included. A reparametrization instead replaces columns in
+place, adding its Jacobian term to the sum but no factor to the product (see
+[the proposal density](#the-proposal-density-forward-and-reverse)).
 
 The generic machinery lives in `dingo.core.factors` (steps, stages, the composer, the
 runner); the gravitational-wave steps, the per-event context, and the chain builders
@@ -163,6 +190,35 @@ views.
   rather than reusing the network-input view, and its reference time is the event
   time. Marginalization settings (time, phase, calibration) are passed per request.
 
+As data flow, including the conditioning columns that a chain feeds back into the
+data preparation:
+
+```{mermaid}
+flowchart TB
+    d[("event data d")]
+    em["event metadata"]
+    md["model metadata"]
+    subgraph ctx ["GWSamplerContext"]
+        direction LR
+        pd["prepared_data"]
+        pr["prior"]
+        lk["likelihood"]
+    end
+    chain["chain steps"]
+    isamp["importance sampling"]
+
+    d --> ctx
+    em --> ctx
+    md --> ctx
+    pd -- "data rows (row-aligned)" --> chain
+    chain -. "conditioning columns (optional)" .-> pd
+    lk -.-> isamp
+    pr -.-> isamp
+
+    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
+    class d,em,md,pd,pr,lk ctxstyle
+```
+
 Alongside the strain data, the context carries the per-event **event metadata**. This
 includes the event time, which sets the likelihood reference time and the sidereal
 right-ascension correction (`RAToEventFrame`), together with any per-event analysis
@@ -189,6 +245,28 @@ stored settings. The chain's torch device is `context.device`. Steps that create
 tensors, such as the pins of a `DeltaFactor`, create them on this device so that their
 outputs can join a chain running on a GPU.
 
+Putting the pieces together:
+
+```{mermaid}
+flowchart TB
+    d[("event data d")] --> ctx["GWSamplerContext"]
+    ctx -- "data views" --> comp["ChainComposer"]
+    comp -. "conditioning parameters" .-> ctx
+    comp --> out(["samples + log_prob"])
+    out --> isamp["importance sampling"]
+    ctx -.-> isamp
+
+    classDef ctxstyle fill:#f4f4f4,stroke:#8c8c8c,color:#1a1a1a
+    classDef step fill:#dbe9f6,stroke:#2980b9,color:#1a1a1a
+    class ctx ctxstyle
+    class comp step
+```
+
+The chain draws its data views through the context and feeds conditioning parameters
+back into the preparation. The likelihood and prior views serve the downstream
+importance sampling in the same way, so a derived context re-targets the whole
+analysis consistently.
+
 ```{note}
 The representation vocabulary here (frequency domains, multibanded decimation,
 base-domain likelihoods) is specific to this domain family. A new domain family

From bc0ac9404141768f5d31b662b916fe21b229bf47 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 18:00:11 +0100
Subject: [PATCH 64/65] Report the phase-maximized SNR under phase
 marginalization

The auxiliary SNR in the phase-marginalized likelihood reused the
marginalized-likelihood term ln_i0(|kappa2C|) in place of the
phase-maximized matched-filter statistic |kappa2C|, biasing every
reported scan SNR low by ~0.5*ln(2*pi*|kappa2C|)/sqrt(rho2opt). Trigger
selection was unaffected (it maximizes the log-likelihood, which was
correct). Found by the cloud review; adds a regression test pinning the
statistic against an independent computation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/gw/likelihood.py      |  5 ++++-
 tests/gw/test_likelihood.py | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/dingo/gw/likelihood.py b/dingo/gw/likelihood.py
index 57d7ec54f..4f7371f97 100644
--- a/dingo/gw/likelihood.py
+++ b/dingo/gw/likelihood.py
@@ -531,7 +531,10 @@ def _log_likelihood_phase_marginalized(self, theta):
             log_likelihood = self.log_Zn + ln_i0(np.abs(kappa2C)) - 1 / 2.0 * rho2opt
 
             if self.return_aux_snr:
-                snr = ln_i0(np.abs(kappa2C)) / rho2opt**0.5
+                # Phase-maximized matched-filter SNR: |kappa2C| / sqrt(rho2opt).
+                # (ln_i0(|kappa2C|) is the phase-marginalized likelihood term
+                # above, not a matched-filter statistic.)
+                snr = np.abs(kappa2C) / rho2opt**0.5
                 return log_likelihood, snr
 
             return log_likelihood
diff --git a/tests/gw/test_likelihood.py b/tests/gw/test_likelihood.py
index 6ecf0dbb0..fd2280595 100644
--- a/tests/gw/test_likelihood.py
+++ b/tests/gw/test_likelihood.py
@@ -355,3 +355,20 @@ def test_multi_rho2opt_consistency(self, likelihood_and_theta):
         )
         assert np.isclose(d_inner_h_single, d_inner_h_multi[0])
         assert np.isclose(rho2opt_single, rho2opt_multi[0])
+
+
+def test_phase_marginalized_aux_snr_is_phase_maximized(domain, event_data):
+    """Under phase marginalization the auxiliary SNR is the phase-maximized
+    matched-filter statistic |<d, h_c>| / sqrt(<h, h>), not the marginalized
+    likelihood term ln_i0(|<d, h_c>|)."""
+    likelihood = make_likelihood(
+        domain,
+        event_data,
+        phase_marginalization_kwargs={"approximation_22_mode": True},
+    )
+    likelihood.return_aux_snr = True
+    theta = dict(THETA)
+    theta.pop("phase", None)
+    _, snr = likelihood.log_likelihood(theta)
+    kappa2C, rho2opt = likelihood._d_inner_h_complex({**theta, "phase": 0.0})
+    assert snr == pytest.approx(np.abs(kappa2C) / rho2opt**0.5)

From e6ac3bc2e54083c32f4a4a3dfe239fd7b8e14ce1 Mon Sep 17 00:00:00 2001
From: Stephen Green <stephenrolandgreen@gmail.com>
Date: Tue, 14 Jul 2026 18:00:11 +0100
Subject: [PATCH 65/65] Fix review leftovers: legacy import and help-text typo

misc_scripts/evaluation.py still imported the removed GWSampler; port
it to GWComposedSampler.from_model. The --log-probs help text
concatenated adjacent string literals without a space.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 dingo/pipe/plot.py         | 2 +-
 misc_scripts/evaluation.py | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/dingo/pipe/plot.py b/dingo/pipe/plot.py
index 1eec0690d..61d71aca8 100644
--- a/dingo/pipe/plot.py
+++ b/dingo/pipe/plot.py
@@ -29,7 +29,7 @@ def create_parser():
     parser.add(
         "--log-probs",
         action="store_true",
-        help="Generate plot of target" "versus proposal log probability",
+        help="Generate plot of target versus proposal log probability",
     )
     # parser.add("--marginal", action="store_true", help="Generate marginal plots")
     # parser.add("--skymap", action="store_true", help="Generate skymap")
diff --git a/misc_scripts/evaluation.py b/misc_scripts/evaluation.py
index 23da2772b..49f2d39b9 100755
--- a/misc_scripts/evaluation.py
+++ b/misc_scripts/evaluation.py
@@ -7,7 +7,7 @@
 
 from dingo.core.posterior_models.build_model import build_model_from_kwargs
 from dingo.gw.result import Result
-from dingo.gw.inference.gw_samplers import GWSampler
+from dingo.gw.inference.sampler import GWComposedSampler
 from dingo.gw.noise.asd_dataset import ASDDataset
 from dingo.gw import injection
 
@@ -63,8 +63,7 @@ def evaluate_model(
 
 
 def run_inference(strain_data, pm_model, injection_dir, num_samples=1000):
-    sampler = GWSampler(model=pm_model)
-    sampler.context = strain_data
+    sampler = GWComposedSampler.from_model(pm_model, strain_data)
     sampler.run_sampler(num_samples=num_samples, batch_size=BATCH_SIZE)
     sampler.to_hdf5(outdir=injection_dir)
     return sampler.samples