Add GNPEChirp to training pipeline

dingo/gw/SVD.py

@@ -256,19 +256,22 @@ def __init__(self, svd_basis: SVDBasis, inverse: bool = False):
         self.svd_basis = svd_basis
         self.inverse = inverse
 
-    def __call__(self, waveform: dict):
+    def __call__(self, input_sample: dict):
         """
         Parameters
         ----------
-        waveform : dict
-            Values should be arrays containing waveforms to be transformed.
+        input_sample : dict
+            input_sample["waveform"] should be a dict with arrays containing waveforms
+            to be transformed.
 
         Returns
         -------
         dict of the same form as the input, but with transformed waveforms.
         """
+        sample = input_sample.copy()
         if not self.inverse:
             func = self.svd_basis.compress
         else:
             func = self.svd_basis.decompress
-        return {k: func(v) for k, v in waveform.items()}
+        sample["waveform"] = {k: func(v) for k, v in sample["waveform"].items()}
+        return sample

dingo/gw/dataset/generate_dataset.py

@@ -18,6 +18,7 @@
 from dingo.gw.prior import build_prior_with_defaults
 from dingo.gw.SVD import ApplySVD, SVDBasis
 from dingo.gw.transforms import WhitenFixedASD
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
 from dingo.gw.waveform_generator import (
     NewInterfaceWaveformGenerator,
     WaveformGenerator,
@@ -183,6 +184,15 @@ def generate_dataset(settings: Dict, num_processes: int) -> WaveformDataset:
                 )
             )
 
+        if "phase_heterodyning" in settings["compression"]:
+            compression_transforms.append(
+                HeterodynePhase(
+                    domain,
+                    inverse=False,
+                    **settings["compression"]["phase_heterodyning"],
+                )
+            )
+
         if "svd" in settings["compression"]:
             svd_settings = settings["compression"]["svd"]
 

dingo/gw/dataset/waveform_dataset.py

@@ -8,6 +8,7 @@
 from dingo.gw.SVD import SVDBasis, ApplySVD
 from dingo.gw.domains import build_domain
 from dingo.gw.transforms import WhitenFixedASD
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
 
 
 class WaveformDataset(DingoDataset, torch.utils.data.Dataset):
@@ -196,6 +197,15 @@ def initialize_decompression(self, svd_size_update: Optional[int] = None):
             svd_basis = SVDBasis(dictionary=self.svd)
             decompression_transform_list.append(ApplySVD(svd_basis, inverse=True))
 
+        if "phase_heterodyning" in self.settings["compression"]:
+            decompression_transform_list.append(
+                HeterodynePhase(
+                    self.domain,
+                    inverse=True,
+                    **self.settings["compression"]["phase_heterodyning"],
+                )
+            )
+
         if "whitening" in self.settings["compression"]:
             decompression_transform_list.append(
                 WhitenFixedASD(
@@ -322,13 +332,12 @@ def __getitems__(
                 for pol, waveforms in self.polarizations.items()
             }
 
-        # Decompression transforms are assumed to apply only to the waveform,
-        # and do not involve parameters.
-        if self.decompression_transform is not None:
-            polarizations = self.decompression_transform(polarizations)
-
-        # Main transforms can depend also on parameters.
+        # Decompression transforms operate on sample dicts since some (e.g.,
+        # HeterodynePhase) need access to parameters.
         data = {"parameters": parameters, "waveform": polarizations}
+        if self.decompression_transform is not None:
+            data = self.decompression_transform(data)
+            polarizations = data["waveform"]
         if self.transform is not None:
             data = self.transform(data)
 

dingo/gw/training/train_builders.py

@@ -19,6 +19,7 @@
     RepackageStrainsAndASDS,
     UnpackDict,
     GNPECoalescenceTimes,
+    GNPEChirp,
     SampleExtrinsicParameters,
     GetDetectorTimes,
     CropMaskStrainRandom,
@@ -125,6 +126,11 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N
         )
         extra_context_parameters += transforms[-1].context_parameters
 
+    if "gnpe_chirp" in data_settings:
+        d = data_settings["gnpe_chirp"]
+        transforms.append(GNPEChirp(d["kernel"], domain, d.get("order", 0)))
+        extra_context_parameters += transforms[-1].context_parameters
+
     # Add the GNPE context to context_parameters the first time the transforms are
     # constructed. We do not want to overwrite the ordering of the parameters in
     # subsequent runs.
@@ -255,6 +261,7 @@ def build_svd_for_embedding_network(
             SelectStandardizeRepackageParameters,
             UnpackDict,
             CropMaskStrainRandom,
+            GNPEChirp,
         ],
     )
 

dingo/gw/transforms/gnpe_transforms.py

@@ -3,6 +3,8 @@
 from bilby.core.prior import PriorDict
 from abc import ABC, abstractmethod
 
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
+
 
 class GNPEBase(ABC):
     """
@@ -227,3 +229,111 @@ def __call__(self, input_sample):
         extrinsic_parameters.update(new_parameters)
         sample["extrinsic_parameters"] = extrinsic_parameters
         return sample
+
+
+class GNPEChirp(GNPEBase):
+    """
+    Relative binning / heterodyning GNPE transform, which factors out the overall chirp
+    from the waveform. This is done based on the proxy parameters chirp_mass_proxy and
+    optionally mass_ratio_proxy. These are defined as blurred version of the parameters
+    chirp_mass and mass_ratio.
+
+    At leading order, the data are transformed by dividing by a fiducial waveform of the
+    form
+
+        exp( - 1j * (3/128) * (pi G chirp_mass_proxy f / c**3)**(-5/3) ) ;
+
+    see 2001.11412, eq. (7.2). This is the leading order chirp due to the emission of
+    quadrupole radiation.
+
+    At next to leading order, this transform also optionally implements 1PN corrections
+    involving the mass ratio. We do not include any amplitude in the fiducial waveform,
+    since at inference time this transform will be applied to noisy data. Multiplying
+    the frequency-domain noise by a complex number of unit norm is allowed because it
+    only changes the phase, not the overall amplitude, which would change the noise PSD.
+    """
+
+    def __init__(self, kernel, domain, order: int = 0, inference: bool = False):
+        """
+        Parameters
+        ----------
+        kernel : dict or str
+            Defines a Bilby prior. If a dict, keys should include chirp_mass,
+            and (possibly) mass_ratio.
+        domain : Domain
+            Only works for a FrequencyDomain at present.
+        order : int
+            Twice the post-Newtonian order for the expansion. Valid orders are 0 and 2.
+        inference : bool = False
+            Whether to use inference or training mode.
+        """
+        # We copy the kernel because the PriorDict constructor modifies the argument.
+        kernel = kernel.copy()
+
+        if order == 0:
+            if "chirp_mass" not in kernel:
+                raise KeyError("Kernel must include chirp_mass key.")
+            if "mass_ratio" in kernel:
+                print(
+                    "Warning: mass_ratio kernel provided, but will be ignored for "
+                    "order 0 GNPE."
+                )
+                kernel.pop("mass_ratio")
+        elif order == 2:
+            if "chirp_mass" not in kernel or "mass_ratio" not in kernel:
+                raise KeyError("Kernel must include chirp_mass and mass_ratio keys.")
+            raise NotImplementedError(
+                "Mass ratio conditioning for 2nd order not set up."
+            )
+        else:
+            raise ValueError(f"Order {order} invalid. Acceptable values are 0 and 2.")
+
+        operators = {"chirp_mass": "+", "mass_ratio": "+"}
+        super().__init__(kernel, operators)
+
+        self.inference = inference
+        self.phase_heterodyning_transform = HeterodynePhase(
+            domain, order, inverse=False
+        )
+
+    def __call__(self, input_sample):
+        sample = input_sample.copy()
+
+        extrinsic_parameters = sample["extrinsic_parameters"].copy()
+
+        # If proxies already exist, use them. Otherwise, sample them. Proxies may
+        # already exist if provided by an unconditional initialization network when
+        # attempting to recover the density from GNPE samples, or when using fixed
+        # initialization parameters.
+        if set(self.proxy_list).issubset(extrinsic_parameters.keys()):
+            proxies = {p: extrinsic_parameters[p] for p in self.proxy_list}
+        else:
+            # The relevant parameters could be in either the intrinsic or extrinsic
+            # parameters list. At inference time, we put all GNPE parameters into the
+            # extrinsic parameters list.
+            parameters = {**sample.get("parameters", {}), **extrinsic_parameters}
+            proxies = self.sample_proxies(parameters)
+            extrinsic_parameters.update(proxies)
+            delta_parameters = {
+                "delta_" + p[: -len("_proxy")]: parameters[p[: -len("_proxy")]]
+                - proxies[p]
+                for p in self.proxy_list
+            }
+            extrinsic_parameters.update(delta_parameters)
+            sample["extrinsic_parameters"] = extrinsic_parameters
+
+        # The only situation where we would expect to not have a waveform to transform
+        # would be when calculating parameter standardizations, since we just want to
+        # draw samples of the parameters at that point, and not prepare any data.
+        if "waveform" in sample:
+            sample["waveform"] = self.phase_heterodyning_transform(
+                {
+                    "waveform": sample["waveform"],
+                    "parameters": {
+                        "chirp_mass": proxies["chirp_mass_proxy"],
+                        "mass_ratio": proxies.get("mass_ratio_proxy"),
+                    },
+                }
+            )["waveform"]
+
+        return sample

dingo/gw/transforms/noise_transforms.py

@@ -104,25 +104,27 @@ def __init__(
 
         self.inverse = inverse
 
-    def __call__(self, sample):
+    def __call__(self, input_sample):
         """
         Parameters
         ----------
-        sample : dict
-            Dictionary of numpy arrays, e.g., with keys corresponding to polarizations.
-            Method whitens each array with the same ASD.
+        input_sample : dict
+            input_sample["waveform"] is a dictionary of numpy arrays, e.g., with keys
+            corresponding to polarizations. Method whitens each array with the same ASD.
 
         Returns
         -------
-        dict of the same form as sample, but with whitened / un-whitened data.
+        dict of the same form as input_sample, but with whitened / un-whitened waveforms.
         """
-        result = {}
-        for k, v in sample.items():
+        sample = input_sample.copy()
+        waveform = {}
+        for k, v in sample["waveform"].items():
             if self.inverse:
-                result[k] = v * self.asd_array
+                waveform[k] = v * self.asd_array
             else:
-                result[k] = v / self.asd_array
-        return result
+                waveform[k] = v / self.asd_array
+        sample["waveform"] = waveform
+        return sample
 
 
 class WhitenAndScaleStrain(object):