Rebin method added

nenupy/__init__.py

@@ -5,7 +5,7 @@
 __copyright__ = "Copyright 2023, nenupy"
 __credits__ = ["Alan Loh"]
 __license__ = "MIT"
-__version__ = "2.4.2"
+__version__ = "2.4.3"
 __maintainer__ = "Alan Loh"
 __email__ = "alan.loh@obspm.fr"
 

nenupy/io/tf.py

@@ -163,33 +163,58 @@ def __init__(self, filename: str):
 
         log.info("\tSetting up default configuration:")
         self.configuration = _ProcessingConfiguration(
-            time_min=Time(self._block_start_unix[0], format="unix", precision=7),
-            time_max=Time(self._block_start_unix[-1] + self._n_time_per_block * self.dt.to_value(u.s), format="unix", precision=7),
-            frequency_min=self._subband_start_hz[0] * u.Hz,
-            frequency_max=self._subband_start_hz[-1] * u.Hz + SUBBAND_WIDTH,
+            time_min=self.time_min,
+            time_max=self.time_max,
+            frequency_min=self.frequency_min,
+            frequency_max=self.frequency_max,
             available_beams=np.array(list(self.beam_indices_dict.keys())).astype(int)
         )
 
     # --------------------------------------------------------- #
     # --------------------- Getter/Setter --------------------- #
+    @property
+    def time_min(self) -> Time:
+        return Time(self._block_start_unix[0], format="unix",precision=7)
+
+    @property
+    def time_max(self) -> Time:
+        return Time(self._block_start_unix[-1] + self._n_time_per_block * self.dt.to_value(u.s), format="unix", precision=7)
+
     @property
     def frequency_min(self) -> u.Quantity:
         freq_mins = []
         for _, boundaries in self.beam_indices_dict.items():
-            freq_mins.append(self.frequency_hz[boundaries[0]])
+            sb_index = int(boundaries[0]/self.n_channels)
+            freq_mins.append(self._subband_start_hz[sb_index])
         return np.min(freq_mins) * u.Hz
 
     @property
     def frequency_max(self) -> u.Quantity:
         freq_maxs = []
         for _, boundaries in self.beam_indices_dict.items():
-            freq_maxs.append(self.frequency_hz[boundaries[1]])
-        return np.max(freq_maxs) * u.Hz
+            sb_index = int((boundaries[1] + 1)/self.n_channels - 1)
+            freq_maxs.append(self._subband_start_hz[sb_index])
+        return np.max(freq_maxs) * u.Hz + SUBBAND_WIDTH
 
     # --------------------------------------------------------- #
     # ------------------------ Methods ------------------------ #
+    def info(self) -> None:
+        """ Display informations about the file. """
+        message = "\n".join([
+            f"{self.filename}",
+            f"time_min: {self.time_min.isot}",
+            f"time_max: {self.time_max.isot}",
+            f"dt: {self.dt.to(u.ms)}",
+            f"frequency_min: {self.frequency_min.to(u.MHz)}",
+            f"frequency_max: {self.frequency_max.to(u.MHz)}",
+            f"df: {self.df.to(u.kHz)}",
+            f"Available beam indices: {list(self.beam_indices_dict.keys())}"
+        ])
+        print(message) 
+
     def get(self, stokes: Union[str, List[str]] = "I"):
 
+        # Select the data in time and frequency (the beam selection is implicit on the frequency idexing)
         frequency_hz, time_unix, data = self._select_data()
 
         # Stop the pipeline if the data is empty
@@ -212,13 +237,28 @@ def get(self, stokes: Union[str, List[str]] = "I"):
                 higher_edge_channels=edge_chans[1] if isinstance(edge_chans, tuple) else edge_chans,
             )
 
-        data = utils.compute_stokes_parameters(data_array=data, stokes=stokes)
+        # Rebin the data
+        if self.configuration.rebin_dt is not None:
+            log.info("\tRebinning in time...")
+            time_unix, data = utils.rebin_along_dimension(
+                data=data,
+                axis_array=time_unix,
+                axis=0,
+                dx=self.dt.to_value(u.s),
+                new_dx=self.configuration.rebin_dt.to_value(u.s)
+            )
+        if self.configuration.rebin_df is not None:
+            log.info("\tRebinning in frequency...")
+            frequency_hz, data = utils.rebin_along_dimension(
+                data=data,
+                axis_array=frequency_hz,
+                axis=1,
+                dx=self.df.to_value(u.Hz),
+                new_dx=self.configuration.rebin_df.to_value(u.Hz)
+            )
 
-        frequency_hz, time_unix, data = self._time_frequency_rebin(
-            data=data,
-            times=time_unix,
-            freqs=frequency_hz,
-        )
+        # Compute the selected Stokes parameters
+        data = utils.compute_stokes_parameters(data_array=data, stokes=stokes)
 
         log.info("Computing the data...")
         with ProgressBar():
@@ -319,7 +359,7 @@ def _assemble_to_tf(self, data: np.ndarray, mask:  np.ndarray) -> da.Array:
 
         return data
 
-    def _select_data(self) -> Tuple[np.ndarray, nd.array, da.Array]:
+    def _select_data(self) -> Tuple[np.ndarray, np.ndarray, da.Array]:
         """ """
 
         log.info("\tComputing the time selection...")
@@ -481,5 +521,5 @@ def _time_frequency_rebin(self, data: da.Array, times: da.Array, freqs: da.Array
             freqs = np.nanmean(freqs, axis=1)
             log.info("Data are frequency-averaged.")
 
-        return freqs, times, data 
+        return freqs, times, data
 

nenupy/io/tf_utils.py

@@ -7,7 +7,7 @@
 import dask.array as da
 import astropy.units as u
 from astropy.time import Time
-from typing import Union, List
+from typing import Union, List, Tuple
 import logging
 log = logging.getLogger(__name__)
 
@@ -17,6 +17,7 @@
     "compute_spectra_frequencies",
     "compute_spectra_time",
     "compute_stokes_parameters",
+    "rebin_along_dimension",
     "sort_beam_edges",
     "spectra_data_to_matrix"
 ]
@@ -144,6 +145,38 @@ def get_bandpass(n_channels: int) -> np.ndarray:
 
     return g**2.0
 
+def rebin_along_dimension(data: np.ndarray, axis_array: np.ndarray, axis: int, dx: float, new_dx: float) -> Tuple[np.ndarray, np.ndarray]:
+    """ """
+
+    # Basic checks to make sure that dimensions are OK
+    if axis_array.ndim != 1:
+        raise IndexError("axis_array should be 1D.")
+    elif data.shape[axis] != axis_array.size:
+        raise IndexError(f"Axis selected ({axis}) dimension {data.shape[axis]} does not match axis_array's shape {axis_array.shape}.")
+    elif dx > new_dx:
+        raise ValueError("Rebin expect a `new_dx` value larger than original `dx`.")
+
+    initial_size = axis_array.size
+    bin_size = int(np.floor(new_dx / dx))
+    final_size = int(np.floor(initial_size / bin_size))
+    leftovers = initial_size % final_size
+
+    d_shape = data.shape
+
+    # Reshape the data and the axis to ease the averaging
+    data = data[tuple([slice(None) if i != axis else slice(None, initial_size - leftovers) for i in range(len(d_shape))])].reshape(
+        d_shape[:axis] + (final_size, int((initial_size - leftovers) / final_size)) + d_shape[axis + 1:]
+    )
+    axis_array = axis_array[: initial_size - leftovers].reshape(
+        (final_size, int((initial_size - leftovers) / final_size))
+    )
+
+    # Average the data and the axis along the right dimension
+    data = np.nanmean(data, axis=axis + 1)
+    axis_array = np.nanmean(axis_array, axis=1)
+
+    return axis_array, data
+
 # ============================================================= #
 # ---------------------- sort_beam_edges ---------------------- #
 def sort_beam_edges(beam_array: np.ndarray, n_channels: int) -> dict: