Merge back PR248 onto develop

cosipy/__init__.py

@@ -12,7 +12,7 @@
 
 from .spacecraftfile import SpacecraftFile
 
-from .ts_map import FastTSMap
+from .ts_map import FastTSMap, MOCTSMap
 
 from .source_injector import SourceInjector
 

cosipy/ts_map/__init__.py

@@ -1,2 +1,3 @@
 from .fast_ts_fit import FastTSMap
-from .fast_norm_fit import FastNormFit
+from .fast_norm_fit import FastNormFit
+from .moc_ts_fit import MOCTSMap

cosipy/ts_map/fast_ts_fit.py

@@ -273,7 +273,7 @@ def get_ei_cds_array(hypothesis_coord, energy_channel, response_path, spectrum,
     def fast_ts_fit(hypothesis_coord, 
                     energy_channel, data_cds_array, bkg_model_cds_array, 
                     orientation, response_path, spectrum, cds_frame,
-                    ts_nside, ts_scheme):
+                    ts_nside, ts_scheme, pixel_idx = None):
 
         """
         Perform a TS fit on a single location at `hypothesis_coord`.
@@ -310,7 +310,10 @@ def fast_ts_fit(hypothesis_coord,
         start_fast_ts_fit = time.time()
 
         # get the indices of the pixels to fit
-        pix = hp.ang2pix(nside = ts_nside, theta = hypothesis_coord.l.deg, phi = hypothesis_coord.b.deg, lonlat = True)
+        if pixel_idx is None:
+            pix = hp.ang2pix(nside = ts_nside, theta = hypothesis_coord.l.deg, phi = hypothesis_coord.b.deg, lonlat = True)
+        else:
+            pix = pixel_idx
 
         # get the expected counts in the flattened cds array
         start_ei_cds_array = time.time()
@@ -331,9 +334,55 @@ def fast_ts_fit(hypothesis_coord,
         time_fast_ts_fit = end_fast_ts_fit - start_fast_ts_fit
 
         return [pix, result[0], result[1], result[2], result[3], result[4], time_ei_cds_array, time_fit, time_fast_ts_fit]
+
+    @staticmethod
+    def zip_comp(*lists):
+
+        """
+        Zip the lists in a way that it expands the lists will one element.
+
+        list1 = [1, 2, 3, 4]
+        list2 = ["a"]
+        list3 = [11, 21, 31, 41]
+
+        zip_comp will produce a tuple like this:
+        ([1, "a", 11],
+         [2, "a", 21], 
+         [3, "a", 31], 
+         [4, "a", 41])
+
+        As you can see, it only allows lists with two length: 1 or the max length.
+
+        Parameters
+        ----------
+        lists : list
+            The input lists
+
+        Returns
+        -------
+        zip :
+            The zippped array. To expand, please use list(returned_object)
+
+        """
+
+        all_lengths = np.unique([len(i) for i in lists])
+
+        if len(all_lengths) > 2:
+            raise ValueError(f"You have input lists with more than two lengths: {all_lengths}. Can't do zip comprehension!")
+
+
+        new_lists = []
+        for i in lists:
+            if len(i) == np.min(all_lengths):
+                new_lists.append(i*np.max(all_lengths))
+            else:
+                new_lists.append(i)
+
+        return zip(*new_lists)
 
 
-    def parallel_ts_fit(self, hypothesis_coords, energy_channel, spectrum, ts_scheme = "RING", start_method = "fork", cpu_cores = None, ts_nside = None):
+    def parallel_ts_fit(self, hypothesis_coords, energy_channel, spectrum, ts_scheme = "RING", start_method = "fork", cpu_cores = None, ts_nside = None,
+                        pixel_idx = [None]):
 
         """
         Perform parallel computation on all the hypothesis coordinates.
@@ -353,7 +402,9 @@ def parallel_ts_fit(self, hypothesis_coords, energy_channel, spectrum, ts_scheme
         cpu_cores : int, optional
             The number of cpu cores you wish to use for the parallel computation (the default is `None`, which implies using all the available number of cores -1 to perform the parallel computation).
         ts_nside : int, optional
-            The nside of the ts map. This must be given if the number of hypothesis_coords isn't equal to the number of pixels of the total ts map, which means that you fit only a portion of the total ts map. (the default is `None`, which means that you fit the full ts map). 
+            The nside of the ts map. This must be given if the number of hypothesis_coords isn't equal to the number of pixels of the total ts map, which means that you fit only a portion of the total ts map. (the default is `None`, which means that you fit the full ts map).
+        pixel_idx : list, optional
+            The pixel indices of the corresponding hypothesis_coords. This parameter is used to match the pixels and the ts values in a regional fit case. 
 
         Returns
         -------
@@ -389,12 +440,12 @@ def parallel_ts_fit(self, hypothesis_coords, energy_channel, spectrum, ts_scheme
             cores = cpu_cores
             logger.info(f"You have total {total_cores} CPU cores, using {cores} CPU cores for parallel computation.")
 
-        start = time.time()
+        start = time.time() 
         multiprocessing.set_start_method(start_method, force = True)
         pool = multiprocessing.Pool(processes = cores)
-        results = pool.starmap(FastTSMap.fast_ts_fit, product(hypothesis_coords, [energy_channel], [data_cds_array], [bkg_model_cds_array], 
-                                                             [self._orientation], [self._response_path], [spectrum], [self._cds_frame], 
-                                                             [ts_nside], [ts_scheme]))
+        results = pool.starmap(FastTSMap.fast_ts_fit, FastTSMap.zip_comp(hypothesis_coords, [energy_channel], [data_cds_array], [bkg_model_cds_array], 
+                                                                         [self._orientation], [self._response_path], [spectrum], [self._cds_frame], 
+                                                                         [ts_nside], [ts_scheme], pixel_idx))
 
         pool.close()
         pool.join()