Realtime waveforms

aframe/pipelines/sandbox/configs/base.cfg

@@ -23,11 +23,11 @@ seed = 1122
 streams_per_gpu = 6
 
 # waveform parameters
-waveform_approximant = IMRPhenomXPHM
+waveform_approximant = IMRPhenomPv2
 waveform_duration = 8
 minimum_frequency = 20
 reference_frequency = 50
-coalescence_time = 6
+right_pad = 2
 
 # training parameters
 kernel_length = 1.5
@@ -70,7 +70,7 @@ waveform_duration = &::luigi_base::waveform_duration
 minimum_frequency = &::luigi_base::minimum_frequency
 reference_frequency = &::luigi_base::reference_frequency
 waveform_approximant = &::luigi_base::waveform_approximant
-coalescence_time = &::luigi_base::coalescence_time
+right_pad = &::luigi_base::right_pad
 prior = &::luigi_base::prior
 request_memory = 64GB
 
@@ -86,7 +86,7 @@ waveform_duration = &::luigi_base::waveform_duration
 minimum_frequency = &::luigi_base::minimum_frequency
 reference_frequency = &::luigi_base::reference_frequency
 waveform_approximant = &::luigi_base::waveform_approximant
-coalescence_time = &::luigi_base::coalescence_time
+right_pad = &::luigi_base::right_pad
 prior = &::luigi_base::prior
 snr_threshold = 4
 request_memory = 16GB
@@ -124,7 +124,7 @@ minimum_frequency = &::luigi_base::minimum_frequency
 reference_frequency = &::luigi_base::reference_frequency
 waveform_duration = &::luigi_base::waveform_duration
 waveform_approximant = &::luigi_base::waveform_approximant
-coalescence_time = &::luigi_base::coalescence_time
+right_pad = &::luigi_base::right_pad
 highpass = &::luigi_base::highpass
 lowpass = &::luigi_base::lowpass
 request_memory = 16GB

aframe/pipelines/sandbox/configs/bbh.cfg

@@ -11,7 +11,7 @@ inherit = $AFRAME_REPO/aframe/pipelines/sandbox/configs/base.cfg
 # override bbh specific parameters
 [luigi_base]
 waveform_duration = 8
-coalescence_time = 4
+right_pad = 2
 kernel_length = 1.5
 prior = priors.priors.end_o3_ratesandpops
 
@@ -29,3 +29,4 @@ lowpass = &::luigi_base::lowpass
 fduration = &::luigi_base::fduration
 seed = &::luigi_base::seed
 use_wandb = true
+precompute_train_waveforms = true

aframe/pipelines/sandbox/configs/bns.cfg

@@ -13,7 +13,7 @@ inherit = $AFRAME_REPO/aframe/pipelines/sandbox/configs/base.cfg
 # override bns specific parameters
 [luigi_base]
 waveform_duration = 90
-coalescence_time = 45
+right_pad = 45
 kernel_length = 8
 prior = priors.priors.end_o3_ratesandpops_bns
 q = 45.6 
@@ -40,3 +40,4 @@ seed = &::luigi_base::seed
 fftlength = &::luigi_base::fftlength
 q = &::luigi_base::q
 use_wandb = true
+precompute_train_waveforms = true

aframe/pipelines/sandbox/configs/review.cfg

@@ -27,7 +27,7 @@ waveform_approximant = IMRPhenomPv2
 waveform_duration = 8
 minimum_frequency = 20
 reference_frequency = 50
-coalescence_time = 6
+right_pad = 2
 
 # training parameters
 kernel_length = 1.5
@@ -79,7 +79,7 @@ minimum_frequency = &::luigi_base::minimum_frequency
 reference_frequency = &::luigi_base::reference_frequency
 waveform_duration = &::luigi_base::waveform_duration
 waveform_approximant = &::luigi_base::waveform_approximant
-coalescence_time = &::luigi_base::coalescence_time
+right_pad = &::luigi_base::right_pad
 highpass = &::luigi_base::highpass
 lowpass = &::luigi_base::lowpass
 request_memory = 6GB

aframe/tasks/data/waveforms/base.py

@@ -29,9 +29,10 @@ class WaveformParams(law.Task):
         default="IMRPhenomXPHM",
         description="Approximant to use for waveform generation",
     )
-    coalescence_time = luigi.FloatParameter(
+    right_pad = luigi.FloatParameter(
         description="Location of the defining point of the signal "
-        "within the generated waveform"
+        "within the generated waveform relative to the right edge "
+        "of the waveform (in seconds)",
     )
 
 

aframe/tasks/data/waveforms/testing.py

@@ -182,7 +182,7 @@ def run(self):
                 self.sample_rate,
                 self.waveform_duration,
                 self.waveform_approximant,
-                self.coalescence_time,
+                self.right_pad,
                 self.highpass,
                 self.lowpass,
                 self.snr_threshold,

aframe/tasks/data/waveforms/training.py

@@ -1,3 +1,6 @@
+from pathlib import Path
+import shutil
+
 import law
 from luigi.util import inherits
 
@@ -10,7 +13,7 @@
 
 
 @inherits(WaveformParams)
-class TrainingWaveforms(
+class DeployTrainingWaveforms(
     AframeDataTask, DeployTask, law.LocalWorkflow, StaticMemoryWorkflow
 ):
     """
@@ -24,14 +27,20 @@ class TrainingWaveforms(
 
     output_dir = PathParameter(
         description="Directory where merged training waveforms will be saved",
-        default=paths().train_datadir / "training_waveforms",
+        default=paths().train_datadir,
+    )
+
+    tmp_dir = PathParameter(
+        description="Directory where temporary "
+        "waveforms will be saved before being merged",
+        default=paths().tmp_dir / "training_waveforms",
     )
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def output(self):
-        return s3_or_local(self.output_dir / f"waveforms-{self.branch}.hdf5")
+        return s3_or_local(self.tmp_dir / f"waveforms-{self.branch}.hdf5")
 
     def run(self):
         from data.waveforms.training import training_waveforms
@@ -47,8 +56,57 @@ def run(self):
             minimum_frequency=self.minimum_frequency,
             reference_frequency=self.reference_frequency,
             waveform_approximant=self.waveform_approximant,
-            coalescence_time=self.coalescence_time,
+            right_pad=self.right_pad,
         )
         chunks = (min(64, num_signals), waveforms.get_waveforms().shape[-1])
         with self.output().open("w") as f:
             waveforms.write(f, chunks=chunks)
+
+
+@inherits(DeployTrainingWaveforms)
+class TrainingWaveforms(AframeDataTask):
+    """
+    Launch condorized generation of training waveforms,
+    and merge results into a single file
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.output_file = self.output_dir / "training_waveforms.hdf5"
+
+    def output(self):
+        return s3_or_local(self.output_file)
+
+    def requires(self):
+        return DeployTrainingWaveforms.req(
+            self,
+            workflow=self.workflow,
+            request_memory=self.request_memory,
+            request_disk=self.request_disk,
+            request_cpus=self.request_cpus,
+        )
+
+    @property
+    def targets(self):
+        return list(self.input().collection.targets.values())
+
+    @property
+    def waveform_files(self):
+        return list(map(Path, [targets.path for targets in self.targets]))
+
+    def run(self):
+        from ledger.injections import (
+            WaveformPolarizationSet,
+            waveform_class_factory,
+        )
+
+        cls = waveform_class_factory(
+            ["cross", "plus"],
+            WaveformPolarizationSet,
+            "WaveformPolarizationSet",
+        )
+        with self.output().open("w") as f:
+            cls.aggregate(self.waveform_files, f, clean=True)
+
+        # clean up temporary directory
+        shutil.rmtree(self.tmp_dir)

aframe/tasks/data/waveforms/validation.py

@@ -124,7 +124,7 @@ def run(self):
             self.sample_rate,
             self.waveform_duration,
             self.waveform_approximant,
-            self.coalescence_time,
+            self.right_pad,
             self.highpass,
             self.lowpass,
             self.snr_threshold,

aframe/tasks/train/base.py

@@ -57,21 +57,30 @@ class TrainBaseParameters(law.Task):
     )
     data_dir = PathParameter(
         description="Directory where training data is stored."
-        "It is expected to contain a `signals.hdf5` file of signals, "
-        "and a `/background` sub-directory containing background "
-        "files used for training",
+        "It is expected to contain a `val_waveforms.hdf5` file of "
+        "signals for validation, a `/background` sub-directory containing "
+        "background, and a `training_waveforms.hdf5` file containing "
+        "training signals if `precompute_train_waveforms` is set to True.",
         default=paths().train_datadir,
     )
+    precompute_train_waveforms = luigi.BoolParameter(
+        default=False,
+        description="Whether to pre-compute the waveforms used "
+        "during training. If True, the training waveforms will be "
+        "read from the `training_waveforms.hdf5` file in the data "
+        "directory. If False, the waveforms will be simulated "
+        "on-the-fly during training.",
+    )
 
 
 @inherits(TrainBaseParameters)
 class TrainBase(law.Task):
     def requires(self):
         reqs = {}
         reqs["strain"] = FetchTrain.req(self)
-        reqs["train_waveforms"] = TrainingWaveforms.req(self)
-
         reqs["val_waveforms"] = ValidationWaveforms.req(self)
+        if self.precompute_train_waveforms:
+            reqs["train_waveforms"] = TrainingWaveforms.req(self)
         return reqs
 
     def configure_wandb(self, args: List[str]) -> None:

docs/projects/data.md

@@ -76,7 +76,7 @@ apptainer run $AFRAME_CONTAINER_ROOT/data.sif \
     --minimum_frequency 20 \
     --reference_frequency 50 \
     --waveform_approximant IMRPhenomXPHM \
-    --coalescence_time 6 \
+    --right_pad 2 \
     --output_file ~/aframe/data/train/train_waveforms.hdf5
 ```
 
@@ -93,7 +93,7 @@ apptainer run $AFRAME_CONTAINER_ROOT/data.sif \
     --sample_rate 2048 \
     --waveform_duration 8 \
     --waveform_approximant IMRPhenomXPHM \
-    --coalescence_time 6 \
+    --right_pad 2 \
     --highpass 32 \
     --snr_threshold 4 \
     --psd ~/aframe/data/train/background/background-1240579783-7829.hdf5