f delay values

radiocosmology · Apr 4, 2023 · 48e3885 · 48e3885
1 parent cbb9608
commit 48e3885
Showing 1 changed file with 18 additions and 5 deletions.
diff --git a/draco/analysis/wavelet.py b/draco/analysis/wavelet.py
@@ -20,7 +20,7 @@ class WaveletSpectrumEstimator(task.SingleTask):
     dataset = config.Property(proptype=str, default="vis")
     average_axis = config.Property(proptype=str)
     ndelay = config.Property(proptype=int, default=128)
-    wavelet = config.Property(proptype=str, default="mexh")
+    wavelet = config.Property(proptype=str, default="morl")
 
     def process(
         self,
@@ -57,20 +57,29 @@ def process(
         nbase = dset_view.global_shape[0]
 
         df = np.abs(data.freq[1] - data.freq[0])
-        delay_scales = np.arange(1, self.ndelay + 1) / (df * self.ndelay)
+        delay_scales = np.arange(1, self.ndelay + 1) / (2 * df * self.ndelay)
+
+        wv_scales = pywt.frequency2scale(self.wavelet, delay_scales * df)
 
         wspec = containers.WaveletSpectrum(
             baseline=nbase,
             axes_from=data,
             attrs_from=data,
             delay=delay_scales,
         )
+        # Copy the index maps for all the flattened axes into the output container, and
+        # write out their order into an attribute so we can reconstruct this easily
+        # when loading in the spectrum
+        for ax in bl_axes:
+            wspec.create_index_map(ax, data.index_map[ax])
+        wspec.attrs["baseline_axes"] = bl_axes
+
         wspec.redistribute("baseline")
         dspec.redistribute("baseline")
-
         ws = wspec.spectrum[:].local_array
         ds = dspec.spectrum[:].local_array
 
+        # Construct the
         F = np.exp(
             -2.0j
             * np.pi
@@ -99,9 +108,13 @@ def process(
                 overwrite_b=True,
             ).T
 
-            scales = self.ndelay / np.arange(1.0, self.ndelay + 1)
             wd, s = pywt.cwt(
-                d_infill, scales=scales, wavelet=self.wavelet, axis=-1, method="fft"
+                d_infill,
+                scales=wv_scales,
+                wavelet=self.wavelet,
+                axis=-1,
+                sampling_period=df,
+                method="fft",
             )
 
             ws[ii] = wd.var(axis=1)