From 589c3753b45b0d0d4839122dc846e3f440e7317b Mon Sep 17 00:00:00 2001 From: Hiroki Yoneda Date: Mon, 22 Jul 2024 17:58:27 +0200 Subject: [PATCH] Modified MAP RL algorithm - updated saving results - modifed parameter definition - output the log prior, posterior - modified TSV prior (bug fix) --- .../image_deconvolution/MAP_RichardsonLucy.py | 153 ++++++++++++------ cosipy/image_deconvolution/prior_base.py | 7 +- cosipy/image_deconvolution/prior_tsv.py | 17 +- 3 files changed, 123 insertions(+), 54 deletions(-) diff --git a/cosipy/image_deconvolution/MAP_RichardsonLucy.py b/cosipy/image_deconvolution/MAP_RichardsonLucy.py index fd588096..bc2b56c9 100644 --- a/cosipy/image_deconvolution/MAP_RichardsonLucy.py +++ b/cosipy/image_deconvolution/MAP_RichardsonLucy.py @@ -20,44 +20,45 @@ class MAP_RichardsonLucy(DeconvolutionAlgorithmBase): (need to fill here in the future) """ - prior_class = {"TSV": PriorTSV} + prior_classes = {"TSV": PriorTSV} def __init__(self, initial_model, dataset, mask, parameter): DeconvolutionAlgorithmBase.__init__(self, initial_model, dataset, mask, parameter) # response weighting filter - self.do_response_weighting = parameter.get('response_weighting', False) + 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) + self.response_weighting_index = parameter.get('response_weighting:index', 0.5) # background normalization optimization - self.do_bkg_norm_optimization = parameter.get('background_normalization_optimization', False) + 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_range', {key: [0.0, 100.0] for key in self.dict_bkg_norm.keys()}) + self.dict_bkg_norm_range = parameter.get('background_normalization_optimization:range', {key: [0.0, 100.0] for key in self.dict_bkg_norm.keys()}) # Prior distribution - self.prior_list = list(parameter.get('prior', {}).keys()) - self.priors = [] + self.prior_key_list = list(parameter.get('prior', {}).keys()) + self.priors = {} ## Gamma distribution - if 'gamma' in self.prior_list: + if 'gamma' in self.prior_key_list: this_prior_parameter = parameter['prior']['gamma'] self.load_gamma_prior(this_prior_parameter) else: self.load_gamma_prior(None) ## other priors - for prior_name in self.prior_list: + for prior_name in self.prior_key_list: if prior_name == 'gamma': continue coefficient = parameter['prior'][prior_name]['coefficient'] - self.priors.append(self.prior_class[prior_name](coefficient, initial_model)) + self.priors[prior_name] = self.prior_classes[prior_name](coefficient, initial_model) # saving results - self.save_results = parameter.get('save_results', False) - self.save_results_directory = parameter.get('save_results_directory', './results') + self.save_results = parameter.get('save_results:activate', False) + self.save_results_directory = parameter.get('save_results:directory', './results') + self.save_results_only_final = parameter.get('save_results:only_final_result', False) if self.save_results is True: if os.path.isdir(self.save_results_directory): @@ -69,13 +70,37 @@ def load_gamma_prior(self, parameter): if parameter is None: self.prior_gamma_model_theta, self.prior_gamma_model_k = np.inf * self.initial_model.unit, 1.0 #flat distribution - self.prior_gamma_background_theta, self.prior_gamma_background_k = np.inf, 1.0 #flat distribution + self.prior_gamma_bkg_theta, self.prior_gamma_bkg_k = np.inf, 1.0 #flat distribution else: self.prior_gamma_model_theta = parameter['model']['theta']['value'] * u.Unit(parameter['model']['theta']['unit']) self.prior_gamma_model_k = parameter['model']['k']['value'] - self.prior_gamma_background_theta = parameter['background']['theta']['value'] - self.prior_gamma_background_k = parameter['background']['k']['value'] + self.prior_gamma_bkg_theta = parameter['background']['theta']['value'] + self.prior_gamma_bkg_k = parameter['background']['k']['value'] + + def log_gamma_prior(self, model): + + eps = np.finfo(model.contents.dtype).eps + + # model + pl_part_model = (self.prior_gamma_model_k - 1.0) * np.sum( np.log(model.contents + eps) ) if model.unit is None else \ + (self.prior_gamma_model_k - 1.0) * np.sum( np.log(model.contents.value + eps) ) + + log_part_model = - np.sum( model.contents / self.prior_gamma_model_theta ) + + # background + pl_part_bkg, log_part_bkg = 0, 0 + + if self.do_bkg_norm_optimization: + for key in self.dict_bkg_norm.keys(): + + bkg_norm = self.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 def initialization(self): """ @@ -134,23 +159,36 @@ def Mstep(self): ratio_list = [ data.event / expectation for data, expectation in zip(self.dataset, self.expectation_list) ] - # model + # model update (EM part) sum_T_product = self.calc_summed_T_product(ratio_list) - self.model = (self.model * sum_T_product + self.prior_gamma_model_k - 1.0) \ - / (self.summed_exposure_map + 1.0 / self.prior_gamma_model_theta) - self.model[:] = np.where(self.model.contents < self.minimum_flux, self.minimum_flux, self.model.contents) - - # prior - sum_grad_log_prior = 0 + 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) + + # model update (prior part) + sum_grad_log_prior = np.zeros_like(self.summed_exposure_map) + + for key in self.priors.keys(): + sum_grad_log_prior += self.priors[key].grad_log_prior(model_EM) + + self.prior_filter = Histogram(self.model.axes, contents = np.exp(-1.0 * sum_grad_log_prior / self.summed_exposure_map.contents) ) + + self.model[:] = self.prior_filter.contents * model_EM.contents - for prior in self.priors: - sum_grad_log_prior += prior.grad_log_prior(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[:] *= np.exp( -1.0 * sum_grad_log_prior / self.summed_exposure_map.contents) + self.model[:] = (self.model.contents - delta_model.contents) + self.response_weighting_filter * delta_model.contents # masking if self.mask is not None: - self.delta_model = self.delta_model.mask_pixels(self.mask) + 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: @@ -158,8 +196,8 @@ def Mstep(self): 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 + self.prior_gamma_background_k - 1.0) \ - / (sum_bkg_model + 1.0 / self.prior_gamma_background_theta) + 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]: @@ -177,20 +215,6 @@ def post_processing(self): #TODO: add acceleration SQUAREM - 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 - delta_model) + delta_model * self.response_weighting_filter - - 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.") @@ -199,24 +223,44 @@ def post_processing(self): self.loglikelihood_list = self.calc_loglikelihood_list(self.expectation_list) logger.debug("The loglikelihood 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) + + for key in self.priors.keys(): + self.log_priors[key] = -1.0 * self.priors[key].log_prior(self.model) + # NOTE: the log_prior is defined as the negative of log of prior function + + # posterior + self.posterior = np.sum(self.loglikelihood_list) + np.sum([self.log_priors[key] for key in self.log_priors.keys()]) + def register_result(self): """ The values below are stored at the end of each iteration. - iteration: iteration number - model: updated image + - prior_filter: prior filter - background_normalization: optimized background normalization - loglikelihood: log-likelihood + - log_prior: log-prior + - posterior: posterior """ - #TODO: add posterior 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), - "loglikelihood": copy.deepcopy(self.loglikelihood_list)} + "loglikelihood": copy.deepcopy(self.loglikelihood_list), + "log_prior": copy.deepcopy(self.log_priors), + "posterior": copy.deepcopy(self.posterior), + } # show intermediate results logger.info(f' background_normalization: {this_result["background_normalization"]}') logger.info(f' loglikelihood: {this_result["loglikelihood"]}') + logger.info(f' log_prior: {this_result["log_prior"]}') + logger.info(f' posterior: {this_result["posterior"]}') # register this_result in self.results self.results.append(this_result) @@ -237,15 +281,21 @@ def finalization(self): """ finalization after running the image deconvolution """ - #TODO: add posterior if self.save_results == True: logger.info('Saving results in {self.save_results_directory}') # model + self.results[-1]["model"].write(f"{self.save_results_directory}/model.hdf5", name = 'result', overwrite = True) + + self.results[-1]["prior_filter"].write(f"{self.save_results_directory}/prior_filter.hdf5", name = 'result', overwrite = True) + for this_result in self.results: + iteration_count = this_result["iteration"] - this_result["model"].write(f"{self.save_results_directory}/model_itr{iteration_count}.hdf5", overwrite = True) + this_result["model"].write(f"{self.save_results_directory}/model.hdf5", name = f'iteration{iteration_count}', overwrite = True) + + this_result["prior_filter"].write(f"{self.save_results_directory}/prior_filter.hdf5", name = f'iteration{iteration_count}', overwrite = True) #fits primary_hdu = fits.PrimaryHDU() @@ -255,6 +305,9 @@ def finalization(self): for key in self.dict_bkg_norm.keys()] cols_loglikelihood = [fits.Column(name=f"{self.dataset[i].name}", array=[float(result['loglikelihood'][i]) for result in self.results], format='D') for i in range(len(self.dataset))] + cols_log_prior = [fits.Column(name=key, array=[float(result['log_prior'][key]) for result in self.results], format='D') + for key in self.results[0]['log_prior'].keys()] + cols_posterior = fits.Column(name = 'posterior', array=[float(result['posterior']) for result in self.results], format='D') table_bkg_norm = fits.BinTableHDU.from_columns([col_iteration] + cols_bkg_norm) table_bkg_norm.name = "bkg_norm" @@ -262,5 +315,11 @@ def finalization(self): table_loglikelihood = fits.BinTableHDU.from_columns([col_iteration] + cols_loglikelihood) table_loglikelihood.name = "loglikelihood" - hdul = fits.HDUList([primary_hdu, table_alpha, table_bkg_norm, table_loglikelihood]) + table_log_prior = fits.BinTableHDU.from_columns([col_iteration] + cols_log_prior) + table_log_prior.name = "log_prior" + + table_posterior = fits.BinTableHDU.from_columns([col_iteration, cols_posterior]) + table_posterior.name = "posterior" + + hdul = fits.HDUList([primary_hdu, table_bkg_norm, table_loglikelihood, table_log_prior, table_posterior]) hdul.writeto(f'{self.save_results_directory}/results.fits', overwrite=True) diff --git a/cosipy/image_deconvolution/prior_base.py b/cosipy/image_deconvolution/prior_base.py index fbf35276..f33fb0cd 100644 --- a/cosipy/image_deconvolution/prior_base.py +++ b/cosipy/image_deconvolution/prior_base.py @@ -1,8 +1,11 @@ class PriorBase: - allowerd_model_class = [] + usable_model_classes = [] def __init__(self, coefficient, model): + + if not self.is_calculable(model): + raise TypeError self.coefficient = coefficient @@ -10,7 +13,7 @@ def __init__(self, coefficient, model): def is_calculable(self, model): - if type(model) in self.allowerd_model_class: + if type(model) in self.usable_model_classes: return True diff --git a/cosipy/image_deconvolution/prior_tsv.py b/cosipy/image_deconvolution/prior_tsv.py index 3e3c4808..73da6093 100644 --- a/cosipy/image_deconvolution/prior_tsv.py +++ b/cosipy/image_deconvolution/prior_tsv.py @@ -7,7 +7,7 @@ class PriorTSV(PriorBase): - allowerd_model_class = [AllSkyImageModel] + usable_model_classes = [AllSkyImageModel] def __init__(self, coefficient, model): @@ -21,22 +21,29 @@ def __init__(self, coefficient, model): theta, phi = hp.pix2ang(nside = nside, ipix = np.arange(npix), nest = nest) - self.neighbour_pixel_index = hp.get_all_neighbours(nside = nside, theta = theta, phi = phi, nest = nest) + self.neighbour_pixel_index = hp.get_all_neighbours(nside = nside, theta = theta, phi = phi, nest = nest) # Its shape is (8, num. of pixels) + + self.num_neighbour_pixels = np.sum(self.neighbour_pixel_index >= 0, axis = 0) # Its shape is (num. of pixels) + + # replace -1 with its pixel index + for idx, ipixel in np.argwhere(self.neighbour_pixel_index == -1): + self.neighbour_pixel_index[idx, ipixel] = ipixel def log_prior(self, model): if self.model_class == AllSkyImageModel: diff = (model[:] - model[self.neighbour_pixel_index]).value - diff[np.isnan(diff)] = 0 + # Its shape is (8, num. of pixels, num. of energies) + diff /= self.num_neighbour_pixels.reshape((1,-1,1)) - return self.coefficient * np.sum(diff**2, axis = 0) + return self.coefficient * np.sum(diff**2) def grad_log_prior(self, model): if self.model_class == AllSkyImageModel: diff = (model[:] - model[self.neighbour_pixel_index]).value - diff[np.isnan(diff)] = 0 + diff /= self.num_neighbour_pixels.reshape((1,-1,1)) return self.coefficient * 4 * np.sum(diff, axis = 0) / model.unit