Piecewise MAD feature for wideband observations

clean_archive.py

@@ -26,7 +26,7 @@ def apply_rcvrstd_cleaner(ar):
     rcvrstd_cleaner = cleaners.load_cleaner('rcvrstd')
     rcvrstd_cleaner.run(ar)
 
-def apply_surgical_cleaner(ar, tmp, cthresh=None, sthresh=None, plot=False, aggressive=False, iterations=1):
+def apply_surgical_cleaner(ar, tmp, cthresh=None, sthresh=None, plot=False, aggressive=False, iterations=1, piecewise_scale=False):
 
     """Apply the surgical cleaner to an archive in place.
 
@@ -42,6 +42,10 @@ def apply_surgical_cleaner(ar, tmp, cthresh=None, sthresh=None, plot=False, aggr
                 and algorithms. (Default: False)
             iterations: Number of times to run the surgical cleaner.
                 (Default: 1)
+            piecewise_scale: If True, force per-segment (piecewise) MAD
+                scaling of the frequency-direction statistics on. If False,
+                the setting is left to the (default or -C) config, so this
+                does not switch off a config that enables it. (Default: False)
 
         Outputs:
             None - The archive is cleaned in place.
@@ -62,6 +66,10 @@ def apply_surgical_cleaner(ar, tmp, cthresh=None, sthresh=None, plot=False, aggr
         param_parts.append("chanthresh={0}".format(cthresh))
     if sthresh is not None:
         param_parts.append("subintthresh={0}".format(sthresh))
+    if piecewise_scale:
+        # Only force it on; when the flag is absent leave the config's value
+        # untouched (so -C UWL, which enables it, is not overridden to off).
+        param_parts.append("piecewise_scale=True")
     surgical_parameters = ",".join(param_parts)
     surgical_cleaner.parse_config_string(surgical_parameters)
     surgical_cleaner.run(ar)
@@ -121,6 +129,12 @@ def apply_bandwagon_cleaner(ar, badchantol=None, badsubtol=None):
               help="Whether to use more aggressive cleaning thresholds and algorithms")
     parser.add_argument("-i", "--iterations", type=int, dest="iterations", default=1,
               help="Number of iterations to run the surgical cleaner [default = 1]")
+    parser.add_argument("-pw", "--piecewise", dest='piecewise_scale', action='store_true', default=False,
+              help="Scale the frequency-direction statistics to sigma using a "
+              "per-segment (piecewise) MAD that matches the piecewise detrend, "
+              "instead of one global MAD across the whole band. Recommended for "
+              "wideband receivers with a strong Tsys gradient (already on in the "
+              "UWL config). [default = off]")
     parser.add_argument("-C", "--config", type=str, dest="config_path", default=None,
               help="Custom config file for misbehaving receivers. Inputting UHF or L will "
               "automatically load the MeerKAT config files for those receivers. Inputting "
@@ -140,6 +154,7 @@ def apply_bandwagon_cleaner(ar, badchantol=None, badsubtol=None):
     output_path = args.output_path
     aggressive = args.aggressive
     iterations = args.iterations
+    piecewise_scale = args.piecewise_scale
     config_path = args.config_path
 
     #raise error if user tries to write multiple input files to one output name
@@ -221,7 +236,7 @@ def apply_bandwagon_cleaner(ar, badchantol=None, badsubtol=None):
                 subint_thresh if subint_thresh is not None else "cfg")
 
         apply_rcvrstd_cleaner(loaded_archive)
-        apply_surgical_cleaner(loaded_archive, template_path, cthresh=chan_thresh, sthresh=subint_thresh, plot=plot, aggressive=aggressive, iterations=iterations)
+        apply_surgical_cleaner(loaded_archive, template_path, cthresh=chan_thresh, sthresh=subint_thresh, plot=plot, aggressive=aggressive, iterations=iterations, piecewise_scale=piecewise_scale)
         apply_bandwagon_cleaner(loaded_archive, badchantol=badchantol, badsubtol=badsubtol)
 
         # Unload the Archive file

coast_guard/clean_utils.py

@@ -112,13 +112,92 @@ def channel_scaler(array2d, **kwargs):
     return scaled
 
 
+def _piecewise_mad_scale(detrended, numpieces, min_points=16):
+    """Scale a 1D detrended residual to robust sigma using a per-segment
+        median/MAD instead of a single global one.
+
+        The residual is split into 'numpieces' roughly-equal segments using
+        the same np.linspace scheme as detrend(), and each segment is
+        centred on its own median and divided by its own MAD (scaled by
+        1.4826). This keeps the definition of "sigma" local in frequency, so
+        a channel is judged against the noise floor of its own sub-band
+        rather than the whole band. On wideband receivers, where the
+        off-pulse RMS varies strongly with Tsys, this avoids inflating the
+        apparent significance of channels in high-Tsys sub-bands.
+
+        Segments too sparse or degenerate to give a stable MAD fall back to
+        the global median/MAD of the whole residual:
+          * fewer than 'min_points' unmasked samples, or
+          * a zero MAD (all-equal or heavily-masked segment).
+
+        Inputs:
+            detrended: A 1D (optionally masked) array of detrended residuals.
+            numpieces: The number of frequency segments to split into.
+            min_points: The minimum number of unmasked samples for a segment
+                to use its own median/MAD. (Default: 16)
+
+        Output:
+            scaled: A 1D masked array of the residual in units of robust
+                sigma. Masked entries of the input remain masked.
+    """
+    detrended = np.ma.masked_array(detrended,
+                                   mask=np.ma.getmaskarray(detrended))
+    # Global fall-back statistics (used for sparse/degenerate segments).
+    global_median = np.ma.median(detrended)
+    global_mad = np.ma.median(np.abs(detrended - global_median))
+
+    scaled = np.ma.masked_array(np.empty(detrended.shape, dtype=float),
+                                mask=np.ma.getmaskarray(detrended).copy())
+    # Same segmentation scheme as detrend()'s numpieces splitting.
+    isplit = np.linspace(0, detrended.size, numpieces + 1, endpoint=1)
+    edges = np.round(isplit).astype(int)
+    for start, stop in zip(edges[:-1], edges[1:]):
+        seg = detrended[start:stop]
+        if np.ma.count(seg) < min_points:
+            # Too few unmasked points for a stable local MAD.
+            median, mad = global_median, global_mad
+        else:
+            median = np.ma.median(seg)
+            mad = np.ma.median(np.abs(seg - median))
+            if mad == 0:
+                # Degenerate segment; local MAD would divide by zero.
+                median, mad = global_median, global_mad
+        scaled[start:stop] = (detrended[start:stop] - median)/(mad * 1.4826)
+    return scaled
+
+
 def subint_scaler(array2d, **kwargs):
     """For each sub-int detrend and scale it.
+
+        By default each sub-int's detrended spectrum is scaled to robust
+        sigma using a single median/MAD taken over the whole band. On
+        wideband receivers the off-pulse RMS varies strongly with frequency
+        (Tsys), so a single global MAD inflates the apparent significance of
+        high-Tsys channels and systematically over-flags them. Passing
+        piecewise_scale=True instead computes the median/MAD per frequency
+        segment, using the same segmentation as the piecewise detrend, so the
+        scaling granularity matches the detrend and each channel is judged
+        against its own sub-band.
+
+        Keyword arguments:
+            subint_order, subint_breakpoints, subint_numpieces: as for the
+                frequency-direction detrend.
+            piecewise_scale: If True, scale to sigma using a per-segment
+                (piecewise) median/MAD rather than a single global MAD.
+                (Default: False, i.e. original global behaviour.)
+            subint_mad_numpieces: Number of frequency segments to use for the
+                piecewise MAD. If None, the finest (last) subint_numpieces
+                value is used so the scaling lines up with the finest detrend
+                pass. Only used when piecewise_scale is True. (Default: None)
     """
     # Grab key-word arguments. If not present use default configs.
     orders = kwargs.pop('subint_order', config.cfg.subint_order)
     breakpoints = kwargs.pop('subint_breakpoints', config.cfg.subint_breakpoints)
     numpieces = kwargs.pop('subint_numpieces', config.cfg.subint_numpieces)
+    # Piecewise-MAD scaling controls. Default to the backward-compatible
+    # global-MAD behaviour when these are absent.
+    piecewise_scale = kwargs.pop('piecewise_scale', False)
+    mad_numpieces = kwargs.pop('subint_mad_numpieces', None)
     if breakpoints is None:
         breakpoints = [[]]*len(orders)
     if numpieces is None:
@@ -133,16 +212,29 @@ def subint_scaler(array2d, **kwargs):
                          min(len(orders), len(breakpoints), len(numpieces))),
                       errors.CoastGuardWarning)
 
+    # Resolve the segmentation for the piecewise MAD. When not given, default
+    # to the finest (last) detrend granularity so the scaling segments line up
+    # with the finest detrend pass. A value <= 1 means "one global segment",
+    # which reduces exactly to the original global-MAD behaviour.
+    if mad_numpieces is None:
+        last_numpieces = numpieces[-1] if len(numpieces) else None
+        mad_numpieces = last_numpieces if last_numpieces is not None else 1
+
     scaled = np.empty_like(array2d)
     nsubs = array2d.shape[0]
     for isub in np.arange(nsubs):
         detrended = array2d[isub,:]
         for order, brkpnts, numpcs in zip(orders, breakpoints, numpieces):
             detrended = iterative_detrend(detrended, order=order, \
                                             bp=brkpnts, numpieces=numpcs)
-        median = np.ma.median(detrended)
-        mad = np.ma.median(np.abs(detrended-median))
-        scaled[isub,:] = (detrended-median)/(mad * 1.4826)  # Scale MAD to be consistent with std for normal distribution
+        if piecewise_scale and mad_numpieces > 1:
+            # Local (per-segment) median/MAD, matching the piecewise detrend.
+            scaled[isub,:] = _piecewise_mad_scale(detrended, mad_numpieces)
+        else:
+            # Single global median/MAD over the whole band (original path).
+            median = np.ma.median(detrended)
+            mad = np.ma.median(np.abs(detrended-median))
+            scaled[isub,:] = (detrended-median)/(mad * 1.4826)  # Scale MAD to be consistent with std for normal distribution
     return scaled
 
 
@@ -415,6 +507,35 @@ def fft_rotate(data, bins):
     # original length (irfft otherwise assumes an even-length signal).
     return np.fft.irfft(phasor*np.fft.rfft(data), n=data.size)
 
+def check_template_nchan(template, data_nchan):
+    """Warn if a 2D (per-channel) template's channel count does not match the
+        data being cleaned.
+
+        remove_profile_inplace() indexes a 2D template by data channel
+        (template[ichan]), so a template with a different number of channels
+        than the data is mis-aligned -- or, if it has more channels, silently
+        uses only the first 'data_nchan' of them. A 1D template is broadcast to
+        every channel and is always fine, so no warning is issued for it.
+
+        Inputs:
+            template: The template array (1D (nbin,) or 2D (nchan x nbin)).
+            data_nchan: The number of channels in the data being cleaned.
+
+        Output:
+            None - Emits a CoastGuardWarning on a 2D channel-count mismatch.
+    """
+    if np.ndim(template) <= 1:
+        return
+    template_nchan = np.shape(template)[0]
+    if template_nchan != data_nchan:
+        warnings.warn(
+            "2D template has %d channels but the data has %d. A 2D template is "
+            "indexed by data channel, so it must have the same number of "
+            "channels as the data; frequency-dependent subtraction may be "
+            "mis-aligned." % (template_nchan, data_nchan),
+            errors.CoastGuardWarning)
+
+
 def remove_profile1d(prof, isub, ichan, template, phs, return_params=False):
     """Fit and subtract a (rotated, scaled) template from a single profile.
 

coast_guard/cleaners/surgical.py

@@ -106,6 +106,32 @@ def _set_config_params(self):
                                aliases=['iterations', 'iter'],
                                nullable=True,
                                help="Number of iterations to run the surgical cleaner [default = 1]")
+        self.configs.add_param('piecewise_scale', config_types.BoolVal,
+                               aliases=['piecewise_scale'],
+                               nullable=True,
+                               help="If True, scale the frequency-direction "
+                                    "statistics to sigma using a per-segment "
+                                    "(piecewise) median/MAD that matches the "
+                                    "piecewise detrend, instead of a single "
+                                    "global MAD across the whole band. This "
+                                    "avoids over-flagging channels in high-Tsys "
+                                    "sub-bands on wideband receivers. "
+                                    "[default = False]")
+        self.configs.add_param('subint_mad_numpieces', config_types.IntVal,
+                               aliases=['subint_mad_numpieces'],
+                               nullable=True,
+                               help="Number of frequency segments to use for "
+                                    "the piecewise MAD scaling (only used when "
+                                    "piecewise_scale is True). If None, the "
+                                    "finest subint_numpieces value is used. "
+                                    "[default = None]")
+        # Establish safe (feature-off) defaults for the optional piecewise
+        # params *before* parsing surgical_default_params. A receiver/custom
+        # config that overrides surgical_default_params need not list these
+        # keys; if it omits them they stay at these defaults (rather than
+        # raising KeyError when read in _clean), and if it sets them the
+        # override below wins.
+        self.parse_config_string('piecewise_scale=False,subint_mad_numpieces=None')
         self.parse_config_string(config.cfg.surgical_default_params)
 
 
@@ -183,6 +209,10 @@ def _clean(self, ar):
                                                   tuple(range(template_data.ndim - 1))).squeeze()
                 template = template_data if template_data.ndim > 1 else template_phs
 
+                # A 2D template is indexed by data channel during subtraction,
+                # so warn if its channel count does not match the data.
+                clean_utils.check_template_nchan(template, patient.get_nchan())
+
 
             print('Estimating template and profile phase offset')
             if self.configs.template is None:
@@ -284,6 +314,8 @@ def _clean(self, ar):
                                         subint_order=self.configs.subint_order, \
                                         subint_breakpoints=self.configs.subint_breakpoints, \
                                         subint_numpieces=self.configs.subint_numpieces, \
+                                        piecewise_scale=self.configs.piecewise_scale, \
+                                        subint_mad_numpieces=self.configs.subint_mad_numpieces, \
                                         aggressive=self.configs.aggressive \
                                         )
             if plot:

coast_guard/configurations/receivers/UWL_3K_Murriyang.cfg

@@ -1,3 +1,7 @@
 #crazy intense RFI mitigation for UWL fold-mode...
-surgical_default_params = 'template=None,chan_breakpoints=None,chan_numpieces=1,chan_order=1,chanthresh=5,subint_breakpoints=None,subint_numpieces=8;16,subint_order=2;1,subintthresh=5,plot=None'
+#piecewise_scale=True: the UWL band spans 704-4032 MHz with a large Tsys
+#gradient, so scale each sub-int's stats to sigma per frequency segment
+#(subint_mad_numpieces=16, matching the finest subint_numpieces) rather than
+#with one global MAD, to avoid over-flagging high-Tsys sub-bands.
+surgical_default_params = 'template=None,chan_breakpoints=None,chan_numpieces=1,chan_order=1,chanthresh=5,subint_breakpoints=None,subint_numpieces=8;16,subint_order=2;1,subintthresh=5,plot=None,piecewise_scale=True,subint_mad_numpieces=16'
 rcvrstd_default_params = 'badchans=None,badfreqs=962:1150;1857:1863;1725:1735;1080:1090;1022:1026;1918:1922;3070:3074,badsubints=None,response=704:4032,trimbw=0,trimfrac=0,trimnum=0'