SSB-47: Refactor find_central_beamlets

CAL/CalibrationCommon/lib/datacontainers/holography_dataset.py

@@ -173,7 +173,8 @@ class HolographyDataset():
                     return False
         return equality
 
-    def find_central_beamlets(_, source, ra_dec, frequencies, beamlets):
+
+    def find_central_beamlets(self, source, ra_dec, frequencies, beamlets):
         '''
         This function finds the central beamlet of a target station for every
         frequency.
@@ -187,68 +188,34 @@ class HolographyDataset():
         @param beamlets:  A list of beamlet numbers. 
         '''
         logger.debug("Finding the central beamlets...")
-        source_ra = source["RA"]
-        source_dec = source["DEC"]
-        # Store each pseudo distance in a dict of dicts: pd[frequency][beamlet]
-        pseudo_distance = dict()
-        # Iterate over all _frequencies...
-        for frequency in frequencies:
-            frequency_string = str(frequency)
-            pseudo_distance[frequency_string] = dict()
-
-            # and iterate over the beamlets.
-            for beamlet in beamlets:
-                beamlet_string = str(beamlet)
-                if beamlet_string not in ra_dec[frequency_string]:
-                    continue
-                ra = ra_dec[frequency_string][beamlet_string]['RA']
-                dec = ra_dec[frequency_string][beamlet_string]['DEC']
-
-                # calculate a pseudo distance that is an indicator if this
-                # beamlet is closer to the source than the ones before.
-                # Note that I am too lazy to calculate the spherical distance.
-                # But keep in mind that this pseudo distance like its cousin
-                # the geometrical distance are not really cutting it since RA
-                # and DEC are coordinates of a spherical coordinate system.
-                # But here the pseudo distance is good enough. 
-                pseudo_distance[frequency_string][beamlet_string] = abs(ra - source_ra) + abs(
-                    dec - source_dec)
-
+        # calculate a pseudo distance that is an indicator if this
+        # beamlet is closer to the source than the ones before.
+        # Note that I am too lazy to calculate the spherical distance.
+        # But keep in mind that this pseudo distance like its cousin
+        # the geometrical distance are not really cutting it since RA
+        # and DEC are coordinates of a spherical coordinate system.
+        # But here the pseudo distance is good enough.
+        pseudo_distance = self._compute_pseudo_distance_per_frequency_beam(ra_dec, source["RA"], source["DEC"])
         # OK.  Done with all the iteration business.  Now check if across all
         # _frequencies the same beamlet is the central one.  It is allowed that
         # the central beamlets are not always the same but it would be better
         # and let the astronomers sleep better if it would always be the same.
         # And also check if there is actually only one central beamlet per
         # frequency.
-        central_beamlet = dict()
-        for (frequency_string, beamlets) in pseudo_distance.items():
-            values = list(beamlets.values())
-            keys = list(beamlets.keys())
-            minimal_distance = min(values)
-            # If I find a second beamlet that is 'closest' to the source for the
-            # same frequency then that is a problem.  The outset was that there
-            # is only a single central beamlet and not a bunch of 'em.
-            if values.count(minimal_distance) == 1:
-                # All is good.  Only one central beamlet.
-                central_beamlet[frequency_string] = keys[values.index(minimal_distance)]
-            else:
-                text = "Found %d beamlets that have the same distance from the source position. Therefore a unique central beamlet does not exist." % (
-                    values.count(minimal_distance))
-                logger.error(text)
-                raise ValueError(text)
+        central_beamlets = self._find_minimal_distance_per_frequency(pseudo_distance)
 
         # Now check if the central beamlet is the same across all _frequencies.
         # I do that by creating a set from the central beamlets of all stations.
         # If there is only one central beamlet for all _frequencies, the size of
         # the set will be 1.
-        central_beamlet_set = set(central_beamlet.values())
+        central_beamlet_set = set(central_beamlets.values())
         if len(central_beamlet_set) == 1:
             logger.debug(
                 "All is good, unicorns everywhere, there is only one central beamlet \"%s\" for all _frequencies.",
                 central_beamlet_set)
         else:
-            logger.warning("Multiple central beamlets have been identified: ", central_beamlet)
-        return central_beamlet
+            logger.warning("Multiple central beamlets have been identified: ", central_beamlets)
+        return central_beamlets
 
     def load_from_beam_specification_and_ms(self, station_name, list_of_hbs_ms_tuples):
         """
@@ -273,6 +240,50 @@ class HolographyDataset():
             logger.exception("Error creating dataset for station \"%s\": %s", station_name, e)
             raise e
 
+    def _compute_pseudo_distance_per_frequency_beam(self, ra_dec, source_ra, source_dec):
+        """
+        Compute the pseudo distance defined as abs(ra-source_ra) + abs(dec-source_dec) for all the ra_dec in
+        ra_dec a function of the frequencies and beamlet
+        :param ra_dec: ra and dec
+        :type ra_dec: Dict[str, Dict[str, numpy.ndarray]]
+        :param source_ra:
+        :param source_dec:
+        :return:
+        """
+        # Store each pseudo distance in a dict of dicts: pd[frequency][beamlet]
+        pseudo_distance = dict()
+        # Iterate over all _frequencies...
+        for frequency_string in ra_dec.keys():
+            pseudo_distance[frequency_string] = dict()
+            # and iterate over the beamlets.
+            for beamlet_string in ra_dec[frequency_string].keys():
+                ra = ra_dec[frequency_string][beamlet_string]['RA']
+                dec = ra_dec[frequency_string][beamlet_string]['DEC']
+
+                pseudo_distance[frequency_string][beamlet_string] = abs(ra - source_ra) + abs(
+                    dec - source_dec)
+        return pseudo_distance
+
+    def _find_minimal_distance_per_frequency(_, distance_per_frequency_beam):
+        minimal_distance_per_frequency = dict()
+        for (frequency_string, beamlets) in distance_per_frequency_beam.items():
+            values = list(beamlets.values())
+            keys = list(beamlets.keys())
+            minimal_distance = min(values)
+            # If I find a second beamlet that is 'closest' to the source for the
+            # same frequency then that is a problem.  The outset was that there
+            # is only a single central beamlet and not a bunch of 'em.
+            if values.count(minimal_distance) == 1:
+                # All is good.  Only one central beamlet.
+                minimal_distance_per_frequency[frequency_string] = keys[values.index(minimal_distance)]
+            else:
+                text = "Found %d beamlets that have the same distance from the source position." \
+                       " Therefore a unique central beamlet does not exist." % (
+                           values.count(minimal_distance))
+                logger.error(text)
+                raise ValueError(text)
+        return minimal_distance_per_frequency
+
     def __read_data(self, station_name, list_of_hbs_ms_tuples):
         """