add fft into generator

ML4GW · Jan 30, 2025 · fb8b47d · fb8b47d
1 parent 7602325
commit fb8b47d
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/ml4gw/waveforms/generator.py b/ml4gw/waveforms/generator.py
@@ -218,4 +218,12 @@ def forward(
         hp_spectrum *= phase_shift
         hc_spectrum *= phase_shift
 
+        hp_spectrum = torch.fft.irfft(hp_spectrum) * self.sample_rate
+        hc_spectrum = torch.fft.irfft(hc_spectrum) * self.sample_rate
+
+        # pad waveforms on left up to duration
+        pad = int((self.duration * self.sample_rate) - hp_spectrum.shape[-1])
+        hp_spectrum = torch.nn.functional.pad(hp_spectrum, (pad, 0))
+        hc_spectrum = torch.nn.functional.pad(hc_spectrum, (pad, 0))
+
         return hc_spectrum, hp_spectrum