SSB-44: refactored code to improve understanding

CAL/CalibrationProcessing/lib/processing/averaging.py

@@ -65,6 +65,17 @@ def average_values_by_sample(data_array, window_size, field_name=None):
     return result
 
 
+def is_datatable_completely_flagged(data_table: numpy.ndarray):
+    """
+    Checks if the datatable is entirely flagged
+    :param data_table: input data
+    :return: if the data is completely flagged
+    :rtype: bool
+    """
+    flags = numpy.array(data_table)['flag']
+    return numpy.product(flags)
+
+
 def _average_datatable_with_function(data_table_in, parameter, expected_array_size, function):
     field_names_mean = []
     formats_mean = []
@@ -100,8 +111,7 @@ def _average_datatable_with_function(data_table_in, parameter, expected_array_si
     fields_to_average = list(data_table_in.dtype.names)
     fields_to_average.remove('flag')
 
-    flags = numpy.array(data_table_in)['flag']
-    is_flagged = numpy.product(flags) == True
+    is_flagged = is_datatable_completely_flagged(data_table_in)
 
     for field_name in fields_to_average:
         datatable_values = numpy.array(data_table_in)
@@ -166,7 +176,6 @@ def average_datatable(data_table_in):
     return result
 
 
-
 def average_datatable_by_time_interval(data_table_in, time_interval):
     """
     Average the datatable with a given sample time interval window
@@ -344,7 +353,6 @@ def average_data(input_data_table):
     return out_data
 
 
-
 def _compute_size_of_new_array(array_size, window_size):
     """
     Round up the array size given a window size.
@@ -364,7 +372,7 @@ def _compute_size_of_new_array(array_size, window_size):
     return new_array_size
 
 
-def _weigthed_average_data_per_polarization(input_dataset, output_dataset,
+def _weighted_average_data_per_polarization(input_dataset, output_dataset,
                                             real_weights, imag_weights, polarization):
     mask = input_dataset['flag'] == False
     try:
@@ -379,6 +387,60 @@ def _weigthed_average_data_per_polarization(input_dataset, output_dataset,
         output_dataset['flag'] = True
 
 
+def __prepare_input_to_weight_dataset(stations, frequency, beam,
+                                      input_data_table):
+    """
+    Prepare the input data for the weighted average
+    :param stations: list of stations to iterate on
+    :type stations: (str)
+    :param frequency: frequency ID
+    :type frequency: str
+    :param beam: beam ID
+    :type beam: str
+    :param input_data_table: data to iterate over
+    :return: the average values,
+     the standard deviation of the real part,
+      the standard deviation of the imaginary part
+    :rtype: list(numpy.ndarray, numpy.ndarray, numpy.ndarray)
+    """
+
+    n_stations = len(stations)
+
+    average_per_beam_station = numpy.zeros(n_stations, dtype=HDS_data_sample_type)
+    std_real_per_beam_station = numpy.zeros(n_stations, dtype=HDS_data_sample_type)
+    std_imag_per_beam_station = numpy.zeros(n_stations, dtype=HDS_data_sample_type)
+
+    for i, station in enumerate(stations):
+        data_per_station_frequency_beam = \
+            input_data_table[station][str(frequency)][beam]
+        average_per_beam_station['XX'][i] = data_per_station_frequency_beam['mean']['XX']
+        average_per_beam_station['XY'][i] = data_per_station_frequency_beam['mean']['XY']
+        average_per_beam_station['YX'][i] = data_per_station_frequency_beam['mean']['YX']
+        average_per_beam_station['YY'][i] = data_per_station_frequency_beam['mean']['YY']
+
+        average_per_beam_station['l'][i] = data_per_station_frequency_beam['mean']['l']
+        average_per_beam_station['m'][i] = data_per_station_frequency_beam['mean']['m']
+
+        std_real_per_beam_station['XX'][i] = \
+            data_per_station_frequency_beam['std']['XX_real']
+        std_real_per_beam_station['XY'][i] = \
+            data_per_station_frequency_beam['std']['XY_real']
+        std_real_per_beam_station['YX'][i] = \
+            data_per_station_frequency_beam['std']['YX_real']
+        std_real_per_beam_station['YY'][i] = \
+            data_per_station_frequency_beam['std']['YY_real']
+
+        std_imag_per_beam_station['XX'][i] = \
+            data_per_station_frequency_beam['std']['XX_imag']
+        std_imag_per_beam_station['XY'][i] = \
+            data_per_station_frequency_beam['std']['XY_imag']
+        std_imag_per_beam_station['YX'][i] = \
+            data_per_station_frequency_beam['std']['YX_imag']
+        std_imag_per_beam_station['YY'][i] = \
+            data_per_station_frequency_beam['std']['YY_imag']
+    return average_per_beam_station, std_real_per_beam_station, std_imag_per_beam_station
+
+
 def weighted_average_dataset_per_station(dataset, input_data_table):
     """
     Perform a weighted average of the dataset using the std deviations in the input data
@@ -389,7 +451,7 @@ def weighted_average_dataset_per_station(dataset, input_data_table):
     :return:
     """
     stations = dataset.reference_stations
-    n_stations = len(stations)
+
     frequencies = dataset.frequencies
     beams = dataset.beamlets
 
@@ -403,41 +465,11 @@ def weighted_average_dataset_per_station(dataset, input_data_table):
             beam_str = str(beam)
             result_per_beam = numpy.zeros(1, dtype=HDS_data_sample_type)
 
-            average_per_beam_station = numpy.zeros(n_stations, dtype=HDS_data_sample_type)
-            std_real_per_beam_station = numpy.zeros(n_stations, dtype=HDS_data_sample_type)
-            std_imag_per_beam_station = numpy.zeros(n_stations, dtype=HDS_data_sample_type)
-
-            for i, station in enumerate(stations):
-                data_per_station_frequency_beam = \
-                    input_data_table[station][str(frequency)][beam_str]
-                average_per_beam_station['XX'][i] = data_per_station_frequency_beam['mean']['XX']
-                average_per_beam_station['XY'][i] = data_per_station_frequency_beam['mean']['XY']
-                average_per_beam_station['YX'][i] = data_per_station_frequency_beam['mean']['YX']
-                average_per_beam_station['YY'][i] = data_per_station_frequency_beam['mean']['YY']
-
-                average_per_beam_station['l'][i] = data_per_station_frequency_beam['mean']['l']
-                average_per_beam_station['m'][i] = data_per_station_frequency_beam['mean']['m']
-
-                std_real_per_beam_station['XX'][i] = \
-                    data_per_station_frequency_beam['std']['XX_real']
-                std_real_per_beam_station['XY'][i] = \
-                    data_per_station_frequency_beam['std']['XY_real']
-                std_real_per_beam_station['YX'][i] = \
-                    data_per_station_frequency_beam['std']['YX_real']
-                std_real_per_beam_station['YY'][i] = \
-                    data_per_station_frequency_beam['std']['YY_real']
-
-                std_imag_per_beam_station['XX'][i] = \
-                    data_per_station_frequency_beam['std']['XX_imag']
-                std_imag_per_beam_station['XY'][i] = \
-                    data_per_station_frequency_beam['std']['XY_imag']
-                std_imag_per_beam_station['YX'][i] = \
-                    data_per_station_frequency_beam['std']['YX_imag']
-                std_imag_per_beam_station['YY'][i] = \
-                    data_per_station_frequency_beam['std']['YY_imag']
+            average_per_beam_station, std_real_per_beam_station, std_imag_per_beam_station = \
+                __prepare_input_to_weight_dataset(stations, frequency, beam_str, input_data_table)
 
             for polarization in ['XX', 'XY', 'YX', 'YY']:
-                _weigthed_average_data_per_polarization(average_per_beam_station,
+                _weighted_average_data_per_polarization(average_per_beam_station,
                                                         result_per_beam,
                                                         std_real_per_beam_station,
                                                         std_imag_per_beam_station,