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/* 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() 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}.") diff --git a/dingo/core/factors.py b/dingo/core/factors.py new file mode 100644 index 000000000..88e517c37 --- /dev/null +++ b/dingo/core/factors.py @@ -0,0 +1,980 @@ +""" +Factorized sampler core: the domain-agnostic spine of the factorized-sampler design +(see vault/Hackathon/Factorized_Sampler_Design.md). + +The posterior is an ordered product of conditional factors, + + q(theta_1, ..., theta_n | d) = prod_i q_i(theta_i | f_i(theta_ 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 quantities + derived from it (the prepared-data representation, the likelihood, event metadata). + Concrete implementations are domain-specific; see + `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 + join a chain whose network factors run on a GPU. + """ + + event_metadata: Optional[dict] + device: Union[torch.device, str] + + def prepared_data(self, conditioning=None) -> torch.Tensor: + """The data representation the factors condition on + (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): + """The likelihood, for likelihood-based factors (synthetic phase) and IS.""" + ... + + +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 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]]: + """Run `run_once` over batches of `total` and concatenate the results. + + 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]] = [] + 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 + + +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.""" + 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 + `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_ (n_rows, n, dim)`. + + Attributes + ---------- + parameters : list[str] + The parameter block this factor produces. + conditioning : list[str] + Earlier-block parameters it conditions on (data dependence is via the context). + """ + + parameters: list[str] + conditioning: list[str] + + @abstractmethod + def sample_and_log_prob( + 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 + `parameters`.""" + + @abstractmethod + def log_prob( + 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).""" + + 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 + 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"] + return metadata + + +class FlowFactor(Factor): + """ + 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__( + self, + model: BasePosteriorModel, + parameters: list[str], + conditioning: Optional[list[str]] = None, + 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 + # 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 [] + 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 + 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 own training metadata (for an unconditional NDE these are its own, + e.g. the GNPE proxies it was trained on). + + 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 = model.metadata["train_settings"]["data"] + context_parameters = data_settings.get("context_parameters") or [] + return cls( + model=model, + parameters=data_settings["inference_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, + ) + + def sample_and_log_prob(self, num_samples, context, given=None): + """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 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 + ) + # 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; 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( + self._network_conditioning(given), self.context_parameters + ) + n_rows = ctx.shape[0] + 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 + ) + 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._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, 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 = _n_rows(theta_net) + z = self.standardization.standardize(theta_net, self._net_parameters) + net_context: tuple[torch.Tensor, ...] + if self.unconditional: + net_context = () + elif not self.context_parameters: + data = context.prepared_data() + data = data.expand(num_samples, *data.shape) + net_context = (data,) + else: + data = context.prepared_data(conditioning=given) + 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} + + +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). + + 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) + 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), device=device) + for p, v in self.values.items() + } + 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 + 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) + + def describe(self) -> dict: + return { + **_describe_default(self), + "values": {k: float(v) for k, v in self.values.items()}, + } + + +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 = _n_rows(self.table) + 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." + ) + # 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( + "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) + 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), + 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] + one_to_one = True + + @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", + given: Optional[dict[str, torch.Tensor]] = None, + ) -> dict[str, torch.Tensor]: + """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" + ) -> torch.Tensor: + """`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, + 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(given, context) + return out, -self.log_det(given, context) + + @property + def consumes(self) -> list[str]: + """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 describe(self) -> dict: + """Provenance descriptor for saved-result metadata.""" + 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 + 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. + + `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] + conditioning: list[str] + one_to_one = True + 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, + 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(given, context) + # 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) + + def describe(self) -> dict: + """Provenance descriptor for saved-result metadata.""" + return _describe_default(self) + + +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, + context: SamplerContext, + given: Optional[dict[str, torch.Tensor]] = None, + ) -> tuple[dict[str, torch.Tensor], Optional[torch.Tensor]]: ... + + +@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`.""" + + step: Step + fan_out: int = 1 + + +class ChainComposer: + """ + 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 + `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` + entries. + """ + + 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 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`). 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] + if missing: + raise ValueError( + f"A step producing {step.parameters} conditions on {missing}, " + f"which no earlier step produces. Check chain order." + ) + 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) + produced.difference_update(getattr(step, "consumes", ())) + + @property + def steps(self) -> list[Step]: + """The stage steps, in order.""" + return [stage.step for stage in self.stages] + + @property + def expansion(self) -> int: + """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( + self, + num_samples: int, + 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.""" + 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], 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 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 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 + base_pending: Optional[int] = base + for i, stage in enumerate(self.stages): + step = stage.step + if base_pending is not None: + 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 + elif step.conditioning: + n = stage.fan_out # fan_out samples per conditioning row + 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 = _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 + # 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. + 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 + 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( + self, samples: dict[str, torch.Tensor], context: SamplerContext + ) -> torch.Tensor: + """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 reversed(self.steps): + if isinstance(step, TargetCorrection): + continue + if isinstance(step, Reparametrization): + params = {k: values[k] for k in step.parameters} + # 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: + 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( + self, + num_samples: int, + context: SamplerContext, + batch_size: Optional[int] = None, + ) -> dict[str, torch.Tensor]: + """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 GibbsBlock: + """ + 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 + (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.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`). + """ + + 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 + # 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_and_log_prob( + 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, 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]: + # 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, context) + state = dict(seed) + for _ in range(self.num_iterations): + for factor in self.factors: + 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, context, given) + state.update(block) + return {p: state[p] for p in self.parameters} + + +class ComposedSampler: + """ + 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): + self.composer = composer + self.context = context + self.samples: Optional[pd.DataFrame] = None + + 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` + 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.samples = pd.DataFrame(merged) + return self.samples 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/dingo/core/result.py b/dingo/core/result.py index 67017271a..f2ada5ed4 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 + # 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, dictionary=dictionary, @@ -97,10 +102,10 @@ def __init__(self, file_name=None, dictionary=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"): - self._rebuild_domain() @property def metadata(self): @@ -145,6 +150,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. @@ -186,11 +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 - def _rebuild_domain(self, verbose=False): - pass + # The old context described the data the samples were drawn from; rebuild it + # 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() + self._build_domain() @property def num_samples(self): @@ -1090,5 +1103,3 @@ def freeze(d): elif isinstance(d, list): return tuple(freeze(value) for value in d) return d - - 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/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/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/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..7db6ce59b 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.sampler 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/context.py b/dingo/gw/inference/context.py new file mode 100644 index 000000000..fc4d65702 --- /dev/null +++ b/dingo/gw/inference/context.py @@ -0,0 +1,639 @@ +"""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 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: + 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) + .cpu() + .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/gw_samplers.py b/dingo/gw/inference/gw_samplers.py deleted file mode 100644 index 2a1b63d89..000000000 --- a/dingo/gw/inference/gw_samplers.py +++ /dev/null @@ -1,682 +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, - build_domain_from_model_metadata, - UniformFrequencyDomain, - Domain, -) -from dingo.gw.domains import build_domain -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) - - -# 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/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/inference/sampler.py b/dingo/gw/inference/sampler.py new file mode 100644 index 000000000..3d992bf26 --- /dev/null +++ b/dingo/gw/inference/sampler.py @@ -0,0 +1,381 @@ +"""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 ( + RAToTrainingFrame, + GNPEFlowFactor, + GNPEKernelCorrection, + GNPEKernelFactor, + 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 `RAToEventFrame` 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_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: + """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 `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, + 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"] + 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( + 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_to_event + steps += _ra_to_event_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 `RAToEventFrame` 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_to_event_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, 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. + + 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_to_event_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 `/