FittingSequence stores the intermediate results of psf iterations

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

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

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)