diff --git a/dingo/gw/dataset/_multibanded_domain_utils.py b/dingo/gw/dataset/_multibanded_domain_utils.py new file mode 100644 index 000000000..dd6219699 --- /dev/null +++ b/dingo/gw/dataset/_multibanded_domain_utils.py @@ -0,0 +1,65 @@ +"""Shared utilities for generating and evaluating MultibandedFrequencyDomain settings.""" + +from copy import deepcopy +from typing import Dict + +import numpy as np + +from dingo.gw.prior import build_prior_with_defaults + + +def build_extreme_prior(settings: dict): + """Build a BBH prior with extreme parameter values to stress-test multibanding. + + Fixes ``chirp_mass`` to its minimum prior value and ``geocent_time`` to 0.12 s + (the typical prior boundary plus the Earth-radius light-crossing time). All other + parameters are sampled from the distributions specified in ``settings``. + + Parameters + ---------- + settings : dict + Dataset settings dict containing an ``'intrinsic_prior'`` key. Not modified. + + Returns + ------- + BBHPriorDict + Prior with extreme fixed values for ``chirp_mass`` and ``geocent_time``. + """ + nominal_prior = build_prior_with_defaults(settings["intrinsic_prior"]) + extreme_settings = deepcopy(settings["intrinsic_prior"]) + extreme_settings["geocent_time"] = 0.12 + # Pin the chirp mass to (essentially) its minimum -- the longest, hardest-to-decimate + # signal. A bare scalar would become a bilby DeltaFunction, which breaks the + # *constrained* sampling required by the mass_1/mass_2 Constraint priors (it raises + # "non-broadcastable output operand with shape ()" in PriorDict.sample). A + # negligibly narrow Uniform samples cleanly while keeping every draw at the minimum. + mc_min = nominal_prior["chirp_mass"].minimum + extreme_settings["chirp_mass"] = ( + f"bilby.core.prior.Uniform(minimum={mc_min}, maximum={mc_min * (1 + 1e-9)})" + ) + return build_prior_with_defaults(extreme_settings) + + +def print_mismatch_stats(mismatches: np.ndarray, num_samples: int) -> None: + """Print a summary of mismatch statistics to stdout. + + Parameters + ---------- + mismatches : np.ndarray + 1D array of mismatch values across all polarisations and samples. + num_samples : int + Number of waveform samples used, reported in the header line. + """ + print("\nMismatches between UFD waveforms and MFD waveforms interpolated to UFD.") + print( + "This is a conservative estimate of the MFD performance when training networks." + ) + print(f"num_samples = {num_samples}") + print(f" Mean mismatch = {np.mean(mismatches)}") + print(f" Standard deviation = {np.std(mismatches)}") + print(f" Max mismatch = {np.max(mismatches)}") + print(f" Median mismatch = {np.median(mismatches)}") + print(" Percentiles:") + print(f" 99 -> {np.percentile(mismatches, 99)}") + print(f" 99.9 -> {np.percentile(mismatches, 99.9)}") + print(f" 99.99 -> {np.percentile(mismatches, 99.99)}") diff --git a/dingo/gw/dataset/evaluate_multibanded_domain.py b/dingo/gw/dataset/evaluate_multibanded_domain.py index b20a2fb43..01e3d6839 100644 --- a/dingo/gw/dataset/evaluate_multibanded_domain.py +++ b/dingo/gw/dataset/evaluate_multibanded_domain.py @@ -5,14 +5,13 @@ from scipy.interpolate import interp1d from dingo.gw.dataset import generate_parameters_and_polarizations -from dingo.gw.domains import build_domain, MultibandedFrequencyDomain +from dingo.gw.dataset._multibanded_domain_utils import (build_extreme_prior, + print_mismatch_stats) +from dingo.gw.domains import MultibandedFrequencyDomain, build_domain from dingo.gw.gwutils import get_mismatch -from dingo.gw.prior import build_prior_with_defaults -from dingo.gw.waveform_generator import ( - NewInterfaceWaveformGenerator, - WaveformGenerator, - generate_waveforms_parallel, -) +from dingo.gw.waveform_generator import (NewInterfaceWaveformGenerator, + WaveformGenerator, + generate_waveforms_parallel) def _evaluate_multibanding_main( @@ -26,15 +25,7 @@ def _evaluate_multibanding_main( if "compression" in settings: del settings["compression"] - # Update prior to challenge the multi-banding: - # - # (a) Set geocent_time = 0.12 s (boundary of usual prior + Earth-radius crossing time) - # (b) Set chirp mass to bottom end of prior. - prior = build_prior_with_defaults(settings["intrinsic_prior"]) - settings["intrinsic_prior"]["geocent_time"] = 0.12 - settings["intrinsic_prior"]["chirp_mass"] = prior["chirp_mass"].minimum - # Rebuild prior with updated settings. - prior = build_prior_with_defaults(settings["intrinsic_prior"]) + prior = build_extreme_prior(settings) print("Prior") for k, v in prior.items(): print(f"{k}: {v}") @@ -88,21 +79,8 @@ def _evaluate_multibanding_main( asd_file="aLIGO_ZERO_DET_high_P_asd.txt", ) - print("\nMismatches between UFD waveforms and MFD waveforms interpolated to MFD.") - print( - "This is a conservative estimate of the MFD performance when training " - "networks." - ) - mismatches = np.concatenate([v for v in mismatches.values()]) - print(f"num_samples = {num_samples}") - print(" Mean mismatch = {}".format(np.mean(mismatches))) - print(" Standard deviation = {}".format(np.std(mismatches))) - print(" Max mismatch = {}".format(np.max(mismatches))) - print(" Median mismatch = {}".format(np.median(mismatches))) - print(" Percentiles:") - print(" 99 -> {}".format(np.percentile(mismatches, 99))) - print(" 99.9 -> {}".format(np.percentile(mismatches, 99.9))) - print(" 99.99 -> {}".format(np.percentile(mismatches, 99.99))) + mismatches = np.concatenate(list(mismatches.values())) + print_mismatch_stats(mismatches, num_samples) def parse_args(): diff --git a/dingo/gw/dataset/generate_multibanded_domain.py b/dingo/gw/dataset/generate_multibanded_domain.py new file mode 100644 index 000000000..a660eb01e --- /dev/null +++ b/dingo/gw/dataset/generate_multibanded_domain.py @@ -0,0 +1,892 @@ +""" +Generate a MultibandedFrequencyDomain (MFD) settings file from a uniform frequency domain +(UFD) settings file by automatically tuning a decimation threshold to meet a target median +waveform mismatch. + +The core idea is to use an extreme prior (minimum chirp mass, boundary geocent_time) to +stress-test the decimation, generate waveforms once, then binary-search over the whitened +waveform difference threshold until the desired mismatch level is reached. + +CLI usage:: + + python -m dingo.gw.dataset.generate_multibanded_domain \\ + --settings_file path/to/settings_wfd_ufd.yaml \\ + --num_samples 1000 \\ + --target_median_mismatch 0.001 \\ + --token_size 16 # for the Dingo-T1 transformer; omit for standard NPE +""" + +import argparse +import os +from copy import deepcopy +from typing import Dict, List, Optional, Tuple + +import numpy as np +import pandas as pd +import torch +import yaml +from bilby.gw.detector import PowerSpectralDensity +from scipy.interpolate import interp1d + +from dingo.gw.dataset._multibanded_domain_utils import (build_extreme_prior, + print_mismatch_stats) +from dingo.gw.dataset.generate_dataset import \ + generate_parameters_and_polarizations +from dingo.gw.domains import (MultibandedFrequencyDomain, + UniformFrequencyDomain, build_domain) +from dingo.gw.gwutils import get_mismatch +from dingo.gw.waveform_generator import (NewInterfaceWaveformGenerator, + WaveformGenerator) + + +def floor_to_power_of_2(x: float) -> float: + """Return the largest power of 2 that is <= x. + + Parameters + ---------- + x : float + Positive input value. + + Returns + ------- + float + Largest power of 2 not exceeding x. + """ + return 2 ** (np.floor(np.log2(x))) + + +def compute_waveform_difference_per_decimation_factor( + decimation_factors: np.ndarray, + waveforms: np.ndarray, + ufd: UniformFrequencyDomain, + waveforms_2x: np.ndarray, + difference_over_full_window: bool = False, +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """Compute the 95th-percentile whitened waveform difference for each decimation factor. + Only the real part of the waveform is considered. + + Two comparison modes are supported: + + - **Center comparison** (default, ``difference_over_full_window=False``): each + decimated bin is compared to the waveform value at the center of the decimation + window, read from the 2x-resolution reference array. This is less conservative + but smoother. + - **Full-window comparison** (``difference_over_full_window=True``): each decimated + bin is compared to all original bins within the decimation window. This is the most + conservative estimate. + + The resulting difference arrays are transformed into monotonically non-increasing + sequences via a right-to-left cumulative maximum, reflecting the physical expectation + that higher-frequency bins are easier to decimate accurately. + + Parameters + ---------- + decimation_factors : np.ndarray + Integer decimation factors to evaluate, e.g. ``2 ** np.arange(8)``. + waveforms : np.ndarray + Whitened real-part waveforms at base UFD resolution, shape + ``(num_samples, num_bins)``. + ufd : UniformFrequencyDomain + Uniform frequency domain corresponding to ``waveforms``. + waveforms_2x : np.ndarray + Whitened real-part waveforms at twice the UFD resolution, used as the + high-resolution reference in center-comparison mode. + difference_over_full_window : bool + If True, use full-window comparison. Default: False. + + Returns + ------- + diffs : List[np.ndarray] + One 1D array per decimation factor, containing the 95th-percentile + whitened difference at each decimated frequency bin, monotonically + non-increasing. + freqs : List[np.ndarray] + One 1D array per decimation factor, containing the corresponding + reference frequencies for each entry in ``diffs``. + """ + data = waveforms.real + data_2x = waveforms_2x.real + freq = ufd.sample_frequencies + assert freq.shape[0] == data.shape[-1] + + freqs_per_decimation_factor: List[np.ndarray] = [] + diffs_per_decimation_factor: List[np.ndarray] = [] + + for decimation_factor in decimation_factors: + kernel = ( + torch.ones((1, 1, decimation_factor), dtype=torch.float64) + / decimation_factor + ) + + def conv( + d: np.ndarray, _k: torch.Tensor = kernel, _df: int = decimation_factor + ) -> torch.Tensor: + return torch.nn.functional.conv1d( + torch.tensor(d, dtype=torch.float64)[:, None, :], + _k, + padding=0, + stride=_df, + ).squeeze() + + data_decimated = conv(data).numpy() + + if difference_over_full_window: + f_ref: np.ndarray = freq + data_ref: np.ndarray = data + data_decimated = np.repeat(data_decimated, decimation_factor, axis=-1) + n_pad = data_ref.shape[1] - data_decimated.shape[1] + data_decimated = np.pad( + data_decimated, ((0, 0), (0, n_pad)), "constant", constant_values=0 + ) + else: + f_ref_tensor = conv(freq[None, :]) + inds_ref = (f_ref_tensor / (ufd.delta_f / 2)).type(torch.int32) + data_ref = data_2x[:, inds_ref] + f_ref = f_ref_tensor.numpy() + + diff = np.abs(data_ref - data_decimated) + diff_perc = np.percentile(diff, 95, axis=0) + diff_perc = np.maximum.accumulate(diff_perc[..., ::-1], axis=-1)[..., ::-1] + diffs_per_decimation_factor.append(diff_perc) + freqs_per_decimation_factor.append(f_ref) + + return diffs_per_decimation_factor, freqs_per_decimation_factor + + +def compute_max_decimation_factor( + decimation_factors: np.ndarray, + diffs_per_decimation_factor: List[np.ndarray], + frequencies_per_decimation_factor: List[np.ndarray], + ufd: UniformFrequencyDomain, + threshold: float, +) -> np.ndarray: + """Determine the maximum permitted decimation factor for each UFD frequency bin. + + Starting from a decimation factor of 1 for all bins, the function iterates over + increasing decimation factors. For each, it finds the lowest frequency above which + the 95th-percentile waveform difference falls below ``threshold``, then assigns that + decimation factor to all bins above that frequency. If the transition frequency for a + higher decimation factor is not strictly above the previous one, the search stops. + + Parameters + ---------- + decimation_factors : np.ndarray + Integer decimation factors in strictly increasing order. + diffs_per_decimation_factor : List[np.ndarray] + Per-decimation-factor whitened difference arrays, as returned by + :func:`compute_waveform_difference_per_decimation_factor`. + frequencies_per_decimation_factor : List[np.ndarray] + Corresponding reference frequency arrays of the same structure. + ufd : UniformFrequencyDomain + Base uniform frequency domain. + threshold : float + Maximum permitted whitened waveform difference. Higher values allow more + aggressive decimation. + + Returns + ------- + max_dec_factor : np.ndarray + Array of shape ``(len(ufd()),)`` with the maximum allowed decimation factor + per UFD bin, monotonically non-decreasing. + """ + max_dec_factor = np.ones(len(ufd())) + f_max_threshold: List[float] = [] + + for decimation_factor, f, diff in zip( + decimation_factors, + frequencies_per_decimation_factor, + diffs_per_decimation_factor, + ): + f_max = f[np.argmax(diff < threshold)] + if f_max_threshold: + if f_max_threshold[-1] >= f_max: + print( + f"Not possible to go to higher decimation factors than " + f"{max_dec_factor.max()} above {f_max_threshold[-1]}" + ) + break + f_max_threshold.append(float(f_max)) + max_dec_factor[int(f_max / ufd.delta_f) :] = decimation_factor + + return max_dec_factor + + +def get_band_nodes_for_adaptive_decimation( + max_dec_factor_array: np.ndarray, + max_dec_factor_global: int = np.inf, + min_mfd_bins_per_band: int = 1, +) -> Tuple[int, List[int]]: + """Convert a per-bin maximum decimation factor array into MFD band nodes. + + Iterates over the domain using the largest power-of-2 decimation factor permitted + by ``max_dec_factor_array``, doubling the decimation factor each time the remaining + bins allow it, until the entire domain is partitioned. + + ``min_mfd_bins_per_band`` controls the granularity of the bands. Each band is grown + in steps of ``dec_factor * min_mfd_bins_per_band`` base-domain bins, so every band + spans an integer number of ``min_mfd_bins_per_band``-bin *tokens*. This is required + by the transformer (Dingo-T1) tokenizer, which partitions the multibanded data into + fixed-size token segments and therefore needs each band to contain a whole number of + tokens. Set ``min_mfd_bins_per_band=16`` to reproduce the T1 banding. The default of + ``1`` recovers the original one-decimated-bin-per-step behaviour used by the standard + (ResNet-embedding) NPE pipeline. + + Parameters + ---------- + max_dec_factor_array : np.ndarray + 1D array of maximum allowed decimation factors per bin, monotonically + non-decreasing. + max_dec_factor_global : int + Global upper bound on the decimation factor. Default: ``np.inf`` (no bound). + min_mfd_bins_per_band : int + Minimum number of multibanded bins per band, i.e. the transformer token size. + Every band width is constrained to be an integer multiple of this value. + Default: ``1`` (no token-alignment constraint). + + Returns + ------- + initial_downsampling : int + Decimation factor of band 0. + band_nodes : List[int] + Bin-index boundaries of the bands. Band *j* spans + ``[band_nodes[j], band_nodes[j+1])``. The first element is always 0. + """ + if len(max_dec_factor_array.shape) != 1: + raise ValueError("max_dec_factor_array needs to be 1D array.") + if not (max_dec_factor_array[1:] >= max_dec_factor_array[:-1]).all(): + raise ValueError("max_dec_factor_array needs to increase monotonically.") + if min_mfd_bins_per_band < 1: + raise ValueError("min_mfd_bins_per_band must be a positive integer.") + + max_dec_factor_array = np.clip(max_dec_factor_array, None, max_dec_factor_global) + N = len(max_dec_factor_array) + dec_factor = int(max(1.0, floor_to_power_of_2(float(max_dec_factor_array[0])))) + band_nodes = [0] + initial_downsampling = dec_factor + # Advance by whole tokens. With min_mfd_bins_per_band=1 each step is a single + # decimated bin and this loop is identical to the original implementation. + upper = dec_factor * min_mfd_bins_per_band + + while upper - 1 < N: + if upper - 1 + dec_factor * min_mfd_bins_per_band >= N: + band_nodes.append(upper) + elif dec_factor * 2 <= max_dec_factor_array[upper]: + band_nodes.append(upper) + # Each band must contain a whole number of tokens. + assert (band_nodes[-1] - band_nodes[-2]) % min_mfd_bins_per_band == 0 + dec_factor *= 2 + upper += dec_factor * min_mfd_bins_per_band + + return initial_downsampling, band_nodes + + +def _get_asd_file(settings: dict) -> str: + """Extract the ASD file path from dataset settings. + + Looks for the ASD under ``settings['compression']['whitening']``. Falls back to the + standard aLIGO design-sensitivity ASD if the key is absent. + + Parameters + ---------- + settings : dict + Dataset settings dict, optionally containing a ``'compression'`` key with a + ``'whitening'`` sub-key. + + Returns + ------- + str + Path or filename of the ASD file. + """ + return settings.get("compression", {}).get( + "whitening", "aLIGO_ZERO_DET_high_P_asd.txt" + ) + + +def _load_asd( + asd_file: str, + ufd_2x: UniformFrequencyDomain, +) -> Tuple[np.ndarray, np.ndarray]: + """Load and interpolate an ASD file to both the UFD and 2x-resolution grids. + + The 1x ASD is obtained by taking every other sample of the 2x ASD, so only + ``ufd_2x`` is needed. A spectral artefact (bump) in the frequency range 477–483 Hz + is suppressed by replacing those values with the ASD value at 477 Hz, preventing it + from artificially inflating the waveform difference in that band. + + Parameters + ---------- + asd_file : str + Path to the ASD file (bilby-compatible format). + ufd_2x : UniformFrequencyDomain + Domain at twice the resolution of the base UFD. + + Returns + ------- + asd : np.ndarray + ASD interpolated to ``ufd.sample_frequencies``. + asd_2x : np.ndarray + ASD interpolated to ``ufd_2x.sample_frequencies`` with bump removed. + """ + psd = PowerSpectralDensity(asd_file=asd_file) + asd_2x = np.interp( + ufd_2x.sample_frequencies, + psd.frequency_array, + psd.asd_array, + left=np.inf, + right=np.inf, + ) + bump_mask = np.where((ufd_2x() > 477) & (ufd_2x() < 483))[0] + asd_2x[bump_mask] = asd_2x[np.argmin(np.abs(ufd_2x() - 477))] + asd = asd_2x[::2] + return asd, asd_2x + + +def _generate_whitened_waveforms( + settings: dict, + prior, + asd_file: str, + num_samples: int, + num_processes: int = 1, +) -> Tuple[ + UniformFrequencyDomain, + pd.DataFrame, + Dict[str, np.ndarray], + Dict[str, np.ndarray], + Dict[str, np.ndarray], + np.ndarray, +]: + """Generate waveforms at 2x resolution, whiten them, and return both 1x and 2x versions. + + Waveforms are generated at twice the base UFD resolution so that, when comparing a + decimated waveform to its reference, the reference value at the center of each + decimation window can be read from the high-resolution array (see + :func:`compute_waveform_difference_per_decimation_factor`). + + Parameters + ---------- + settings : dict + Dataset settings containing ``'domain'`` and ``'waveform_generator'`` keys. + prior + Prior object used to draw waveform parameters. + asd_file : str + Path to the ASD file used for whitening. + num_samples : int + Number of waveforms to generate. + num_processes : int + Number of parallel processes for waveform generation. Default: 1. + + Returns + ------- + ufd : UniformFrequencyDomain + Base uniform frequency domain. + parameters : pd.DataFrame + Sampled waveform parameters. + polarizations : Dict[str, np.ndarray] + Raw (non-whitened) waveforms at 1x resolution, shape + ``(num_samples, len(ufd()))``. Used for mismatch computation. + data_whitened : Dict[str, np.ndarray] + Whitened waveforms at 1x resolution. Used for band-node computation. + data_whitened_2x : Dict[str, np.ndarray] + Whitened waveforms at 2x resolution. Used as high-resolution reference. + asd : np.ndarray + ASD values at ``ufd.sample_frequencies``, used for pre-whitening in + mismatch evaluation. + """ + ufd = build_domain(settings["domain"]) + ufd_2x = build_domain({**ufd.domain_dict, **{"delta_f": ufd.delta_f / 2}}) + + # Select the waveform generator from the settings, mirroring generate_dataset.py, so + # the domain is built with the same generator the dataset will use. The standard + # WaveformGenerator is the default (Dingo-T1 / IMRPhenomXPHM); only EOB-style models + # set new_interface: true. (For Phenom approximants the two interfaces are expected to + # agree; honouring the flag keeps this script consistent with generate_dataset.py.) + if settings["waveform_generator"].get("new_interface", False): + waveform_generator = NewInterfaceWaveformGenerator( + domain=ufd_2x, **settings["waveform_generator"] + ) + else: + waveform_generator = WaveformGenerator( + domain=ufd_2x, **settings["waveform_generator"] + ) + parameters, polarizations_2x = generate_parameters_and_polarizations( + waveform_generator=waveform_generator, + prior=prior, + num_samples=num_samples, + num_processes=num_processes, + ) + + polarizations = {k: v[..., ::2] for k, v in polarizations_2x.items()} + assert polarizations["h_plus"][0].shape == ufd().shape + + asd, asd_2x = _load_asd(asd_file, ufd_2x) + data_whitened = {k: v / asd / ufd.noise_std for k, v in polarizations.items()} + data_whitened_2x = { + k: v / asd_2x / ufd.noise_std for k, v in polarizations_2x.items() + } + + return ufd, parameters, polarizations, data_whitened, data_whitened_2x, asd + + +def _build_mfd_for_threshold( + diffs: List[np.ndarray], + freqs: List[np.ndarray], + decimation_factors: np.ndarray, + ufd: UniformFrequencyDomain, + threshold: float, + delta_f_max_time_shift: float, + min_mfd_bins_per_band: int = 1, +) -> MultibandedFrequencyDomain: + """Construct a MultibandedFrequencyDomain from pre-computed differences and a threshold. + + Parameters + ---------- + diffs : List[np.ndarray] + Per-decimation-factor difference arrays from + :func:`compute_waveform_difference_per_decimation_factor`. + freqs : List[np.ndarray] + Corresponding reference frequency arrays. + decimation_factors : np.ndarray + Decimation factors corresponding to entries in ``diffs`` and ``freqs``. + ufd : UniformFrequencyDomain + Base uniform frequency domain. + threshold : float + Whitened waveform difference threshold; higher values allow more aggressive + decimation. + delta_f_max_time_shift : float + Maximum permitted frequency bin width (Hz) imposed by time-shift considerations. + Sets ``max_dec_factor_global = delta_f_max_time_shift / ufd.delta_f``. + min_mfd_bins_per_band : int + Transformer token size; constrains each band to an integer number of tokens. + Default: ``1`` (no token-alignment constraint). + + Returns + ------- + MultibandedFrequencyDomain + Multibanded domain constructed for the given threshold. + """ + max_dec_factor = compute_max_decimation_factor( + decimation_factors, diffs, freqs, ufd, threshold + ) + initial_downsampling, band_nodes_indices = get_band_nodes_for_adaptive_decimation( + max_dec_factor[ufd.min_idx :], + max_dec_factor_global=int(delta_f_max_time_shift / ufd.delta_f), + min_mfd_bins_per_band=min_mfd_bins_per_band, + ) + delta_f_initial = ufd.delta_f * initial_downsampling + mfd_nodes = ufd()[ufd.min_idx :][np.array(band_nodes_indices)] + return MultibandedFrequencyDomain( + nodes=mfd_nodes, delta_f_initial=delta_f_initial, base_domain=ufd + ) + + +def _compute_mismatches( + polarizations_ufd: Dict[str, np.ndarray], + ufd: UniformFrequencyDomain, + mfd: MultibandedFrequencyDomain, + asd_ufd: np.ndarray, +) -> np.ndarray: + """Compute mismatches between UFD waveforms and MFD-decimated versions interpolated back. + + For each waveform, the procedure is: + + 1. Decimate the UFD waveform to the MFD via ``mfd.decimate``. + 2. Interpolate the MFD waveform back to UFD sample frequencies. + 3. Whiten both the original and interpolated waveforms with ``asd_ufd``. + 4. Compute the mismatch (1 - overlap) using inner products summed from ``ufd.min_idx``. + + This mirrors the evaluation in :func:`evaluate_multibanded_domain._evaluate_multibanding_main` + while reusing pre-generated waveforms to avoid redundant waveform generation during + threshold search. + + Parameters + ---------- + polarizations_ufd : Dict[str, np.ndarray] + Raw waveforms in the UFD, shape ``(num_samples, len(ufd()))``. + ufd : UniformFrequencyDomain + Base uniform frequency domain. + mfd : MultibandedFrequencyDomain + Candidate multibanded domain to evaluate. + asd_ufd : np.ndarray + ASD values at ``ufd.sample_frequencies``, shape ``(len(ufd()),)``, used + for whitening prior to inner-product computation. + + Returns + ------- + mismatches : np.ndarray + 1D array of mismatch values concatenated across all polarisations and samples. + """ + ufd_freqs = ufd() + mfd_freqs = mfd() + + all_mismatches: List[np.ndarray] = [] + for waveforms in polarizations_ufd.values(): + pol_mismatches = np.empty(len(waveforms)) + for i, wf_ufd in enumerate(waveforms): + wf_mfd = mfd.decimate(wf_ufd) + wf_mfd_interp = interp1d(mfd_freqs, wf_mfd, fill_value="extrapolate")( + ufd_freqs + ) + wf_ufd_white = wf_ufd / asd_ufd + wf_mfd_interp_white = wf_mfd_interp / asd_ufd + pol_mismatches[i] = get_mismatch(wf_ufd_white, wf_mfd_interp_white, ufd) + all_mismatches.append(pol_mismatches) + + return np.concatenate(all_mismatches) + + +def _output_settings_path(settings_file: str) -> str: + """Derive the output MFD settings file path from the input UFD settings path. + + Replaces the first occurrence of ``'_ufd'`` in the filename stem with ``'_mfd'``. + If ``'_ufd'`` is not present, appends ``'_mfd'`` before the ``.yaml`` extension. + The output is placed in the same directory as ``settings_file``. + + Parameters + ---------- + settings_file : str + Path to the input UFD settings YAML file. + + Returns + ------- + str + Absolute path for the output MFD settings YAML file. + """ + directory = os.path.dirname(os.path.abspath(settings_file)) + basename = os.path.basename(settings_file) + name, ext = os.path.splitext(basename) + if "_ufd" in name: + out_name = name.replace("_ufd", "_mfd", 1) + ext + else: + out_name = name + "_mfd" + ext + return os.path.join(directory, out_name) + + +def generate_multibanded_domain_settings( + settings_file: str, + num_samples: int, + target_median_mismatch: float, + num_processes: int = 1, + delta_f_max_time_shift: float = 2.0, + decimation_factors: Optional[np.ndarray] = None, + initial_threshold: float = 5e-3, + max_iterations: int = 20, + token_size: Optional[int] = None, +) -> str: + """Generate a MultibandedFrequencyDomain settings file targeting a given median mismatch. + + Loads a uniform frequency domain (UFD) settings file, generates waveforms once using + an extreme prior (minimum chirp mass, boundary geocent time), then searches over the + whitened waveform difference threshold until the MFD achieves the desired median + mismatch. The resulting MFD settings are saved next to the input file, and mismatch + statistics are printed to stdout. + + The search uses two phases: + + 1. **Bracketing**: starting from ``initial_threshold``, walk geometrically outward + (multiplying or dividing by a step factor that doubles each step) until the target + mismatch is bracketed. This focuses evaluations near the likely solution rather than + at extreme values that are unlikely to be close to the answer. + 2. **Bisection**: refine within the bracket until the MFD nodes converge (discrete + structure) or the bracket width drops below 0.1%. + + Waveforms are generated only once and reused across all iterations, keeping the runtime + cost proportional to a single waveform generation plus + :math:`O(N_{\\text{iter}} \\cdot N_{\\text{samples}})` cheap operations. + + Parameters + ---------- + settings_file : str + Path to the UFD waveform dataset settings YAML. Must contain ``'domain'``, + ``'waveform_generator'``, and ``'intrinsic_prior'`` keys. + num_samples : int + Number of waveforms used to determine the threshold and evaluate the mismatch. + target_median_mismatch : float + Desired upper bound on the median mismatch between UFD and MFD waveforms. + num_processes : int + Number of parallel processes for waveform generation. Default: 1. + delta_f_max_time_shift : float + Maximum permitted frequency bin width (Hz) set by time-shift requirements. + Controls the global upper bound on the decimation factor. Default: 2.0. + decimation_factors : np.ndarray, optional + Decimation factors to evaluate. Default: ``2 ** np.arange(8)``. + initial_threshold : float + Starting point for the threshold search. Should be a reasonable central estimate; + the search walks outward from here. Default: 5e-3. + max_iterations : int + Maximum number of iterations for each search phase. Default: 20. + token_size : int, optional + Transformer token size (multibanded bins per token). When set, every band is + constrained to contain a whole number of tokens, as required by the Dingo-T1 + transformer tokenizer; use ``16`` to reproduce T1. When ``None`` (default), bands + are not token-aligned, matching the standard ResNet-embedding NPE pipeline. + + Returns + ------- + output_path : str + Path of the saved MFD settings YAML file. + + Raises + ------ + RuntimeError + If ``target_median_mismatch`` cannot be achieved even with the most conservative + decimation reached during the downward walk. + """ + if decimation_factors is None: + decimation_factors = 2 ** np.arange(8) + + min_mfd_bins_per_band = token_size if token_size is not None else 1 + if min_mfd_bins_per_band < 1: + raise ValueError("token_size must be a positive integer.") + if token_size is not None: + print( + f"Token-aligned banding enabled: bands constrained to multiples of " + f"{token_size} multibanded bins (transformer token size)." + ) + + with open(settings_file, "r") as f: + settings = yaml.safe_load(f) + + asd_file = _get_asd_file(settings) + + prior = build_extreme_prior(settings) + print("Prior (extreme settings for stress-testing multibanding):") + for k, v in prior.items(): + print(f" {k}: {v}") + + print(f"\nGenerating {num_samples} waveforms at 2x resolution...") + ( + ufd, + _parameters, + polarizations, + data_whitened, + data_whitened_2x, + asd, + ) = _generate_whitened_waveforms( + settings, prior, asd_file, num_samples, num_processes + ) + + print("Computing waveform differences per decimation factor...") + diffs, freqs = compute_waveform_difference_per_decimation_factor( + decimation_factors=decimation_factors, + waveforms=data_whitened["h_cross"], + ufd=ufd, + waveforms_2x=data_whitened_2x["h_cross"], + ) + + # --- Two-phase threshold search --- + # + # Phase 1 (bracketing): start at initial_threshold and walk geometrically outward, + # doubling the step factor each iteration, until the target is bracketed. This avoids + # evaluating at extreme values far from the solution. + # + # Phase 2 (bisection): refine within the bracket. Terminate early when both endpoints + # map to identical MFD nodes (the threshold→nodes mapping is a step function). + + best_mfd: Optional[MultibandedFrequencyDomain] = None + best_mismatches: Optional[np.ndarray] = None + + def _eval(t: float) -> Tuple[MultibandedFrequencyDomain, np.ndarray, float]: + m = _build_mfd_for_threshold( + diffs, + freqs, + decimation_factors, + ufd, + t, + delta_f_max_time_shift, + min_mfd_bins_per_band=min_mfd_bins_per_band, + ) + mis = _compute_mismatches(polarizations, ufd, m, asd) + return m, mis, float(np.median(mis)) + + def _log(label: str, t: float, med: float, m: MultibandedFrequencyDomain) -> None: + comp = len(ufd()[ufd.frequency_mask]) / len(m()) + print( + f" {label:>12} {t:>12.3e} {med:>16.4e} {m.num_bands:>6} {comp:>11.1f}x" + ) + + print( + f"\nSearching for threshold targeting median mismatch ≤ {target_median_mismatch:.2e}:" + ) + print( + f" {'Step':>12} {'Threshold':>12} {'Median mismatch':>16} " + f"{'Bands':>6} {'Compression':>12}" + ) + + # Evaluate starting point + mfd_init, mis_init, med_init = _eval(initial_threshold) + _log("init", initial_threshold, med_init, mfd_init) + + # Phase 1: bracketing walk + threshold_low: Optional[float] = None + threshold_high: Optional[float] = None + + if med_init <= target_median_mismatch: + # Initial threshold is feasible; walk upward to find where target is exceeded. + best_mfd, best_mismatches = mfd_init, mis_init + threshold_low = initial_threshold + factor = 2.0 + for i in range(max_iterations): + t = threshold_low * factor + mfd_t, mis_t, med_t = _eval(t) + _log(f"up {i + 1}", t, med_t, mfd_t) + if med_t > target_median_mismatch: + threshold_high = t + break + best_mfd, best_mismatches = mfd_t, mis_t + threshold_low = t + factor *= 2.0 + if threshold_high is None: + print( + " Target achievable at all tested thresholds; using most aggressive." + ) + else: + # Initial threshold is too aggressive; walk downward to find where target is met. + threshold_high = initial_threshold + factor = 2.0 + for i in range(max_iterations): + t = threshold_high / factor + mfd_t, mis_t, med_t = _eval(t) + _log(f"down {i + 1}", t, med_t, mfd_t) + if med_t <= target_median_mismatch: + threshold_low = t + best_mfd, best_mismatches = mfd_t, mis_t + break + threshold_high = t + factor *= 2.0 + if threshold_low is None: + raise RuntimeError( + f"Target median mismatch {target_median_mismatch:.2e} is too stringent: " + f"even the smallest tested threshold yields median mismatch > target. " + "Consider increasing the target or using more waveform samples." + ) + + # Phase 2: bisection within [threshold_low, threshold_high] + if threshold_low is not None and threshold_high is not None: + lo_nodes = _build_mfd_for_threshold( + diffs, + freqs, + decimation_factors, + ufd, + threshold_low, + delta_f_max_time_shift, + min_mfd_bins_per_band=min_mfd_bins_per_band, + ).nodes + hi_nodes = _build_mfd_for_threshold( + diffs, + freqs, + decimation_factors, + ufd, + threshold_high, + delta_f_max_time_shift, + min_mfd_bins_per_band=min_mfd_bins_per_band, + ).nodes + if np.allclose(lo_nodes, hi_nodes): + print("\n Bracket maps to identical MFD nodes; no bisection needed.") + else: + print(f"\n Refining bracket [{threshold_low:.3e}, {threshold_high:.3e}]:") + previous_nodes: Optional[np.ndarray] = None + for iteration in range(max_iterations): + threshold = np.sqrt(threshold_low * threshold_high) + mfd, mismatches, median_mismatch = _eval(threshold) + _log(f"bisect {iteration + 1}", threshold, median_mismatch, mfd) + + if median_mismatch <= target_median_mismatch: + best_mfd, best_mismatches = mfd, mismatches + threshold_low = threshold + else: + threshold_high = threshold + + if previous_nodes is not None and np.allclose( + mfd.nodes, previous_nodes + ): + print(" Nodes unchanged; bisection converged.") + break + previous_nodes = mfd.nodes.copy() + + if threshold_high / threshold_low < 1.001: + print(" Bracket width < 0.1%; bisection converged.") + break + + # --- Report and save --- + mfd_final = best_mfd + mismatches_final = best_mismatches + + print(f"\nFinal MFD:") + print(f" Nodes : {mfd_final.nodes.tolist()}") + print(f" delta_f_initial: {mfd_final._delta_f_bands[0].item()}") + print(f" Num bands : {mfd_final.num_bands}") + print( + f" Compression : {len(ufd()[ufd.frequency_mask]) / len(mfd_final()):.2f}x" + ) + print_mismatch_stats(mismatches_final, num_samples) + + output_settings = deepcopy(settings) + output_settings["domain"] = { + "type": "MultibandedFrequencyDomain", + "nodes": [float(f) for f in mfd_final.nodes], + "delta_f_initial": float(mfd_final._delta_f_bands[0]), + "base_domain": settings["domain"], + } + + output_path = _output_settings_path(settings_file) + with open(output_path, "w") as f: + yaml.dump(output_settings, f) + print(f"\nSaved MFD settings to: {output_path}") + + return output_path + + +def parse_args(): + parser = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, + description="Generate a waveform dataset settings file with a MultibandedFrequencyDomain " + "from a waveform dataset settings file with a UniformFrequencyDomain.", + ) + parser.add_argument( + "--settings_file", + type=str, + required=True, + help="Path to the UFD waveform dataset settings YAML file.", + ) + parser.add_argument( + "--num_samples", + type=int, + default=1000, + help="Number of waveforms used for threshold tuning and mismatch evaluation.", + ) + parser.add_argument( + "--target_median_mismatch", + type=float, + required=True, + help="Desired upper bound on the median mismatch between UFD and MFD waveforms.", + ) + parser.add_argument( + "--num_processes", + type=int, + default=1, + help="Number of parallel processes for waveform generation.", + ) + parser.add_argument( + "--token_size", + type=int, + default=None, + help="Transformer token size (multibanded bins per token). Set to 16 for " + "Dingo-T1 token-aligned banding; omit for standard NPE.", + ) + return parser.parse_args() + + +def main() -> None: + """Entry point for the generate_multibanded_domain CLI.""" + args = parse_args() + generate_multibanded_domain_settings( + settings_file=args.settings_file, + num_samples=args.num_samples, + target_median_mismatch=args.target_median_mismatch, + num_processes=args.num_processes, + token_size=args.token_size, + ) + + +if __name__ == "__main__": + main() diff --git a/dingo/gw/domains/multibanded_frequency_domain.py b/dingo/gw/domains/multibanded_frequency_domain.py index 04189e008..97c281660 100644 --- a/dingo/gw/domains/multibanded_frequency_domain.py +++ b/dingo/gw/domains/multibanded_frequency_domain.py @@ -337,11 +337,12 @@ def sample_frequencies(self) -> np.ndarray: @property def frequency_mask(self) -> np.ndarray: - """Array of len(self) consisting of ones. + """Boolean array of len(self) consisting of True values. As the MultibandedFrequencyDomain starts from f_min, no masking is generally - required.""" - return np.ones_like(self.sample_frequencies) + required. The mask is boolean so it can be used directly for index selection + (e.g. ``frequencies[domain.frequency_mask]``).""" + return np.ones_like(self.sample_frequencies, dtype=bool) @property def frequency_mask_length(self) -> int: diff --git a/pyproject.toml b/pyproject.toml index fa510cd77..39be50938 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,7 @@ dingo_build_svd = "dingo.gw.dataset.utils:build_svd_cli" dingo_estimate_psds = "dingo.gw.noise.asd_estimation:download_and_estimate_cli" dingo_evaluate_multibanded_domain = "dingo.gw.dataset.evaluate_multibanded_domain:main" dingo_generate_asd_dataset = "dingo.gw.noise.generate_dataset:generate_dataset" +dingo_generate_multibanded_domain = "dingo.gw.dataset.generate_multibanded_domain:main" dingo_generate_dataset = "dingo.gw.dataset.generate_dataset:main" dingo_generate_dataset_dag = "dingo.gw.dataset.generate_dataset_dag:main" dingo_generate_synthetic_asd_dataset = "dingo.gw.noise.synthetic.generate_dataset:main" diff --git a/tests/gw/test_multibanded_domain_generation.py b/tests/gw/test_multibanded_domain_generation.py new file mode 100644 index 000000000..9947c0484 --- /dev/null +++ b/tests/gw/test_multibanded_domain_generation.py @@ -0,0 +1,440 @@ +""" +Tests for generate_multibanded_domain.py, evaluate_multibanded_domain.py, and +_multibanded_domain_utils.py. + +Tests for the top-level generate_multibanded_domain_settings and +_evaluate_multibanding_main functions (which require waveform generation) are not +included here; those are covered by integration tests. +""" + +import os +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest +from bilby.core.prior import DeltaFunction + +from dingo.gw.dataset._multibanded_domain_utils import build_extreme_prior +from dingo.gw.dataset.evaluate_multibanded_domain import \ + _evaluate_multibanding_main +from dingo.gw.dataset.generate_multibanded_domain import ( + _build_mfd_for_threshold, _compute_mismatches, _output_settings_path, + compute_max_decimation_factor, + compute_waveform_difference_per_decimation_factor, floor_to_power_of_2, + get_band_nodes_for_adaptive_decimation) +from dingo.gw.domains import MultibandedFrequencyDomain, UniformFrequencyDomain +from dingo.gw.prior import default_intrinsic_dict + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + +INTRINSIC_PRIOR = default_intrinsic_dict + + +@pytest.fixture +def settings(): + return {"intrinsic_prior": INTRINSIC_PRIOR.copy()} + + +@pytest.fixture +def ufd(): + return UniformFrequencyDomain(f_min=20.0, f_max=256.0, delta_f=1.0) + + +@pytest.fixture +def ufd_2x(ufd): + return UniformFrequencyDomain( + f_min=ufd.f_min, f_max=ufd.f_max, delta_f=ufd.delta_f / 2 + ) + + +@pytest.fixture +def mfd(ufd): + # Two bands: [20, 64) with dec=1, [64, 256) with dec=2 + return MultibandedFrequencyDomain( + nodes=[20.0, 64.0, 256.0], + delta_f_initial=1.0, + base_domain=ufd, + ) + + +# --------------------------------------------------------------------------- +# floor_to_power_of_2 +# --------------------------------------------------------------------------- + + +class TestFloorToPowerOf2: + @pytest.mark.parametrize( + "x, expected", + [ + (1.0, 1.0), + (2.0, 2.0), + (3.0, 2.0), + (4.0, 4.0), + (7.9, 4.0), + (8.0, 8.0), + (128.0, 128.0), + (129.0, 128.0), + ], + ) + def test_values(self, x, expected): + assert floor_to_power_of_2(x) == expected + + +# --------------------------------------------------------------------------- +# get_band_nodes_for_adaptive_decimation +# --------------------------------------------------------------------------- + + +class TestGetBandNodesForAdaptiveDecimation: + def test_all_ones_gives_single_band(self): + max_dec = np.ones(100) + initial_ds, nodes = get_band_nodes_for_adaptive_decimation(max_dec) + assert initial_ds == 1 + assert nodes[0] == 0 + assert nodes[-1] == len(max_dec) + + def test_nodes_start_at_zero(self): + max_dec = np.concatenate([np.ones(50), 4 * np.ones(50)]) + _, nodes = get_band_nodes_for_adaptive_decimation(max_dec) + assert nodes[0] == 0 + + def test_nodes_monotonically_increasing(self): + max_dec = np.concatenate([np.ones(32), 2 * np.ones(32), 4 * np.ones(32)]) + _, nodes = get_band_nodes_for_adaptive_decimation(max_dec) + assert all(nodes[i] < nodes[i + 1] for i in range(len(nodes) - 1)) + + def test_global_max_dec_factor_clipping(self): + max_dec = np.concatenate([np.ones(16), 2 * np.ones(16), 4 * np.ones(32)]) + initial_ds, _ = get_band_nodes_for_adaptive_decimation( + max_dec, max_dec_factor_global=2 + ) + assert initial_ds <= 2 + + def test_raises_for_non_monotonic_input(self): + max_dec = np.array([1.0, 2.0, 1.0]) + with pytest.raises(ValueError): + get_band_nodes_for_adaptive_decimation(max_dec) + + def test_raises_for_2d_input(self): + max_dec = np.ones((10, 10)) + with pytest.raises(ValueError): + get_band_nodes_for_adaptive_decimation(max_dec) + + +# --------------------------------------------------------------------------- +# compute_max_decimation_factor +# --------------------------------------------------------------------------- + + +class TestComputeMaxDecimationFactor: + def test_shape(self, ufd): + decimation_factors = np.array([2, 4]) + n = len(ufd()) // 2 + diffs = [np.zeros(n), np.zeros(n)] + freqs = [np.linspace(0, ufd.f_max, n), np.linspace(0, ufd.f_max, n)] + result = compute_max_decimation_factor( + decimation_factors, diffs, freqs, ufd, threshold=0.1 + ) + assert result.shape == (len(ufd()),) + + def test_result_monotonically_nondecreasing(self, ufd): + # Monotonically decreasing diff → decimation factor increases with frequency + decimation_factors = np.array([2, 4]) + n = len(ufd()) // 2 + diffs = [np.linspace(1.0, 0.0, n), np.linspace(1.0, 0.0, n)] + freqs = [np.linspace(0, ufd.f_max, n), np.linspace(0, ufd.f_max, n)] + result = compute_max_decimation_factor( + decimation_factors, diffs, freqs, ufd, threshold=0.5 + ) + assert np.all(result[1:] >= result[:-1]) + + def test_all_values_among_valid_decimation_factors(self, ufd): + decimation_factors = np.array([2, 4, 8]) + n = len(ufd()) // 2 + diffs = [np.zeros(n)] * 3 + freqs = [np.linspace(0, ufd.f_max, n)] * 3 + result = compute_max_decimation_factor( + decimation_factors, diffs, freqs, ufd, threshold=0.1 + ) + valid = np.concatenate([[1], decimation_factors]) + assert all(v in valid for v in np.unique(result)) + + +# --------------------------------------------------------------------------- +# compute_waveform_difference_per_decimation_factor +# --------------------------------------------------------------------------- + + +class TestComputeWaveformDifferencePerdecimationFactor: + def test_flat_waveforms_give_zero_difference(self, ufd, ufd_2x): + n_samples = 5 + waveforms = np.ones((n_samples, len(ufd()))) + waveforms_2x = np.ones((n_samples, len(ufd_2x()))) + decimation_factors = np.array([2, 4]) + diffs, freqs = compute_waveform_difference_per_decimation_factor( + decimation_factors, waveforms, ufd, waveforms_2x + ) + for d in diffs: + assert np.allclose(d, 0.0, atol=1e-10) + + def test_output_length_matches_decimation_factors(self, ufd, ufd_2x): + n_samples = 5 + waveforms = np.random.randn(n_samples, len(ufd())) + waveforms_2x = np.random.randn(n_samples, len(ufd_2x())) + decimation_factors = np.array([2, 4, 8]) + diffs, freqs = compute_waveform_difference_per_decimation_factor( + decimation_factors, waveforms, ufd, waveforms_2x + ) + assert len(diffs) == len(decimation_factors) + assert len(freqs) == len(decimation_factors) + + def test_diffs_are_nonnegative(self, ufd, ufd_2x): + n_samples = 10 + waveforms = np.random.randn(n_samples, len(ufd())) + waveforms_2x = np.random.randn(n_samples, len(ufd_2x())) + decimation_factors = np.array([2]) + diffs, _ = compute_waveform_difference_per_decimation_factor( + decimation_factors, waveforms, ufd, waveforms_2x + ) + assert np.all(diffs[0] >= 0) + + def test_diffs_are_monotonically_nonincreasing(self, ufd, ufd_2x): + # The cumulative max ensures the diff array is non-increasing + n_samples = 10 + waveforms = np.random.randn(n_samples, len(ufd())) + waveforms_2x = np.random.randn(n_samples, len(ufd_2x())) + decimation_factors = np.array([2]) + diffs, _ = compute_waveform_difference_per_decimation_factor( + decimation_factors, waveforms, ufd, waveforms_2x + ) + assert np.all(diffs[0][:-1] >= diffs[0][1:]) + + +# --------------------------------------------------------------------------- +# _build_mfd_for_threshold +# --------------------------------------------------------------------------- + + +class TestBuildMfdForThreshold: + def test_returns_valid_mfd(self, ufd, ufd_2x): + n_samples = 5 + waveforms = np.ones((n_samples, len(ufd()))) + waveforms_2x = np.ones((n_samples, len(ufd_2x()))) + decimation_factors = np.array([2, 4]) + diffs, freqs = compute_waveform_difference_per_decimation_factor( + decimation_factors, waveforms, ufd, waveforms_2x + ) + mfd = _build_mfd_for_threshold(diffs, freqs, decimation_factors, ufd, 0.01, 2.0) + assert isinstance(mfd, MultibandedFrequencyDomain) + assert mfd.num_bands >= 1 + + def test_higher_threshold_gives_equal_or_fewer_bins(self, ufd, ufd_2x): + # Higher threshold → more aggressive decimation → fewer MFD bins + n_samples = 10 + np.random.seed(0) + waveforms = np.random.randn(n_samples, len(ufd())) + waveforms_2x = np.random.randn(n_samples, len(ufd_2x())) + decimation_factors = np.array([2, 4, 8]) + diffs, freqs = compute_waveform_difference_per_decimation_factor( + decimation_factors, waveforms, ufd, waveforms_2x + ) + mfd_conservative = _build_mfd_for_threshold( + diffs, freqs, decimation_factors, ufd, 1e-4, 2.0 + ) + mfd_aggressive = _build_mfd_for_threshold( + diffs, freqs, decimation_factors, ufd, 1.0, 2.0 + ) + assert len(mfd_aggressive()) <= len(mfd_conservative()) + + +# --------------------------------------------------------------------------- +# _compute_mismatches +# --------------------------------------------------------------------------- + + +class TestComputeMismatches: + def test_mismatches_in_unit_interval(self, ufd, mfd): + n_samples = 5 + rng = np.random.default_rng(0) + waveforms = rng.standard_normal( + (n_samples, len(ufd())) + ) + 1j * rng.standard_normal((n_samples, len(ufd()))) + polarizations = {"h_plus": waveforms} + asd = np.ones(len(ufd())) + mismatches = _compute_mismatches(polarizations, ufd, mfd, asd) + assert mismatches.shape == (n_samples,) + assert np.all(mismatches >= 0) + assert np.all(mismatches <= 1) + + def test_constant_waveform_gives_near_zero_mismatch(self, ufd, mfd): + # A constant complex waveform is perfectly preserved by decimation + interpolation + n_samples = 3 + constant = np.ones(len(ufd()), dtype=complex) + polarizations = {"h_plus": np.stack([constant] * n_samples)} + asd = np.ones(len(ufd())) + mismatches = _compute_mismatches(polarizations, ufd, mfd, asd) + assert np.allclose(mismatches, 0.0, atol=1e-10) + + +# --------------------------------------------------------------------------- +# _output_settings_path +# --------------------------------------------------------------------------- + + +class TestOutputSettingsPath: + def test_replaces_ufd_with_mfd(self): + path = _output_settings_path("/some/dir/settings_wfd_ufd.yaml") + assert os.path.basename(path) == "settings_wfd_mfd.yaml" + assert os.path.dirname(path) == "/some/dir" + + def test_replaces_ufd_suffix_in_longer_name(self): + path = _output_settings_path("/data/settings_wfd_ufd_32s.yaml") + assert os.path.basename(path) == "settings_wfd_mfd_32s.yaml" + + def test_appends_mfd_when_ufd_absent(self): + path = _output_settings_path("/some/dir/settings_wfd.yaml") + assert os.path.basename(path) == "settings_wfd_mfd.yaml" + + def test_output_in_same_directory_as_input(self): + path = _output_settings_path("/some/dir/settings_wfd_ufd.yaml") + assert os.path.dirname(path) == "/some/dir" + + +# --------------------------------------------------------------------------- +# build_extreme_prior +# --------------------------------------------------------------------------- + + +class TestBuildExtremePrior: + def test_chirp_mass_fixed_at_nominal_minimum(self, settings): + from dingo.gw.prior import build_prior_with_defaults + + nominal_minimum = build_prior_with_defaults(settings["intrinsic_prior"])[ + "chirp_mass" + ].minimum + prior = build_extreme_prior(settings) + assert isinstance(prior["chirp_mass"], DeltaFunction) + assert prior["chirp_mass"].peak == nominal_minimum + + def test_geocent_time_fixed_at_boundary(self, settings): + prior = build_extreme_prior(settings) + assert isinstance(prior["geocent_time"], DeltaFunction) + assert prior["geocent_time"].peak == 0.12 + + def test_input_settings_not_mutated(self, settings): + original_chirp_mass = settings["intrinsic_prior"]["chirp_mass"] + original_geocent_time = settings["intrinsic_prior"]["geocent_time"] + build_extreme_prior(settings) + assert settings["intrinsic_prior"]["chirp_mass"] == original_chirp_mass + assert settings["intrinsic_prior"]["geocent_time"] == original_geocent_time + + def test_other_parameters_unchanged(self, settings): + from dingo.gw.prior import build_prior_with_defaults + + nominal = build_prior_with_defaults(settings["intrinsic_prior"]) + extreme = build_extreme_prior(settings) + for key in nominal: + if key in ("chirp_mass", "geocent_time"): + continue + assert str(extreme[key]) == str(nominal[key]) + + +# --------------------------------------------------------------------------- +# _evaluate_multibanding_main +# --------------------------------------------------------------------------- + +MFD_SETTINGS = { + "domain": { + "type": "MultibandedFrequencyDomain", + "nodes": [20.0, 64.0, 256.0], + "delta_f_initial": 1.0, + "base_domain": { + "type": "UniformFrequencyDomain", + "f_min": 20.0, + "f_max": 256.0, + "delta_f": 1.0, + }, + }, + "waveform_generator": { + "approximant": "IMRPhenomPv2", + "f_ref": 20.0, + "new_interface": False, + }, + "intrinsic_prior": INTRINSIC_PRIOR, +} + + +@pytest.fixture +def mfd_settings_file(tmp_path): + import yaml + + path = tmp_path / "settings_wfd_mfd.yaml" + with open(path, "w") as f: + yaml.dump(MFD_SETTINGS, f) + return str(path) + + +class TestEvaluateMultibandingMain: + def test_raises_for_ufd_settings(self, tmp_path): + import yaml + + ufd_settings = dict(MFD_SETTINGS) + ufd_settings["domain"] = { + "type": "UniformFrequencyDomain", + "f_min": 20.0, + "f_max": 256.0, + "delta_f": 1.0, + } + path = str(tmp_path / "settings_ufd.yaml") + with open(path, "w") as f: + yaml.dump(ufd_settings, f) + with pytest.raises(ValueError, match="MultibandedFrequencyDomain"): + _evaluate_multibanding_main(path, num_samples=2) + + def test_uses_extreme_prior(self, mfd_settings_file): + """Verify that the extreme prior (min chirp mass, geocent_time=0.12) is applied.""" + captured_prior = {} + + def mock_generate(waveform_generator, prior, num_samples, num_processes): + captured_prior["prior"] = prior + ufd = waveform_generator.domain + n = len(ufd()) + pols = { + "h_plus": np.ones((num_samples, n), dtype=complex), + "h_cross": np.ones((num_samples, n), dtype=complex), + } + return pd.DataFrame({"dummy": range(num_samples)}), pols + + def mock_parallel(waveform_generator, parameters): + ufd = waveform_generator.domain + n = len(ufd()) + return { + "h_plus": np.ones((len(parameters), n), dtype=complex), + "h_cross": np.ones((len(parameters), n), dtype=complex), + } + + with ( + patch( + "dingo.gw.dataset.evaluate_multibanded_domain.generate_parameters_and_polarizations", + side_effect=mock_generate, + ), + patch( + "dingo.gw.dataset.evaluate_multibanded_domain.generate_waveforms_parallel", + side_effect=mock_parallel, + ), + ): + _evaluate_multibanding_main(mfd_settings_file, num_samples=2) + + from dingo.gw.prior import build_prior_with_defaults + + nominal_minimum = build_prior_with_defaults(INTRINSIC_PRIOR)[ + "chirp_mass" + ].minimum + prior = captured_prior["prior"] + assert isinstance(prior["chirp_mass"], DeltaFunction) + assert prior["chirp_mass"].peak == nominal_minimum + assert isinstance(prior["geocent_time"], DeltaFunction) + assert prior["geocent_time"].peak == 0.12