Merge pull request #124 from hiyoneda/image_deconvolution

Changed the axis order of coordsys conversion matrix / added doctring to image deconvolution classes

cosipy/image_deconvolution/RichardsonLucy.py

@@ -1,14 +1,16 @@
 import copy
 import numpy as np
-import astropy.units as u
 from tqdm.autonotebook import tqdm
-import gc
 
 from histpy import Histogram
 
 from .deconvolution_algorithm_base import DeconvolutionAlgorithmBase
 
 class RichardsonLucy(DeconvolutionAlgorithmBase):
+    """
+    A class for the RichardsonLucy algorithm. 
+    The algorithm here is based on Knoedlseder+99, Knoedlseder+05, Siegert+20.
+    """
 
     def __init__(self, initial_model_map, data, parameter):
         DeconvolutionAlgorithmBase.__init__(self, initial_model_map, data, parameter)
@@ -43,37 +45,61 @@ def __init__(self, initial_model_map, data, parameter):
 
             self.smoothing_sigma = parameter['smoothing_FWHM'] / 2.354820 # degree
 
-            self.smoothing_max_sigma = parameter.get('smoothing_max_sigma', default = 5.0)
-            self.gaussian_filter = self.calc_gaussian_filter(self.smoothing_sigma, self.smoothing_max_sigma)
+            self.gaussian_filter = self.calc_gaussian_filter(self.smoothing_sigma)
 
     def pre_processing(self):
         pass
 
     def Estep(self):
-#        self.expectation = self.calc_expectation(self.model_map, self.data, self.use_sparse)
+        """
+        Notes
+        -----
+        Expect count histogram is calculated in the post processing.
+        """
         print("... skip E-step ...")
 
     def Mstep(self):
+        """
+        M-step in RL algorithm.
+
+        Notes
+        -----
+        Background normalization is also optimized based on a generalized RL algirithm.
+        Currenly we use a signle normalization parameter. 
+        In the future, the normalization will be optimized for each background group defined in some file.
+        """
+        # Currenly (2024-01-12) this method can work for both local coordinate CDS and in galactic coordinate CDS.
+        # This is just because in DC2 the rotate response for galactic coordinate CDS does not have an axis for time/scatt binning.
+        # However it is likely that it will have such an axis in the future in order to consider background variability depending on time and pointign direction etc.
+        # Then, the implementation here will not work. Thus, keep in mind that we need to modify it once the response format is fixed.
+
         diff = self.data.event_dense / self.expectation - 1
 
         delta_map_part1 = self.model_map / self.data.image_response_dense_projected
         delta_map_part2 = Histogram(self.model_map.axes, unit = self.data.image_response_dense_projected.unit)
 
         if self.data.response_on_memory == True:
-            diff_x_response_this_pix = np.tensordot(diff.contents, self.data.image_response_dense.contents, axes = ([1,2,3], [2,3,4])) 
+            diff_x_response = np.tensordot(diff.contents, self.data.image_response_dense.contents, axes = ([1,2,3], [2,3,4])) 
             # [Time/ScAtt, Em, Phi, PsiChi] x [NuLambda, Ei, Em, Phi, PsiChi] -> [Time/ScAtt, NuLambda, Ei]
 
-            delta_map_part2[:] = np.tensordot(self.data.coordsys_conv_matrix.contents, diff_x_response_this_pix, axes = ([1,2], [0,1])) * diff_x_response_this_pix.unit * self.data.coordsys_conv_matrix.unit #lb, Ei
-            # [lb, Time/ScAtt, NuLambda] x [Time/ScAtt, NuLambda, Ei] -> [lb, Ei]
+            delta_map_part2[:] = np.tensordot(self.data.coordsys_conv_matrix.contents, diff_x_response, axes = ([0,2], [0,1])) \
+                                 * diff_x_response.unit * self.data.coordsys_conv_matrix.unit
+            # [Time/ScAtt, lb, NuLambda] x [Time/ScAtt, NuLambda, Ei] -> [lb, Ei]
             # note that coordsys_conv_matrix is the sparse, so the unit should be recovered.
 
         else:
             for ipix in tqdm(range(self.npix_local)):
-                response_this_pix = np.sum(self.data.full_detector_response[ipix].to_dense(), axis = (4,5)) # may not work with the DC2 response format
+                if self.data.is_miniDC2_format == True:
+                    response_this_pix = np.sum(self.data.full_detector_response[ipix].to_dense(), axis = (4,5)) # [Ei, Em, Phi, PsiChi]
+                else:
+                    response_this_pix = self.data.full_detector_response[ipix].to_dense() # [Ei, Em, Phi, PsiChi]
 
-                diff_x_response_this_pix = np.tensordot(diff.contents, response_this_pix, axes = ([1,2,3], [1,2,3])) # Ti, Ei
+                diff_x_response_this_pix = np.tensordot(diff.contents, response_this_pix, axes = ([1,2,3], [1,2,3]))
+                # [Time/ScAtt, Em, Phi, PsiChi] x [Ei, Em, Phi, PsiChi] -> [Time/ScAtt, Ei]
 
-                delta_map_part2 += np.tensordot(self.data.coordsys_conv_matrix[:,:,ipix], diff_x_response_this_pix, axes = ([1],[0])) * diff_x_response_this_pix.unit * self.data.coordsys_conv_matrix.unit #lb, Ei
+                delta_map_part2 += np.tensordot(self.data.coordsys_conv_matrix[:,:,ipix], diff_x_response_this_pix, axes = ([0],[0])) \
+                                   * diff_x_response_this_pix.unit * self.data.coordsys_conv_matrix.unit #lb, Ei
+                # [Time/ScAtt, lb] x [Time/ScAtt, Ei] -> [lb, Ei]
 
         self.delta_map = delta_map_part1 * delta_map_part2
 
@@ -86,6 +112,12 @@ def Mstep(self):
                 self.bkg_norm = self.bkg_norm_range[1]
 
     def post_processing(self):
+        """
+        Here three processes will be performed.
+        - response weighting filter: the delta map is renormalized as pixels with large exposure times will have more feedback.
+        - gaussian smoothing filter: the delta map is blurred with a Gaussian function.
+        - acceleration of RL algirithm: the normalization of delta map is increased as long as the updated image has no non-negative components.
+        """
 
         if self.do_response_weighting:
             self.delta_map[:,:] *= self.response_weighting_filter
@@ -103,11 +135,28 @@ def post_processing(self):
         self.expectation = self.calc_expectation(self.model_map, self.data)
 
     def check_stopping_criteria(self, i_iteration):
+        """
+        If i_iteration is smaller than iteration_max, the iterative process will continue.
+
+        Returns
+        -------
+        bool
+        """
         if i_iteration < self.iteration_max:
             return False
         return True
 
     def register_result(self, i_iteration):
+        """
+        The values below are stored at the end of each iteration.
+        - iteration: iteration number
+        - model_map: updated image
+        - delta_map: delta map after M-step 
+        - processed_delta_map: delta map after post-processing
+        - alpha: acceleration parameter in RL algirithm
+        - background_normalization: optimized background normalization
+        - loglikelihood: log-likelihood
+        """
         loglikelihood = self.calc_loglikelihood(self.data, self.model_map, self.expectation)
 
         this_result = {"iteration": i_iteration, 
@@ -136,6 +185,14 @@ def show_result(self, i_iteration):
         print(f'    background_normalization: {self.result["background_normalization"]}')
 
     def calc_alpha(self, delta, model_map, almost_zero = 1e-5): #almost_zero is needed to prevent producing a flux below zero
+        """
+        Calculate the acceleration parameter in RL algorithm.
+
+        Returns
+        -------
+        float
+            Acceleration parameter
+        """
         alpha = -1.0 / np.min( delta / model_map ) * (1 - almost_zero)
         alpha = min(alpha, self.alpha_max)
         if alpha < 1.0:

cosipy/image_deconvolution/coordsys_conversion_matrix.py

@@ -1,16 +1,18 @@
-import astropy.units as u
-from astropy.time import Time
-from astropy.coordinates import SkyCoord, cartesian_to_spherical, Galactic
 import numpy as np
 import healpy as hp
 from tqdm.autonotebook import tqdm
+import sparse
+import astropy.units as u
+from astropy.time import Time
+from astropy.coordinates import SkyCoord, cartesian_to_spherical, Galactic
 
 from scoords import Attitude, SpacecraftFrame
 from histpy import Histogram, Axes, Axis, HealpixAxis
 
-import sparse
-
 class CoordsysConversionMatrix(Histogram):
+    """
+    A class for coordinate conversion matrix (ccm).
+    """
 
     def __init__(self, edges, contents = None, sumw2 = None,
                  labels=None, axis_scale = None, sparse = None, unit = None,
@@ -24,40 +26,50 @@ def __init__(self, edges, contents = None, sumw2 = None,
     @classmethod
     def time_binning_ccm(cls, full_detector_response, orientation, time_intervals, nside_model = None, is_nest_model = False):
         """
+        Calculate a ccm from a given orientation.
+        
         Parameters
         ----------
-        full_detector_response: 
-        orientation:
-        time_intervals: 2d np.array. it is the same format of binned_data.axes['Time'].edges
-        nside_model: If it is None, it will be the same as the NSIDE in the response.
+        full_detector_response : :py:class:`cosipy.response.FullDetectorResponse`
+            Response
+        orientation : :py:class:`cosipy.spacecraftfile.SpacecraftFile`
+            Orientation
+        time_intervals : :py:class:`np.array`
+            The same format of binned_data.axes['Time'].edges
+        nside_model : int or None, default None
+            If it is None, it will be the same as the NSIDE in the response.
+        is_nest_model : bool, default False
+            If scheme of the model map is nested, it should be False while it is rare.
 
         Returns
         -------
-        coordsys_conv_matrix: Axes [ "lb", "Time", "NuLambda" ]
+        :py:class:`cosipy.image_deconvolution.CoordsysConversionMatrix`
+            Its axes are [ "Time", "lb", "NuLambda" ].
         """
 
         if nside_model is None:
             nside_model = full_detector_response.nside
 
-        axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", label = "lb")
         axis_time = Axis(edges = time_intervals, label = "Time")
+        axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", label = "lb")
         axis_local_map = full_detector_response.axes["NuLambda"]
 
-        axis_coordsys_conv_matrix = [ axis_model_map, axis_time, axis_local_map ] #lb, Time, NuLambda
+        axis_coordsys_conv_matrix = [ axis_time, axis_model_map, axis_local_map ] #Time, lb, NuLambda
 
         contents = []
 
-        for ipix in tqdm(range(hp.nside2npix(nside_model))):
-            l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
-            pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')
-
-            ccm_thispix = np.zeros((axis_time.nbins, axis_local_map.nbins)) # without unit
+        for i_time, [init_time, end_time] in tqdm(enumerate(axis_time.bounds), total = len(axis_time.bounds)):
+            ccm_thispix = np.zeros((axis_model_map.nbins, axis_local_map.nbins)) # without unit
 
-            for i_time, [init_time, end_time] in enumerate(axis_time.bounds):
-                init_time = Time(init_time, format = 'unix')
-                end_time = Time(end_time, format = 'unix')
+            init_time = Time(init_time, format = 'unix')
+            end_time = Time(end_time, format = 'unix')
 
-                filtered_orientation = orientation.source_interval(init_time, end_time)
+            filtered_orientation = orientation.source_interval(init_time, end_time)
+
+            for ipix in range(hp.nside2npix(nside_model)):
+                l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
+                pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')
+
                 pixel_movement = filtered_orientation.get_target_in_sc_frame(target_name = f"pixel_{ipix}_{i_time}",
                                                                              target_coord = pixel_coord,
                                                                              quiet = True,
@@ -70,10 +82,10 @@ def time_binning_ccm(cls, full_detector_response, orientation, time_intervals, n
                                                                     src_path = pixel_movement,
                                                                     save = False)
 
-                ccm_thispix[i_time] = dwell_time_map.data 
+                ccm_thispix[ipix] = dwell_time_map.data 
                 # (HealpixMap).data returns the numpy array without its unit. dwell_time_map.unit is u.s.
 
-            ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_time.nbins, axis_local_map.nbins)) )
+            ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_model_map.nbins, axis_local_map.nbins)) )
 
             contents.append(ccm_thispix_sparse)
 
@@ -86,15 +98,28 @@ def time_binning_ccm(cls, full_detector_response, orientation, time_intervals, n
     @classmethod
     def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table, nside_model = None, use_averaged_pointing = False):
         """
+        Calculate a ccm from a given exposure_table.
+
         Parameters
         ----------
-        full_detector_response: 
-        exposure_table:
-        use_averaged_pointing: if this is True, the ccm loses accuracy but the calculatiion gets much faster.
+        full_detector_response : :py:class:`cosipy.response.FullDetectorResponse`
+            Response
+        exposure_table : :py:class:`cosipy.image_deconvolution.SpacecraftAttitudeExposureTable`
+            Scatt exposure table
+        nside_model : int or None, default None
+            If it is None, it will be the same as the NSIDE in the response.
+        use_averaged_pointing : bool, default False
+            If it is True, first the averaged Z- and X-pointings are calculated.
+            Then the dwell time map is calculated once for ach model pixel and each scatt_binning_index.
+            If it is False, the dwell time map is calculated for each attitude in zpointing and xpointing in the exposure table.
+            Then the calculated dwell time maps are summed up. 
+            In the former case, the computation is fast but may lose the angular resolution. 
+            In the latter case, the conversion matrix is more accurate but it takes a long time to calculate it.
 
         Returns
         -------
-        coordsys_conv_matrix: Axes [ "lb", "ScAtt", "NuLambda" ]
+        :py:class:`cosipy.image_deconvolution.CoordsysConversionMatrix'
+            Its axes are [ "ScAtt", "lb", "NuLambda" ].
         """
 
         if nside_model is None:
@@ -104,41 +129,41 @@ def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table,
 
         n_scatt_bins = len(exposure_table)
 
-        axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", scheme = exposure_table.scheme, label = "lb")
         axis_scatt = Axis(edges = np.arange(n_scatt_bins+1), label = "ScAtt")
+        axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", scheme = exposure_table.scheme, label = "lb")
         axis_local_map = full_detector_response.axes["NuLambda"]
 
-        axis_coordsys_conv_matrix = [ axis_model_map, axis_scatt, axis_local_map ] #lb, ScAtt, NuLambda
+        axis_coordsys_conv_matrix = [ axis_scatt, axis_model_map, axis_local_map ] #lb, ScAtt, NuLambda
 
         contents = []
 
-        for ipix in tqdm(range(hp.nside2npix(nside_model))):
-            l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
-            pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')
-
-            ccm_thispix = np.zeros((axis_scatt.nbins, axis_local_map.nbins)) # without unit
+        for i_scatt_bin in tqdm(range(n_scatt_bins)):
+            ccm_thispix = np.zeros((axis_model_map.nbins, axis_local_map.nbins)) # without unit
 
-            for idx in range(n_scatt_bins):
-                row = exposure_table.iloc[idx]
-
-                scatt_binning_index = row['scatt_binning_index']
-                num_pointings = row['num_pointings']
-                #healpix_index = row['healpix_index']
-                zpointing = row['zpointing']
-                xpointing = row['xpointing']
-                zpointing_averaged = row['zpointing_averaged']
-                xpointing_averaged = row['xpointing_averaged']
-                delta_time = row['delta_time']
-                exposure = row['exposure']
-
-                if use_averaged_pointing:
-                    z = SkyCoord([zpointing_averaged[0]], [zpointing_averaged[1]], frame="galactic", unit="deg")
-                    x = SkyCoord([xpointing_averaged[0]], [xpointing_averaged[1]], frame="galactic", unit="deg")
-                else:
-                    z = SkyCoord(zpointing.T[0], zpointing.T[1], frame="galactic", unit="deg")
-                    x = SkyCoord(xpointing.T[0], xpointing.T[1], frame="galactic", unit="deg")
+            row = exposure_table.iloc[i_scatt_bin]
+
+            scatt_binning_index = row['scatt_binning_index']
+            num_pointings = row['num_pointings']
+            #healpix_index = row['healpix_index']
+            zpointing = row['zpointing']
+            xpointing = row['xpointing']
+            zpointing_averaged = row['zpointing_averaged']
+            xpointing_averaged = row['xpointing_averaged']
+            delta_time = row['delta_time']
+            exposure = row['exposure']
 
-                attitude = Attitude.from_axes(x = x, z = z, frame = 'galactic')
+            if use_averaged_pointing:
+                z = SkyCoord([zpointing_averaged[0]], [zpointing_averaged[1]], frame="galactic", unit="deg")
+                x = SkyCoord([xpointing_averaged[0]], [xpointing_averaged[1]], frame="galactic", unit="deg")
+            else:
+                z = SkyCoord(zpointing.T[0], zpointing.T[1], frame="galactic", unit="deg")
+                x = SkyCoord(xpointing.T[0], xpointing.T[1], frame="galactic", unit="deg")
+
+            attitude = Attitude.from_axes(x = x, z = z, frame = 'galactic')
+
+            for ipix in range(hp.nside2npix(nside_model)):
+                l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
+                pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')
 
                 src_path_cartesian = SkyCoord(np.dot(attitude.rot.inv().as_matrix(), pixel_coord.cartesian.xyz.value),
                                               representation_type = 'cartesian', frame = SpacecraftFrame())
@@ -159,9 +184,9 @@ def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table,
 
                 hist, bins = np.histogram(pixels, bins = axis_local_map.edges, weights = weights)
 
-                ccm_thispix[idx] = hist
+                ccm_thispix[ipix] = hist
 
-            ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_scatt.nbins, axis_local_map.nbins)) )
+            ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_model_map.nbins, axis_local_map.nbins)) )
 
             contents.append(ccm_thispix_sparse)
 
@@ -173,12 +198,27 @@ def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table,
 
     @classmethod
     def open(cls, filename, name = 'hist'):
+        """
+        Open a ccm from a file.
+
+        Parameters
+        ----------
+        filename : str
+            Path to file.
+        name : str, default 'hist'
+            Name of group where the histogram was saved.
+
+        Returns
+        -------
+        :py:class:`cosipy.image_deconvolution.CoordsysConversionMatrix'
+            Its axes are [ "lb", "Time" or "ScAtt", "NuLambda" ].
+        """
 
         new = super().open(filename, name)
 
         new = cls(new.axes, contents = new.contents, sumw2 = new.contents, unit = new.unit) 
 
-        new.binning_method = new.axes.labels[1] # 'Time' or 'ScAtt'
+        new.binning_method = new.axes.labels[0] # 'Time' or 'ScAtt'
 
         return new