From f0ab40816e31e4624f852ea0ee9f937e30c3167c Mon Sep 17 00:00:00 2001 From: Simon Birrer Date: Fri, 19 Jul 2024 12:20:48 +0200 Subject: [PATCH] FittingSequence stores the intermediate results of psf iterations --- lenstronomy/Util/class_creator.py | 2 +- lenstronomy/Workflow/fitting_sequence.py | 22 +++++++++++++++++++++ test/test_Workflow/test_fitting_sequence.py | 7 +++++++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/lenstronomy/Util/class_creator.py b/lenstronomy/Util/class_creator.py index e1d457ae7..d0c0e3b36 100644 --- a/lenstronomy/Util/class_creator.py +++ b/lenstronomy/Util/class_creator.py @@ -89,7 +89,7 @@ def create_class_instances( :param index_lens_model_list: :param index_source_light_model_list: :param index_lens_light_model_list: optional, list of list of all model indexes for each modeled band - :param index_point_source_model_list: + :param index_point_source_model_list: optional, list of list of all model indexes for each modeled band :param optical_depth_model_list: list of strings indicating the optical depth model to compute (differential) extinctions from the source :param index_optical_depth_model_list: :param band_index: int, index of band to consider. Has an effect if only partial models are considered for a specific band diff --git a/lenstronomy/Workflow/fitting_sequence.py b/lenstronomy/Workflow/fitting_sequence.py index 7f469b18f..f83162fbd 100644 --- a/lenstronomy/Workflow/fitting_sequence.py +++ b/lenstronomy/Workflow/fitting_sequence.py @@ -1,3 +1,5 @@ +import copy + from lenstronomy.Workflow.psf_fitting import PsfFitting from lenstronomy.Workflow.alignment_matching import AlignmentFitting from lenstronomy.Workflow.flux_calibration import FluxCalibration @@ -67,6 +69,8 @@ def __init__( num_bands=len(self.multi_band_list), ) self._mcmc_init_samples = None + self._psf_iteration_memory = [] + self._psf_iteration_index = 0 # index of the sequence of the PSF iteration (how many times it is being run) @property def kwargs_fixed(self): @@ -583,6 +587,7 @@ def psf_iteration(self, compute_bands=None, **kwargs_psf_iter): for band_index in range(len(self.multi_band_list)): if compute_bands[band_index] is True: kwargs_psf = self.multi_band_list[band_index][1] + kwargs_psf_before = copy.deepcopy(kwargs_psf) image_model = SingleBandMultiModel( self.multi_band_list, kwargs_model, @@ -595,6 +600,9 @@ def psf_iteration(self, compute_bands=None, **kwargs_psf_iter): kwargs_psf, kwargs_params=kwargs_temp, **kwargs_psf_iter ) self.multi_band_list[band_index][1] = kwargs_psf + self._psf_iteration_memory.append({'sequence': self._psf_iteration_index, 'band': band_index, + "psf_before": kwargs_psf_before, "psf_after": kwargs_psf}) + self._psf_iteration_index += 1 return 0 def align_images( @@ -866,3 +874,17 @@ def best_fit_from_samples(self, samples, logl): # get corresponding kwargs kwargs_result = self.param_class.args2kwargs(best_fit_result, bijective=True) return kwargs_result + + @property + def psf_iteration_memory(self): + """ + returns all the psf iterations performed in the FittingSequence + It stores in a list of dictionaries: + "sequence": what PSF sequence it is (0, 1 etc) + "band": index of the imaging band that is being corrected + "psf_before" kwargs_psf prior to the iteration + "psf_after" kwargs_psf as a result of the iteration + + :return: list of all psf corrections + """ + return self._psf_iteration_memory \ No newline at end of file diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 932a956b5..13ec0a2a9 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -377,6 +377,13 @@ def test_fitting_sequence(self): fitting_list_three.append(["emcee", kwargs_test]) fittingSequence.fit_sequence(fitting_list_three) + psf_iteration_list = fittingSequence.psf_iteration_memory + assert len(psf_iteration_list) == 1 + assert "sequence" in psf_iteration_list[0] + assert "band" in psf_iteration_list[0] + assert "psf_before" in psf_iteration_list[0] + assert "psf_after" in psf_iteration_list[0] + def test_cobaya(self): np.random.seed(42)