sk_sg_filter can now work if one of the sigma is zero

Signed-off-by: Devansh Agarwal <[email protected]>
thepetabyteproject · Jan 31, 2021 · a0bcf90 · a0bcf90
1 parent 0b82d9e
commit a0bcf90
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 61 deletions.
diff --git a/bin/your_rfimask.py b/bin/your_rfimask.py
@@ -98,43 +98,16 @@
                 handlers=[RichHandler(rich_tracebacks=True)],
             )
 
-    if values.savgol_sigma == 0 and values.spectral_kurtosis_sigma == 0:
-        raise ValueError("All sigma values cannot be zero.")
+    your_object = Your(file=values.files)
+    bandpass = your_object.bandpass(values.nspectra)
 
-    else:
-        your_object = Your(file=values.files)
-        bandpass = your_object.bandpass(values.nspectra)
-
-    if values.savgol_sigma > 0 and values.spectral_kurtosis_sigma == 0:
-        logging.info("Applying Savgol")
-        mask = savgol_filter(
-            bandpass=bandpass,
-            channel_bandwidth=your_object.your_header.foff,
-            frequency_window=values.savgol_frequency_window,
-            sigma=values.savgol_sigma,
-        )
-
-    elif values.savgol_sigma == 0 and values.spectral_kurtosis_sigma > 0:
-        logging.info("Applying Spectral Kurtosis")
-        mask = sk_filter(
-            data=your_object.get_data(0, values.nspectra),
-            channel_bandwidth=your_object.your_header.foff,
-            tsamp=your_object.your_header.tsamp,
-            N=None,
-            d=None,
-            sigma=values.spectral_kurtosis_sigma,
-        )
-
-    else:
-        logging.info("Applying both Spectral Kurtosis and Savgol")
-        mask = sk_sg_filter(
-            data=your_object.get_data(0, values.nspectra),
-            your_object=your_object,
-            nchans=your_object.your_header.nchans,
-            spectral_kurtosis_sigma=values.spectral_kurtosis_sigma,
-            savgol_frequency_window=values.savgol_frequency_window,
-            savgol_sigma=values.savgol_sigma,
-        )
+    mask = sk_sg_filter(
+        data=your_object.get_data(0, values.nspectra),
+        your_object=your_object,
+        spectral_kurtosis_sigma=values.spectral_kurtosis_sigma,
+        savgol_frequency_window=values.savgol_frequency_window,
+        savgol_sigma=values.savgol_sigma,
+    )
 
     basename = f"{values.output_dir}/{your_object.your_header.basename}_your_rfi_mask"
     chan_nos = np.array(range(your_object.your_header.nchans), dtype=int)

diff --git a/your/utils/rfi.py b/your/utils/rfi.py
@@ -42,7 +42,8 @@ def spectral_kurtosis(data, N=1, d=None):
          numpy.ndarray: Spectral Kurtosis along frequency axis
 
     """
-    data = data.astype("float32")
+    zero_mask = data == 0
+    data = np.ma.array(data.astype(float), mask=zero_mask)
     S1 = data.sum(0)
     S2 = (data ** 2).sum(0)
     M = data.shape[0]
@@ -101,7 +102,6 @@ def calc_N(channel_bandwidth, tsamp):
 def sk_sg_filter(
     data,
     your_object,
-    nchans,
     spectral_kurtosis_sigma=6,
     savgol_frequency_window=15,
     savgol_sigma=5,
@@ -113,7 +113,6 @@ def sk_sg_filter(
     Args:
         data (numpy.ndarray): 2D frequency time data
         your_object: Your object
-        nchans (int): number of channels
         spectral_kurtosis_sigma (float): sigma value to apply cutoff on for SK filter
         savgol_frequency_window (float): frequency window for savgol filter(MHz)
         savgol_sigma (float): sigma value to apply cutoff on for savgol filter
@@ -123,27 +122,32 @@ def sk_sg_filter(
          numpy.ndarray: mask for channels
 
     """
-
-    logger.info(
-        f"Applying spectral kurtosis filter with sigma={spectral_kurtosis_sigma}"
-    )
-    sk_mask = sk_filter(
-        data=data,
-        channel_bandwidth=your_object.your_header.foff,
-        tsamp=your_object.your_header.tsamp,
-        sigma=spectral_kurtosis_sigma,
-    )
-    bp = data.sum(0)[~sk_mask]
-    logger.info(
-        f"Applying savgol filter with frequency_window={savgol_frequency_window} and sigma={savgol_sigma}"
-    )
-    sg_mask = savgol_filter(
-        bandpass=bp,
-        channel_bandwidth=your_object.your_header.foff,
-        frequency_window=savgol_frequency_window,
-        sigma=savgol_sigma,
-    )
     mask = np.zeros(data.shape[1], dtype=np.bool)
-    mask[sk_mask] = True
-    mask[np.where(mask == False)[0][sg_mask]] = True
+    if spectral_kurtosis_sigma > 0:
+        logger.debug(
+            f"Applying spectral kurtosis filter with sigma={spectral_kurtosis_sigma}"
+        )
+        sk_mask = sk_filter(
+            data=data,
+            channel_bandwidth=your_object.your_header.foff,
+            tsamp=your_object.your_header.tsamp,
+            sigma=spectral_kurtosis_sigma,
+        )
+        mask[sk_mask] = True
+    elif savgol_sigma > 0:
+        bp = data.sum(0)[~sk_mask]
+        logger.debug(
+            f"Applying savgol filter with frequency_window={savgol_frequency_window} and sigma={savgol_sigma}"
+        )
+        sg_mask = savgol_filter(
+            bandpass=bp,
+            channel_bandwidth=your_object.your_header.foff,
+            frequency_window=savgol_frequency_window,
+            sigma=savgol_sigma,
+        )
+        mask[np.where(mask == False)[0][sg_mask]] = True
+    else:
+        raise ValueError(
+            "Both savgol_sigma and spectral_kurtosis_sigma cannot be zero."
+        )
     return mask