Masked Sensitivity Mapping

[Jammy2211]

The following example script runs sensitivity mapping for a strong lens:

https://github.com/Jammy2211/autolens_workspace_test/blob/main/slam/integration/source_lp/mass_total/sensitivity.py

It performs a 2 x 2 sensitivity grid, in order to run fast.

At the end, I have added a new input the sensitivity_mask, which signifies that only one of the four grid cells should run (in autolens convention True means to ignore).

    subhalo_result = slam.subhalo.sensitivity_imaging_lp.run(
        settings_search=settings_search,
        mask=mask,
        psf=dataset.psf,
        adapt_images=al.AdaptImages.from_result(result=source_lp_result),
        mass_result=mass_result,
        subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),
        grid_dimension_arcsec=3.0,
        number_of_steps=2,
        sensitivity_mask=[[True, True], [True, False]],
    )

On the following PyAutoFit branch I got the sensitivity map to successfully only run 1 in 4 times:

https://github.com/rhayes777/PyAutoFit/tree/feature/sensitivity_mask

The main edit was here:

https://github.com/rhayes777/PyAutoFit/blob/feature/sensitivity_mask/autofit/non_linear/grid/sensitivity/__init__.py

    def _make_jobs(self) -> Generator[Job, None, None]:
        """
        Create a list of jobs to be run on separate processes.

        Each job fits a perturb image with the original model
        and a model which includes a perturbation.
        """
        for number, (perturb_instance, perturb_model, label) in enumerate(
            zip(self._perturb_instances, self._perturb_models, self._labels)
        ):

            shape_index = self.shape_index_from_number(number=number)

            should_bypass = False

            if self.mask is not None:
                should_bypass = np.asarray(self.mask)[shape_index]

            if not should_bypass:

                if self.perturb_model_prior_func is not None:
                    perturb_model = self.perturb_model_prior_func(
                        perturb_instance=perturb_instance, perturb_model=perturb_model
                    )

                simulate_instance = copy(self.instance)
                simulate_instance.perturb = perturb_instance

                paths = self.paths.for_sub_analysis(
                    label,
                )

                yield self.job_cls(
                    simulate_instance=simulate_instance,
                    model=self.model,
                    perturb_model=perturb_model,
                    base_instance=self.instance,
                    simulate_cls=self.simulate_cls,
                    base_fit_cls=self.base_fit_cls,
                    perturb_fit_cls=self.perturb_fit_cls,
                    paths=paths,
                    number=number,
                )

The code is set up to retain the overall grid information of the sensitivity map, for example, with this mask the job that does run is still job number 3, as seen in results.csv:

index,centre_0,centre_1,log_evidence_increase,log_likelihood_increase
    3,     1.5,     1.5,  -1613.7677295368567,    -1613.7677295368567

This is important so that I could change the mask and run more cells in the sensitivity map without issue, and to in general retain the sensitivity map structure.

The problem is that the SensitivityResult currently does not support a mask input, and calling any of its methods after a mask is used raises an exception:

Traceback (most recent call last):
  File "/mnt/c/Users/Jammy/Code/PyAuto/autolens_workspace_test/slam/integration/source_lp/mass_total/sensitivity.py", line 356, in <module>
    fit()
  File "/mnt/c/Users/Jammy/Code/PyAuto/autolens_workspace_test/slam/integration/source_lp/mass_total/sensitivity.py", line 342, in fit
    subhalo_result = slam.subhalo.sensitivity_imaging_lp.run(
  File "/mnt/c/Users/Jammy/Code/PyAuto/autolens_workspace_test/slam/subhalo/sensitivity_imaging_lp.py", line 626, in run
    subhalo_util.visualize_sensitivity(
  File "/mnt/c/Users/Jammy/Code/PyAuto/autolens_workspace_test/slam/subhalo/subhalo_util.py", line 142, in visualize_sensitivity
    plotter.subplot_figures_of_merit_grid()
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoLens/autolens/lens/sensitivity.py", line 367, in subplot_figures_of_merit_grid
    figures_of_merit = self.result.figure_of_merit_array(
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoLens/autolens/lens/sensitivity.py", line 133, in figure_of_merit_array
    return self._array_2d_from(values=figures_of_merits)
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoLens/autolens/lens/sensitivity.py", line 89, in _array_2d_from
    values_reshaped = [value for values in values.native for value in values]
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoFit/autofit/non_linear/grid/grid_list.py", line 75, in native
    return np.reshape(np.array(self), self.shape)
  File "/home/jammy/venvs/PyAuto/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 285, in reshape
    return _wrapfunc(a, 'reshape', newshape, order=order)
  File "/home/jammy/venvs/PyAuto/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 59, in _wrapfunc
    return bound(*args, **kwds)
ValueError: cannot reshape array of size 1 into shape (2,2)

Can you do the following:

1. Make sure you agree that my implementation of the mask is the best way to do it, and refactor as you please.
2. Update SensitvityResult so that its attributes still work when a mask is input.

An example SensitivityResult attribute is the log evidence of the base run of every grid on the sensitivity map:

    @property
    @as_grid_list
    def log_evidences_base(self) -> GridList:
        """
        The log evidences of the base model for each sensitivity fit
        """
        return [sample.log_evidence for sample in self.samples]

The problem is that self.samples is no longer a list of 4 entries but fewer than 4 entries, depending on how the mask sets up the shape.

I think we basically need the GridList object to account for the mask when doing this mapping, and to do so in such a way that all masked entries in log_evidences_base become None.

https://github.com/rhayes777/PyAutoFit/blob/feature/sensitivity_mask/autofit/non_linear/grid/grid_list.py