Add frequency flatten task

nenupy/__init__.py

@@ -5,7 +5,7 @@
 __copyright__ = "Copyright 2023, nenupy"
 __credits__ = ["Alan Loh"]
 __license__ = "MIT"
-__version__ = "2.6.8"
+__version__ = "2.6.9"
 __maintainer__ = "Alan Loh"
 __email__ = "[email protected]"
 

nenupy/io/tf.py

@@ -79,6 +79,7 @@ class TFTask:
     .. autosummary::
 
         ~nenupy.io.tf.TFTask.correct_bandpass
+        ~nenupy.io.tf.TFTask.flatten_subband
         ~nenupy.io.tf.TFTask.remove_channels
         ~nenupy.io.tf.TFTask.correct_polarization
         ~nenupy.io.tf.TFTask.correct_faraday_rotation
@@ -114,16 +115,46 @@ def _func_call(self) -> Callable:
 
     @classmethod
     def correct_bandpass(cls):
-        """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.correct_bandpass` to correct the polyphase-filter bandpass reponse."""
+        """ :class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.correct_bandpass` to correct the polyphase-filter bandpass reponse.
 
+            .. figure:: ./_images/io_images/tf_bandpass_correction.png
+                :width: 650
+                :align: center
+
+        """
         def wrapper_task(data, channels):
             return utils.correct_bandpass(data=data, n_channels=channels)
 
         return cls("Correct bandpass", wrapper_task, ["channels"], repeatable=False)
 
+    @classmethod
+    def flatten_subband(cls):
+        """_summary_
+
+            .. figure:: ./_images/io_images/tf_sb_flatten.png
+                :width: 650
+                :align: center
+
+            Warning
+            -------
+            This is  a warning
+
+
+        """
+        def wrapper_task(data, channels):
+            return utils.flatten_subband(data=data, channels=channels)
+
+        return cls("Flatten subband", wrapper_task, ["channels"], repeatable=False)
+
     @classmethod
     def remove_channels(cls):
-        """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.remove_channels_per_subband` to set a list of sub-band channels to `NaN` values."""
+        """:class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.remove_channels_per_subband` to set a list of sub-band channels to `NaN` values.
+
+            .. figure:: ./_images/io_images/tf_remove_channels.png
+                :width: 650
+                :align: center
+
+        """
 
         def wrapper_task(data, channels, remove_channels):
             if (remove_channels is None) or (len(remove_channels) == 0):
@@ -242,6 +273,13 @@ def wrapper_task(
 
     @classmethod
     def time_rebin(cls):
+        """_summary_
+
+            .. figure:: ./_images/io_images/tf_time_rebin.png
+                :width: 650
+                :align: center
+    
+        """
         def rebin_time(time_unix, data, dt, rebin_dt):
             if rebin_dt is None:
                 return time_unix, data
@@ -258,6 +296,13 @@ def rebin_time(time_unix, data, dt, rebin_dt):
 
     @classmethod
     def frequency_rebin(cls):
+        """_summary_
+
+            .. figure:: ./_images/io_images/tf_frequency_rebin.png
+                :width: 650
+                :align: center
+
+        """
         def rebin_freq(frequency_hz, data, df, rebin_df):
             if rebin_df is None:
                 return frequency_hz, data
@@ -274,6 +319,14 @@ def rebin_freq(frequency_hz, data, df, rebin_df):
 
     @classmethod
     def get_stokes(cls):
+        """_summary_
+
+            .. figure:: ./_images/io_images/tf_stokes.png
+                :width: 650
+                :align: center
+
+        """
+
         def compute_stokes(data, stokes):
             if (stokes is None) or (stokes == "") or (stokes == []):
                 return data
@@ -438,7 +491,6 @@ def set_default(self) -> None:
         self.tasks = [
             TFTask.correct_bandpass(),
             TFTask.remove_channels(),
-            # TFTask.correct_faraday_rotation(),
             TFTask.time_rebin(),
             TFTask.frequency_rebin(),
             TFTask.get_stokes(),

nenupy/io/tf_utils.py

@@ -36,6 +36,7 @@
     "crop_subband_edges",
     "de_disperse_array",
     "de_faraday_data",
+    "flatten_subband",
     "get_bandpass",
     "polarization_angle",
     "rebin_along_dimension",
@@ -46,7 +47,6 @@
     "TFPipelineParameters"
 ]
 
-
 # ============================================================= #
 # ---------------- apply_dreambeam_corrections ---------------- #
 def apply_dreambeam_corrections(
@@ -479,6 +479,50 @@ def get_bandpass(n_channels: int) -> np.ndarray:
     return g**2.0
 
 
+# ============================================================= #
+# ---------------------- flatten_subband ---------------------- #
+def flatten_subband(data: np.ndarray, channels: int) -> np.ndarray:
+
+    # Check that data has not been altered, i.e. dimension 1 should be a multiple of channels
+    if data.shape[1] % channels != 0:
+        raise ValueError(
+            f"data's frequency dimension (of size {data.shape[1]}) is "
+            f"not a multiple of channels={channels}. data's second "
+            "dimension should be of size number_of_subbands*number_of_channels."
+        )
+    n_subbands = int(data.shape[1]/channels)
+    pol_dims = data.ndim - 2 # the first two dimensions should be time and frequency
+
+    # Compute the median spectral profile (along the time axis)
+    median_frequency_profile = np.nanmedian(data, axis=0)
+    subband_shape = (n_subbands, channels) + data.shape[2:]
+    # Reshape to have the spectral profile as (subbands, channels, (polarizations...))
+    median_subband_profile = median_frequency_profile.reshape(subband_shape)
+
+    # # Select two data points (away from subband edges) that will be
+    # # used to compute the affine function that approximates each subband.
+    # ind1, ind2 = int(np.round(channels*1/3)), int(np.round(channels*2/3))
+    # # Get the y-values corresponding to these two indices, each y is of shape (subbands, (polarizations...))
+    # y1, y2 = median_subband_profile[:, ind1, ...], median_subband_profile[:, ind2, ...]
+
+    # Split the subband in two and compute the 2 medians that will be
+    # used to compute the affine function that approximates each subband.
+    ind1, ind2 = int(np.floor(channels/2))/2, channels - int(np.floor(channels/2))/2
+    y1 = np.nanmedian(median_subband_profile[:, :int(np.floor(channels/2)), ...], axis=1)
+    y2 = np.nanmedian(median_subband_profile[:, int(np.ceil(channels/2)):, ...], axis=1)
+
+    # Compute the linear approximations of each subbands, linear_subbands's shape is (channels, subbands, (polarizations...))
+    x_values = np.arange(channels)[(...,) + (np.newaxis,) * (pol_dims + 1)] # +1 --> subbands
+    linear_subbands = (x_values - ind1) * (y2 - y1) / (ind2 - ind1) + y1 # linear equation
+
+    # Compute the subband mean value and the normalised linear subbands
+    subband_mean_values = np.nanmedian(linear_subbands, axis=0) # shape (subbands, (polarizations))
+    normalised_linear_subbands = np.swapaxes(linear_subbands / subband_mean_values[None, ...], 0, 1).reshape(data.shape[1:])
+
+    # Correct the data by the normalised linear subbands to flatten them
+    return data / normalised_linear_subbands[None, ...]
+
+
 # ============================================================= #
 # -------------------- polarization_angle --------------------- #
 def polarization_angle(stokes_u: np.ndarray, stokes_q: np.ndarray) -> np.ndarray: