Using rfi mitigation in candmaker

thepetabyteproject · Feb 3, 2021 · 4f143f1 · 4f143f1
1 parent 8ea65e2
commit 4f143f1
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/bin/your_candmaker.py b/bin/your_candmaker.py
@@ -111,7 +111,17 @@ def cand2h5(cand_val):
 
     logger.debug(f"Source file list: {files}")
     cand = Candidate(
-        files, snr=snr, width=width, dm=dm, label=label, tcand=tcand, device=gpu_id
+        files,
+        snr=snr,
+        width=width,
+        dm=dm,
+        label=label,
+        tcand=tcand,
+        device=gpu_id,
+        spectral_kurtosis_sigma=args.spectral_kurtosis_sigma,
+        savgol_frequency_window=args.savgol_frequency_window,
+        savgol_sigma=args.savgol_sigma,
+        flag_rfi=args.flag_rfi,
     )
     if os.path.exists(str(kill_mask_path)):
         kill_mask = np.zeros(cand.nchans, dtype=np.bool)
@@ -215,6 +225,37 @@ def cand2h5(cand_val):
         type=str,
         default=".",
     )
+    parser.add_argument(
+        "-r",
+        "--flag_rfi",
+        help="Turn on RFI flagging",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "-sksig",
+        "--spectral_kurtosis_sigma",
+        help="Sigma for spectral kurtosis filter",
+        type=float,
+        default=4,
+        required=False,
+    )
+    parser.add_argument(
+        "-sgsig",
+        "--savgol_sigma",
+        help="Sigma for savgol filter",
+        type=float,
+        default=4,
+        required=False,
+    )
+    parser.add_argument(
+        "-sgfw",
+        "--savgol_frequency_window",
+        help="Filter window for savgol filter (MHz)",
+        type=float,
+        default=15,
+        required=False,
+    )
     parser.add_argument(
         "-opt",
         "--opt_dm",

diff --git a/tests/test_candidate.py b/tests/test_candidate.py
@@ -149,6 +149,7 @@ def test_resize(cand):
     except AttributeError:
         pass
 
+
 def test_rfi_mask():
     fil_file = os.path.join(_install_dir, "data/28.fil")
     cand = Candidate(
@@ -170,9 +171,10 @@ def test_rfi_mask():
     assert cand.data[:, 172:177].sum() == 0
     assert cand.data[:, ~cand.rfi_mask].sum() != 0
 
+
 def test_kill_mask():
     fil_file = os.path.join(_install_dir, "data/28.fil")
-    km = np.zeros(336, dtype='bool')
+    km = np.zeros(336, dtype="bool")
     km[[10, 12, 25, 100, 300]] = True
     cand = Candidate(
         fp=fil_file,
@@ -184,7 +186,7 @@ def test_kill_mask():
         min_samp=256,
         device=0,
         flag_rfi=False,
-        kill_mask=km
+        kill_mask=km,
     )
     cand.get_chunk()
     assert cand.data[:, cand.kill_mask].sum() == 0