dingo.gw.dataset package

Submodules

dingo.gw.dataset.evaluate_multibanded_domain module

dingo.gw.dataset.evaluate_multibanded_domain.main() → None

dingo.gw.dataset.evaluate_multibanded_domain.parse_args()

dingo.gw.dataset.generate_dataset module

dingo.gw.dataset.generate_dataset.generate_dataset(settings: Dict, num_processes: int) → WaveformDataset

Generate a waveform dataset.

Parameters:

settings (dict) – Dictionary of settings to configure the dataset
num_processes (int)

Return type:

A WaveformDataset based on the settings.

dingo.gw.dataset.generate_dataset.generate_parameters_and_polarizations(waveform_generator: WaveformGenerator, prior: BBHPriorDict, num_samples: int, num_processes: int) → Tuple[DataFrame, Dict[str, ndarray]]

Generate a dataset of waveforms based on parameters drawn from the prior.

Parameters:

waveform_generator (WaveformGenerator)
prior (Prior)
num_samples (int)
num_processes (int)

Returns:

pandas DataFrame of parameters
dictionary of numpy arrays corresponding to waveform polarizations

dingo.gw.dataset.generate_dataset.main() → None

dingo.gw.dataset.generate_dataset.parse_args()

dingo.gw.dataset.generate_dataset.train_svd_basis(dataset: WaveformDataset, size: int, n_train: int)

Train (and optionally validate) an SVD basis.

Parameters:

dataset (WaveformDataset) – Contains waveforms to be used for building SVD.
size (int) – Number of elements to keep for the SVD basis.
n_train (int) – Number of training waveforms to use. Remaining are used for validation. Note that the actual number of training waveforms is n_train * len(polarizations), since there is one waveform used for each polarization.

Returns:

Since EOB waveforms can fail to generate, provide also the number used in training and validation.

Return type:

SVDBasis, n_train, n_test

dingo.gw.dataset.generate_dataset_dag module

dingo.gw.dataset.generate_dataset_dag.configure_runs(settings, num_jobs, temp_dir)

Prepare and save settings .yaml files for generating subsets of the dataset. Generally this will produce two .yaml files, one for generating the main dataset, one for the SVD training.

Parameters:

settings (dict) – Settings for full dataset configuration.
num_jobs (int) – Number of jobs over which to split the run.
temp_dir (str) – Name of (temporary) directory in which to place temporary output files.

dingo.gw.dataset.generate_dataset_dag.create_args_string(args_dict: Dict): Generate argument string from dictionary of argument names and arguments.

dingo.gw.dataset.generate_dataset_dag.create_dag(args, settings): Create a Condor DAG from command line arguments to carry out the five steps in the workflow.

dingo.gw.dataset.generate_dataset_dag.main()

dingo.gw.dataset.generate_dataset_dag.modulus_check(a: int, b: int, a_label: str, b_label: str): Raise error if a % b != 0.

dingo.gw.dataset.generate_dataset_dag.parse_args()

dingo.gw.dataset.generate_multibanded_domain module

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:

dingo_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

which is equivalent to:

python -m dingo.gw.dataset.generate_multibanded_domain \
    --settings_file path/to/settings_wfd_ufd.yaml \
    ...

dingo.gw.dataset.generate_multibanded_domain.compute_max_decimation_factor(decimation_factors: ndarray, diffs_per_decimation_factor: List[ndarray], frequencies_per_decimation_factor: List[ndarray], ufd: UniformFrequencyDomain, threshold: float) → 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. The search also stops if the difference never falls below threshold, in which case that decimation factor is not applied anywhere.

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 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 – Array of shape (len(ufd()),) with the maximum allowed decimation factor per UFD bin, monotonically non-decreasing.

Return type:

np.ndarray

dingo.gw.dataset.generate_multibanded_domain.compute_waveform_difference_per_decimation_factor(decimation_factors: ndarray, waveforms: ndarray, ufd: UniformFrequencyDomain, waveforms_2x: ndarray, difference_over_full_window: bool = False) → Tuple[List[ndarray], List[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.

dingo.gw.dataset.generate_multibanded_domain.floor_to_power_of_2(x: float) → float

Return the largest power of 2 that is <= x.

Parameters:: x (float) – Positive input value.
Returns:: Largest power of 2 not exceeding x.
Return type:: float

dingo.gw.dataset.generate_multibanded_domain.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: ndarray | None = None, initial_threshold: float = 0.005, max_iterations: int = 20, token_size: int | None = None, difference_over_full_window: bool = False) → 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:

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.
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 \(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.
difference_over_full_window (bool) – Comparison mode used when computing the whitened waveform differences, see compute_waveform_difference_per_decimation_factor(). If True, each decimated bin is compared against all original bins in its decimation window (most conservative); if False (default), it is compared against the window center at 2x resolution.

Returns:

output_path – Path of the saved MFD settings YAML file.

Return type:

str

Raises:

RuntimeError – If target_median_mismatch cannot be achieved even with the most conservative decimation reached during the downward walk.

dingo.gw.dataset.generate_multibanded_domain.get_band_nodes_for_adaptive_decimation(max_dec_factor_array: ndarray, max_dec_factor_global: int = 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.

dingo.gw.dataset.generate_multibanded_domain.main() → None: Entry point for the generate_multibanded_domain CLI.

dingo.gw.dataset.generate_multibanded_domain.parse_args()

dingo.gw.dataset.utils module

dingo.gw.dataset.utils.build_svd_cli(): Command-line function to build an SVD based on an uncompressed dataset file.

dingo.gw.dataset.utils.merge_datasets(dataset_list: List[WaveformDataset]) → WaveformDataset

Merge a collection of datasets into one.

Parameters:: dataset_list (list[WaveformDataset]) – A list of WaveformDatasets. Each item should be a dictionary containing parameters and polarizations.
Return type:: WaveformDataset containing the merged data.

dingo.gw.dataset.utils.merge_datasets_cli(): Command-line function to combine a collection of datasets into one. Used for parallelized waveform generation.

dingo.gw.dataset.waveform_dataset module

class dingo.gw.dataset.waveform_dataset.WaveformDataset(file_name: str | None = None, dictionary: dict | None = None, transform=None, precision: Literal['single', 'double'] | None = None, domain_update: dict | None = None, svd_size_update: int | None = None, leave_waveforms_on_disk: bool | None = False)

Bases: DingoDataset, Dataset

This class stores a dataset of waveforms (polarizations) and corresponding parameters.

It can load the dataset either from an HDF5 file or suitable dictionary.

It is possible to either load the entire dataset into memory or to load the dataset during training (leave_waveforms_on_disk=True) to reduce the memory footprint. At the moment, it is only possible to load the waveforms on-demand since the standardization dict for all parameters in the dataset has to be computed at the beginning of training.

The waveform data is consumed through a __getitem__() or __getitems__() call which optionally loads the polarizations and applies a chain of transformations, which are classes that implement a __call__() method.

For constructing, provide either file_name, or dictionary containing data and settings entries, or neither.

Parameters:

file_name (str) – HDF5 file containing a dataset
dictionary (dict) – Contains settings and data entries. The dictionary keys should be ‘settings’, ‘parameters’, and ‘polarizations’.
transform (Transform) – Transform to be applied to dataset samples when accessed through __getitem__
precision (str ('single', 'double')) – If provided, changes precision of loaded dataset.
domain_update (dict) – If provided, update domain from existing domain using new settings.
svd_size_update (int) – If provided, reduces the SVD size when decompressing (for speed).
leave_waveforms_on_disk (bool) – If True, the values for the waveforms are not loaded into RAM when initializing the waveform dataset. Instead, they are loaded lazily in __getitem__().

dataset_type = 'waveform_dataset'

property dtype_map: Mapping[str, DTypeLike | DTypeMap] | None

Mapping from group names to target dtypes for HDF5 loading.

This enables direct dtype conversion during HDF5 read, avoiding intermediate memory allocation when changing precision.

initialize_decompression(svd_size_update: int | None = None)

Sets up decompression transforms. These are applied to the raw dataset before self.transform. E.g., SVD decompression.

Parameters:: svd_size_update (int) – If provided, reduces the SVD size when decompressing (for speed).

load_supplemental(domain_update: dict | None = None, svd_size_update: int | None = None)

Method called immediately after loading a dataset.

Creates (and possibly updates) domain, updates dtypes, and initializes any decompression transform. Also zeros data below f_min, and truncates above f_max.

Parameters:

domain_update (dict) – If provided, update domain from existing domain using new settings.
svd_size_update (int) – If provided, reduces the SVD size when decompressing (for speed).

parameter_mean_std()

update_domain(domain_update: dict | None = None)

Update the domain based on new configuration.

The waveform dataset provides waveform polarizations in a particular domain. In Frequency domain, this is [0, domain._f_max]. Furthermore, data is set to 0 below domain._f_min. In practice one may want to train a network based on slightly different domain settings, which corresponds to truncating the likelihood integral.

This method provides functionality for that. It truncates and/or zeroes the dataset to the range specified by the domain, by calling domain.update_data.

Parameters:: domain_update (dict) – Settings dictionary. Must contain a subset of the keys contained in domain_dict.

dingo.gw.dataset package

Submodules

dingo.gw.dataset.evaluate_multibanded_domain module

dingo.gw.dataset.generate_dataset module

dingo.gw.dataset.generate_dataset_dag module

dingo.gw.dataset.generate_multibanded_domain module

dingo.gw.dataset.utils module

dingo.gw.dataset.waveform_dataset module

Module contents