diff --git a/cosipy/image_deconvolution/RLparallelscript.py b/cosipy/image_deconvolution/RLparallelscript.py index ef6030ea3..acc327129 100644 --- a/cosipy/image_deconvolution/RLparallelscript.py +++ b/cosipy/image_deconvolution/RLparallelscript.py @@ -163,4 +163,4 @@ def plot_reconstructed_image(result, source_position=None): # source_position s MPI.Finalize() if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/cosipy/image_deconvolution/RichardsonLucy.py b/cosipy/image_deconvolution/RichardsonLucy.py deleted file mode 100644 index 235b3d393..000000000 --- a/cosipy/image_deconvolution/RichardsonLucy.py +++ /dev/null @@ -1,242 +0,0 @@ -import os -import numpy as np -import astropy.units as u -import astropy.io.fits as fits -import logging -logger = logging.getLogger(__name__) - -from histpy import Histogram - -from .RichardsonLucySimple import RichardsonLucySimple - -class RichardsonLucy(RichardsonLucySimple): - """ - A class for the RichardsonLucy algorithm. - The algorithm here is based on Knoedlseder+99, Knoedlseder+05, Siegert+20. - - An example of parameter is as follows. - - iteration_max: 100 - minimum_flux: - value: 0.0 - unit: "cm-2 s-1 sr-1" - acceleration: - activate: True - alpha_max: 10.0 - response_weighting: - activate: True - index: 0.5 - smoothing: - activate: True - FWHM: - value: 2.0 - unit: "deg" - stopping_criteria: - statistics: "log-likelihood" - threshold: 1e-2 - background_normalization_optimization: - activate: True - range: {"albedo": [0.9, 1.1]} - save_results: - activate: True - directory: "/results" - only_final_result: True - """ - - def __init__(self, initial_model, dataset, mask, parameter): - - super().__init__(initial_model, dataset, mask, parameter) - - # acceleration - self.do_acceleration = parameter.get('acceleration:activate', False) - if self.do_acceleration == True: - self.alpha_max = parameter.get('acceleration:alpha_max', 1.0) - - # smoothing - self.do_smoothing = parameter.get('smoothing:activate', False) - if self.do_smoothing: - self.smoothing_fwhm = parameter.get('smoothing:FWHM:value') * u.Unit(parameter.get('smoothing:FWHM:unit')) - logger.info(f"Gaussian filter with FWHM of {self.smoothing_fwhm} will be applied to delta images ...") - - # stopping criteria - self.stopping_criteria_statistics = parameter.get('stopping_criteria:statistics', "log-likelihood") - self.stopping_criteria_threshold = parameter.get('stopping_criteria:threshold', 1e-2) - - if not self.stopping_criteria_statistics in ["log-likelihood"]: - raise ValueError - - def initialization(self): - """ - initialization before running the image deconvolution - """ - - super().initialization() - - # expected count histograms - self.expectation_list = self.calc_expectation_list(model = self.initial_model, dict_bkg_norm = self.dict_bkg_norm) - logger.info("The expected count histograms were calculated with the initial model map.") - - def pre_processing(self): - """ - pre-processing for each iteration - """ - pass - - def Estep(self): - """ - E-step (but it will be skipped). - Note that self.expectation_list is updated in self.post_processing(). - """ - pass - - def Mstep(self): - """ - M-step in RL algorithm. - """ - super().Mstep() - - 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. - """ - - self.processed_delta_model = self.delta_model.copy() - - if self.do_response_weighting: - self.processed_delta_model *= self.response_weighting_filter - - if self.do_smoothing: - self.processed_delta_model = self.processed_delta_model.smoothing(fwhm = self.smoothing_fwhm) - - if self.do_acceleration: - self.alpha = self.calc_alpha(self.processed_delta_model, self.model) - else: - self.alpha = 1.0 - - self.model += self.processed_delta_model * self.alpha - self.model[:] = np.where(self.model.contents < self.minimum_flux, self.minimum_flux, self.model.contents) - - if self.mask is not None: - self.model = self.model.mask_pixels(self.mask) - - # update expectation_list - self.expectation_list = self.calc_expectation_list(self.model, dict_bkg_norm = self.dict_bkg_norm) - logger.debug("The expected count histograms were updated with the new model map.") - - # update log_likelihood_list - self.log_likelihood_list = self.calc_log_likelihood_list(self.expectation_list) - logger.debug("The log-likelihood list was updated with the new expected count histograms.") - - def register_result(self): - """ - The values below are stored at the end of each iteration. - - iteration: iteration number - - model: updated image - - delta_model: delta map after M-step - - processed_delta_model: delta map after post-processing - - alpha: acceleration parameter in RL algirithm - - background_normalization: optimized background normalization - - log-likelihood: log-likelihood - """ - - this_result = {"iteration": self.iteration_count, - "model": self.model.copy(), - "delta_model": self.delta_model, - "processed_delta_model": self.processed_delta_model, - "background_normalization": self.dict_bkg_norm.copy(), - "alpha": self.alpha, - "log-likelihood": self.log_likelihood_list} - - # show intermediate results - logger.info(f' alpha: {this_result["alpha"]}') - logger.info(f' background_normalization: {this_result["background_normalization"]}') - logger.info(f' log-likelihood: {this_result["log-likelihood"]}') - - # register this_result in self.results - self.results.append(this_result) - - def check_stopping_criteria(self): - """ - If iteration_count is smaller than iteration_max, the iterative process will continue. - - Returns - ------- - bool - """ - if self.iteration_count == 1: - return False - elif self.iteration_count == self.iteration_max: - return True - - if self.stopping_criteria_statistics == "log-likelihood": - - log_likelihood = np.sum(self.results[-1]["log-likelihood"]) - log_likelihood_before = np.sum(self.results[-2]["log-likelihood"]) - - logger.debug(f'Delta log-likelihood: {log_likelihood - log_likelihood_before}') - - if log_likelihood - log_likelihood_before < 0: - - logger.warning("The likelihood is not increased in this iteration. The image reconstruction may be unstable.") - return False - - elif log_likelihood - log_likelihood_before < self.stopping_criteria_threshold: - return True - - return False - - def finalization(self): - """ - finalization after running the image deconvolution - """ - if self.save_results == True: - logger.info(f'Saving results in {self.save_results_directory}') - - counter_name = "iteration" - - # model - histkey_filename = [("model", f"{self.save_results_directory}/model.hdf5"), - ("delta_model", f"{self.save_results_directory}/delta_model.hdf5"), - ("processed_delta_model", f"{self.save_results_directory}/processed_delta_model.hdf5")] - - for key, filename in histkey_filename: - - self.save_histogram(filename = filename, - counter_name = counter_name, - histogram_key = key, - only_final_result = self.save_only_final_result) - - #fits - fits_filename = f'{self.save_results_directory}/results.fits' - - self.save_results_as_fits(filename = fits_filename, - counter_name = counter_name, - values_key_name_format = [("alpha", "ALPHA", "D")], - dicts_key_name_format = [("background_normalization", "BKG_NORM", "D")], - lists_key_name_format = [("log-likelihood", "LOG-LIKELIHOOD", "D")]) - - def calc_alpha(self, delta_model, model): - """ - Calculate the acceleration parameter in RL algorithm. - - Returns - ------- - float - Acceleration parameter - """ - diff = -1 * (model / delta_model).contents - - diff[(diff <= 0) | (delta_model.contents == 0)] = np.inf - - if self.mask is not None: - diff[np.invert(self.mask.contents)] = np.inf - - alpha = min(np.min(diff), self.alpha_max) - - if alpha < 1.0: - alpha = 1.0 - - return alpha diff --git a/cosipy/image_deconvolution/RichardsonLucySimple.py b/cosipy/image_deconvolution/RichardsonLucySimple.py deleted file mode 100644 index 22e4239c9..000000000 --- a/cosipy/image_deconvolution/RichardsonLucySimple.py +++ /dev/null @@ -1,211 +0,0 @@ -import os -import numpy as np -import logging -logger = logging.getLogger(__name__) - -from histpy import Histogram - -from .deconvolution_algorithm_base import DeconvolutionAlgorithmBase - -class RichardsonLucySimple(DeconvolutionAlgorithmBase): - """ - A class for the original RichardsonLucy algorithm. - Basically, this class can be used for testing codes. - - An example of parameter is as follows. - - iteration_max: 100 - minimum_flux: - value: 0.0 - unit: "cm-2 s-1 sr-1" - background_normalization_optimization: - activate: True - range: {"albedo": [0.9, 1.1]} - response_weighting: - activate: True - index: 0.5 - save_results: - activate: True - directory: "/results" - only_final_result: True - """ - - def __init__(self, initial_model, dataset, mask, parameter): - - super().__init__(initial_model, dataset, mask, parameter) - - # background normalization optimization - self.do_bkg_norm_optimization = parameter.get('background_normalization_optimization:activate', False) - if self.do_bkg_norm_optimization: - self.dict_bkg_norm_range = parameter.get('background_normalization_optimization:range', {key: [0.0, 100.0] for key in self.dict_bkg_norm.keys()}) - - # response_weighting - self.do_response_weighting = parameter.get('response_weighting:activate', False) - if self.do_response_weighting: - self.response_weighting_index = parameter.get('response_weighting:index', 0.5) - - # saving results - self.save_results = parameter.get('save_results:activate', False) - self.save_results_directory = parameter.get('save_results:directory', './results') - self.save_only_final_result = parameter.get('save_results:only_final_result', False) - - if self.save_results is True: - if os.path.isdir(self.save_results_directory): - logger.warning(f"A directory {self.save_results_directory} already exists. Files in {self.save_results_directory} may be overwritten. Make sure that is not a problem.") - else: - os.makedirs(self.save_results_directory) - - def initialization(self): - """ - initialization before running the image deconvolution - """ - # clear counter - self.iteration_count = 0 - - # clear results - self.results.clear() - - # copy model - self.model = self.initial_model.copy() - - # calculate exposure map - self.summed_exposure_map = self.calc_summed_exposure_map() - - # mask setting - if self.mask is None and np.any(self.summed_exposure_map.contents == 0): - self.mask = Histogram(self.model.axes, - contents = self.summed_exposure_map.contents > 0, - copy_contents = False) - self.model = self.model.mask_pixels(self.mask) - logger.info("There are zero-exposure pixels. A mask to ignore them was set.") - - # response-weighting filter - if self.do_response_weighting: - self.response_weighting_filter = (self.summed_exposure_map.contents / np.max(self.summed_exposure_map.contents))**self.response_weighting_index - logger.info("The response weighting filter was calculated.") - - # calculate summed background models for M-step - if self.do_bkg_norm_optimization: - self.dict_summed_bkg_model = {} - for key in self.dict_bkg_norm.keys(): - self.dict_summed_bkg_model[key] = self.calc_summed_bkg_model(key) - - def pre_processing(self): - """ - pre-processing for each iteration - """ - pass - - def Estep(self): - """ - E-step in RL algoritm - """ - - self.expectation_list = self.calc_expectation_list(self.model, dict_bkg_norm = self.dict_bkg_norm) - logger.debug("The expected count histograms were updated with the new model map.") - - def Mstep(self): - """ - M-step in RL algorithm. - """ - - ratio_list = [ data.event / expectation for data, expectation in zip(self.dataset, self.expectation_list) ] - - # delta model - sum_T_product = self.calc_summed_T_product(ratio_list) - self.delta_model = self.model * (sum_T_product/self.summed_exposure_map - 1) - - # masking - if self.mask is not None: - self.delta_model = self.delta_model.mask_pixels(self.mask) - - # background normalization optimization - if self.do_bkg_norm_optimization: - for key in self.dict_bkg_norm.keys(): - - sum_bkg_T_product = self.calc_summed_bkg_model_product(key, ratio_list) - sum_bkg_model = self.dict_summed_bkg_model[key] - bkg_norm = self.dict_bkg_norm[key] * (sum_bkg_T_product / sum_bkg_model) - - bkg_range = self.dict_bkg_norm_range[key] - if bkg_norm < bkg_range[0]: - bkg_norm = bkg_range[0] - elif bkg_norm > bkg_range[1]: - bkg_norm = bkg_range[1] - - self.dict_bkg_norm[key] = bkg_norm - - def post_processing(self): - """ - Post-processing. - """ - - # response_weighting - if self.do_response_weighting: - self.model[:] += self.delta_model.contents * self.response_weighting_filter - else: - self.model[:] += self.delta_model.contents - - self.model[:] = np.where(self.model.contents < self.minimum_flux, self.minimum_flux, self.model.contents) - - # masking again - if self.mask is not None: - self.model = self.model.mask_pixels(self.mask) - - def register_result(self): - """ - Register results at the end of each iteration. - """ - - this_result = {"iteration": self.iteration_count, - "model": self.model.copy(), - "delta_model": self.delta_model, - "background_normalization": self.dict_bkg_norm.copy()} - - # show intermediate results - logger.info(f' background_normalization: {this_result["background_normalization"]}') - - # register this_result in self.results - self.results.append(this_result) - - def check_stopping_criteria(self): - """ - If iteration_count is smaller than iteration_max, the iterative process will continue. - - Returns - ------- - bool - """ - if self.iteration_count < self.iteration_max: - return False - return True - - def finalization(self): - """ - finalization after running the image deconvolution - """ - - if self.save_results == True: - logger.info(f'Saving results in {self.save_results_directory}') - - counter_name = "iteration" - - # model - histkey_filename = [("model", f"{self.save_results_directory}/model.hdf5"), - ("delta_model", f"{self.save_results_directory}/delta_model.hdf5")] - - for key, filename in histkey_filename: - - self.save_histogram(filename = filename, - counter_name = counter_name, - histogram_key = key, - only_final_result = self.save_only_final_result) - - #fits - fits_filename = f'{self.save_results_directory}/results.fits' - - self.save_results_as_fits(filename = fits_filename, - counter_name = counter_name, - values_key_name_format = [], - dicts_key_name_format = [("background_normalization", "BKG_NORM", "D")], - lists_key_name_format = []) diff --git a/cosipy/image_deconvolution/__init__.py b/cosipy/image_deconvolution/__init__.py index ad11db1e8..b97384aae 100644 --- a/cosipy/image_deconvolution/__init__.py +++ b/cosipy/image_deconvolution/__init__.py @@ -1,17 +1,19 @@ from .image_deconvolution import ImageDeconvolution, ParallelImageDeconvolution -from .image_deconvolution_data_interface_base import ImageDeconvolutionDataInterfaceBase -from .dataIF_COSI_DC2 import DataIF_COSI_DC2 -from .dataIF_Parallel import DataIF_Parallel +from .data_interfaces.image_deconvolution_data_interface_base import ImageDeconvolutionDataInterfaceBase +from .data_interfaces.dataIF_COSI_DC2 import DataIF_COSI_DC2 +from .data_interfaces.dataIF_Parallel import DataIF_Parallel +from .data_interfaces.data_interface_collection import DataInterfaceCollection -from .model_base import ModelBase -from .allskyimage import AllSkyImageModel +from .models.model_base import ModelBase +from .models.allskyimage import AllSkyImageModel -from .deconvolution_algorithm_base import DeconvolutionAlgorithmBase -from .RichardsonLucy import RichardsonLucy -from .RichardsonLucySimple import RichardsonLucySimple -from .MAP_RichardsonLucy import MAP_RichardsonLucy +from .algorithms.deconvolution_algorithm_base import DeconvolutionAlgorithmBase +from .algorithms.RichardsonLucyBasic import RichardsonLucyBasic +from .algorithms.RichardsonLucy import RichardsonLucy +from .algorithms.RichardsonLucyAdvanced import RichardsonLucyAdvanced +from .algorithms.MAP_RichardsonLucy import MAP_RichardsonLucy -from .scatt_exposure_table import SpacecraftAttitudeExposureTable -from .time_binned_exposure_table import TimeBinnedExposureTable -from .coordsys_conversion_matrix import CoordsysConversionMatrix +from .exposure_tables.scatt_exposure_table import SpacecraftAttitudeExposureTable +from .exposure_tables.time_binned_exposure_table import TimeBinnedExposureTable +from .exposure_tables.coordsys_conversion_matrix import CoordsysConversionMatrix diff --git a/cosipy/image_deconvolution/MAP_RichardsonLucy.py b/cosipy/image_deconvolution/algorithms/MAP_RichardsonLucy.py similarity index 63% rename from cosipy/image_deconvolution/MAP_RichardsonLucy.py rename to cosipy/image_deconvolution/algorithms/MAP_RichardsonLucy.py index 36586195d..8f475a85f 100644 --- a/cosipy/image_deconvolution/MAP_RichardsonLucy.py +++ b/cosipy/image_deconvolution/algorithms/MAP_RichardsonLucy.py @@ -1,19 +1,21 @@ -import os -import copy import numpy as np import astropy.units as u -import astropy.io.fits as fits import logging logger = logging.getLogger(__name__) from histpy import Histogram -from .RichardsonLucySimple import RichardsonLucySimple +from ..utils import _to_float +from .RichardsonLucy import RichardsonLucy from .prior_tsv import PriorTSV from .prior_entropy import PriorEntropy -class MAP_RichardsonLucy(RichardsonLucySimple): +from .response_weighting_filter import ResponseWeightingFilter + +from ..constants import DEFAULT_STOPPING_THRESHOLD, DEFAULT_RESPONSE_WEIGHTING_INDEX + +class MAP_RichardsonLucy(RichardsonLucy): """ A class for the RichardsonLucy algorithm using prior distributions. @@ -78,13 +80,36 @@ def __init__(self, initial_model, dataset, mask, parameter): this_prior_parameter = parameter['prior'][prior_name] self.priors[prior_name] = self.prior_classes[prior_name](this_prior_parameter, initial_model) + # response_weighting + self.response_weighting_enabled = parameter.get('response_weighting:activate', False) + if self.response_weighting_enabled: + self.response_weighting_index = parameter.get('response_weighting:index', DEFAULT_RESPONSE_WEIGHTING_INDEX) + # stopping criteria self.stopping_criteria_statistics = parameter.get('stopping_criteria:statistics', "log-posterior") - self.stopping_criteria_threshold = parameter.get('stopping_criteria:threshold', 1e-2) + self.stopping_criteria_threshold = parameter.get('stopping_criteria:threshold', DEFAULT_STOPPING_THRESHOLD) if not self.stopping_criteria_statistics in ["log-likelihood", "log-posterior"]: raise ValueError + # update parameter summary + self._parameter_summary += [ + ("prior", parameter['prior']), + ] + + self._parameter_summary += [ + ("response_weighting_enabled", self.response_weighting_enabled), + ] + if self.response_weighting_enabled: + self._parameter_summary += [ + ("response_weighting_index", self.response_weighting_index), + ] + + self._parameter_summary += [ + ("stopping_criteria_statistics", self.stopping_criteria_statistics), + ("stopping_criteria_threshold", self.stopping_criteria_threshold), + ] + def load_gamma_prior(self, parameter): if parameter is None: @@ -97,7 +122,7 @@ def load_gamma_prior(self, parameter): self.prior_gamma_bkg_theta = parameter['background']['theta']['value'] self.prior_gamma_bkg_k = parameter['background']['k']['value'] - def log_gamma_prior(self, model): + def log_gamma_prior(self, model, dict_bkg_norm): eps = np.finfo(model.contents.dtype).eps @@ -110,40 +135,46 @@ def log_gamma_prior(self, model): # background pl_part_bkg, log_part_bkg = 0, 0 - if self.do_bkg_norm_optimization: - for key in self.dict_bkg_norm.keys(): + if self.bkg_norm_optimization_enabled: + for key in dict_bkg_norm.keys(): - bkg_norm = self.dict_bkg_norm[key] + bkg_norm = dict_bkg_norm[key] pl_part_bkg += (self.prior_gamma_bkg_k - 1.0) * np.log(bkg_norm) log_part_bkg += -1.0 * np.sum( bkg_norm / self.prior_gamma_bkg_theta ) - return pl_part_model + log_part_model, pl_part_bkg + log_part_bkg + return _to_float(pl_part_model + log_part_model), _to_float(pl_part_bkg + log_part_bkg) def initialization(self): """ - initialization before running the image deconvolution + Initialize before running image deconvolution. + + This method sets up response weighting filter based on the exposure map. """ super().initialization() - # expected count histograms - self.expectation_list = self.calc_expectation_list(model = self.initial_model, dict_bkg_norm = self.dict_bkg_norm) - logger.info("The expected count histograms were calculated with the initial model map.") + # response-weighting filter + if self.response_weighting_enabled: + self.response_weighting_filter = ResponseWeightingFilter(self.summed_exposure_map, self.response_weighting_index) def pre_processing(self): """ pre-processing for each iteration """ - pass - def Estep(self): + if self.iteration_count == 1: + self.Estep() + logger.info("The expected count histograms were calculated with the initial model map.") + + def processing_core(self): """ - E-step (but it will be skipped). - Note that self.expectation_list is updated in self.post_processing(). + Core processing for each iteration. """ - pass + + # Note that Estep() is performed in self.post_processing(). + self.Mstep() def Mstep(self): """ @@ -153,7 +184,7 @@ def Mstep(self): ratio_list = [ data.event / expectation for data, expectation in zip(self.dataset, self.expectation_list) ] # model update (EM part) - sum_T_product = self.calc_summed_T_product(ratio_list) + sum_T_product = self.dataset.calc_summed_T_product(ratio_list) model_EM = (self.model * sum_T_product + self.prior_gamma_model_k - 1.0) \ / (self.summed_exposure_map + 1.0 / self.prior_gamma_model_theta) model_EM[:] = np.where( model_EM.contents < self.minimum_flux, self.minimum_flux, model_EM.contents) @@ -166,39 +197,22 @@ def Mstep(self): self.prior_filter = Histogram(self.model.axes, contents = np.exp( sum_grad_log_prior / (self.summed_exposure_map.contents + 1.0 / self.prior_gamma_model_theta))) - self.model[:] = self.prior_filter.contents * model_EM.contents + self.delta_model = self.prior_filter * model_EM.contents - self.model - # applying response_weighting_filter - if self.do_response_weighting: - if self.iteration_count == 1: - delta_model = self.model - self.initial_model - else: - delta_model = self.model - self.results[-1]['model'] - - self.model[:] = (self.model.contents - delta_model.contents) + self.response_weighting_filter * delta_model.contents - - # masking - if self.mask is not None: - self.model = self.model.mask_pixels(self.mask) - - self.model[:] = np.where( self.model.contents < self.minimum_flux, self.minimum_flux, self.model.contents) - # background normalization optimization - if self.do_bkg_norm_optimization: + if self.bkg_norm_optimization_enabled: for key in self.dict_bkg_norm.keys(): - sum_bkg_T_product = self.calc_summed_bkg_model_product(key, ratio_list) + sum_bkg_T_product = self.dataset.calc_summed_bkg_model_product(key, ratio_list) sum_bkg_model = self.dict_summed_bkg_model[key] bkg_norm = (self.dict_bkg_norm[key] * sum_bkg_T_product + self.prior_gamma_bkg_k - 1.0) \ / (sum_bkg_model + 1.0 / self.prior_gamma_bkg_theta) - bkg_range = self.dict_bkg_norm_range[key] - if bkg_norm < bkg_range[0]: - bkg_norm = bkg_range[0] - elif bkg_norm > bkg_range[1]: - bkg_norm = bkg_range[1] + self.dict_delta_bkg_norm[key] = bkg_norm - self.dict_bkg_norm[key] - self.dict_bkg_norm[key] = bkg_norm + # apply response_weighting_filter + if self.response_weighting_enabled: + self.delta_model = self.response_weighting_filter.apply(self.delta_model) def post_processing(self): """ @@ -206,23 +220,31 @@ def post_processing(self): - response weighting filter: the delta map is renormalized as pixels with large exposure times will have more feedback. """ - #TODO: add acceleration SQUAREM + # update model + self.model += self.delta_model + self._ensure_model_constraints() + + # update background normalization + if self.bkg_norm_optimization_enabled: + for key in self.dict_bkg_norm.keys(): + self.dict_bkg_norm[key] += self.dict_delta_bkg_norm[key] + self._ensure_bkg_norm_range() # update expectation_list - self.expectation_list = self.calc_expectation_list(self.model, dict_bkg_norm = self.dict_bkg_norm) + self.Estep() logger.debug("The expected count histograms were updated with the new model map.") # update log_likelihood_list - self.log_likelihood_list = self.calc_log_likelihood_list(self.expectation_list) + self.log_likelihood_list = self.dataset.calc_log_likelihood_list(self.expectation_list) logger.debug("The log-likelihood list was updated with the new expected count histograms.") # update log priors self.log_priors = {} - self.log_priors['gamma_model'], self.log_priors['gamma_bkg'] = self.log_gamma_prior(self.model) + self.log_priors['gamma_model'], self.log_priors['gamma_bkg'] = self.log_gamma_prior(self.model, self.dict_bkg_norm) for key in self.priors.keys(): - self.log_priors[key] = self.priors[key].log_prior(self.model) + self.log_priors[key] = _to_float(self.priors[key].log_prior(self.model)) # log-posterior self.log_posterior = np.sum(self.log_likelihood_list) + np.sum([self.log_priors[key] for key in self.log_priors.keys()]) @@ -240,19 +262,17 @@ def register_result(self): """ this_result = {"iteration": self.iteration_count, - "model": copy.deepcopy(self.model), - "prior_filter": copy.deepcopy(self.prior_filter), - "background_normalization": copy.deepcopy(self.dict_bkg_norm), - "log-likelihood": copy.deepcopy(self.log_likelihood_list), - "log-prior": copy.deepcopy(self.log_priors), - "log-posterior": copy.deepcopy(self.log_posterior), + "model": self.model.copy(), + "prior_filter": self.prior_filter.copy(), + "background_normalization": self.dict_bkg_norm.copy(), + "log-likelihood": self.log_likelihood_list.copy(), + "log-prior": self.log_priors.copy(), + "log-posterior": self.log_posterior, } # show intermediate results - logger.info(f' background_normalization: {this_result["background_normalization"]}') - logger.info(f' log-likelihood: {this_result["log-likelihood"]}') - logger.info(f' log-prior: {this_result["log-prior"]}') - logger.info(f' log-posterior: {this_result["log-posterior"]}') + for key in ["background_normalization", "log-likelihood", "log-prior", "log-posterior"]: + logger.info(f"{key}: {this_result[key]}") # register this_result in self.results self.results.append(this_result) @@ -266,11 +286,12 @@ def check_stopping_criteria(self): bool """ - if self.iteration_count == 1: - return False - elif self.iteration_count == self.iteration_max: + if self.iteration_count >= self.iteration_max: return True + if self.iteration_count == 1: + return False # need at least 2 results to compute delta + if self.stopping_criteria_statistics == "log-likelihood": log_likelihood = np.sum(self.results[-1]["log-likelihood"]) @@ -308,27 +329,17 @@ def finalization(self): finalization after running the image deconvolution """ - if self.save_results == True: - logger.info(f'Saving results in {self.save_results_directory}') + if not self.save_results: + return - counter_name = "iteration" - - # model - histkey_filename = [("model", f"{self.save_results_directory}/model.hdf5"), - ("prior_filter", f"{self.save_results_directory}/prior_filter.hdf5")] + logger.info(f"Saving results in {self.save_results_directory}") - for key, filename in histkey_filename: - - self.save_histogram(filename = filename, - counter_name = counter_name, - histogram_key = key, - only_final_result = self.save_only_final_result) - - #fits - fits_filename = f'{self.save_results_directory}/results.fits' - - self.save_results_as_fits(filename = fits_filename, - counter_name = counter_name, - values_key_name_format = [("log-posterior", "LOG-POSTERIOR", "D")], - dicts_key_name_format = [("background_normalization", "BKG_NORM", "D"), ("log-prior", "LOG-PRIOR", "D")], - lists_key_name_format = [("log-likelihood", "LOG-LIKELIHOOD", "D")]) + self._save_standard_results( + counter_name = "iteration", + histogram_keys = [("model", f"{self.save_results_directory}/model.hdf5", self.save_only_final_result), + ("prior_filter", f"{self.save_results_directory}/prior_filter.hdf5", self.save_only_final_result)], + fits_filename = f"{self.save_results_directory}/results.fits", + values_key_name_format = [("log-posterior", "LOG-POSTERIOR", "D")], + dicts_key_name_format = [("background_normalization", "BKG_NORM", "D"), ("log-prior", "LOG-PRIOR", "D")], + lists_key_name_format = [("log-likelihood", "LOG-LIKELIHOOD", "D")], + ) diff --git a/cosipy/image_deconvolution/algorithms/RichardsonLucy.py b/cosipy/image_deconvolution/algorithms/RichardsonLucy.py new file mode 100644 index 000000000..bfabce458 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/RichardsonLucy.py @@ -0,0 +1,170 @@ +import numpy as np +import logging +logger = logging.getLogger(__name__) + +from histpy import Histogram + +from ..utils import _to_float +from .RichardsonLucyBasic import RichardsonLucyBasic + +from ..constants import DEFAULT_BKG_NORM_RANGE + +class RichardsonLucy(RichardsonLucyBasic): + """ + Standard Richardson-Lucy algorithm with background optimization. + + This class extends RichardsonLucyBasic with background normalization + optimization, making it suitable for practical image reconstruction + with real data. + + Features: + - E-step + M-step (from RichardsonLucyBasic) + - Background normalization optimization + + For advanced features (response weighting, acceleration, smoothing), + use RichardsonLucyAdvanced. + For MAP estimation with priors, use MAP_RichardsonLucy. + + An example of parameter is as follows. + + iteration_max: 100 + minimum_flux: + value: 0.0 + unit: "cm-2 s-1 sr-1" + background_normalization_optimization: + activate: True + range: {"albedo": [0.9, 1.1]} + save_results: + activate: True + directory: "./results" + only_final_result: True + """ + + def __init__(self, initial_model, dataset, mask, parameter): + + super().__init__(initial_model, dataset, mask, parameter) + + # background normalization optimization + self.bkg_norm_optimization_enabled = parameter.get('background_normalization_optimization:activate', False) + if self.bkg_norm_optimization_enabled: + self.dict_delta_bkg_norm = {} + self.dict_bkg_norm_range = parameter.get('background_normalization_optimization:range', {key: DEFAULT_BKG_NORM_RANGE for key in self.dict_bkg_norm.keys()}) + + # update parameter summary + self._parameter_summary += [ + ("bkg_norm_optimization_enabled", self.bkg_norm_optimization_enabled), + ] + if self.bkg_norm_optimization_enabled: + self._parameter_summary += [ + ("dict_bkg_norm_range", self.dict_bkg_norm_range), + ] + + def initialization(self): + """ + initialization before running the image deconvolution + """ + + super().initialization() + + # calculate summed background models for M-step + if self.bkg_norm_optimization_enabled: + self.dict_summed_bkg_model = {} + for key in self.dict_bkg_norm.keys(): + self.dict_summed_bkg_model[key] = self.dataset.calc_summed_bkg_model(key) + + def Mstep(self): + """ + M-step in RL algorithm. + In this step, self.delta_model and self.delta_bkg_norm will be updated. + """ + + ratio_list = [ data.event / expectation for data, expectation in zip(self.dataset, self.expectation_list) ] + + # delta model + sum_T_product = self.dataset.calc_summed_T_product(ratio_list) + self.delta_model = self.model * (sum_T_product/self.summed_exposure_map - 1) + + # masking + if self.mask is not None: + self.delta_model = self.delta_model.mask_pixels(self.mask) + + logger.debug("The delta model was updated.") + + # background normalization optimization + if self.bkg_norm_optimization_enabled: + for key in self.dict_bkg_norm.keys(): + + sum_bkg_T_product = self.dataset.calc_summed_bkg_model_product(key, ratio_list) + sum_bkg_model = self.dict_summed_bkg_model[key] + + self.dict_delta_bkg_norm[key] = self.dict_bkg_norm[key] * (sum_bkg_T_product / sum_bkg_model - 1) + + def _ensure_bkg_norm_range(self): + """ + Ensure background normalization is within allowed range. + + This method clips background normalization values to their + allowed ranges. If a value is outside the range, it is + set to the nearest boundary value. + """ + + for key in self.dict_bkg_norm.keys(): + bkg_norm = self.dict_bkg_norm[key] + bkg_range = self.dict_bkg_norm_range[key] + + if bkg_norm < bkg_range[0]: + bkg_norm = bkg_range[0] + elif bkg_norm > bkg_range[1]: + bkg_norm = bkg_range[1] + + self.dict_bkg_norm[key] = _to_float(bkg_norm) + + def post_processing(self): + """ + Post-processing. + """ + + # updating model + self.model[:] += self.delta_model.contents + self._ensure_model_constraints() + + # update background normalization + if self.bkg_norm_optimization_enabled: + for key in self.dict_bkg_norm.keys(): + self.dict_bkg_norm[key] += self.dict_delta_bkg_norm[key] + self._ensure_bkg_norm_range() + + def register_result(self): + """ + Register results at the end of each iteration. + """ + + this_result = {"iteration": self.iteration_count, + "model": self.model.copy(), + "background_normalization": self.dict_bkg_norm.copy()} + + # show intermediate results + for key in ["background_normalization"]: + logger.info(f"{key}: {this_result[key]}") + + # register this_result in self.results + self.results.append(this_result) + + def finalization(self): + """ + finalization after running the image deconvolution + """ + + if not self.save_results: + return + + logger.info(f"Saving results in {self.save_results_directory}") + + self._save_standard_results( + counter_name = "iteration", + histogram_keys = [("model", f"{self.save_results_directory}/model.hdf5", self.save_only_final_result)], + fits_filename = f"{self.save_results_directory}/results.fits", + values_key_name_format = [], + dicts_key_name_format = [("background_normalization", "BKG_NORM", "D")], + lists_key_name_format = [], + ) diff --git a/cosipy/image_deconvolution/algorithms/RichardsonLucyAdvanced.py b/cosipy/image_deconvolution/algorithms/RichardsonLucyAdvanced.py new file mode 100644 index 000000000..0adc5317d --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/RichardsonLucyAdvanced.py @@ -0,0 +1,307 @@ +import numpy as np +import astropy.units as u +import logging +logger = logging.getLogger(__name__) + +from yayc import Configurator +from histpy import Histogram +from copy import deepcopy + +from .RichardsonLucy import RichardsonLucy +from .accelerators.build_accelerator import build_accelerator +from .accelerators.accelerator_base import EMStepResult + +from .response_weighting_filter import ResponseWeightingFilter + +from ..constants import DEFAULT_STOPPING_THRESHOLD, DEFAULT_RESPONSE_WEIGHTING_INDEX + +class RichardsonLucyAdvanced(RichardsonLucy): + """ + A class for the RichardsonLucy algorithm. + The algorithm here is based on Knoedlseder+99, Knoedlseder+05, Siegert+20. + + An example of parameter is as follows. + + iteration_max: 100 + minimum_flux: + value: 0.0 + unit: "cm-2 s-1 sr-1" + acceleration: + activate: True + accel_factor_max: 10.0 + response_weighting: + activate: True + index: 0.5 + smoothing: + activate: True + FWHM: + value: 2.0 + unit: "deg" + stopping_criteria: + statistics: "log-likelihood" + threshold: 1e-2 + background_normalization_optimization: + activate: True + range: {"albedo": [0.9, 1.1]} + save_results: + activate: True + directory: "./results" + only_final_result: True + """ + + def __init__(self, initial_model, dataset, mask, parameter): + + super().__init__(initial_model, dataset, mask, parameter) + + # acceleration + self.acceleration_enabled = parameter.get('acceleration:activate', False) + if self.acceleration_enabled: + self.accelerator_parameter = Configurator(parameter['acceleration']) + + # response_weighting + self.response_weighting_enabled = parameter.get('response_weighting:activate', False) + if self.response_weighting_enabled: + self.response_weighting_index = parameter.get('response_weighting:index', DEFAULT_RESPONSE_WEIGHTING_INDEX) + + # smoothing + self.smoothing_enabled = parameter.get('smoothing:activate', False) + if self.smoothing_enabled: + self.smoothing_fwhm = parameter.get('smoothing:FWHM:value') * u.Unit(parameter.get('smoothing:FWHM:unit')) + + # stopping criteria + self.stopping_criteria_statistics = parameter.get('stopping_criteria:statistics', "log-likelihood") + self.stopping_criteria_threshold = parameter.get('stopping_criteria:threshold', DEFAULT_STOPPING_THRESHOLD) + + if not self.stopping_criteria_statistics in ["log-likelihood"]: + raise ValueError + + # update parameter summary + self._parameter_summary += [ + ("acceleration_enabled", self.acceleration_enabled), + ] + if self.acceleration_enabled: + self._parameter_summary += [ + ("accelerator_parameter", parameter['acceleration']), + ] + + self._parameter_summary += [ + ("response_weighting_enabled", self.response_weighting_enabled), + ] + if self.response_weighting_enabled: + self._parameter_summary += [ + ("response_weighting_index", self.response_weighting_index), + ] + + self._parameter_summary += [ + ("smoothing_enabled", self.smoothing_enabled), + ] + if self.smoothing_enabled: + self._parameter_summary += [ + ("smoothing_fwhm", self.smoothing_fwhm), + ] + + self._parameter_summary += [ + ("stopping_criteria_statistics", self.stopping_criteria_statistics), + ("stopping_criteria_threshold", self.stopping_criteria_threshold), + ] + + def initialization(self): + """ + Initialize before running image deconvolution. + + This method sets up response weighting filter based on the exposure map. + """ + + super().initialization() + + # response-weighting filter + if self.response_weighting_enabled: + self.response_weighting_filter = ResponseWeightingFilter(self.summed_exposure_map, self.response_weighting_index) + + # build accelerator + if self.acceleration_enabled: + self.accelerator = build_accelerator(self.accelerator_parameter) + + def pre_processing(self): + """ + pre-processing for each iteration + """ + + if self.iteration_count == 1: + self.Estep() + logger.info("The expected count histograms were calculated with the initial model map.") + + def processing_core(self): + """ + Core processing for each iteration. + """ + + self.Mstep() + # Note that Estep() is performed in self.post_processing(). + + def Mstep(self): + """ + Mstep + """ + super().Mstep() + + # apply response_weighting_filter + if self.response_weighting_enabled: + self.delta_model = self.response_weighting_filter.apply(self.delta_model) + + # apply smoothing + if self.smoothing_enabled: + self.delta_model = self.delta_model.smoothing(fwhm = self.smoothing_fwhm) + + def post_processing(self): + """ + perform the acceleration of RL algirithm + """ + + if self.acceleration_enabled: + + # em_results[0]: "before" state (before any update) + em_results = [ + EMStepResult( + model = self.model.copy(), + dict_bkg_norm = self.dict_bkg_norm.copy(), + expectation_list = deepcopy(self.expectation_list), + source_expectation_list = deepcopy(self.source_expectation_list), + bkg_expectation_list = deepcopy(self.bkg_expectation_list), + ), + ] + + # Run n_em_steps_required EM steps, appending each result + for _ in range(self.accelerator.n_em_steps_required): + self.model += self.delta_model + if self.bkg_norm_optimization_enabled: + self.dict_bkg_norm = { + key: self.dict_bkg_norm[key] + self.dict_delta_bkg_norm[key] + for key in self.dict_bkg_norm + } + self.Estep() # updates self.source_expectation_list / bkg_expectation_list + if _ < self.accelerator.n_em_steps_required - 1: + self.Mstep() # need delta_model for next iteration (not needed on last step) + em_results.append(EMStepResult( + model = self.model.copy(), + dict_bkg_norm = self.dict_bkg_norm.copy(), + expectation_list = deepcopy(self.expectation_list), + source_expectation_list = deepcopy(self.source_expectation_list), + bkg_expectation_list = deepcopy(self.bkg_expectation_list), + )) + + result = self.accelerator.compute( + em_results = em_results, + dataset = self.dataset, + mask = self.mask, + ) + self.model = result.model + self.dict_bkg_norm = result.dict_bkg_norm + self._accel_result = result + + else: + self.model += self.delta_model + if self.bkg_norm_optimization_enabled: + self.dict_bkg_norm = { + key: self.dict_bkg_norm[key] + self.dict_delta_bkg_norm[key] + for key in self.dict_bkg_norm + } + self._accel_result = None + + self._ensure_model_constraints() + if self.bkg_norm_optimization_enabled: + self._ensure_bkg_norm_range() + + # always recompute expectation and LH after _ensure_model_constraints + self.Estep() + logger.debug("Expected count histograms updated.") + self.log_likelihood_list = self.dataset.calc_log_likelihood_list(self.expectation_list) + logger.debug("Log-likelihood list updated.") + + def register_result(self): + """ + The values below are stored at the end of each iteration. + - iteration: iteration number + - model: updated image + - delta_model: delta map after M-step + - processed_delta_model: delta map after post-processing + - accel_factor: acceleration parameter in RL algirithm + - background_normalization: optimized background normalization + - log-likelihood: log-likelihood + """ + + this_result = {"iteration": self.iteration_count, + "model": self.model.copy(), + "background_normalization": self.dict_bkg_norm.copy(), + "log-likelihood": self.log_likelihood_list.copy()} + + if self._accel_result is not None and self._accel_result.extras is not None: + this_result.update(self._accel_result.extras) + + for key in ["background_normalization", "log-likelihood"]: + logger.info(f"{key}: {this_result[key]}") + + if self.acceleration_enabled: + for key, _ in self.accelerator.logged_result_fields: + if key in this_result: + logger.info(f"{key}: {this_result[key]}") + + # register this_result in self.results + self.results.append(this_result) + + def check_stopping_criteria(self): + """ + If iteration_count is smaller than iteration_max, the iterative process will continue. + + Returns + ------- + bool + """ + + if self.iteration_count >= self.iteration_max: + return True + + if self.iteration_count == 1: + return False # need at least 2 results to compute delta + + if self.stopping_criteria_statistics == "log-likelihood": + + log_likelihood = np.sum(self.results[-1]["log-likelihood"]) + log_likelihood_before = np.sum(self.results[-2]["log-likelihood"]) + delta_log_likelihood = log_likelihood - log_likelihood_before + + logger.debug(f'Delta log-likelihood: {delta_log_likelihood}') + + if delta_log_likelihood < 0: + logger.warning(f"Log-likelihood decreased {delta_log_likelihood}. Reconstruction may be unstable.") + return False + + elif log_likelihood - log_likelihood_before < self.stopping_criteria_threshold: + return True + + return False + + def finalization(self): + """ + finalization after running the image deconvolution + """ + + if not self.save_results: + return + + logger.info(f"Saving results in {self.save_results_directory}") + + values_key_name_format = [] + if self.acceleration_enabled: + for key, fits_fmt in self.accelerator.logged_result_fields: + if key in self.results[0]: + values_key_name_format.append((key, key.upper(), fits_fmt)) + + self._save_standard_results( + counter_name = "iteration", + histogram_keys = [("model", f"{self.save_results_directory}/model.hdf5", self.save_only_final_result)], + fits_filename = f"{self.save_results_directory}/results.fits", + values_key_name_format = values_key_name_format, + dicts_key_name_format = [("background_normalization", "BKG_NORM", "D")], + lists_key_name_format = [("log-likelihood", "LOG-LIKELIHOOD", "D")], + ) diff --git a/cosipy/image_deconvolution/algorithms/RichardsonLucyBasic.py b/cosipy/image_deconvolution/algorithms/RichardsonLucyBasic.py new file mode 100644 index 000000000..afe5370c0 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/RichardsonLucyBasic.py @@ -0,0 +1,210 @@ +import os +import numpy as np +import astropy.units as u +import logging +logger = logging.getLogger(__name__) + +from histpy import Histogram + +from .deconvolution_algorithm_base import DeconvolutionAlgorithmBase +from ..constants import DEFAULT_MINIMUM_FLUX + +class RichardsonLucyBasic(DeconvolutionAlgorithmBase): + """ + Basic Richardson-Lucy algorithm. + + This class implements the core RL algorithm with minimal parameters. + It is designed as an entry point to understand the fundamental EM structure of the Richardson-Lucy algorithm. + + Features: + - E-step + M-step + - Minimum flux constraint + - Masking + + An example of parameter is as follows. + + iteration_max: 100 + minimum_flux: + value: 0.0 + unit: "cm-2 s-1 sr-1" + save_results: + activate: True + directory: "./results" + only_final_result: True + """ + + def __init__(self, initial_model, dataset, mask, parameter): + + super().__init__(initial_model, dataset, mask, parameter) + + # minimum flux + self.minimum_flux = parameter.get('minimum_flux:value', DEFAULT_MINIMUM_FLUX) + + minimum_flux_unit = parameter.get('minimum_flux:unit', initial_model.unit) + if minimum_flux_unit is not None: + self.minimum_flux = self.minimum_flux*u.Unit(minimum_flux_unit) + + # saving results + self.save_results = parameter.get('save_results:activate', False) + self.save_results_directory = parameter.get('save_results:directory', './results') + self.save_only_final_result = parameter.get('save_results:only_final_result', False) + + if self.save_results is True: + if os.path.isdir(self.save_results_directory): + logger.warning(f"A directory {self.save_results_directory} already exists. Files in {self.save_results_directory} may be overwritten. Make sure that is not a problem.") + else: + os.makedirs(self.save_results_directory) + + # update parameter summary + self._parameter_summary += [ + ("minimum_flux", self.minimum_flux), + ("save_results", self.save_results), + ] + if self.save_results: + self._parameter_summary += [ + ("save_results_directory", self.save_results_directory), + ("save_only_final_result", self.save_only_final_result), + ] + + def initialization(self): + """ + initialization before running the image deconvolution + """ + + # show parameters + self._show_parameters() + + # clear counter + self.iteration_count = 0 + + # clear results + self.results.clear() + + # copy model + self.model = self.initial_model.copy() + + # calculate exposure map + self.summed_exposure_map = self.dataset.calc_summed_exposure_map() + + # mask setting + if self.mask is None and np.any(self.summed_exposure_map.contents == 0): + self.mask = Histogram(self.model.axes, + contents = self.summed_exposure_map.contents > 0, + copy_contents = False) + self.model = self.model.mask_pixels(self.mask) + logger.info("There are zero-exposure pixels. A mask to ignore them was set.") + + def pre_processing(self): + """ + pre-processing for each iteration + """ + + pass + + def processing_core(self): + """ + Core processing for each iteration. + """ + + self.Estep() + self.Mstep() + + def Estep(self): + """ + E-step. + In this step, self.expectation_list will be updated. + """ + + self.source_expectation_list = self.dataset.calc_source_expectation_list(self.model) + self.bkg_expectation_list = self.dataset.calc_bkg_expectation_list(self.dict_bkg_norm) + self.expectation_list = self.dataset.combine_expectation_list(self.source_expectation_list, self.bkg_expectation_list) + + logger.debug("The expected count histograms were updated.") + + def Mstep(self): + """ + M-step. + In this step, self.delta_model will be updated. + """ + + ratio_list = [ data.event / expectation for data, expectation in zip(self.dataset, self.expectation_list) ] + + # delta model + sum_T_product = self.dataset.calc_summed_T_product(ratio_list) + self.delta_model = self.model * (sum_T_product/self.summed_exposure_map - 1) + + # masking + if self.mask is not None: + self.delta_model = self.delta_model.mask_pixels(self.mask) + + logger.debug("The delta model was updated.") + + def _ensure_model_constraints(self): + """ + Ensure model satisfies physical constraints. + + This method enforces: + 1. Minimum flux constraint (non-negative flux) + 2. Masking (zero-exposure pixels) + """ + + # checking minimum flux + self.model[:] = np.where(self.model.contents < self.minimum_flux, self.minimum_flux, self.model.contents) + + # masking again + if self.mask is not None: + self.model = self.model.mask_pixels(self.mask) + + def post_processing(self): + """ + Post-processing. + """ + + # updating model + self.model[:] += self.delta_model.contents + + self._ensure_model_constraints() + + def register_result(self): + """ + Register results at the end of each iteration. + """ + + this_result = {"iteration": self.iteration_count, + "model": self.model.copy()} + + # register this_result in self.results + self.results.append(this_result) + + def check_stopping_criteria(self): + """ + If iteration_count is smaller than iteration_max, the iterative process will continue. + + Returns + ------- + bool + """ + + if self.iteration_count >= self.iteration_max: + return True + + return False + + def finalization(self): + """ + finalization after running the image deconvolution + """ + + if not self.save_results: + return + + logger.info(f"Saving results in {self.save_results_directory}") + + self._save_standard_results( + counter_name = "iteration", + histogram_keys = [("model", f"{self.save_results_directory}/model.hdf5", self.save_only_final_result)], + fits_filename = f"{self.save_results_directory}/results.fits", + values_key_name_format = [], + dicts_key_name_format = [], + lists_key_name_format = [], + ) diff --git a/cosipy/image_deconvolution/algorithms/__init__.py b/cosipy/image_deconvolution/algorithms/__init__.py new file mode 100644 index 000000000..70ce1aa76 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/__init__.py @@ -0,0 +1,9 @@ +from .deconvolution_algorithm_base import DeconvolutionAlgorithmBase +from .RichardsonLucyBasic import RichardsonLucyBasic +from .RichardsonLucy import RichardsonLucy +from .RichardsonLucyAdvanced import RichardsonLucyAdvanced +from .MAP_RichardsonLucy import MAP_RichardsonLucy +from .prior_base import PriorBase +from .prior_entropy import PriorEntropy +from .prior_tsv import PriorTSV +from .response_weighting_filter import ResponseWeightingFilter diff --git a/cosipy/image_deconvolution/algorithms/accelerators/__init__.py b/cosipy/image_deconvolution/algorithms/accelerators/__init__.py new file mode 100644 index 000000000..ed62c96b9 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/accelerators/__init__.py @@ -0,0 +1,4 @@ +from .accelerator_base import AcceleratorBase, AcceleratorResult, EMStepResult +from .build_accelerator import build_accelerator +from .max_step_accelerator import MaxStepAccelerator +from .line_search_accelerator import LineSearchAccelerator diff --git a/cosipy/image_deconvolution/algorithms/accelerators/accelerator_base.py b/cosipy/image_deconvolution/algorithms/accelerators/accelerator_base.py new file mode 100644 index 000000000..ee3120006 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/accelerators/accelerator_base.py @@ -0,0 +1,126 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Union +import numpy as np +import logging +logger = logging.getLogger(__name__) + +from histpy import Histogram + +from ...utils import _to_float + +@dataclass +class EMStepResult: + """ + Holds the state and results of one EM step (or the initial "before" state). + + em_results passed to AcceleratorBase.compute() is a list of EMStepResult: + - em_results[0] : "before" state + - em_results[1] : result of 1st EM step + - em_results[2] : result of 2nd EM step (e.g. SQUAREM) + - ... + """ + + model : Histogram + dict_bkg_norm : dict + expectation_list : Union[list, None] = None + source_expectation_list : Union[list, None] = None + bkg_expectation_list : Union[list, None] = None + + +@dataclass +class AcceleratorResult: + """ + Return value of AcceleratorBase.compute(). + """ + + model : Histogram + dict_bkg_norm : dict + extras : Union[dict, None] = None + + +class AcceleratorBase(ABC): + + n_em_steps_required : int = 1 + + # Fields from AcceleratorResult.extras to be saved in results and logged. + # Each entry is a (field_name, fits_format) tuple, e.g. [("accel_factor", "D")]. + # Algorithm classes (e.g. RichardsonLucyAdvanced) iterate over this list in + # register_result and finalization, so accelerator-specific keys never need + # to be hardcoded in the algorithm. + logged_result_fields: list = [] + + def __init__(self, parameter): + self._parameter = parameter + + @abstractmethod + def compute( + self, + em_results : list, + dataset, + mask, + ) -> AcceleratorResult: + raise NotImplementedError + + def _compute_accel_factor_max(self, delta_model, model, mask) -> float: + """ + Compute the maximum allowed accel_factor such that + model + accel_factor * delta_model >= 0 element-wise. + + Parameters + ---------- + delta_model : Histogram + delta of model. + model : Histogram + model before the EM step. + mask : Histogram or None + + Returns + ------- + float + Maximum safe accel_factor + """ + + diff = -1 * (model / delta_model).contents + diff[(diff <= 0) | (delta_model.contents == 0)] = np.inf + + if mask is not None: + diff[np.invert(mask.contents)] = np.inf + + accel_factor = _to_float(np.min(diff)) + + if accel_factor < 1.0: + accel_factor = 1.0 + + return accel_factor + + def _compute_accel_factor_bkg_max(self, delta_bkg_norm, bkg_norm) -> float: + """ + Compute the maximum allowed accel_factor_bkg such that + bkg_norm + accel_factor_bkg * delta_bkg_norm >= 0 for all keys. + + For each key, the constraint is active only when delta_bkg_norm < 0. + The returned value is the minimum across all keys, so a single + accel_factor_bkg is guaranteed to be safe for all background components. + + Parameters + ---------- + delta_bkg_norm : dict + Per-key delta of background normalization. + bkg_norm : dict + Per-key background normalization before the EM step. + + Returns + ------- + float + Maximum safe accel_factor_bkg + """ + + accel_factor_bkg_max = np.inf + + for key in delta_bkg_norm: + d = delta_bkg_norm[key] + if d < 0: + accel_factor_bkg_max = min(accel_factor_bkg_max, -bkg_norm[key] / d) + + return max(1.0, _to_float(accel_factor_bkg_max)) diff --git a/cosipy/image_deconvolution/algorithms/accelerators/build_accelerator.py b/cosipy/image_deconvolution/algorithms/accelerators/build_accelerator.py new file mode 100644 index 000000000..2dcf1ab0e --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/accelerators/build_accelerator.py @@ -0,0 +1,61 @@ +""" +Factory that maps algorithm names to AcceleratorBase subclasses. + +Adding a new accelerator +------------------------ +1. Implement a subclass of AcceleratorBase. +2. Import it here. +3. Add the name -> class mapping to accelerator_classes. +""" + +import logging + +from .accelerator_base import AcceleratorBase +from .max_step_accelerator import MaxStepAccelerator +from .line_search_accelerator import LineSearchAccelerator + +logger = logging.getLogger(__name__) + +ACCELERATOR_CLASSES : dict = { + "MaxStep" : MaxStepAccelerator, + "LineSearch": LineSearchAccelerator, +} + +DEFAULT_ACCELERATOR = "MaxStep" + + +def build_accelerator(parameter) -> AcceleratorBase | None: + """ + Construct an AcceleratorBase from the ``acceleration`` block. + + Parameters + ---------- + parameter : dict-like + activate: true + algorithm: MaxStep # optional, default MaxStep + accel_factor_max: 10.0 + + Returns + ------- + AcceleratorBase + + Raises + ------ + ValueError + When the requested algorithm name is not in ACCELERATOR_CLASSES. + """ + + algorithm_name = parameter.get("algorithm", DEFAULT_ACCELERATOR) + + if algorithm_name not in ACCELERATOR_CLASSES: + available = ", ".join(ACCELERATOR_CLASSES.keys()) + raise ValueError( + f'Unknown accelerator "{algorithm_name}". Available: {available}' + ) + + cls = ACCELERATOR_CLASSES[algorithm_name] + accelerator = cls(parameter) + + logger.info(f"[Accelerator '{algorithm_name}' created]") + return accelerator + diff --git a/cosipy/image_deconvolution/algorithms/accelerators/line_search_accelerator.py b/cosipy/image_deconvolution/algorithms/accelerators/line_search_accelerator.py new file mode 100644 index 000000000..ad218a246 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/accelerators/line_search_accelerator.py @@ -0,0 +1,291 @@ +""" +Acceleration via line search: finds the accel_factor in [1, accel_factor_max] +that maximises the log-likelihood. + +For the model-only (1-D) case, Brent's method (scipy.optimize.minimize_scalar) +is used. When background normalization is optimised independently (2-D), either +a gradient-based method (L-BFGS-B) or a grid search can be selected via the +``search_method`` parameter. + +YAML configuration example +--------------------------- +acceleration: + activate: true + algorithm: LineSearch + accel_factor_max: 10.0 # hard upper bound for model accel_factor + accel_bkg_norm: + activate: false # whether to accelerate bkg_norm at all + independent: false # if true, bkg gets its own accel_factor + max: 10.0 # hard upper bound for bkg accel_factor (independent mode only) + search_method: gradient # "gradient" or "grid" (independent mode only) + grid_n: 10 # grid points per axis (grid mode only) +""" + +import numpy as np +import logging +from scipy.optimize import minimize_scalar, minimize + +from .accelerator_base import AcceleratorBase, AcceleratorResult + +logger = logging.getLogger(__name__) + +DEFAULT_ACCEL_FACTOR_MAX = 10.0 +DEFAULT_GRID_N = 10 + + +class LineSearchAccelerator(AcceleratorBase): + """ + RL acceleration by line search over accel_factor. + + Three bkg_norm modes are supported: + + accel_bkg_norm.activate=False (default) + bkg_norm is NOT accelerated; only the model is scaled. + Line search is 1-D (Brent's method). + + accel_bkg_norm.activate=True, independent=False + bkg_norm is scaled by the same accel_factor as the model. + Line search is still 1-D (Brent's method). + + accel_bkg_norm.activate=True, independent=True + bkg_norm gets its own accel_factor optimised independently. + Line search becomes 2-D. Two search methods are available: + - "gradient": L-BFGS-B (fast, may miss global optimum) + - "grid" : exhaustive grid search over grid_n x grid_n points + (slower but more robust) + + In all cases the search is bounded by accel_factor_max (and + accel_bkg_norm.max for the independent bkg case). + """ + + n_em_steps_required = 1 + logged_result_fields = [("accel_factor", "D"), ("accel_factor_bkg", "D")] + + def __init__(self, parameter): + super().__init__(parameter) + + self.accel_factor_max = float(parameter.get("accel_factor_max", DEFAULT_ACCEL_FACTOR_MAX)) + self.accel_bkg_norm = bool(parameter.get("accel_bkg_norm:activate", False)) + self.accel_bkg_norm_independent = bool(parameter.get("accel_bkg_norm:independent", False)) + self.accel_factor_bkg_max = float(parameter.get("accel_bkg_norm:max", DEFAULT_ACCEL_FACTOR_MAX)) + self.bkg_search_method = parameter.get("accel_bkg_norm:search_method", "gradient") + self.bkg_grid_n = int(parameter.get("accel_bkg_norm:grid_n", DEFAULT_GRID_N)) + + if self.bkg_search_method not in ("gradient", "grid"): + raise ValueError( + f'Unknown search_method "{self.bkg_search_method}". ' + f'Available: "gradient", "grid"' + ) + + logger.info( + f"[LineSearchAccelerator]" + f"\n accel_factor_max: {self.accel_factor_max}" + f"\n accel_bkg_norm: {self.accel_bkg_norm}" + f"\n accel_bkg_norm_independent: {self.accel_bkg_norm_independent}" + f"\n accel_factor_bkg_max: {self.accel_factor_bkg_max}" + + (f"\n search_method: {self.bkg_search_method}" + f"\n grid_n: {self.bkg_grid_n}" + if self.accel_bkg_norm and self.accel_bkg_norm_independent else "") + ) + + def compute(self, em_results, dataset, mask) -> AcceleratorResult: + + before = em_results[0] + after = em_results[1] + + delta_model = after.model - before.model + delta_bkg_norm = {key: after.dict_bkg_norm[key] - before.dict_bkg_norm[key] + for key in before.dict_bkg_norm} + + # Upper bound from non-negativity constraint on model + alpha_max = min(self._compute_accel_factor_max(delta_model, before.model, mask), + self.accel_factor_max) + + # Upper bound from non-negativity constraint on bkg_norm + alpha_bkg_max = min(self._compute_accel_factor_bkg_max(delta_bkg_norm, before.dict_bkg_norm), + self.accel_factor_bkg_max) + + if alpha_max <= 1.0: + # No room to accelerate + return AcceleratorResult( + model = after.model, + dict_bkg_norm = after.dict_bkg_norm, + extras = {"accel_factor": 1.0, "accel_factor_bkg": 1.0}, + ) + + # --- define log-likelihood as a function of accel factors --- + + def _ll_1d(alpha): + """LH as a function of a single accel_factor (model, and optionally bkg).""" + new_model = before.model + delta_model * alpha + + if self.accel_bkg_norm and not self.accel_bkg_norm_independent: + new_dict_bkg_norm = {key: before.dict_bkg_norm[key] + delta_bkg_norm[key] * alpha + for key in before.dict_bkg_norm} + else: + new_dict_bkg_norm = after.dict_bkg_norm + + src_exp_list = [ + src_before + (src_after - src_before) * alpha + for src_before, src_after + in zip(before.source_expectation_list, after.source_expectation_list) + ] + bkg_exp_list = dataset.calc_bkg_expectation_list(new_dict_bkg_norm) + exp_list = dataset.combine_expectation_list(src_exp_list, bkg_exp_list) + return float(np.sum(dataset.calc_log_likelihood_list(exp_list))) + + def _ll_2d(alpha_m, alpha_b): + """LH as a function of independent model and bkg accel_factors.""" + new_model = before.model + delta_model * alpha_m + new_dict_bkg_norm = {key: before.dict_bkg_norm[key] + delta_bkg_norm[key] * alpha_b + for key in before.dict_bkg_norm} + src_exp_list = [ + src_before + (src_after - src_before) * alpha_m + for src_before, src_after + in zip(before.source_expectation_list, after.source_expectation_list) + ] + bkg_exp_list = dataset.calc_bkg_expectation_list(new_dict_bkg_norm) + exp_list = dataset.combine_expectation_list(src_exp_list, bkg_exp_list) + return float(np.sum(dataset.calc_log_likelihood_list(exp_list))) + + # --- run line search --- + + ll_after = _ll_1d(1.0) + + if self.accel_bkg_norm and self.accel_bkg_norm_independent: + # 2-D search + accel_factor, accel_factor_bkg = self._search_2d(_ll_2d, alpha_max, alpha_bkg_max) + ll_opt = _ll_2d(accel_factor, accel_factor_bkg) + new_dict_bkg_norm = {key: before.dict_bkg_norm[key] + delta_bkg_norm[key] * accel_factor_bkg + for key in before.dict_bkg_norm} + else: + # 1-D optimisation (Brent's method) + # In same-factor mode, bkg is also constrained, so use the tighter bound + alpha_1d_max = min(alpha_max, alpha_bkg_max) if self.accel_bkg_norm else alpha_max + + opt = minimize_scalar( + lambda a: -_ll_1d(a), + bounds=(1.0, alpha_1d_max), + method="bounded", + ) + accel_factor = float(opt.x) + accel_factor_bkg = accel_factor if (self.accel_bkg_norm and not self.accel_bkg_norm_independent) else 1.0 + ll_opt = -float(opt.fun) + + if self.accel_bkg_norm and not self.accel_bkg_norm_independent: + new_dict_bkg_norm = {key: before.dict_bkg_norm[key] + delta_bkg_norm[key] * accel_factor + for key in before.dict_bkg_norm} + else: + new_dict_bkg_norm = after.dict_bkg_norm + + # Fall back to accel_factor=1 if line search did not improve LH + if ll_opt < ll_after: + logger.debug( + f"[LineSearchAccelerator] line search (alpha={accel_factor:.3f}, alpha_bkg={accel_factor_bkg:.3f}) " + f"did not improve LH ({ll_opt:.6f} < {ll_after:.6f}). Falling back to accel_factor=1." + ) + accel_factor = 1.0 + accel_factor_bkg = 1.0 + new_model = after.model + new_dict_bkg_norm = after.dict_bkg_norm + else: + new_model = before.model + delta_model * accel_factor + + logger.debug(f"[LineSearchAccelerator] accel_factor={accel_factor:.4f}, accel_factor_bkg={accel_factor_bkg:.4f}") + + return AcceleratorResult( + model = new_model, + dict_bkg_norm = new_dict_bkg_norm, + extras = {"accel_factor": accel_factor, "accel_factor_bkg": accel_factor_bkg}, + ) + + def _search_2d(self, ll_2d_func, alpha_max, alpha_bkg_max): + """ + 2-D search for the optimal (accel_factor, accel_factor_bkg). + + Parameters + ---------- + ll_2d_func : callable + Function of (alpha_m, alpha_b) returning log-likelihood. + alpha_max : float + Upper bound for model accel_factor. + alpha_bkg_max : float + Upper bound for bkg accel_factor. + + Returns + ------- + (float, float) + Optimal (accel_factor, accel_factor_bkg). + """ + + if self.bkg_search_method == "grid": + return self._grid_search_2d(ll_2d_func, alpha_max, alpha_bkg_max) + else: + # gradient (L-BFGS-B) + + # 1. Optimize source to get a good initial point + opt_s = minimize_scalar( + lambda s: -ll_2d_func(s, 1.0), + bounds=(1.0, alpha_max), + method="bounded", + ) + alpha_init = float(opt_s.x) + + # 2. Optimize bkg to get a good initial point + opt_b = minimize_scalar( + lambda b: -ll_2d_func(1.0, b), + bounds=(1.0, alpha_bkg_max), + method="bounded", + ) + alpha_bkg_init = float(opt_b.x) + + # 3. Optimize both with both initial points + opt = minimize( + lambda x: -ll_2d_func(x[0], x[1]), + x0=[alpha_init, alpha_bkg_init], + bounds=[(1.0, alpha_max), (1.0, alpha_bkg_max)], + method="L-BFGS-B", + ) + + accel_factor = float(opt.x[0]) + accel_factor_bkg = float(opt.x[1]) + + return float(opt.x[0]), float(opt.x[1]) + + def _grid_search_2d(self, ll_2d_func, alpha_max, alpha_bkg_max): + """ + Exhaustive grid search over [1, alpha_max] x [1, alpha_bkg_max]. + + Parameters + ---------- + ll_2d_func : callable + alpha_max : float + alpha_bkg_max : float + + Returns + ------- + (float, float) + Optimal (accel_factor, accel_factor_bkg). + """ + + alphas_model = np.linspace(1.0, alpha_max, self.bkg_grid_n) + alphas_bkg = np.linspace(1.0, alpha_bkg_max, self.bkg_grid_n) + + best_ll = -np.inf + best_alpha_m = 1.0 + best_alpha_b = 1.0 + + for am in alphas_model: + for ab in alphas_bkg: + ll = ll_2d_func(am, ab) + if ll > best_ll: + best_ll = ll + best_alpha_m = am + best_alpha_b = ab + + logger.debug( + f"[LineSearchAccelerator] grid search best: " + f"alpha={best_alpha_m:.4f}, alpha_bkg={best_alpha_b:.4f}, ll={best_ll:.6f}" + ) + + return best_alpha_m, best_alpha_b diff --git a/cosipy/image_deconvolution/algorithms/accelerators/max_step_accelerator.py b/cosipy/image_deconvolution/algorithms/accelerators/max_step_accelerator.py new file mode 100644 index 000000000..370fe1756 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/accelerators/max_step_accelerator.py @@ -0,0 +1,115 @@ +""" +Acceleration via the maximum safe step: finds the largest accel_factor in +[1, accel_factor_max] such that the updated model remains non-negative +element-wise, then accepts it only if it improves the log-likelihood. + +YAML configuration example +--------------------------- +acceleration: + activate: true + algorithm: MaxStep + accel_factor_max: 10.0 # hard upper bound for accel_factor + accel_bkg_norm: false # if true, bkg_norm is scaled by the same accel_factor +""" + +import numpy as np +import logging + +from .accelerator_base import AcceleratorBase, AcceleratorResult + +from ...utils import _to_float + +logger = logging.getLogger(__name__) + +DEFAULT_ACCEL_FACTOR_MAX = 10.0 + + +class MaxStepAccelerator(AcceleratorBase): + """ + Finds the largest accel_factor in [1, accel_factor_max] such that + + model_before + accel_factor * delta_model >= 0 (element-wise) + + If accel_factor does not improve LH, falls back to accel_factor=1. + accel_factor is stored in extras={"accel_factor": accel_factor}. + + algorithm: MaxStep + accel_factor_max: 10.0 + accel_bkg_norm: False + """ + + n_em_steps_required = 1 + logged_result_fields = [("accel_factor", "D")] + + def __init__(self, parameter): + super().__init__(parameter) + self.accel_factor_max = float(parameter.get("accel_factor_max", DEFAULT_ACCEL_FACTOR_MAX)) + self.accel_bkg_norm = bool(parameter.get("accel_bkg_norm", False)) + logger.info( + f"[MaxStepAccelerator]" + f"\n accel_factor_max: {self.accel_factor_max}" + f"\n accel_bkg_norm: {self.accel_bkg_norm}" + ) + + def compute( + self, + em_results : list, + dataset, + mask, + ) -> AcceleratorResult: + + before = em_results[0] + after = em_results[1] + + delta_model = after.model - before.model + + accel_factor = self._compute_accel_factor(delta_model, before.model, mask) + + if accel_factor > 1.0: + new_model = before.model + delta_model * accel_factor + if self.accel_bkg_norm: + new_dict_bkg_norm = { + key: before.dict_bkg_norm[key] + (after.dict_bkg_norm[key] - before.dict_bkg_norm[key]) * accel_factor + for key in before.dict_bkg_norm + } + else: + new_dict_bkg_norm = after.dict_bkg_norm + + # LH check: reuse source expectation, recompute bkg only + source_expectation_list = [ + src_before + (src_after - src_before) * accel_factor + for src_before, src_after in zip(before.source_expectation_list, after.source_expectation_list) + ] + bkg_expectation_list = dataset.calc_bkg_expectation_list(new_dict_bkg_norm) + expectation_list = dataset.combine_expectation_list(source_expectation_list, bkg_expectation_list) + + ll_accel = np.sum(dataset.calc_log_likelihood_list(expectation_list)) + ll_after = np.sum(dataset.calc_log_likelihood_list(after.expectation_list)) + + if ll_accel < ll_after: + logger.debug( + f"[MaxStepAccelerator] accel_factor={accel_factor:.3f} did not improve LH " + f"({ll_accel:.6f} < {ll_after:.6f}). Falling back to accel_factor=1." + ) + accel_factor = 1.0 + new_model = after.model + new_dict_bkg_norm = after.dict_bkg_norm + + else: + new_model = after.model + new_dict_bkg_norm = after.dict_bkg_norm + + logger.debug(f"[MaxStepAccelerator] accel_factor={accel_factor}") + + return AcceleratorResult( + model = new_model, + dict_bkg_norm = new_dict_bkg_norm, + extras = {"accel_factor": accel_factor}, + ) + + def _compute_accel_factor(self, delta_model, model, mask) -> float: + + accel_factor = min(self._compute_accel_factor_max(delta_model, model, mask), \ + self.accel_factor_max) + + return accel_factor diff --git a/cosipy/image_deconvolution/deconvolution_algorithm_base.py b/cosipy/image_deconvolution/algorithms/deconvolution_algorithm_base.py similarity index 50% rename from cosipy/image_deconvolution/deconvolution_algorithm_base.py rename to cosipy/image_deconvolution/algorithms/deconvolution_algorithm_base.py index 5c0574131..e416c2829 100644 --- a/cosipy/image_deconvolution/deconvolution_algorithm_base.py +++ b/cosipy/image_deconvolution/algorithms/deconvolution_algorithm_base.py @@ -1,11 +1,13 @@ import numpy as np -import astropy.units as u import astropy.io.fits as fits import functools from abc import ABC, abstractmethod import logging logger = logging.getLogger(__name__) +from ..constants import CHUNK_SIZE_FITS, DEFAULT_ITERATION_MAX +from ..utils import _to_float + class DeconvolutionAlgorithmBase(ABC): """ A base class for image deconvolution algorithms. @@ -13,8 +15,7 @@ class DeconvolutionAlgorithmBase(ABC): - initialization - pre_processing - - Estep - - Mstep + - processing_core - post_processing - register_result - check_stopping_criteria @@ -26,11 +27,10 @@ class DeconvolutionAlgorithmBase(ABC): Attributes ---------- initial_model: :py:class:`cosipy.image_deconvolution.ModelBase` or its subclass - dataset: list of :py:class:`cosipy.image_deconvolution.ImageDeconvolutionDataInterfaceBase` or its subclass + dataset: :py:class:`cosipy.image_deconvolution.DataInterfaceCollection` parameter : py:class:`yayc.Configurator` results: list of results dict_bkg_norm: the dictionary of background normalizations - dict_dataset_indexlist_for_bkg_models: the indices of data corresponding to each background model in the dataset """ def __init__(self, initial_model, dataset, mask, parameter): @@ -42,32 +42,18 @@ def __init__(self, initial_model, dataset, mask, parameter): self.results = [] # background normalization - self.dict_bkg_norm = {} - self.dict_dataset_indexlist_for_bkg_models = {} - for data in self.dataset: - for key in data.keys_bkg_models(): - if not key in self.dict_bkg_norm.keys(): - self.dict_bkg_norm[key] = 1.0 - self.dict_dataset_indexlist_for_bkg_models[key] = [] - - for key in self.dict_dataset_indexlist_for_bkg_models.keys(): - for index, data in enumerate(self.dataset): - if key in data.keys_bkg_models(): - self.dict_dataset_indexlist_for_bkg_models[key].append(index) + self.dict_bkg_norm = {key: 1.0 for key in dataset.keys_bkg_models()} logger.debug(f'dict_bkg_norm: {self.dict_bkg_norm}') - logger.debug(f'dict_dataset_indexlist_for_bkg_models: {self.dict_dataset_indexlist_for_bkg_models}') - - # minimum flux - self.minimum_flux = parameter.get('minimum_flux:value', 0.0) - - minimum_flux_unit = parameter.get('minimum_flux:unit', initial_model.unit) - if minimum_flux_unit is not None: - self.minimum_flux = self.minimum_flux*u.Unit(minimum_flux_unit) # parameters of the iteration self.iteration_count = 0 - self.iteration_max = parameter.get('iteration_max', 1) + self.iteration_max = parameter.get('iteration_max', DEFAULT_ITERATION_MAX) + + # parameter summary (subclasses can append to this list) + self._parameter_summary = [ + ("iteration_max", self.iteration_max), + ] @abstractmethod def initialization(self): @@ -83,23 +69,31 @@ def pre_processing(self): """ raise NotImplementedError + @abstractmethod + def processing_core(self): + """ + Core processing for each iteration. + + This method should implement the main algorithm logic. + For EM-based algorithms, this typically includes: + - E-step: Expectation calculation + - M-step: Maximization/update + + Subclasses must implement this method. + """ + raise NotImplementedError + @abstractmethod def Estep(self): """ - E-step. - In this step, only self.expectation_list should be updated. - If it will be performed in other step, typically post_processing() or check_stopping_criteria(), - this step can be skipped. + E-step: Expectation calculation """ raise NotImplementedError @abstractmethod def Mstep(self): """ - M-step. - In this step, only self.delta_model should be updated. - If you want to apply some operations to self.delta_model, - these should be performed in post_processing(). + M-step: Maximization/update """ raise NotImplementedError @@ -136,8 +130,6 @@ def finalization(self): """ raise NotImplementedError -### A subclass should not override the methods below. ### - def iteration(self): """ Perform one iteration of image deconvolution. @@ -150,11 +142,8 @@ def iteration(self): logger.info("<< Pre-processing >>") self.pre_processing() - logger.info("<< E-step >>") - self.Estep() - - logger.info("<< M-step >>") - self.Mstep() + logger.info("<< Processing Core>>") + self.processing_core() logger.info("<< Post-processing >>") self.post_processing() @@ -164,131 +153,10 @@ def iteration(self): logger.info("<< Checking Stopping Criteria >>") stop_iteration = self.check_stopping_criteria() - logger.info("--> {}".format("Stop" if stop_iteration else "Continue")) + logger.info("-> {}".format("Stop" if stop_iteration else "Continue")) return stop_iteration - def calc_expectation_list(self, model, dict_bkg_norm = None, almost_zero = 1e-12): - """ - Calculate a list of expected count histograms corresponding to each data in the registered dataset. - - Parameters - ---------- - model: :py:class:`cosipy.image_deconvolution.ModelBase` or its subclass - Model - dict_bkg_norm : dict, default None - background normalization for each background model, e.g, {'albedo': 0.95, 'activation': 1.05} - almost_zero : float, default 1e-12 - In order to avoid zero components in extended count histogram, a tiny offset is introduced. - It should be small enough not to effect statistics. - - Returns - ------- - list of :py:class:`histpy.Histogram` - List of expected count histograms - """ - - return [data.calc_expectation(model, dict_bkg_norm = dict_bkg_norm, almost_zero = almost_zero) for data in self.dataset] - - def calc_log_likelihood_list(self, expectation_list): - """ - Calculate a list of log-likelihood from each data in the registered dataset and the corresponding given expected count histogram. - - Parameters - ---------- - expectation_list : list of :py:class:`histpy.Histogram` - List of expected count histograms - - Returns - ------- - list of float - List of Log-likelihood - """ - - return [data.calc_log_likelihood(expectation) for data, expectation in zip(self.dataset, expectation_list)] - - def calc_summed_exposure_map(self): - """ - Calculate a list of exposure maps from the registered dataset. - - Returns - ------- - :py:class:`histpy.Histogram` - """ - - return self._histogram_sum([ data.exposure_map for data in self.dataset ]) - - def calc_summed_bkg_model(self, key): - """ - Calculate the sum of histograms for a given background model in the registered dataset. - - Parameters - ---------- - key: str - Background model name - - Returns - ------- - float - """ - - indexlist = self.dict_dataset_indexlist_for_bkg_models[key] - - return sum([self.dataset[i].summed_bkg_model(key) for i in indexlist]) - - def calc_summed_T_product(self, dataspace_histogram_list): # dataspace_histogram_list = ratio_list = d_i/E_i - """ - For each data in the registered dataset, the product of the corresponding input histogram with the transpose of the response function is computed. - Then, this method returns the sum of all of the products. - - Parameters - ---------- - dataspace_histogram_list: list of :py:class:`histpy.Histogram` - - Returns - ------- - :py:class:`histpy.Histogram` - """ - - return self._histogram_sum([data.calc_T_product(hist) - for data, hist in zip(self.dataset, dataspace_histogram_list)]) - - def calc_summed_bkg_model_product(self, key, dataspace_histogram_list): - """ - For each data in the registered dataset, the product of the corresponding input histogram with the specified background model is computed. - Then, this method returns the sum of all of the products. - - Parameters - ---------- - key: str - Background model name - dataspace_histogram_list: list of :py:class:`histpy.Histogram` - - Returns - ------- - flaot - """ - - indexlist = self.dict_dataset_indexlist_for_bkg_models[key] - - return sum( - self.dataset[i].calc_bkg_model_product(key = key, dataspace_histogram = dataspace_histogram_list[i]) - for i in indexlist - ) - - @staticmethod - def _histogram_sum(hlist): - """ - Sum a list of Histograms. If only one input, just return it. - """ - if len(hlist) == 1: - return hlist[0] - else: - result = hlist[0].copy() - for h in hlist[1:]: - result += h - return result - def save_histogram(self, filename, counter_name, histogram_key, only_final_result = False): # save last result @@ -331,7 +199,7 @@ def save_results_as_fits(self, filename, counter_name, values_key_name_format, d dict_keys = list(self.results[0][key].keys()) - chunk_size = 998 # when the number of columns >= 1000, the fits file may not be saved. + chunk_size = CHUNK_SIZE_FITS # when the number of columns >= 1000, the fits file may not be saved. for i_chunk, chunked_dict_keys in enumerate([dict_keys[i:i+chunk_size] for i in range(0, len(dict_keys), chunk_size)]): cols_dict = [fits.Column(name=dict_key, array=[result[key][dict_key] for result in self.results], format=fits_format) for dict_key in chunked_dict_keys] @@ -358,3 +226,61 @@ def save_results_as_fits(self, filename, counter_name, values_key_name_format, d # write fits.HDUList(hdu_list).writeto(filename, overwrite=True) + + def _save_standard_results(self, counter_name, histogram_keys, fits_filename, + values_key_name_format=None, dicts_key_name_format=None, lists_key_name_format=None): + """ + Save standard results including histograms and FITS files. + + Parameters + ---------- + counter_name : str + Name of the counter (e.g., "iteration") + histogram_keys : list of tuple + List of (key, filename, only_final_result) for histograms to save. + fits_filename : str + Path to the FITS file. + values_key_name_format : list of tuple, optional + List of (key, name, fits_format) for single values to save in FITS. + dicts_key_name_format : list of tuple, optional + List of (key, name, fits_format) for dictionaries to save in FITS. + lists_key_name_format : list of tuple, optional + List of (key, name, fits_format) for lists to save in FITS. + """ + # Save histograms + for key, filename, only_final_result in histogram_keys: + self.save_histogram( + filename = filename, + counter_name = counter_name, + histogram_key = key, + only_final_result = only_final_result + ) + + # Save FITS file (use default if not specified) + self.save_results_as_fits( + filename = fits_filename, + counter_name = counter_name, + values_key_name_format = values_key_name_format if values_key_name_format is not None else [], + dicts_key_name_format = dicts_key_name_format if dicts_key_name_format is not None else [], + lists_key_name_format = lists_key_name_format if lists_key_name_format is not None else [] + ) + + def _show_parameters(self): + """ + Log all parameters registered in _parameter_summary. + Subclasses should append their parameters to self._parameter_summary + in their __init__, then call this method at the end of __init__. + + Example + ------- + # In a subclass __init__: + self._parameter_summary += [ + ("bkg_norm_optimization_enabled", self.bkg_norm_optimization_enabled), + ("minimum_flux", self.minimum_flux), + ] + self._show_parameters() + """ + + logger.info(f"[{self.__class__.__name__}]") + for name, value in self._parameter_summary: + logger.info(f" {name}: {value}") diff --git a/cosipy/image_deconvolution/prior_base.py b/cosipy/image_deconvolution/algorithms/prior_base.py similarity index 73% rename from cosipy/image_deconvolution/prior_base.py rename to cosipy/image_deconvolution/algorithms/prior_base.py index 8859c67fb..8bd5f6a98 100644 --- a/cosipy/image_deconvolution/prior_base.py +++ b/cosipy/image_deconvolution/algorithms/prior_base.py @@ -10,7 +10,7 @@ class PriorBase: Parameters ---------- parameter : dict - parameters for the prior probability. + Parameters for the prior probability. model : object of a subclass of :py:class:`cosipy.image_deconvolution.ModelBase` Model object to which the prior will be applied. @@ -22,6 +22,28 @@ class PriorBase: Scaling coefficient for the prior probability. model_class : type Type of the model being used. + + Notes + ----- + All prior subclasses must implement: + + - ``log_prior(model)`` : returns the scalar log prior value + - ``grad_log_prior(model)`` : returns the gradient of the log prior + with respect to the model, with units + inverse to the model + + The ``coefficient`` (:math:`c`) controls the strength of the prior. + A larger value enforces the prior more strongly relative to the likelihood. + Setting ``coefficient = 0`` effectively disables the prior. + + **YAML parameter block (common to all priors)** + + .. code-block:: yaml + + prior: + : + coefficient: 1.0 # regularization strength c (dimensionless) + # ... prior-specific parameters follow """ usable_model_classes = [] @@ -90,6 +112,7 @@ def grad_log_prior(self, model): numpy.ndarray Gradient of the log prior. Its shape must be the same as the input model. + Its unit must be the inverse of the model unit. """ raise NotImplementedError diff --git a/cosipy/image_deconvolution/algorithms/prior_entropy.py b/cosipy/image_deconvolution/algorithms/prior_entropy.py new file mode 100644 index 000000000..1c2b4d177 --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/prior_entropy.py @@ -0,0 +1,112 @@ +import astropy.units as u +import healpy as hp +import numpy as np + +from .prior_base import PriorBase + +from ..models.allskyimage import AllSkyImageModel + +class PriorEntropy(PriorBase): + """ + Maximum Entropy prior for all-sky image models. + + This prior encodes a preference for solutions that are as smooth and + featureless as possible relative to a reference map, penalizing + deviation of the reconstructed image from the reference. + + Parameters + ---------- + parameter : dict + Parameters for the entropy prior. + model : AllSkyImageModel + All-sky image model to which the prior will be applied. + + Attributes + ---------- + reference_map : astropy.units.Quantity + The reference (default) image :math:`m` against which entropy is measured. + Typically set to a flat map representing the prior expectation of the image. + + Notes + ----- + **Mathematical definition** + + The log maximum entropy prior is defined as: + + .. math:: + + \\log P_{\\mathrm{ME}}(\\lambda) = + c_{\\mathrm{ME}} \\sum_{i} \\lambda_i \\left(1 - \\log\\frac{\\lambda_i}{m_i}\\right) + + where :math:`\lambda_i` is the flux in pixel :math:`i`, + :math:`m_i` is the reference map value in pixel :math:`i`, + and :math:`c_{\mathrm{ME}}` is the regularization coefficient. + + The gradient with respect to :math:`\lambda_i` is: + + .. math:: + + \\frac{\\partial \\log P_{\\mathrm{ME}}}{\\partial \\lambda_i} = + -c_{\\mathrm{ME}} \\log\\frac{\\lambda_i}{m_i} + + **YAML parameter block** + + .. code-block:: yaml + + prior: + entropy: + coefficient: 1.0 # regularization strength lambda (dimensionless) + reference_map: + value: 1.0e-4 # reference map value (uniform map) + unit: "cm-2 s-1 sr-1" + """ + + usable_model_classes = [AllSkyImageModel] + + def __init__(self, parameter, model): + + super().__init__(parameter, model) + + self.reference_map = self.parameter['reference_map']['value'] * u.Unit(self.parameter['reference_map']['unit']) + + def log_prior(self, model): + """ + Calculate the logarithm of the entropy prior probability. + + Parameters + ---------- + model : AllSkyImageModel + Model for which to calculate the log prior. + + Returns + ------- + float + The logarithm of the entropy prior probability. + """ + + if self.model_class == AllSkyImageModel: + + image_ratio = (model/self.reference_map).contents.to('').value + + return self.coefficient * np.sum(model.contents.value * (1 - np.log(image_ratio))) + + def grad_log_prior(self, model): + """ + Calculate the gradient of the log entropy prior. + + Parameters + ---------- + model : AllSkyImageModel + Model for which to calculate the gradient. + + Returns + ------- + numpy.ndarray + Gradient of the log prior, in units inverse to the model. + """ + + if self.model_class == AllSkyImageModel: + + image_ratio = (model/self.reference_map).contents.to('').value + + return -1 * self.coefficient * np.log(image_ratio) / model.unit diff --git a/cosipy/image_deconvolution/prior_tsv.py b/cosipy/image_deconvolution/algorithms/prior_tsv.py similarity index 66% rename from cosipy/image_deconvolution/prior_tsv.py rename to cosipy/image_deconvolution/algorithms/prior_tsv.py index 47745e94a..afe72ea77 100644 --- a/cosipy/image_deconvolution/prior_tsv.py +++ b/cosipy/image_deconvolution/algorithms/prior_tsv.py @@ -3,18 +3,19 @@ from .prior_base import PriorBase -from .allskyimage import AllSkyImageModel +from ..models.allskyimage import AllSkyImageModel class PriorTSV(PriorBase): """ Total Squared Variation (TSV) prior for all-sky image models. - This prior implements a smoothness constraint by penalizing differences between neighboring pixels. + This prior implements a smoothness constraint by penalizing squared + differences between neighboring pixels. Parameters ---------- - parameter: dict - parameters for the TSV prior. + parameter : dict + Parameters for the TSV prior. model : AllSkyImageModel All-sky image model to which the prior will be applied. @@ -23,12 +24,43 @@ class PriorTSV(PriorBase): usable_model_classes : list List containing AllSkyImageModel as the only compatible model class. neighbour_pixel_index : numpy.ndarray - Array of shape (8, npix) containing indices of neighboring pixels. - Some pixels have only 7 neighboring pixels. In this case, healpy returns -1 - as the index of a neighboring pixel, but it can cause calculation errors in - this code. So, such a pixel index is replaced with its own pixel index. + Array of shape (-1, npix) containing indices of neighboring pixels. + For the case of HealPix, some pixels have only 7 neighboring pixels. + In this case, healpy returns -1 as the index of a neighboring pixel, + but it can cause calculation errors in this code. So, such a pixel + index is replaced with its own pixel index. num_neighbour_pixels : numpy.ndarray Array of shape (npix,) containing the number of valid neighbors for each pixel. + + Notes + ----- + **Mathematical definition** + + The log TSV prior is defined as: + + .. math:: + + \\log P_{\\mathrm{TSV}}(\\lambda) = + -c_{\\mathrm{TSV}} \\sum_{i} \\sum_{j \\in \\sigma(i)} (\\lambda_i - \\lambda_j)^2 + + where :math:`\\lambda_i` is the flux in pixel :math:`i`, + :math:`\\sigma(i)` is the set of neighboring pixels of pixel :math:`i`, + and :math:`c_{\mathrm{TSV}}` is the regularization coefficient (``coefficient`` in YAML). + + The gradient with respect to :math:`\\lambda_i` is: + + .. math:: + + \\frac{\\partial \\log P_{\\mathrm{TSV}}}{\\partial \\lambda_i} = + -4c_{\\mathrm{TSV}} \\sum_{j \\in \\sigma(i)} (\\lambda_i - \\lambda_j) + + **YAML parameter block** + + .. code-block:: yaml + + prior: + TSV: + coefficient: 1.0e+6 # regularization strength lambda (dimensionless) """ usable_model_classes = [AllSkyImageModel] @@ -67,6 +99,7 @@ def log_prior(self, model): float The logarithm of the TSV prior probability. """ + if self.model_class == AllSkyImageModel: diff = (model[:] - model[self.neighbour_pixel_index]).value @@ -88,6 +121,7 @@ def grad_log_prior(self, model): numpy.ndarray Gradient of the log prior, in units inverse to the model. """ + if self.model_class == AllSkyImageModel: diff = (model[:] - model[self.neighbour_pixel_index]).value diff --git a/cosipy/image_deconvolution/algorithms/response_weighting_filter.py b/cosipy/image_deconvolution/algorithms/response_weighting_filter.py new file mode 100644 index 000000000..f3adc363d --- /dev/null +++ b/cosipy/image_deconvolution/algorithms/response_weighting_filter.py @@ -0,0 +1,60 @@ +import numpy as np + +import logging +logger = logging.getLogger(__name__) + +class ResponseWeightingFilter: + """ + Response weighting filter. + + This class calculates and stores a response weighting filter + based on the exposure map. The filter renormalizes the delta map + so that pixels with large exposure times have more feedback. + """ + + def __init__(self, exposure_map, index=0.5): + """ + Initialize response weighting filter. + + Parameters + ---------- + exposure_map : Histogram + Summed exposure map + index : float, optional + Response weighting index (default: 0.5) + """ + + self.exposure_map = exposure_map + self.index = index + self.filter = self._calculate_filter() + + logger.info(f"[Response weighting filter created (index={index})]") + + def _calculate_filter(self): + """ + Calculate the filter array. + """ + + max_exposure = np.max(self.exposure_map.contents) + if max_exposure == 0: + logger.warning("Maximum exposure is zero") + return np.ones_like(self.exposure_map.contents) + + return (self.exposure_map.contents / max_exposure)**self.index + + def apply(self, delta_model): + """ + Apply filter to delta model. + + Parameters + ---------- + delta_model : Histogram + Delta model to be weighted + + Returns + ------- + Histogram + Weighted delta model (delta_model * filter) + """ + + return delta_model * self.filter diff --git a/cosipy/image_deconvolution/constants.py b/cosipy/image_deconvolution/constants.py new file mode 100644 index 000000000..8a811ed5b --- /dev/null +++ b/cosipy/image_deconvolution/constants.py @@ -0,0 +1,23 @@ +""" +Constants for the image deconvolution module. +This module centralizes magic numbers and default values. +""" +import astropy.units as u + +# Numerical Constants +NUMERICAL_ZERO = 1e-12 # Small value to avoid division by zero in expectation calculations. +CHUNK_SIZE_FITS = 998 # Maximum columns in FITS table (FITS limit is 1000, using 998 for safety). + +# Physical Constants +EARTH_RADIUS_KM = 6378.0 # Earth radius in km + +# Default Values - General Deconvolution Algorithm Parameters +DEFAULT_MINIMUM_FLUX = 0.0 # Default minimum flux to enforce non-negativity. +DEFAULT_ITERATION_MAX = 1 # Default maximum number of iterations. +DEFAULT_STOPPING_THRESHOLD = 1e-2 # Default convergence threshold for log-likelihood change. + +# Default Values - Background Normalization +DEFAULT_BKG_NORM_RANGE = [0.0, 100.0] # Default allowed range [min, max] for background normalization factors. + +# Default Values - Response Weighting +DEFAULT_RESPONSE_WEIGHTING_INDEX = 0.5 # Default power index for response weighting: filter = (exposure/max)^index. diff --git a/cosipy/image_deconvolution/data_interfaces/__init__.py b/cosipy/image_deconvolution/data_interfaces/__init__.py new file mode 100644 index 000000000..0b4c70fd5 --- /dev/null +++ b/cosipy/image_deconvolution/data_interfaces/__init__.py @@ -0,0 +1,4 @@ +from .image_deconvolution_data_interface_base import ImageDeconvolutionDataInterfaceBase +from .data_interface_collection import DataInterfaceCollection +from .dataIF_COSI_DC2 import DataIF_COSI_DC2 +from .dataIF_Parallel import DataIF_Parallel diff --git a/cosipy/image_deconvolution/dataIF_COSI_DC2.py b/cosipy/image_deconvolution/data_interfaces/dataIF_COSI_DC2.py similarity index 93% rename from cosipy/image_deconvolution/dataIF_COSI_DC2.py rename to cosipy/image_deconvolution/data_interfaces/dataIF_COSI_DC2.py index dd942afdb..542bcf881 100644 --- a/cosipy/image_deconvolution/dataIF_COSI_DC2.py +++ b/cosipy/image_deconvolution/data_interfaces/dataIF_COSI_DC2.py @@ -12,16 +12,8 @@ from cosipy.response import FullDetectorResponse from .image_deconvolution_data_interface_base import ImageDeconvolutionDataInterfaceBase - -def tensordot_sparse(A, A_unit, B, axes): - """ - perform a tensordot operation on A and B. A is sparse - and so does not carry a unit; rather it must be passed - as a separate argument. B is a normal Quantity. Return - a proper Quantity as the result. - """ - dotprod = np.tensordot(A, B.value, axes=axes) - return u.Quantity(dotprod, unit= A_unit * B.unit, copy=False) +from .utils import tensordot_sparse +from ..constants import NUMERICAL_ZERO class DataIF_COSI_DC2(ImageDeconvolutionDataInterfaceBase): @@ -214,7 +206,7 @@ def _calc_exposure_map(self): logger.info("Finished...") - def calc_expectation(self, model, dict_bkg_norm = None, almost_zero = 1e-12): + def calc_source_expectation(self, model): """ Calculate expected counts from a given model. @@ -222,11 +214,6 @@ def calc_expectation(self, model, dict_bkg_norm = None, almost_zero = 1e-12): ---------- model : :py:class:`cosipy.image_deconvolution.AllSkyImageModel` Model map - dict_bkg_norm : dict, default None - background normalization for each background model, e.g, {'albedo': 0.95, 'activation': 1.05} - almost_zero : float, default 1e-12 - In order to avoid zero components in extended count histogram, a tiny offset is introduced. - It should be small enough not to effect statistics. Returns ------- @@ -257,11 +244,30 @@ def calc_expectation(self, model, dict_bkg_norm = None, almost_zero = 1e-12): # [Time/ScAtt, NuLambda, Ei] x [NuLambda, Ei, Em, Phi, PsiChi] -> [Time/ScAtt, Em, Phi, PsiChi] expectation *= model.axes['lb'].pixarea() - expectation += almost_zero + expectation += NUMERICAL_ZERO + + return Histogram(self.data_axes, contents = expectation, copy_contents = False) + + def calc_bkg_expectation(self, dict_bkg_norm): + """ + Calculate expected counts from a given background normalizations. + + Parameters + ---------- + dict_bkg_norm : dict, default None + background normalization for each background model, e.g, {'albedo': 0.95, 'activation': 1.05} - if dict_bkg_norm is not None: - for key in self.keys_bkg_models(): - expectation += self.bkg_model(key).contents * dict_bkg_norm[key] + Returns + ------- + :py:class:`histpy.Histogram` + Expected count histogram + + Notes + ----- + This method should be implemented in a more general class, for example, extended source response class in the future. + """ + + expectation = sum(self.bkg_model(key).contents * dict_bkg_norm[key] for key in self.keys_bkg_models()) return Histogram(self.data_axes, contents = expectation, copy_contents = False) @@ -301,9 +307,6 @@ def calc_T_product(self, dataspace_histogram): return Histogram(self.model_axes, contents = tprod, copy_contents = False) - - return hist - def calc_bkg_model_product(self, key, dataspace_histogram): """ Calculate the product of the input histogram with the background model. diff --git a/cosipy/image_deconvolution/dataIF_Parallel.py b/cosipy/image_deconvolution/data_interfaces/dataIF_Parallel.py similarity index 84% rename from cosipy/image_deconvolution/dataIF_Parallel.py rename to cosipy/image_deconvolution/data_interfaces/dataIF_Parallel.py index 91b593721..637f16c71 100644 --- a/cosipy/image_deconvolution/dataIF_Parallel.py +++ b/cosipy/image_deconvolution/data_interfaces/dataIF_Parallel.py @@ -18,7 +18,8 @@ from histpy import Histogram, Axes, Axis, HealpixAxis from cosipy.response import FullDetectorResponse -from cosipy.image_deconvolution import ImageDeconvolutionDataInterfaceBase +from .image_deconvolution_data_interface_base import ImageDeconvolutionDataInterfaceBase +from ..constants import NUMERICAL_ZERO def load_response_matrix(comm, start_col, end_col, filename): ''' @@ -255,33 +256,47 @@ def _calc_exposure_map(self): logger.info("Finished...") - def calc_expectation(self, model, dict_bkg_norm = None, almost_zero = 1e-12): + def _allgatherv_slice(self, expectation_slice): """ - Calculate expected counts from a given model. + Gather expectation slices from all MPI processes into a full Histogram. Parameters ---------- - model : :py:class:`cosipy.image_deconvolution.AllSkyImageModel` - Model map - dict_bkg_norm : dict, default None - background normalization for each background model, e.g, {'albedo': 0.95, 'activation': 1.05} - almost_zero : float, default 1e-12 - In order to avoid zero components in extended count histogram, a tiny offset is introduced. - It should be small enough not to effect statistics. + expectation_slice : Histogram + Local slice computed by this process. Returns ------- - :py:class:`histpy.Histogram` - Expected count histogram + Histogram + Full expectation over all processes (serial: same as input). + """ + if not self.parallel: + return expectation_slice + + all_sizes = [(self.averow) * self.row_size] * (self.numtasks - 1) \ + + [(self.averow + self.extra_rows) * self.row_size] + displacements = np.insert(np.cumsum(all_sizes), 0, 0)[:-1] + total_size = int(np.sum(all_sizes)) + + recvbuf = np.empty(total_size, dtype=np.float64) + self._comm.Allgatherv(expectation_slice.contents.flatten(), [recvbuf, all_sizes, displacements, MPI.DOUBLE]) + + epsilon = np.concatenate([recvbuf[displacements[i]:displacements[i] + all_sizes[i]].reshape(expectation_slice.contents.shape[:-1] + (-1,)) for i in range(self.numtasks)], axis=-1) + + return Histogram(self.event.axes, contents=epsilon, unit=self.event.unit) + + def calc_source_expectation(self, model): + """ + Calculate expected counts from the source model only. - Notes - ----- - This method should be implemented in a more general class, for example, extended source response class in the future. + Parameters + ---------- + model : AllSkyImageModel + + Returns + ------- + histpy.Histogram """ - # 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. expectation_slice = Histogram(self._data_axes_slice) @@ -298,40 +313,29 @@ def calc_expectation(self, model, dict_bkg_norm = None, almost_zero = 1e-12): expectation_slice[:] = np.tensordot( map_rotated, self._image_response.contents, axes = ([1,2], [0,1])) # [Time/ScAtt, NuLambda, Ei] x [NuLambda, Ei, Em, Phi, PsiChi] -> [Time/ScAtt, Em, Phi, PsiChi] - if dict_bkg_norm is not None: - for key in self.keys_bkg_models(): - expectation_slice += self.bkg_model_slice(key) * dict_bkg_norm[key] + expectation_slice += NUMERICAL_ZERO - expectation_slice += almost_zero + return self._allgatherv_slice(expectation_slice) - # Parallel - if self.parallel: - ''' - Synchronization Barrier 1 - ''' - - # Calculate all_sizes and displacements - all_sizes = [(self.averow) * self.row_size] * (self.numtasks-1) + [(self.averow + self.extra_rows) * self.row_size] - displacements = np.insert(np.cumsum(all_sizes), 0, 0)[0:-1] - - # Create a buffer to receive the gathered data - total_size = int(np.sum(all_sizes)) - recvbuf = np.empty(total_size, dtype=np.float64) # Receive buffer + def calc_bkg_expectation(self, dict_bkg_norm): + """ + Calculate expected counts from background models only. - # Gather all arrays into recvbuf - self._comm.Allgatherv(expectation_slice.contents.flatten(), [recvbuf, all_sizes, displacements, MPI.DOUBLE]) # For multiple MPI processes, full = [slice1, ... sliceN] + Parameters + ---------- + dict_bkg_norm : dict - # Reshape the received buffer back into the original 3D array shape - epsilon = np.concatenate([ recvbuf[displacements[i]:displacements[i] + all_sizes[i]].reshape(expectation_slice.contents.shape[:-1] + (-1,)) for i in range(self.numtasks) ], axis=-1) + Returns + ------- + histpy.Histogram + """ - # Add to list that manages multiple datasets - expectation = Histogram(self.event.axes, contents=epsilon, unit=self.event.unit) + expectation_slice = Histogram(self._data_axes_slice) - # Serial - else: - expectation = expectation_slice # If single process, then full = slice + for key in self.keys_bkg_models(): + expectation_slice += self.bkg_model_slice(key) * dict_bkg_norm[key] - return expectation + return self._allgatherv_slice(expectation_slice) def calc_T_product(self, dataspace_histogram): """ diff --git a/cosipy/image_deconvolution/data_interfaces/data_interface_collection.py b/cosipy/image_deconvolution/data_interfaces/data_interface_collection.py new file mode 100644 index 000000000..de601f179 --- /dev/null +++ b/cosipy/image_deconvolution/data_interfaces/data_interface_collection.py @@ -0,0 +1,236 @@ +""" +Provides DataInterfaceCollection: a class that encapsulates a list of +ImageDeconvolutionDataInterfaceBase objects and exposes the collective +operations (expectation, log-likelihood, T-product, background operations) +that were previously scattered as methods of DeconvolutionAlgorithmBase. +""" + +from collections.abc import Sequence +import numpy as np +import logging +from typing import Union + +logger = logging.getLogger(__name__) + + +from ..utils import _to_float + + +class DataInterfaceCollection(Sequence): + """ + Wraps a list of :py:class:`ImageDeconvolutionDataInterfaceBase` objects + and provides collective statistical/response operations. + + Parameters + ---------- + dataset : list of ImageDeconvolutionDataInterfaceBase + The registered dataset. + + Attributes + ---------- + dataset : list + The underlying list of data interface objects. + """ + + def __init__(self, dataset: list): + self._dataset = dataset + + self._dict_dataset_indexlist_for_bkg_models = {} + + for idx, data in enumerate(self._dataset): + for key in data.keys_bkg_models(): + if key in self._dict_dataset_indexlist_for_bkg_models: + self._dict_dataset_indexlist_for_bkg_models[key].append(idx) + else: + self._dict_dataset_indexlist_for_bkg_models[key] = [idx] + + # Basic access + @property + def dataset(self) -> list: + return self._dataset + + def keys_bkg_models(self) -> list: + return list(self._dict_dataset_indexlist_for_bkg_models.keys()) + + def __len__(self) -> int: + return len(self._dataset) + + def __getitem__(self, idx): + return self._dataset[idx] + + + # Expectation + def calc_expectation_list(self, model, dict_bkg_norm = Union[dict,None]) -> list: + """ + Calculate a list of expected count histograms corresponding to each data in the registered dataset. + + Parameters + ---------- + model: :py:class:`cosipy.image_deconvolution.ModelBase` or its subclass + Model + dict_bkg_norm : dict, default None + background normalization for each background model, e.g, {'albedo': 0.95, 'activation': 1.05} + + Returns + ------- + list of :py:class:`histpy.Histogram` + List of expected count histograms + """ + + return [ + data.calc_expectation(model, dict_bkg_norm) + for data in self._dataset + ] + + def calc_source_expectation_list(self, model) -> list: + """ + Compute source-only expected counts for every dataset entry. + + Returns + ------- + list of histpy.Histogram + """ + return [ + data.calc_source_expectation(model) + for data in self._dataset + ] + + def calc_bkg_expectation_list(self, dict_bkg_norm) -> list: + """ + Compute background-only expected counts for every dataset entry. + + Returns + ------- + list of histpy.Histogram + """ + return [ + data.calc_bkg_expectation(dict_bkg_norm) + for data in self._dataset + ] + + def combine_expectation_list(self, expectation_list_src, expectation_list_bkg) -> list: + """ + Sum source and background expectation lists element-wise. + + Returns + ------- + list of histpy.Histogram + """ + return [ + src + bkg + for src, bkg in zip(expectation_list_src, expectation_list_bkg) + ] + + # Log-likelihood + def calc_log_likelihood_list(self, expectation_list: list) -> list: + """ + Calculate a list of log-likelihood from each data in the registered dataset and the corresponding given expected count histogram. + + Parameters + ---------- + expectation_list : list of :py:class:`histpy.Histogram` + List of expected count histograms + + Returns + ------- + list of float + List of Log-likelihood + """ + + return [_to_float(data.calc_log_likelihood(expectation)) for data, expectation in zip(self._dataset, expectation_list)] + + def calc_total_log_likelihood(self, expectation_list: list) -> float: + """ + Convenience: sum of all per-dataset log-likelihoods. + """ + + return float(np.sum(self.calc_log_likelihood_list(expectation_list))) + + # Response-transpose products + def calc_summed_T_product(self, dataspace_histogram_list: list): + """ + For each data in the registered dataset, the product of the corresponding input histogram with the transpose of the response function is computed. + Then, this method returns the sum of all of the products. + + Parameters + ---------- + dataspace_histogram_list: list of :py:class:`histpy.Histogram` + + Returns + ------- + :py:class:`histpy.Histogram` + """ + + return self._histogram_sum([data.calc_T_product(hist) + for data, hist in zip(self._dataset, dataspace_histogram_list)]) + + # Exposure map + def calc_summed_exposure_map(self): + """ + Calculate a list of exposure maps from the registered dataset. + + Returns + ------- + :py:class:`histpy.Histogram` + """ + + return self._histogram_sum([data.exposure_map for data in self._dataset]) + + def calc_summed_bkg_model(self, key: str) -> float: + """ + Calculate the sum of histograms for a given background model in the registered dataset. + + Parameters + ---------- + key: str + Background model name + + Returns + ------- + float + """ + + indexlist = self._dict_dataset_indexlist_for_bkg_models[key] + + return sum([self._dataset[i].summed_bkg_model(key) for i in indexlist]) + + def calc_summed_bkg_model_product(self, key: str, dataspace_histogram_list: list) -> float: + """ + For each data in the registered dataset, the product of the corresponding input histogram with the specified background model is computed. + Then, this method returns the sum of all of the products. + + Parameters + ---------- + key: str + Background model name + dataspace_histogram_list: list of :py:class:`histpy.Histogram` + + Returns + ------- + flaot + """ + + if len(dataspace_histogram_list) != len(self._dataset): + logger.error(f"The length of the input histogram list ({len(dataspace_histogram_list)}) is not equal to that of the dataset ({len(self._dataset)}).") + raise ValueError + + indexlist = self._dict_dataset_indexlist_for_bkg_models[key] + + return sum( + self._dataset[i].calc_bkg_model_product(key = key, dataspace_histogram = dataspace_histogram_list[i]) + for i in indexlist + ) + + # Utility + @staticmethod + def _histogram_sum(hlist): + """ + Sum a list of Histograms. If only one input, just return it. + """ + if len(hlist) == 1: + return hlist[0] + else: + result = hlist[0].copy() + for h in hlist[1:]: + result += h + return result diff --git a/cosipy/image_deconvolution/image_deconvolution_data_interface_base.py b/cosipy/image_deconvolution/data_interfaces/image_deconvolution_data_interface_base.py similarity index 78% rename from cosipy/image_deconvolution/image_deconvolution_data_interface_base.py rename to cosipy/image_deconvolution/data_interfaces/image_deconvolution_data_interface_base.py index 64e4f12cf..a6a7cc847 100644 --- a/cosipy/image_deconvolution/image_deconvolution_data_interface_base.py +++ b/cosipy/image_deconvolution/data_interfaces/image_deconvolution_data_interface_base.py @@ -32,8 +32,12 @@ class ImageDeconvolutionDataInterfaceBase(ABC): It returns a binned histogram of a background model with the given key. - summed_bkg_model(key) It returns the summed value of the background histogram with the given key. - - calc_expectation(model) - It returns a histogram of expected counts from the given model. + - calc_source_expectation(model) + It returns a histogram of expected counts from the given source model. + - calc_bkg_expectation(dict_bkg_norm) + It returns a histogram of expected counts from the given background normalizations + - calc_expectation(model, dict_bkg_norm) + It returns a histogram of expected counts from the given source model and background normalizations - calc_T_product(dataspace_histogram) It returns the product of the input histogram with the transonse matrix of the response function. - calc_bkg_model_product(key, dataspace_histogram) @@ -90,27 +94,57 @@ def summed_bkg_model(self, key): return self._summed_bkg_models[key] @abstractmethod - def calc_expectation(self, model, dict_bkg_norm = None, almost_zero = 1e-12): + def calc_source_expectation(self, model): """ - Calculate expected counts from a given model map. + Calculate expected counts from the source model only (R · model). Parameters ---------- - model : :py:class:`cosipy.image_deconvolution.ModelBase` or its subclass - Model - dict_bkg_norm : dict, default None - background normalization for each background model, e.g, {'albedo': 0.95, 'activation': 1.05} - almost_zero : float, default 1e-12 - In order to avoid zero components in extended count histogram, a tiny offset is introduced. - It should be small enough not to effect statistics. + model : ModelBase or subclass Returns ------- - :py:class:`histpy.Histogram` - Expected count histogram + histpy.Histogram + """ + raise NotImplementedError + + @abstractmethod + def calc_bkg_expectation(self, dict_bkg_norm): + """ + Calculate expected counts from background models only. + + Parameters + ---------- + dict_bkg_norm : dict + Background normalization, e.g. {'albedo': 0.95}. + + Returns + ------- + histpy.Histogram """ raise NotImplementedError + def calc_expectation(self, model, dict_bkg_norm): + """ + Calculate total expected counts (source + background). + + Default implementation sums calc_source_expectation and + calc_bkg_expectation. Subclasses may override for efficiency. + + Parameters + ---------- + model : ModelBase or subclass + dict_bkg_norm : dict + + Returns + ------- + histpy.Histogram + """ + + expectation = self.calc_source_expectation(model) + self.calc_bkg_expectation(dict_bkg_norm) + + return expectation + @abstractmethod def calc_T_product(self, dataspace_histogram): """ diff --git a/cosipy/image_deconvolution/data_interfaces/utils.py b/cosipy/image_deconvolution/data_interfaces/utils.py new file mode 100644 index 000000000..2b85ce450 --- /dev/null +++ b/cosipy/image_deconvolution/data_interfaces/utils.py @@ -0,0 +1,16 @@ +""" +Utility functions for data interfaces. +""" +import numpy as np +import astropy.units as u + + +def tensordot_sparse(A, A_unit, B, axes): + """ + Perform a tensordot operation on A and B. A is sparse + and so does not carry a unit; rather it must be passed + as a separate argument. B is a normal Quantity. Return + a proper Quantity as the result. + """ + dotprod = np.tensordot(A, B.value, axes=axes) + return u.Quantity(dotprod, unit=A_unit * B.unit, copy=False) diff --git a/cosipy/image_deconvolution/exposure_tables/__init__.py b/cosipy/image_deconvolution/exposure_tables/__init__.py new file mode 100644 index 000000000..9466fbf9f --- /dev/null +++ b/cosipy/image_deconvolution/exposure_tables/__init__.py @@ -0,0 +1,4 @@ +from .exposure_table_base import ExposureTableBase +from .scatt_exposure_table import SpacecraftAttitudeExposureTable +from .time_binned_exposure_table import TimeBinnedExposureTable +from .coordsys_conversion_matrix import CoordsysConversionMatrix diff --git a/cosipy/image_deconvolution/coordsys_conversion_matrix.py b/cosipy/image_deconvolution/exposure_tables/coordsys_conversion_matrix.py similarity index 96% rename from cosipy/image_deconvolution/coordsys_conversion_matrix.py rename to cosipy/image_deconvolution/exposure_tables/coordsys_conversion_matrix.py index 359809e5c..1fdd38a27 100644 --- a/cosipy/image_deconvolution/coordsys_conversion_matrix.py +++ b/cosipy/image_deconvolution/exposure_tables/coordsys_conversion_matrix.py @@ -9,7 +9,8 @@ from scoords import Attitude, SpacecraftFrame from histpy import Histogram, Axes, Axis, HealpixAxis -from .dataIF_COSI_DC2 import tensordot_sparse +from ..data_interfaces.utils import tensordot_sparse +from ..constants import EARTH_RADIUS_KM class CoordsysConversionMatrix(Histogram): """ @@ -32,7 +33,7 @@ def copy(self): return new @classmethod - def from_exposure_table(cls, exposure_table, full_detector_response, nside_model = None, scheme_model = 'ring', use_averaged_pointing = False, earth_occ = True, r_earth = 6378.0): + def from_exposure_table(cls, exposure_table, full_detector_response, nside_model = None, scheme_model = 'ring', use_averaged_pointing = False, earth_occ = True, r_earth = EARTH_RADIUS_KM): """ Calculate a ccm from a given exposure_table. @@ -54,7 +55,7 @@ def from_exposure_table(cls, exposure_table, full_detector_response, nside_model In the latter case, the conversion matrix is more accurate but it takes a long time to calculate it. earth_occ: bool, default True If it is True, the earth occultation is considered. - r_earth : float, default 6378.0 + r_earth : float, default EARTH_RADIUS_KM Earth's radius in kilometers. Returns @@ -147,7 +148,7 @@ def from_exposure_table(cls, exposure_table, full_detector_response, nside_model return coordsys_conv_matrix @classmethod - def _calc_exposure_time_map(cls, nside_model, num_pointings, earth_zenith, altitude, livetime, is_nest_model = False, earth_occ = True, r_earth = 6378.0): + def _calc_exposure_time_map(cls, nside_model, num_pointings, earth_zenith, altitude, livetime, is_nest_model, earth_occ, r_earth): """ Calculate exposure time map considering Earth occultation. @@ -169,11 +170,11 @@ def _calc_exposure_time_map(cls, nside_model, num_pointings, earth_zenith, altit Array of spacecraft altitudes in kilometers for each pointing. livetime : numpy.ndarray Array of livetimes in seconds for each pointing. - is_nest_model : bool, default False + is_nest_model : bool If True, use nested HEALPix pixel ordering scheme. If False, use ring ordering. - earth_occ: bool, default True + earth_occ: bool If it is True, the earth occultation is considered. - r_earth : float, default 6378.0 + r_earth : float Earth's radius in kilometers. Returns @@ -184,6 +185,7 @@ def _calc_exposure_time_map(cls, nside_model, num_pointings, earth_zenith, altit exposure time in seconds for pointing i and pixel j that is within the Earth occultation region. """ + npix_model = hp.nside2npix(nside_model) exposure_time_map = np.zeros((num_pointings, npix_model)) @@ -240,6 +242,7 @@ def calc_exposure_map(self, full_detector_response): Exposure map with axes ["ScAtt", "lb", "Ei"] representing the effective area x time for each attitude bin, sky pixel, and energy bin. """ + effective_area = full_detector_response.to_dr().project(['NuLambda', 'Ei']) exposure_map_contents = tensordot_sparse(self.contents, self.unit, diff --git a/cosipy/image_deconvolution/exposure_table_base.py b/cosipy/image_deconvolution/exposure_tables/exposure_table_base.py similarity index 99% rename from cosipy/image_deconvolution/exposure_table_base.py rename to cosipy/image_deconvolution/exposure_tables/exposure_table_base.py index d84e98d81..f3585ceb4 100644 --- a/cosipy/image_deconvolution/exposure_table_base.py +++ b/cosipy/image_deconvolution/exposure_tables/exposure_table_base.py @@ -345,6 +345,7 @@ def _get_averaged_pointing(cls, pointing, livetime): :py:class:`np.array` Averaged pointing in degrees, as np.array([l, b]) """ + if np.all(livetime == 0) == True: averaged_vector = np.sum(hp.ang2vec(pointing.T[0], pointing.T[1], lonlat = True).T, axis = (1)) logger.warning("Livetime is all zero") diff --git a/cosipy/image_deconvolution/scatt_exposure_table.py b/cosipy/image_deconvolution/exposure_tables/scatt_exposure_table.py similarity index 99% rename from cosipy/image_deconvolution/scatt_exposure_table.py rename to cosipy/image_deconvolution/exposure_tables/scatt_exposure_table.py index 5d361ae33..5686c0af1 100644 --- a/cosipy/image_deconvolution/scatt_exposure_table.py +++ b/cosipy/image_deconvolution/exposure_tables/scatt_exposure_table.py @@ -63,6 +63,7 @@ def __init__(self, df, nside, scheme='ring'): scheme : str, default 'ring' Healpix scheme ('ring' or 'nested') """ + super().__init__(df, nside=nside, scheme=scheme) @classmethod @@ -91,6 +92,7 @@ def from_orientation(cls, orientation, nside, scheme = 'ring', start = None, sto ------- :py:class:`cosipy.spacecraftfile.SpacecraftAttitudeExposureTable` """ + df = cls.analyze_orientation(orientation, nside, scheme, start, stop, min_livetime, min_num_pointings) # nside and scheme are no longer stored in df @@ -291,6 +293,7 @@ def get_binned_data(self, unbinned_event, psichi_binning = 'local', sparse = Fal :py:class:`histpy.Histogram` Binned data with axes ["ScAtt", "Em", "Phi", "PsiChi"]. """ + exposure_dict = {(row['healpix_index_z_pointing'], row['healpix_index_x_pointing']): row['scatt_binning_index'] for _, row in self.iterrows()} # from BinnedData.py diff --git a/cosipy/image_deconvolution/time_binned_exposure_table.py b/cosipy/image_deconvolution/exposure_tables/time_binned_exposure_table.py similarity index 99% rename from cosipy/image_deconvolution/time_binned_exposure_table.py rename to cosipy/image_deconvolution/exposure_tables/time_binned_exposure_table.py index f54dc0754..f31eb46b3 100644 --- a/cosipy/image_deconvolution/time_binned_exposure_table.py +++ b/cosipy/image_deconvolution/exposure_tables/time_binned_exposure_table.py @@ -43,6 +43,7 @@ class TimeBinnedExposureTable(ExposureTableBase): time_format: str time_scale: str """ + binning_index_name = 'time_binning_index' binning_method = 'Time' additional_column_scaler = [('tstart', 'D', ''), @@ -59,6 +60,7 @@ def __init__(self, df, format, scale): time_format: str time_scale: str """ + super().__init__(df, format=format, scale=scale) @classmethod @@ -79,6 +81,7 @@ def from_orientation(cls, orientation, tstart_list, tstop_list, **kwargs): ------- :py:class:`cosipy.spacecraftfile.SpacecraftAttitudeExposureTable` """ + if tstart_list.isscalar == True: tstart_list = Time([tstart_list]) @@ -180,6 +183,7 @@ def get_binned_data(self, unbinned_event, psichi_binning = 'local', sparse = Fal :py:class:`histpy.Histogram` Binned data with axes ["Time", "Em", "Phi", "PsiChi"]. """ + # Get TimeTags from unbinned event time_tags = unbinned_event.cosi_dataset['TimeTags'] diff --git a/cosipy/image_deconvolution/image_deconvolution.py b/cosipy/image_deconvolution/image_deconvolution.py index 1d771cf8e..eb7f42fea 100644 --- a/cosipy/image_deconvolution/image_deconvolution.py +++ b/cosipy/image_deconvolution/image_deconvolution.py @@ -5,19 +5,24 @@ from yayc import Configurator from pathlib import Path +from typing import Union -from .allskyimage import AllSkyImageModel +from .data_interfaces.data_interface_collection import DataInterfaceCollection -from .RichardsonLucy import RichardsonLucy -from .RichardsonLucySimple import RichardsonLucySimple -from .MAP_RichardsonLucy import MAP_RichardsonLucy +from .models.allskyimage import AllSkyImageModel + +from .algorithms.RichardsonLucyBasic import RichardsonLucyBasic +from .algorithms.RichardsonLucy import RichardsonLucy +from .algorithms.RichardsonLucyAdvanced import RichardsonLucyAdvanced +from .algorithms.MAP_RichardsonLucy import MAP_RichardsonLucy class ImageDeconvolution: """ A class to reconstruct all-sky images from COSI data based on image deconvolution methods. """ model_classes = {"AllSkyImage": AllSkyImageModel} - deconvolution_algorithm_classes = {"RL": RichardsonLucy, "RLsimple": RichardsonLucySimple, "MAP_RL": MAP_RichardsonLucy} + deconvolution_algorithm_classes = {"RLbasic": RichardsonLucyBasic, "RL": RichardsonLucy, + "RLadvanced": RichardsonLucyAdvanced, "MAP_RL": MAP_RichardsonLucy} def __init__(self): self._dataset = None @@ -27,17 +32,19 @@ def __init__(self): self._model_class = None self._deconvolution_class = None - def set_dataset(self, dataset: list): + def set_dataset(self, dataset: Union[list,DataInterfaceCollection]): """ Set dataset as a list. Parameters ---------- - dataset : list of :py:class:`cosipy.image_deconvolution.ImageDeconvolutionDataInterfaceBase` or its subclass + dataset : DataInterfaceCollection or list of :py:class:`cosipy.image_deconvolution.ImageDeconvolutionDataInterfaceBase` Each component contaning an event histogram, a background model, a response matrix, and a coordsys_conversion_matrix. """ - - self._dataset = dataset + if isinstance(dataset, list): + self._dataset = DataInterfaceCollection(dataset) + else: + self._dataset = dataset logger.debug(f"dataset for image deconvolution was set -> {self._dataset}") @@ -62,6 +69,7 @@ def read_parameterfile(self, parameter_filepath: str | Path): parameter_filepath : str or pathlib.Path Path of parameter file. """ + self._parameter = Configurator.open(parameter_filepath) logger.debug(f"parameter file for image deconvolution was set -> {parameter_filepath}") @@ -71,6 +79,7 @@ def dataset(self): """ Return the dataset. """ + return self._dataset @property @@ -78,6 +87,7 @@ def parameter(self): """ Return the registered parameter. """ + return self._parameter def override_parameter(self, *args): @@ -91,8 +101,9 @@ def override_parameter(self, *args): Examples -------- - >>> image_deconvolution.override_parameter("deconvolution:parameter_RL:iteration = 30") + >>> image_deconvolution.override_parameter("deconvolution:parameter:iteration_max = 30") """ + self._parameter.override(args) @property @@ -100,6 +111,7 @@ def initial_model(self): """ Return the initial model. """ + if self._initial_model is None: logger.warning("Need to initialize model in the image_deconvolution instance!") @@ -110,6 +122,7 @@ def mask(self): """ Return the mask. """ + return self._mask @property @@ -117,6 +130,7 @@ def results(self): """ Return the results. """ + return self._deconvolution.results def initialize(self): @@ -142,12 +156,18 @@ def model_initialization(self): bool whether the instantiation and initialization are successfully done. """ + # set self._model_class model_name = self.parameter['model_definition']['class'] # Options include "AllSkyImage", etc. - if not model_name in self.model_classes.keys(): # See model_classes dictionary declared above - logger.error(f'The model class "{model_name}" does not exist!') - raise ValueError + if model_name not in self.model_classes.keys(): # See model_classes dictionary declared above + available_models = ', '.join(self.model_classes.keys()) + error_msg = ( + f'Unknown model class "{model_name}". ' + f'Available models: {available_models}' + ) + logger.error(error_msg) + raise ValueError(error_msg) self._model_class = self.model_classes[model_name] @@ -156,7 +176,6 @@ def model_initialization(self): parameter_model_property = Configurator(self.parameter['model_definition']['property']) self._initial_model = self._model_class.instantiate_from_parameters(parameter_model_property) - logger.info("---- parameters ----") logger.info(parameter_model_property.dump()) # setting initial values @@ -173,7 +192,6 @@ def model_initialization(self): logger.error("The model axes mismatches with the reponse in the dataset!") raise ValueError - logger.info("---- parameters ----") logger.info(parameter_model_initialization.dump()) def register_deconvolution_algorithm(self): @@ -185,15 +203,21 @@ def register_deconvolution_algorithm(self): bool whether the deconvolution algorithm is successfully registered. """ + logger.info("<< Registering the deconvolution algorithm >>") parameter_deconvolution = Configurator(self.parameter['deconvolution']) algorithm_name = parameter_deconvolution['algorithm'] algorithm_parameter = Configurator(parameter_deconvolution['parameter']) - if not algorithm_name in self.deconvolution_algorithm_classes.keys(): - logger.error(f'The algorithm "{algorithm_name}" does not exist!') - raise ValueError + if algorithm_name not in self.deconvolution_algorithm_classes.keys(): + available_algorithms = ', '.join(self.deconvolution_algorithm_classes.keys()) + error_msg = ( + f'Unknown deconvolution algorithm "{algorithm_name}". ' + f'Available algorithms: {available_algorithms}' + ) + logger.error(error_msg) + raise ValueError(error_msg) self._deconvolution_class = self.deconvolution_algorithm_classes[algorithm_name] # Alias to class constructor self._deconvolution = self._deconvolution_class(initial_model = self.initial_model, # Initialize object for relevant class @@ -201,7 +225,6 @@ def register_deconvolution_algorithm(self): mask = self.mask, parameter = algorithm_parameter) - logger.info("---- parameters ----") logger.info(parameter_deconvolution.dump()) def run_deconvolution(self): @@ -213,6 +236,7 @@ def run_deconvolution(self): list List containing results (reconstructed image, likelihood etc) at each iteration. """ + logger.info("#### Image Deconvolution Starts ####") logger.info(f"<< Initialization >>") @@ -272,4 +296,4 @@ def _finalize(self): # only in the master node if self.is_master_node: - super()._finalize() \ No newline at end of file + super()._finalize() diff --git a/cosipy/image_deconvolution/models/__init__.py b/cosipy/image_deconvolution/models/__init__.py new file mode 100644 index 000000000..5d2154608 --- /dev/null +++ b/cosipy/image_deconvolution/models/__init__.py @@ -0,0 +1,2 @@ +from .model_base import ModelBase +from .allskyimage import AllSkyImageModel diff --git a/cosipy/image_deconvolution/allskyimage.py b/cosipy/image_deconvolution/models/allskyimage.py similarity index 99% rename from cosipy/image_deconvolution/allskyimage.py rename to cosipy/image_deconvolution/models/allskyimage.py index 51172bb6b..93b2db363 100644 --- a/cosipy/image_deconvolution/allskyimage.py +++ b/cosipy/image_deconvolution/models/allskyimage.py @@ -176,4 +176,5 @@ def total_flux(self): """ Calculate the total flux at each energy bin. """ + return np.sum(self, axis = 0) * self.axes['lb'].pixarea() diff --git a/cosipy/image_deconvolution/model_base.py b/cosipy/image_deconvolution/models/model_base.py similarity index 99% rename from cosipy/image_deconvolution/model_base.py rename to cosipy/image_deconvolution/models/model_base.py index 3197c70c4..6f7524bed 100644 --- a/cosipy/image_deconvolution/model_base.py +++ b/cosipy/image_deconvolution/models/model_base.py @@ -31,7 +31,6 @@ def instantiate_from_parameters(cls, parameter): ------- py:class:`ModelBase` """ - raise NotImplementedError @abstractmethod @@ -44,7 +43,6 @@ def set_values_from_parameters(self, parameter): parameter : py:class:`yayc.Configurator` Parameters for the specified algorithm. """ - raise NotImplementedError def mask_pixels(self, mask, fill_value = 0): diff --git a/cosipy/image_deconvolution/prior_entropy.py b/cosipy/image_deconvolution/prior_entropy.py deleted file mode 100644 index d329bd464..000000000 --- a/cosipy/image_deconvolution/prior_entropy.py +++ /dev/null @@ -1,67 +0,0 @@ -import astropy.units as u -import healpy as hp -import numpy as np - -from .prior_base import PriorBase - -from .allskyimage import AllSkyImageModel - -class PriorEntropy(PriorBase): - """ - Entropy prior for all-sky image models. - - Parameters - ---------- - parameter : dict - Parameters for the entropy prior. - model : AllSkyImageModel - All-sky image model to which the prior will be applied. - """ - - usable_model_classes = [AllSkyImageModel] - - def __init__(self, parameter, model): - - super().__init__(parameter, model) - - self.prior_map = self.parameter['prior_map']['value'] * u.Unit(self.parameter['prior_map']['unit']) - - def log_prior(self, model): - """ - Calculate the logarithm of the entropy prior probability. - - Parameters - ---------- - model : AllSkyImageModel - Model for which to calculate the log prior. - - Returns - ------- - float - The logarithm of the entropy prior probability. - """ - if self.model_class == AllSkyImageModel: - - image_ratio = (model/self.prior_map).contents.to('').value - - return self.coefficient * np.sum(model.contents.value * (1 - np.log(image_ratio))) - - def grad_log_prior(self, model): - """ - Calculate the gradient of the log entropy prior. - - Parameters - ---------- - model : AllSkyImageModel - Model for which to calculate the gradient. - - Returns - ------- - numpy.ndarray - Gradient of the log prior, in units inverse to the model. - """ - if self.model_class == AllSkyImageModel: - - image_ratio = (model/self.prior_map).contents.to('').value - - return -1 * self.coefficient * np.log(image_ratio) / model.unit diff --git a/cosipy/image_deconvolution/utils.py b/cosipy/image_deconvolution/utils.py new file mode 100644 index 000000000..8e88b0014 --- /dev/null +++ b/cosipy/image_deconvolution/utils.py @@ -0,0 +1,11 @@ +""" +Shared utility functions for the image_deconvolution module. +""" + +def _to_float(x) -> float: + """ + Convert to float, handling astropy Quantity. + """ + if hasattr(x, 'unit'): + return float(x.to('')) + return float(x) diff --git a/cosipy/test_data/image_deconvolution/imagedeconvolution_parfile_test.yml b/cosipy/test_data/image_deconvolution/imagedeconvolution_parfile_test.yml index 39e9a58ab..a70a9fe16 100644 --- a/cosipy/test_data/image_deconvolution/imagedeconvolution_parfile_test.yml +++ b/cosipy/test_data/image_deconvolution/imagedeconvolution_parfile_test.yml @@ -16,7 +16,7 @@ model_definition: value: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] #the number of these values should be the same as "the number of energy_edges - 1". unit: "cm-2 s-1 sr-1" # do not change it as for now deconvolution: - algorithm: "RLsimple" + algorithm: "RLbasic" parameter: iteration_max: 2 background_normalization_optimization: diff --git a/docs/tutorials/image_deconvolution/511keV-Galactic-ImageDeconvolution.ipynb b/docs/tutorials/image_deconvolution/511keV-Galactic-ImageDeconvolution.ipynb index 76d5cabe8..548f37af5 100644 --- a/docs/tutorials/image_deconvolution/511keV-Galactic-ImageDeconvolution.ipynb +++ b/docs/tutorials/image_deconvolution/511keV-Galactic-ImageDeconvolution.ipynb @@ -503,7 +503,7 @@ "|||||\n", "|parameter_RL:iteration_max | int | The maximum number of the iteration | |\n", "|parameter_RL:acceleration:activate | bool | whether the accelerated ML-EM algorithm (Knoedlseder+99) is used | |\n", - "|parameter_RL:acceleration:alpha_max | float | the maximum value for the acceleration parameter | |\n", + "|parameter_RL:acceleration:accel_factor_max | float | the maximum value for the acceleration parameter | |\n", "|parameter_RL:response_weighting:activate | bool | whether a delta map is renormalized based on the exposure time on each pixel, namely $w_j = (\\sum_{i} R_{ij})^{\\beta}$ (see Knoedlseder+05, Siegert+20) | |\n", "|parameter_RL:response_weighting:index | float | $\\beta$ in the above equation | |\n", "|parameter_RL:smoothing:activate | bool | whether a Gaussian filter is used (see Knoedlseder+05, Siegert+20) | |\n", @@ -547,7 +547,7 @@ "source": [ "image_deconvolution.override_parameter(\"deconvolution:parameter:iteration_max = 50\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:background_normalization_optimization:activate = True\")\n", - "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:alpha_max = 2.0\")\n", + "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:accel_factor_max = 2.0\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:smoothing:activate = False\")\n", "#image_deconvolution.override_parameter(\"deconvolution:parameter:smoothing:activate = True\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:response_weighting:activate = False\")\n", @@ -788,7 +788,7 @@ " iteration_max: 50\n", " acceleration:\n", " activate: True\n", - " alpha_max: 2.0\n", + " accel_factor_max: 2.0\n", " response_weighting:\n", " activate: False\n", " index: 0.5\n", diff --git a/docs/tutorials/image_deconvolution/imagedeconvolution_parfile_gal_511keV.yml b/docs/tutorials/image_deconvolution/imagedeconvolution_parfile_gal_511keV.yml index ae61d1efe..36c39ebb9 100644 --- a/docs/tutorials/image_deconvolution/imagedeconvolution_parfile_gal_511keV.yml +++ b/docs/tutorials/image_deconvolution/imagedeconvolution_parfile_gal_511keV.yml @@ -16,12 +16,14 @@ model_definition: value: [1e-4] #the number of these values should be the same as "the number of energy_edges - 1". unit: "cm-2 s-1 sr-1" # do not change it as for now deconvolution: - algorithm: "RL" + algorithm: "RLadvanced" parameter: iteration_max: 10 acceleration: activate: True - alpha_max: 10.0 + algorithm: "MaxStep" + accel_factor_max: 10.0 + accel_bkg_norm: False response_weighting: activate: True index: 0.5 diff --git a/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/511keV-ScAtt-ImageDeconvolution.ipynb b/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/511keV-ScAtt-ImageDeconvolution.ipynb index cc006be4d..30caa236e 100644 --- a/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/511keV-ScAtt-ImageDeconvolution.ipynb +++ b/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/511keV-ScAtt-ImageDeconvolution.ipynb @@ -694,7 +694,7 @@ "|||||\n", "|parameter_RL:iteration_max | int | The maximum number of the iteration | |\n", "|parameter_RL:acceleration:activate | bool | whether the accelerated ML-EM algorithm (Knoedlseder+99) is used | |\n", - "|parameter_RL:acceleration:alpha_max | float | the maximum value for the acceleration parameter | |\n", + "|parameter_RL:acceleration:accel_factor_max | float | the maximum value for the acceleration parameter | |\n", "|parameter_RL:response_weighting:activate | bool | whether a delta map is renormalized based on the exposure time on each pixel, namely $w_j = (\\sum_{i} R_{ij})^{\\beta}$ (see Knoedlseder+05, Siegert+20) | |\n", "|parameter_RL:response_weighting:index | float | $\\beta$ in the above equation | |\n", "|parameter_RL:smoothing:activate | bool | whether a Gaussian filter is used (see Knoedlseder+05, Siegert+20) | |\n", @@ -747,7 +747,7 @@ "parameter:\n", " acceleration:\n", " activate: true\n", - " alpha_max: 10.0\n", + " accel_factor_max: 10.0\n", " background_normalization_optimization:\n", " activate: true\n", " range:\n", @@ -827,7 +827,7 @@ "parameter:\n", " acceleration:\n", " activate: true\n", - " alpha_max: 10\n", + " accel_factor_max: 10\n", " background_normalization_optimization:\n", " activate: true\n", " range:\n", @@ -858,7 +858,7 @@ "source": [ "image_deconvolution.override_parameter(\"deconvolution:parameter:iteration_max = 30\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:background_normalization_optimization:activate = True\")\n", - "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:alpha_max = 10\")\n", + "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:accel_factor_max = 10\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:smoothing:FWHM:value = 3.0\")\n", "\n", "image_deconvolution.initialize()" diff --git a/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_511keV.yml b/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_511keV.yml index ae61d1efe..36c39ebb9 100644 --- a/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_511keV.yml +++ b/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_511keV.yml @@ -16,12 +16,14 @@ model_definition: value: [1e-4] #the number of these values should be the same as "the number of energy_edges - 1". unit: "cm-2 s-1 sr-1" # do not change it as for now deconvolution: - algorithm: "RL" + algorithm: "RLadvanced" parameter: iteration_max: 10 acceleration: activate: True - alpha_max: 10.0 + algorithm: "MaxStep" + accel_factor_max: 10.0 + accel_bkg_norm: False response_weighting: activate: True index: 0.5 diff --git a/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/GRB-ScAtt-ImageDeconvolution.ipynb b/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/GRB-ScAtt-ImageDeconvolution.ipynb index 32c256aa9..3161e545c 100644 --- a/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/GRB-ScAtt-ImageDeconvolution.ipynb +++ b/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/GRB-ScAtt-ImageDeconvolution.ipynb @@ -1293,7 +1293,7 @@ "|||||\n", "|parameter_RL:iteration_max | int | The maximum number of the iteration | |\n", "|parameter_RL:acceleration:activate | bool | whether the accelerated ML-EM algorithm (Knoedlseder+99) is used | |\n", - "|parameter_RL:acceleration:alpha_max | float | the maximum value for the acceleration parameter | |\n", + "|parameter_RL:acceleration:accel_factor_max | float | the maximum value for the acceleration parameter | |\n", "|parameter_RL:response_weighting:activate | bool | whether a delta map is renormalized based on the exposure time on each pixel, namely $w_j = (\\sum_{i} R_{ij})^{\\beta}$ (see Knoedlseder+05, Siegert+20) | |\n", "|parameter_RL:response_weighting:index | float | $\\beta$ in the above equation | |\n", "|parameter_RL:smoothing:activate | bool | whether a Gaussian filter is used (see Knoedlseder+05, Siegert+20) | |\n", @@ -1364,7 +1364,7 @@ "parameter:\n", " acceleration:\n", " activate: true\n", - " alpha_max: 10.0\n", + " accel_factor_max: 10.0\n", " background_normalization_optimization:\n", " activate: true\n", " range:\n", @@ -1462,7 +1462,7 @@ "parameter:\n", " acceleration:\n", " activate: true\n", - " alpha_max: 10\n", + " accel_factor_max: 10\n", " background_normalization_optimization:\n", " activate: false\n", " range:\n", @@ -1493,7 +1493,7 @@ "source": [ "image_deconvolution.override_parameter(\"deconvolution:parameter:iteration_max = 200\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:background_normalization_optimization:activate = False\") # in this case, the total counts of the background model is zero. So, here we don't fit the bkg normalization.\n", - "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:alpha_max = 10\")\n", + "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:accel_factor_max = 10\")\n", "image_deconvolution.override_parameter(\"deconvolution:parameter:smoothing:FWHM:value = 10.0\")\n", "\n", "image_deconvolution.initialize()" diff --git a/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_GRB.yml b/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_GRB.yml index 415494212..90a443760 100644 --- a/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_GRB.yml +++ b/docs/tutorials/image_deconvolution/optional/time_binned_data_w_local_coordinate/imagedeconvolution_parfile_scatt_GRB.yml @@ -16,12 +16,14 @@ model_definition: value: [1e-5,1e-5,1e-5,1e-5,1e-5,1e-5,1e-5,1e-5,1e-5,1e-5] #the number of these values should be the same as "the number of energy_edges - 1". unit: "cm-2 s-1 sr-1" # do not change it as for now deconvolution: - algorithm: "RL" + algorithm: "RLadvanced" parameter: iteration_max: 10 acceleration: activate: True - alpha_max: 10.0 + algorithm: "MaxStep" + accel_factor_max: 10.0 + accel_bkg_norm: False response_weighting: activate: True index: 0.5 diff --git a/tests/image_deconvolution/conftest.py b/tests/image_deconvolution/conftest.py index 3eafd4bec..15cfa9fd5 100644 --- a/tests/image_deconvolution/conftest.py +++ b/tests/image_deconvolution/conftest.py @@ -4,7 +4,7 @@ from histpy import Histogram, Axis, Axes from cosipy import test_data -from cosipy.image_deconvolution import DataIF_COSI_DC2, AllSkyImageModel +from cosipy.image_deconvolution import DataIF_COSI_DC2, AllSkyImageModel, DataInterfaceCollection @pytest.fixture def dataset(): @@ -19,7 +19,7 @@ def dataset(): rsp = precomputed_response, coordsys_conv_matrix = None) - return [data] + return DataInterfaceCollection([data]) @pytest.fixture def model(dataset): diff --git a/tests/image_deconvolution/test_accelerators.py b/tests/image_deconvolution/test_accelerators.py new file mode 100644 index 000000000..e91c55c43 --- /dev/null +++ b/tests/image_deconvolution/test_accelerators.py @@ -0,0 +1,136 @@ +import pytest +import numpy as np + +from cosipy.image_deconvolution.algorithms.accelerators.accelerator_base import EMStepResult +from cosipy.image_deconvolution.algorithms.accelerators.max_step_accelerator import MaxStepAccelerator +from cosipy.image_deconvolution.algorithms.accelerators.line_search_accelerator import LineSearchAccelerator + +def make_model(dataset, value=1.0): + from cosipy.image_deconvolution import AllSkyImageModel + model = AllSkyImageModel( + dataset[0].model_axes['lb'].nside, + dataset[0].model_axes['Ei'].edges, + ) + model[:] = value * model.unit + return model + + +def make_em_results(dataset, model_value=1.0, delta_value=0.1, bkg_norm_value=1.0): + """Return a (before, after) pair of EMStepResult.""" + model_before = make_model(dataset, value=model_value) + model_after = make_model(dataset, value=model_value + delta_value) + + bkg_before = {"bkg": bkg_norm_value} + bkg_after = {"bkg": bkg_norm_value + 0.01} + + src_exp_before = dataset.calc_source_expectation_list(model_before) + src_exp_after = dataset.calc_source_expectation_list(model_after) + + before = EMStepResult( + model = model_before, + dict_bkg_norm = bkg_before, + expectation_list = dataset.combine_expectation_list(src_exp_before, dataset.calc_bkg_expectation_list(bkg_before)), + source_expectation_list = src_exp_before, + bkg_expectation_list = dataset.calc_bkg_expectation_list(bkg_before), + ) + after = EMStepResult( + model = model_after, + dict_bkg_norm = bkg_after, + expectation_list = dataset.combine_expectation_list(src_exp_after, dataset.calc_bkg_expectation_list(bkg_after)), + source_expectation_list = src_exp_after, + bkg_expectation_list = dataset.calc_bkg_expectation_list(bkg_after), + ) + return [before, after] + + +def check_result(result, extra_keys): + """Common assertions for AcceleratorResult.""" + assert np.all(result.model.contents >= 0), "model must be non-negative" + for key in extra_keys: + assert key in result.extras, f"extras must contain '{key}'" + if key in ["accel_factor", "accel_factor_bkg"]: + assert result.extras[key] >= 1.0, f"{key} must be >= 1.0" + + +# --------------------------------------------------------------------------- +# MaxStepAccelerator +# --------------------------------------------------------------------------- + +class TestMaxStepAccelerator: + + @pytest.fixture + def accelerator(self): + from yayc import Configurator + return MaxStepAccelerator(Configurator({"accel_factor_max": 10.0, "accel_bkg_norm": False})) + + def test_basic(self, accelerator, dataset, mask): + result = accelerator.compute(make_em_results(dataset), dataset, mask) + print(f"accel_factor = {result.extras['accel_factor']}") + check_result(result, ["accel_factor"]) + + def test_with_bkg_norm(self, dataset, mask): + from yayc import Configurator + acc = MaxStepAccelerator(Configurator({"accel_factor_max": 10.0, "accel_bkg_norm": True})) + result = acc.compute(make_em_results(dataset), dataset, mask) + print(f"accel_factor = {result.extras['accel_factor']}") + check_result(result, ["accel_factor"]) + + +# --------------------------------------------------------------------------- +# LineSearchAccelerator +# --------------------------------------------------------------------------- + +class TestLineSearchAccelerator: + + def make_accelerator(self, params): + from yayc import Configurator + return LineSearchAccelerator(Configurator(params)) + + def test_1d_no_bkg(self, dataset, mask): + """Model only, 1-D line search.""" + acc = self.make_accelerator({"accel_factor_max": 10.0}) + result = acc.compute(make_em_results(dataset), dataset, mask) + print(f"accel_factor = {result.extras['accel_factor']}") + check_result(result, ["accel_factor", "accel_factor_bkg"]) + + def test_1d_with_bkg_same_factor(self, dataset, mask): + """Model and bkg scaled by the same accel_factor.""" + acc = self.make_accelerator({ + "accel_factor_max": 10.0, + "accel_bkg_norm": {"activate": True, "independent": False, "max": 10.0}, + }) + result = acc.compute(make_em_results(dataset), dataset, mask) + print(f"accel_factor = {result.extras['accel_factor']}, accel_factor_bkg = {result.extras['accel_factor_bkg']}") + check_result(result, ["accel_factor", "accel_factor_bkg"]) + assert result.extras["accel_factor"] == pytest.approx(result.extras["accel_factor_bkg"]) + + def test_2d_independent_gradient(self, dataset, mask): + """Independent bkg accel_factor, gradient search.""" + acc = self.make_accelerator({ + "accel_factor_max": 10.0, + "accel_bkg_norm": {"activate": True, "independent": True, "max": 10.0, + "search_method": "gradient", "grid_n": 5}, + }) + result = acc.compute(make_em_results(dataset), dataset, mask) + print(f"accel_factor = {result.extras['accel_factor']}, accel_factor_bkg = {result.extras['accel_factor_bkg']}") + check_result(result, ["accel_factor", "accel_factor_bkg"]) + + def test_2d_independent_grid(self, dataset, mask): + """Independent bkg accel_factor, grid search.""" + acc = self.make_accelerator({ + "accel_factor_max": 10.0, + "accel_bkg_norm": {"activate": True, "independent": True, "max": 10.0, + "search_method": "grid", "grid_n": 5}, + }) + result = acc.compute(make_em_results(dataset), dataset, mask) + print(f"accel_factor = {result.extras['accel_factor']}, accel_factor_bkg = {result.extras['accel_factor_bkg']}") + check_result(result, ["accel_factor", "accel_factor_bkg"]) + + def test_invalid_search_method_raises(self): + from yayc import Configurator + with pytest.raises(ValueError): + LineSearchAccelerator(Configurator({ + "accel_factor_max": 10.0, + "accel_bkg_norm": {"activate": True, "independent": True, "max": 10.0, + "search_method": "invalid"}, + })) diff --git a/tests/image_deconvolution/test_algorithm.py b/tests/image_deconvolution/test_algorithm.py index be978bbe5..af0604d6e 100644 --- a/tests/image_deconvolution/test_algorithm.py +++ b/tests/image_deconvolution/test_algorithm.py @@ -1,204 +1,183 @@ +""" +Integration tests for image deconvolution algorithms. +Each class groups tests for one algorithm. Within a class, tests share +a common parameter template via a fixture, and each test method exercises +a specific scenario (e.g. with/without background normalization, different +stopping criteria). Numerical values are regression checks: they should +only be updated intentionally when the algorithm changes. +""" + import pytest import numpy as np from yayc import Configurator -from cosipy.image_deconvolution import RichardsonLucySimple, RichardsonLucy, MAP_RichardsonLucy - -def test_RicharsonLucySimple(dataset, model, mask, tmp_path): - - num_iteration = 3 - - parameter = Configurator({"iteration_max": num_iteration, - "response_weighting": {"activate": True, "index": 0.5}, - "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, - "background_normalization_optimization": {"activate": True, - "range": {"bkg": [0.9, 1.1]}}, - "save_results": {"activate": True, "directory": f"{str(tmp_path)}", "only_final_result": True} - }) - - algorithm = RichardsonLucySimple(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - -def test_RicharsonLucy(dataset, model, mask, tmp_path): - - num_iteration = 3 - - parameter = Configurator({"iteration_max": num_iteration, - "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, - "acceleration": {"activate": True, "alpha_max": 10.0}, - "response_weighting": {"activate": True, "index": 0.5}, - "smoothing": {"activate": True, "FWHM": {"value": 2.0, "unit": "deg"}}, - "background_normalization_optimization": {"activate": True, - "range": {"bkg": [0.9, 1.1]}}, - "save_results": {"activate": True, "directory": f"{str(tmp_path)}", "only_final_result": True} - }) - - # w/ acceleration - algorithm = RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - - assert np.isclose(algorithm.results[-1]['log-likelihood'][0], 5495.120521335304) - - # wo/ acceleration and overwrite the directory - parameter["acceleration:activate"] = False - - algorithm = RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - - assert np.isclose(algorithm.results[-1]['log-likelihood'][0], 5270.562770130176) - -def test_MAP_RichardsonLucy(dataset, model, mask, tmp_path): - - num_iteration = 10 - - parameter = Configurator({"iteration_max": num_iteration, - "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, - "response_weighting": {"activate": True, "index": 0.5}, - "background_normalization_optimization": {"activate": True, - "range": {"bkg": [0.9, 1.1]}}, - "stopping_criteria": {"statistics": "log-posterior", - "threshold": 1e-2}, - "prior": {"TSV" :{"coefficient": 1.e-10}, - "gamma":{"model":{"theta": {"value": np.inf, "unit": "cm-2 s-1 sr-1"}, - "k": {"value": 0.999}}, - "background": {"theta": {"value": np.inf}, "k": {"value": 1.0}} - } - }, - "save_results": {"activate": True, "directory": f"{str(tmp_path)}", "only_final_result": True} - }) - - # first run - algorithm = MAP_RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - - assert np.isclose(algorithm.results[-1]['log-posterior'], 6567.857548203495) - - # background fixed - parameter["background_normalization_optimization:activate"] = False - - algorithm = MAP_RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - - assert np.isclose(algorithm.results[-1]['log-posterior'], 6202.336733778631) +from cosipy.image_deconvolution import ( + RichardsonLucyBasic, + RichardsonLucy, + RichardsonLucyAdvanced, + MAP_RichardsonLucy, +) - # with large threshold - parameter["stopping_criteria:threshold"] = 1e10 - algorithm = MAP_RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) +# --------------------------------------------------------------------------- +# Helper +# --------------------------------------------------------------------------- +def run(algorithm, num_iteration): + """Initialize and iterate an algorithm. Returns the algorithm instance.""" algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: + for _ in range(num_iteration): + if algorithm.iteration(): break - algorithm.finalization() - - assert len(algorithm.results) == 2 - - # with log-likelihood - parameter["stopping_criteria:statistics"] = "log-likelihood" - - algorithm = MAP_RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - - assert np.isclose(algorithm.results[-1]['log-likelihood'][0], 3931.3528198012773) - - # without gamma prior - parameter["stopping_criteria:statistics"] = 'log-posterior' - parameter["prior"] = {"TSV":{"coefficient": 1.e-10}} - - algorithm = MAP_RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - algorithm.initialization() - - for i in range(num_iteration): - stop = algorithm.iteration() - if stop: - break - - algorithm.finalization() - - assert np.isclose(algorithm.results[-1]['log-posterior'], 3931.29811442966) - - # wrong statistics - parameter["stopping_criteria:statistics"] = "likelihooooooooood!!!" - - with pytest.raises(ValueError) as e_info: - algorithm = MAP_RichardsonLucy(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) + return algorithm + + +# --------------------------------------------------------------------------- +# RichardsonLucyBasic +# --------------------------------------------------------------------------- + +class TestRichardsonLucyBasic: + + @pytest.fixture + def parameter(self, tmp_path): + return Configurator({ + "iteration_max": 3, + "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, + "save_results": {"activate": True, "directory": str(tmp_path), "only_final_result": True}, + }) + + def test_basic(self, dataset, model, mask, parameter): + alg = RichardsonLucyBasic(model, dataset, mask, parameter) + run(alg, parameter["iteration_max"]) + + +# --------------------------------------------------------------------------- +# RichardsonLucy +# --------------------------------------------------------------------------- + +class TestRichardsonLucy: + + @pytest.fixture + def parameter(self, tmp_path): + return Configurator({ + "iteration_max": 3, + "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, + "background_normalization_optimization": {"activate": True, "range": {"bkg": [0.9, 1.1]}}, + "save_results": {"activate": True, "directory": str(tmp_path), "only_final_result": True}, + }) + + def test_basic(self, dataset, model, mask, parameter): + alg = RichardsonLucy(model, dataset, mask, parameter) + run(alg, parameter["iteration_max"]) + + +# --------------------------------------------------------------------------- +# RichardsonLucyAdvanced +# --------------------------------------------------------------------------- + +class TestRichardsonLucyAdvanced: + """ + Tests for RichardsonLucyAdvanced covering acceleration on/off. + Numerical assertions are regression checks. + """ + + NUM_ITERATION = 3 + + @pytest.fixture + def base_parameter(self, tmp_path): + return Configurator({ + "iteration_max": self.NUM_ITERATION, + "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, + "acceleration": {"activate": True, "alpha_max": 10.0}, + "response_weighting": {"activate": True, "index": 0.5}, + "smoothing": {"activate": True, "FWHM": {"value": 2.0, "unit": "deg"}}, + "background_normalization_optimization": {"activate": True, "range": {"bkg": [0.9, 1.1]}}, + "save_results": {"activate": True, "directory": str(tmp_path), "only_final_result": True}, + }) + + def test_with_acceleration(self, dataset, model, mask, base_parameter): + alg = run(RichardsonLucyAdvanced(model, dataset, mask, base_parameter), self.NUM_ITERATION) + # Regression check + assert np.isclose(alg.results[-1]['log-likelihood'][0], 5495.120521335304) + + def test_without_acceleration(self, dataset, model, mask, base_parameter): + base_parameter["acceleration:activate"] = False + alg = run(RichardsonLucyAdvanced(model, dataset, mask, base_parameter), self.NUM_ITERATION) + # Regression check + assert np.isclose(alg.results[-1]['log-likelihood'][0], 5270.562770130176) + + +# --------------------------------------------------------------------------- +# MAP_RichardsonLucy +# --------------------------------------------------------------------------- + +class TestMAPRichardsonLucy: + """ + Tests for MAP_RichardsonLucy covering different prior configurations + and stopping criteria. Numerical assertions are regression checks. + """ + + NUM_ITERATION = 10 + + @pytest.fixture + def base_parameter(self, tmp_path): + return Configurator({ + "iteration_max": self.NUM_ITERATION, + "minimum_flux": {"value": 0.0, "unit": "cm-2 s-1 sr-1"}, + "response_weighting": {"activate": True, "index": 0.5}, + "background_normalization_optimization": {"activate": True, "range": {"bkg": [0.9, 1.1]}}, + "stopping_criteria": {"statistics": "log-posterior", "threshold": 1e-2}, + "prior": { + "TSV": {"coefficient": 1e-10}, + "gamma": { + "model": {"theta": {"value": np.inf, "unit": "cm-2 s-1 sr-1"}, "k": {"value": 0.999}}, + "background": {"theta": {"value": np.inf}, "k": {"value": 1.0}}, + }, + }, + "save_results": {"activate": True, "directory": str(tmp_path), "only_final_result": True}, + }) + + def test_with_gamma_and_tsv_prior(self, dataset, model, mask, base_parameter): + """Full setup: TSV + gamma prior, bkg optimization enabled.""" + alg = run(MAP_RichardsonLucy(model, dataset, mask, base_parameter), self.NUM_ITERATION) + # Regression check + assert np.isclose(alg.results[-1]['log-posterior'], 6567.857548203495) + + def test_without_bkg_optimization(self, dataset, model, mask, base_parameter): + """Background normalization fixed at 1.0.""" + base_parameter["background_normalization_optimization:activate"] = False + alg = run(MAP_RichardsonLucy(model, dataset, mask, base_parameter), self.NUM_ITERATION) + # Regression check + assert np.isclose(alg.results[-1]['log-posterior'], 6202.336733778631) + + def test_stopping_criteria_threshold(self, dataset, model, mask, base_parameter): + """Large threshold causes early stopping after 2 iterations.""" + base_parameter["stopping_criteria:threshold"] = 1e10 + alg = run(MAP_RichardsonLucy(model, dataset, mask, base_parameter), self.NUM_ITERATION) + assert len(alg.results) == 2 + + def test_stopping_criteria_log_likelihood(self, dataset, model, mask, base_parameter): + """Use log-likelihood instead of log-posterior as stopping criterion.""" + base_parameter["background_normalization_optimization:activate"] = False + base_parameter["stopping_criteria:threshold"] = 1e10 + base_parameter["stopping_criteria:statistics"] = "log-likelihood" + alg = run(MAP_RichardsonLucy(model, dataset, mask, base_parameter), self.NUM_ITERATION) + # Regression check + assert np.isclose(alg.results[-1]['log-likelihood'][0], 3931.3528198012773) + + def test_without_gamma_prior(self, dataset, model, mask, base_parameter): + """TSV prior only, no gamma prior.""" + base_parameter["background_normalization_optimization:activate"] = False + base_parameter["stopping_criteria:statistics"] = "log-posterior" + base_parameter["stopping_criteria:threshold"] = 1e10 + base_parameter["prior"] = {"TSV": {"coefficient": 1e-10}} + alg = run(MAP_RichardsonLucy(model, dataset, mask, base_parameter), self.NUM_ITERATION) + # Regression check + assert np.isclose(alg.results[-1]['log-posterior'], 3931.29811442966) + + def test_invalid_stopping_statistics_raises(self, dataset, model, mask, base_parameter): + """Invalid stopping_criteria:statistics should raise ValueError.""" + base_parameter["stopping_criteria:statistics"] = "likelihooooooooood!!!" + with pytest.raises(ValueError): + MAP_RichardsonLucy(model, dataset, mask, base_parameter) diff --git a/tests/image_deconvolution/test_algorithm_base.py b/tests/image_deconvolution/test_algorithm_base.py index 5d90317f1..f75fb58d3 100644 --- a/tests/image_deconvolution/test_algorithm_base.py +++ b/tests/image_deconvolution/test_algorithm_base.py @@ -3,45 +3,27 @@ from cosipy.image_deconvolution import DeconvolutionAlgorithmBase + def test_deconvolution_algorithm_base(dataset, model, mask): + """All abstract methods should raise NotImplementedError.""" DeconvolutionAlgorithmBase.__abstractmethods__ = set() - parameter = Configurator({}) - - algorithm = DeconvolutionAlgorithmBase(initial_model = model, - dataset = dataset, - mask = mask, - parameter = parameter) - - with pytest.raises(RuntimeError) as e_info: - algorithm.initialization() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.pre_processing() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.Estep() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.Mstep() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.post_processing() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.register_result() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.check_stopping_criteria() - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - algorithm.finalization() - assert e_info.type is NotImplementedError + algorithm = DeconvolutionAlgorithmBase( + initial_model = model, + dataset = dataset, + mask = mask, + parameter = Configurator({}), + ) + + for method in [ + algorithm.initialization, + algorithm.pre_processing, + algorithm.processing_core, + algorithm.post_processing, + algorithm.register_result, + algorithm.check_stopping_criteria, + algorithm.finalization, + ]: + with pytest.raises(NotImplementedError): + method() diff --git a/tests/image_deconvolution/test_dataIF_collection.py b/tests/image_deconvolution/test_dataIF_collection.py new file mode 100644 index 000000000..0d107526e --- /dev/null +++ b/tests/image_deconvolution/test_dataIF_collection.py @@ -0,0 +1,43 @@ +import pytest +from histpy import Histogram +import numpy as np +import astropy.units as u + +from cosipy.image_deconvolution import DataIF_COSI_DC2, AllSkyImageModel, DataInterfaceCollection +from cosipy.response import FullDetectorResponse +from cosipy import test_data + +def test_dataIF_collection(): + + event_binned_data = Histogram.open(test_data.path / "test_event_histogram_galacticCDS.hdf5").project(["Em", "Phi", "PsiChi"]) + dict_bkg_binned_data = {"bkg": Histogram.open(test_data.path / "test_event_histogram_galacticCDS.hdf5").project(["Em", "Phi", "PsiChi"])} + precomputed_response = Histogram.open(test_data.path / "test_precomputed_response.h5") + + data = DataIF_COSI_DC2.load(name = "testdata_galacticCDS", + event_binned_data = event_binned_data, + dict_bkg_binned_data = dict_bkg_binned_data, + rsp = precomputed_response, + coordsys_conv_matrix = None) + + data_collection = DataInterfaceCollection([data]) + + model = AllSkyImageModel(precomputed_response.axes['NuLambda'].nside, precomputed_response.axes['Ei'].edges) + model[:] = 1.0 * model.unit + + expectation_list = data_collection.calc_expectation_list(model = model, dict_bkg_norm = {"bkg": 1.0}) + + expectation_list_src = data_collection.calc_source_expectation_list(model = model) + + expectation_list_bkg = data_collection.calc_bkg_expectation_list(dict_bkg_norm = {"bkg": 1.0}) + + assert np.all(expectation_list[0].contents == data_collection.combine_expectation_list(expectation_list_src, expectation_list_bkg)[0].contents) + + log_likelihood_list = data_collection.calc_total_log_likelihood(expectation_list) + + assert np.isclose(log_likelihood_list, -8330.6845369209) + + assert np.isclose(data_collection.calc_summed_bkg_model("bkg"), 4676.0) + + assert np.isclose(data_collection.calc_summed_bkg_model_product("bkg", expectation_list), 319612.9093017483) + + assert np.isclose(np.sum(data_collection.calc_summed_exposure_map()), 22580.773319674616 * u.Unit("cm2 s sr")) diff --git a/tests/image_deconvolution/test_prior_base.py b/tests/image_deconvolution/test_prior_base.py index b3a8b9264..077614115 100644 --- a/tests/image_deconvolution/test_prior_base.py +++ b/tests/image_deconvolution/test_prior_base.py @@ -1,29 +1,25 @@ import pytest import numpy as np -from cosipy.image_deconvolution.prior_base import PriorBase +from cosipy.image_deconvolution.algorithms.prior_base import PriorBase + def test_PriorBase(): + PriorBase.__abstractmethods__ = set() - - # no class is allowered - with pytest.raises(TypeError) as e_info: - parameter = {'coefficient': 10} - test_model = np.zeros(2) - prior = PriorBase(parameter, test_model) - - # As a test, np.ndarray is added + + # Instantiation should fail when model class is not in usable_model_classes + with pytest.raises(TypeError): + PriorBase({'coefficient': 10}, np.zeros(2)) + + # Allow np.ndarray for testing PriorBase.usable_model_classes.append(np.ndarray) - parameter = {'coefficient': 10} - test_model = np.zeros(2) - prior = PriorBase(parameter, test_model) - - # other function tests - with pytest.raises(RuntimeError) as e_info: - prior.log_prior(test_model) - assert e_info.type is NotImplementedError - - with pytest.raises(RuntimeError) as e_info: - prior.grad_log_prior(test_model) - assert e_info.type is NotImplementedError + prior = PriorBase({'coefficient': 10}, np.zeros(2)) + + # Abstract methods should raise NotImplementedError + with pytest.raises(NotImplementedError): + prior.log_prior(np.zeros(2)) + + with pytest.raises(NotImplementedError): + prior.grad_log_prior(np.zeros(2)) diff --git a/tests/image_deconvolution/test_priors.py b/tests/image_deconvolution/test_priors.py index 1e2a1d7ef..5a4690f89 100644 --- a/tests/image_deconvolution/test_priors.py +++ b/tests/image_deconvolution/test_priors.py @@ -4,8 +4,8 @@ import numpy as np import healpy as hp -from cosipy.image_deconvolution.prior_tsv import PriorTSV -from cosipy.image_deconvolution.prior_entropy import PriorEntropy +from cosipy.image_deconvolution.algorithms.prior_tsv import PriorTSV +from cosipy.image_deconvolution.algorithms.prior_entropy import PriorEntropy from cosipy.image_deconvolution import AllSkyImageModel def test_PriorTSV(): @@ -38,7 +38,7 @@ def test_PriorTSV(): def test_PriorEntropy(): - parameter = {'coefficient': 1.0, 'prior_map': {'value': 1.0, 'unit': 'cm-2 s-1 sr-1'}} + parameter = {'coefficient': 1.0, 'reference_map': {'value': 1.0, 'unit': 'cm-2 s-1 sr-1'}} nside = 1 allskyimage_model = AllSkyImageModel(nside = nside,