Multibanding (#292)

* Implement MultibandedFrequencyDomain.

Lots of code taken from Max's implementation. Not sure this works with the new waveform interface.

* Enable inference for MultibandedFrequencyDomain.

* Add brute force phase grid to likelihood (for MultibandedFrequencyDomain).

* Add toggle for uniform vs multi-banded domain for Result, Injection classes.

* Raise NotImplementedError if calling torch_sample_frequencies in MFD.

* Update MFD for batched time translation.

* Fix domain_update test for ASD dataset.

* Fix domain_update test for ASD dataset.

* Whiten data before decimating at inference time.

* Improve usage of MultibandedFrequencyDomain / storage of event data.

No longer use context setter in Sampler, instead always store event data as it was generated. Use transform in Sampler to decimate data.

Ensure DecimateWaveformsAndASDS does nothing to data already in MFD.

* Update dingo-pipe for MFD.

By default use the base domain. Can be changed with option `importance-sampling-settings = {use_base_domain: False}`.

* Implement torch frequency arrays for MFD, update numpy arrays to float32.

This is needed for GNPE sampling.

* Extend waveform generator to multi-banded frequency domain for all relevant waveform approximants.

For TD-native waveforms, generate in FD then decimate. Also extend methods that generate waveform modes.

Add relevant unit tests.

* Remove IrregularFrequencyDomain.

* Add tests for MFD.

* Add additional test.

* Restructure UniformFrequencyDomain and MultibandedFrequencyDomain to inherit from common BaseFrequencyDomain superclass.

* Fix bug in build_domain().

* Resolve several comments in PR.

* Resolve several more comments in PR.

* Address several additional comments, fix some docstrings.

* Updated link in documentation to point to public notebook.

* Improve check for TD waveforms for decimation method.

* Fix bug in recursive_hdf5_load and other small changes.

* Remove adaptive MFD generation.

* Remove adaptive MFD generation.

* Fix test.

* Add diagnostic script to test MFD performance.

* Improve MFD diagnostic script.

* fix bug in error message

* fix error when checking time domain type

* Update stripping of window factor from ASD dataset for MFD.

* elif -> if in Injection.asd setter

---------

Co-authored-by: Annalena Kofler <annalena.kofler@tuebingen.mpg.de>
Co-authored-by: Maximilian Dax <maximilian.dax@tuebingen.mpg.de>

Author: 3 people

dingo/core/dataset.py

@@ -162,7 +162,7 @@ def from_file(self, file_name: str):
             # Set the keys that the class expects
             for k, v in loaded_dict.items():
                 assert k in self._data_keys
-                vars(self)[k] = v
+                setattr(self, k, v)
             try:
                 self.settings = ast.literal_eval(f.attrs["settings"])
             except KeyError:
@@ -179,6 +179,6 @@ def to_dictionary(self):
     def from_dictionary(self, dictionary: dict):
         for k, v in dictionary.items():
             if k in self._data_keys or k == "settings":
-                vars(self)[k] = v
+                setattr(self, k, v)
         if "version" not in dictionary:
             self.version = f"dingo={get_version()}"

dingo/core/samplers.py

@@ -75,12 +75,12 @@ def __init__(
             # 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._context = self.model.context
             self.event_metadata = self.model.event_metadata
             self.base_model_metadata = self.metadata["base"]
         else:
             self.unconditional_model = False
-            self.context = None
+            self._context = None
             self.event_metadata = None
             self.base_model_metadata = self.metadata
 

dingo/gw/data/data_preparation.py

@@ -7,7 +7,8 @@
 from dingo.core.utils.misc import recursive_check_dicts_are_equal
 from dingo.gw.data.data_download import download_raw_data
 from dingo.gw.gwutils import get_window
-from dingo.gw.domains import build_domain_from_model_metadata, FrequencyDomain
+from dingo.gw.domains import UniformFrequencyDomain
+from dingo.gw.domains import build_domain_from_model_metadata
 
 
 def load_raw_data(time_event, settings, event_dataset=None):
@@ -35,7 +36,6 @@ def load_raw_data(time_event, settings, event_dataset=None):
 
     # first try to load the event data from the saved dataset
     if event_dataset is not None:
-
         if isfile(event_dataset):
             dataset = DingoDataset(file_name=event_dataset, data_keys=[event])
             if settings is not None:
@@ -66,7 +66,7 @@ def load_raw_data(time_event, settings, event_dataset=None):
 def parse_settings_for_raw_data(model_metadata, time_psd, time_buffer):
     domain_type = model_metadata["dataset_settings"]["domain"]["type"]
 
-    if domain_type == "FrequencyDomain":
+    if domain_type == "UniformFrequencyDomain":
         data_settings = model_metadata["train_settings"]["data"]
         settings = {
             "window": data_settings["window"],
@@ -98,7 +98,7 @@ def data_to_domain(raw_data, settings_raw_data, domain, **kwargs):
 
     """
 
-    if type(domain) == FrequencyDomain:
+    if isinstance(domain, UniformFrequencyDomain):
         window = get_window(kwargs["window"])
         data = {"waveform": {}, "asds": {}}
         # convert event strains to frequency domain

dingo/gw/dataset/evaluate_multibanded_domain.py

@@ -0,0 +1,130 @@
+import argparse
+
+import numpy as np
+import yaml
+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.gwutils import get_mismatch
+from dingo.gw.prior import build_prior_with_defaults
+from dingo.gw.waveform_generator import (
+    NewInterfaceWaveformGenerator,
+    WaveformGenerator,
+    generate_waveforms_parallel,
+)
+
+
+def _evaluate_multibanding_main(
+    settings_file: str,
+    num_samples: int,
+):
+    with open(settings_file, "r") as f:
+        settings = yaml.safe_load(f)
+
+    # Ignore any compression settings
+    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"])
+    print("Prior")
+    for k, v in prior.items():
+        print(f"{k}: {v}")
+
+    domain = build_domain(settings["domain"])
+    print("\nDomain")
+    print(domain.domain_dict)
+
+    if not isinstance(domain, MultibandedFrequencyDomain):
+        raise ValueError("Waveform dataset domain not a MultibandedFrequencyDomain.")
+
+    if settings["waveform_generator"].get("new_interface", False):
+        waveform_generator_mfd = NewInterfaceWaveformGenerator(
+            domain=domain,
+            **settings["waveform_generator"],
+        )
+        waveform_generator_ufd = NewInterfaceWaveformGenerator(
+            domain=domain.base_domain,
+            **settings["waveform_generator"],
+        )
+    else:
+        waveform_generator_mfd = WaveformGenerator(
+            domain=domain,
+            **settings["waveform_generator"],
+        )
+        waveform_generator_ufd = WaveformGenerator(
+            domain=domain.base_domain,
+            **settings["waveform_generator"],
+        )
+
+    # Generate MFD waveforms.
+    parameters, polarizations_mfd = generate_parameters_and_polarizations(
+        waveform_generator_mfd, prior, num_samples, 1
+    )
+
+    # Generate UFD waveforms, re-using the parameter choices from before.
+    polarizations_ufd = generate_waveforms_parallel(waveform_generator_ufd, parameters)
+
+    # Compare UFD waveforms against MFD waveforms interpolated to MFD.
+    mismatches = {}
+    ufd = domain.base_domain
+    mfd = domain
+    for pol, d in polarizations_mfd.items():
+        mismatches[pol] = np.empty(len(d))
+        for i in range(len(d)):
+            mfd_interpolated = interp1d(mfd(), d[i], fill_value="extrapolate")(ufd())
+            mismatches[pol][i] = get_mismatch(
+                polarizations_ufd[pol][i],
+                mfd_interpolated,
+                ufd,
+                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)))
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="Evaluate performance of multibanding on waveform dataset.",
+    )
+    parser.add_argument(
+        "--settings-file",
+        type=str,
+        required=True,
+        help="YAML file containing database settings",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=5000,
+        help="Number of waveform evaluations for comparison.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    _evaluate_multibanding_main(args.settings_file, args.num_samples)

dingo/gw/dataset/generate_dataset.py

@@ -25,6 +25,7 @@
 )
 from dingo.core.utils.misc import call_func_strict_output_dim
 
+
 def generate_parameters_and_polarizations(
     waveform_generator: WaveformGenerator,
     prior: BBHPriorDict,
@@ -198,11 +199,14 @@ def generate_dataset(settings: Dict, num_processes: int) -> WaveformDataset:
                 n_test = svd_settings.get("num_validation_samples", 0)
 
                 func = partial(
-                    generate_parameters_and_polarizations, 
-                    waveform_generator, 
+                    generate_parameters_and_polarizations,
+                    waveform_generator,
                     prior,
-                    num_processes=num_processes)
-                parameters, polarizations = call_func_strict_output_dim(func, n_train + n_test)
+                    num_processes=num_processes,
+                )
+                parameters, polarizations = call_func_strict_output_dim(
+                    func, n_train + n_test
+                )
                 svd_dataset_settings = copy.deepcopy(settings)
                 svd_dataset_settings["num_samples"] = len(parameters)
                 del svd_dataset_settings["compression"]["svd"]
@@ -231,11 +235,14 @@ def generate_dataset(settings: Dict, num_processes: int) -> WaveformDataset:
         waveform_generator.transform = Compose(compression_transforms)
 
     func = partial(
-        generate_parameters_and_polarizations, 
-        waveform_generator, 
+        generate_parameters_and_polarizations,
+        waveform_generator,
         prior,
-        num_processes=num_processes)
-    parameters, polarizations = call_func_strict_output_dim(func, settings["num_samples"])
+        num_processes=num_processes,
+    )
+    parameters, polarizations = call_func_strict_output_dim(
+        func, settings["num_samples"]
+    )
     dataset_dict["parameters"] = parameters
     dataset_dict["polarizations"] = polarizations
 
@@ -277,7 +284,6 @@ def parse_args():
 def _generate_dataset_main(
     settings_file: str, out_file: str, num_processes: int
 ) -> None:
-
     if not Path(settings_file).is_file():
         raise FileNotFoundError(f"dataset generation, failed to find {settings_file}")
     if not Path(out_file).parent.is_dir():