Method to remove subband edge channels

nenupy/__init__.py

@@ -5,7 +5,7 @@
 __copyright__ = 'Copyright 2021, nenupy'
 __credits__ = ['Alan Loh']
 __license__ = 'MIT'
-__version__ = '2.2.9'
+__version__ = '2.3.0'
 __maintainer__ = 'Alan Loh'
 __email__ = '[email protected]'
 

nenupy/undysputed/dynspec.py

@@ -531,7 +531,9 @@ def fmax(self):
     # --------------------------------------------------------- #
     # ------------------------ Methods ------------------------ #
     def get_stokes(self, stokes='I'):
-        """
+        """ fft0[..., :] = [XX, YY]
+            fft1[..., :] = [Re(XY*), Im(XY*)]
+            XX=I+Q YY=I-Q XY=U+iV YX=U-iV
         """
         stokes_data = {
             'i': np.sum(self.fft0, axis=4),
@@ -671,6 +673,7 @@ def _to2d(self, data):
                 int(self.fftlen/2)
             )
         )
+        # data = data[:, :, ::-1, :].reshape((n_times, n_freqs-nfb))
         data = data[:, :, ::-1, :].reshape((n_times, n_freqs))
         return data
 
@@ -733,7 +736,21 @@ def _to2d(self, data):
     #         data *= broadband[np.newaxis, :, np.newaxis]
     #         return data.reshape((ntimes, nfreqs)) / spectrum
     #     elif method.lower() == 'none':
-    #         return data 
+    #         return data
+
+    def _polarization_correction(fft0, fft1, jones):
+        """
+            fft0[..., :] = [XX, YY]
+            fft1[..., :] = [Re(XY*), Im(XY*)]
+            XX=I+Q YY=I-Q XY=U+iV YX=U-iV
+        """
+        # fft0.shape = (time_blocks, subbands, time_per_blocks, channels, 2)
+        # jones.shape = (time, frequency, 2, 2)
+        xx = fft0[..., 0]
+        yy = fft1[..., 1]
+        xy = fft1[..., 0] + 1j*fft1[..., 1]
+        yx = fft1[..., 0] - 1j*fft1[..., 1]
+
 # ============================================================= #
 
 
@@ -766,6 +783,7 @@ def __init__(self, lanefiles=[]):
         self.rebin_df = None
         self.jump_correction = False
         self.bp_correction = 'none'
+        self.edge_channels_to_remove = 0
         self.freq_flat = False
         self.clean_rfi = False
 
@@ -1077,6 +1095,44 @@ def bp_correction(self, b):
         return    
 
 
+    @property
+    def edge_channels_to_remove(self) -> int:
+        """ Number of channels to remove at each subband edge.
+            Generally, the 2 extreme channels are not taken into account.
+
+            :type: `int`
+        """
+        return self._edge_channels_to_remove
+    @edge_channels_to_remove.setter
+    def edge_channels_to_remove(self, n: int) -> None:
+        default_val = 0
+        if not isinstance(n, int):
+            log.warning(
+                "The number of channels to remove at each subband "
+                f"edge needs to be an integer. Setting default value {default_val}."
+            )
+            n = default_val
+        elif n*2 > self.fftlen:
+            log.warning(
+                f"Each subband is divided in {self.fftlen} channels "
+                f"in this dataset. Removing {n} channels at both "
+                "edge would not leave any data left. Setting default "
+                f"value {default_val}."
+            )
+            n = default_val
+        elif n < 0:
+            log.warning(
+                "The number of channels to remove should be positive. "
+                f"Setting default value {default_val}."
+            )
+            n = default_val
+        if n > 0:
+            log.info(
+                f"{n} channels will be removed at subband edges."
+            )
+        self._edge_channels_to_remove = n
+
+
     @property
     def jump_correction(self):
         """ Correct or not the known 6-minute gain jumps
@@ -1302,6 +1358,11 @@ def get(self, stokes='I'):
                 fftlen=li.fftlen
             )
             selfreqs = li._freqs[beam_start:beam_stop][fmin_idx:fmax_idx].compute()*u.Hz
+            # Remove subband edge channels
+            data, selfreqs = self._remove_edge_channels(data=data, frequencies=selfreqs)
+            # mask = np.ones(a.size, dtype=bool)
+            # mask::4] = 0
+            # data = data[:, ]
             # RFI mitigation
             data = self._clean(
                 data=data
@@ -1348,7 +1409,8 @@ def get(self, stokes='I'):
 
     # --------------------------------------------------------- #
     # ----------------------- Internal ------------------------ #
-    def _bandpass(self, fftlen):
+    @staticmethod
+    def _bandpass(fftlen):
         """ Computes the bandpass correction for a beamlet.
         """
         kaiser_file = join(
@@ -1424,6 +1486,30 @@ def _bp_correct(self, data, fftlen):
             return data
 
 
+    def _remove_edge_channels(self, data, frequencies):
+        """ Remove the N channels at each subband edge.
+        """
+        if self.edge_channels_to_remove > 0:
+            ntimes, nfreqs = data.shape
+            nsubbands = int(nfreqs / self.fftlen)
+            data = data.reshape(
+                (ntimes, nsubbands, self.fftlen)
+            )[:, :, self.edge_channels_to_remove:self.fftlen - self.edge_channels_to_remove]
+            new_nfreqs = nsubbands*(self.fftlen - self.edge_channels_to_remove*2)
+            data = data.reshape((ntimes, new_nfreqs))
+            frequencies = frequencies.reshape(
+                (nsubbands, self.fftlen)
+            )[:, :, self.edge_channels_to_remove:self.fftlen - self.edge_channels_to_remove]
+            frequencies = frequencies.reshape((new_nfreqs,))
+            log.info(
+                f"{self.edge_channels_to_remove} channels have been "
+                "removed at the subband edges. Each subband now "
+                f"contains {self.fftlen - 2*self.edge_channels_to_remove}"
+                f"/{self.fftlen} channels."
+            )
+        return data, frequencies
+
+
     def _freqFlattening(self, data):
         """ Flatten the sub-band response
         """