Repository: rhayes777/PyAutoFit (python)

Issue:
Bug where sensitivity mapping no longer supports tuple step_size
In a previous issue, we made it so that the `number_of_steps` in the `Sensitivity` class could be a tuple so that different dimensions had
different steps:

```
sensitivity = af.Sensitivity(
    paths=paths,
    simulation_instance=simulation_instance,
    base_model=base_model,
    perturb_model=perturb_model,
    simulate_cls=SimulateImaging(mask=mask, psf=dataset.psf),
    base_fit_cls=BaseFit(adapt_result=result),
    perturb_fit_cls=PerturbFit(adapt_result=result),
    perturb_model_prior_func=perturb_model_prior_func,
  #  number_of_steps=2,
    number_of_steps=(4, 3, 2),
    number_of_cores=2,
)
```

We have since added a function which customized the prior on sensitivity mapping:

```
    @property
    def _perturb_models(self) -> Generator[AbstractPriorModel, None, None]:
        """
        A list of models representing a perturbation at each grid square.

        By default models have priors with limits at the edges of a grid square.
        These limits can be scaled using the limit_scale variable. If the variable
        is 2 then the priors will have width twice the step size.
        """
        half_step = self.limit_scale * self.step_size / 2
        for list_ in self._lists:
            limits = [
                (
                    prior.value_for(max(0.0, centre - half_step)),
                    prior.value_for(min(1.0, centre + half_step)),
                )
                for centre, prior in zip(list_, self.perturb_model.priors_ordered_by_id)
            ]
            yield self.perturb_model.with_limits(limits)
```

When a tuple is used the following error is raised:

```
Traceback (most recent call last):
  File "/mnt/c/Users/Jammy/Code/PyAuto/autolens_workspace/scripts/imaging/advanced/subhalo/sensitivity/start_here.py", line 618, in <module>
    sensitivity_result = sensitivity.run()
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoFit/autofit/non_linear/grid/sensitivity.py", line 311, in run
    for result in process_class.run_jobs(
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoFit/autofit/non_linear/parallel/process.py", line 162, in run_jobs
    for job in jobs:
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoFit/autofit/non_linear/grid/sensitivity.py", line 430, in _make_jobs
    for number, (perturb_instance, perturb_model, label) in enumerate(
  File "/mnt/c/Users/Jammy/Code/PyAuto/PyAutoFit/autofit/non_linear/grid/sensitivity.py", line 412, in _perturb_models
    half_step = self.limit_scale * self.step_size / 2
TypeError: unsupported operand type(s) for /: 'tuple' and 'int'
```

Pull request description:
feature/sensitivity tuples
Resolve #883


Patch:
diff --git a/autofit/non_linear/grid/sensitivity/__init__.py b/autofit/non_linear/grid/sensitivity/__init__.py
index fa19768c0..d59fedc54 100644
--- a/autofit/non_linear/grid/sensitivity/__init__.py
+++ b/autofit/non_linear/grid/sensitivity/__init__.py
@@ -15,6 +15,7 @@ from autofit.non_linear.grid.sensitivity.result import SensitivityResult
 from autofit.non_linear.parallel import Process
 from autofit.text.text_util import padding
 
+
 class Sensitivity:
     def __init__(
         self,
@@ -26,7 +27,7 @@ class Sensitivity:
         base_fit_cls: Callable,
         perturb_fit_cls: Callable,
         job_cls: ClassVar = Job,
-        perturb_model_prior_func : Optional[Callable] = None,
+        perturb_model_prior_func: Optional[Callable] = None,
         number_of_steps: Union[Tuple[int, ...], int] = 4,
         number_of_cores: int = 2,
         limit_scale: int = 1,
@@ -130,7 +131,6 @@ class Sensitivity:
                 writer = csv.writer(f)
                 writer.writerow(headers)
                 for result_ in results:
-
                     values = physical_values[result_.number]
                     writer.writerow(
                         padding(item)
@@ -144,8 +144,10 @@ class Sensitivity:
 
         result = SensitivityResult(
             samples=[result.result.samples.summary() for result in results],
-            perturb_samples=[result.perturb_result.samples.summary() for result in results],
-            shape=self.shape
+            perturb_samples=[
+                result.perturb_result.samples.summary() for result in results
+            ],
+            shape=self.shape,
         )
 
         self.paths.save_json("result", to_dict(result))
@@ -170,7 +172,9 @@ class Sensitivity:
         if isinstance(self.number_of_steps, tuple):
             return self.number_of_steps
 
-        return tuple(self.number_of_steps for _ in range(self.perturb_model.prior_count))
+        return tuple(
+            self.number_of_steps for _ in range(self.perturb_model.prior_count)
+        )
 
     @property
     def step_size(self) -> Union[float, Tuple]:
@@ -181,7 +185,9 @@ class Sensitivity:
             The size of a step in any given dimension in hyper space.
         """
         if isinstance(self.number_of_steps, tuple):
-            return tuple(1 / number_of_steps for number_of_steps in self.number_of_steps)
+            return tuple(
+                1 / number_of_steps for number_of_steps in self.number_of_steps
+            )
         return 1 / self.number_of_steps
 
     @property
@@ -253,14 +259,23 @@ class Sensitivity:
         These limits can be scaled using the limit_scale variable. If the variable
         is 2 then the priors will have width twice the step size.
         """
-        half_step = self.limit_scale * self.step_size / 2
+        if isinstance(self.step_size, tuple):
+            step_sizes = self.step_size
+        else:
+            step_sizes = (self.step_size,) * self.perturb_model.prior_count
+
+        half_steps = [self.limit_scale * step_size / 2 for step_size in step_sizes]
         for list_ in self._lists:
             limits = [
                 (
                     prior.value_for(max(0.0, centre - half_step)),
                     prior.value_for(min(1.0, centre + half_step)),
                 )
-                for centre, prior in zip(list_, self.perturb_model.priors_ordered_by_id)
+                for centre, prior, half_step in zip(
+                    list_,
+                    self.perturb_model.priors_ordered_by_id,
+                    half_steps,
+                )
             ]
             yield self.perturb_model.with_limits(limits)
 
@@ -274,11 +289,9 @@ class Sensitivity:
         for number, (perturb_instance, perturb_model, label) in enumerate(
             zip(self._perturb_instances, self._perturb_models, self._labels)
         ):
-
             if self.perturb_model_prior_func is not None:
                 perturb_model = self.perturb_model_prior_func(
-                    perturb_instance=perturb_instance,
-                    perturb_model=perturb_model
+                    perturb_instance=perturb_instance, perturb_model=perturb_model
                 )
 
             simulate_instance = copy(self.instance)
@@ -298,4 +311,4 @@ class Sensitivity:
                 perturb_fit_cls=self.perturb_fit_cls,
                 paths=paths,
                 number=number,
-            )
\ No newline at end of file
+            )


Test patch:
diff --git a/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py b/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
index 5fb46770d..4d1330d1c 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
@@ -28,7 +28,6 @@ def test_labels(sensitivity):
     ]
 
 
-
 def test_perform_job(job):
     result = job.perform()
     assert isinstance(result, s.JobResult)
@@ -40,23 +39,33 @@ class TestPerturbationModels:
     @pytest.mark.parametrize(
         "limit_scale, fl, fu, sl, su",
         [
-            (1.0, 0.0, 0.5, 0.5, 1.0,),
-            (2.0, 0.0, 0.75, 0.25, 1.0,),
-            (4.0, 0.0, 1.0, 0.0, 1.0,),
-        ]
+            (
+                1.0,
+                0.0,
+                0.5,
+                0.5,
+                1.0,
+            ),
+            (
+                2.0,
+                0.0,
+                0.75,
+                0.25,
+                1.0,
+            ),
+            (
+                4.0,
+                0.0,
+                1.0,
+                0.0,
+                1.0,
+            ),
+        ],
     )
-    def test_perturb_models(
-            self,
-            sensitivity,
-            limit_scale,
-            fl, fu, sl, su
-    ):
+    def test_perturb_models(self, sensitivity, limit_scale, fl, fu, sl, su):
         sensitivity.limit_scale = limit_scale
         jobs = sensitivity._make_jobs()
-        models = [
-            job.perturb_model
-            for job in jobs
-        ]
+        models = [job.perturb_model for job in jobs]
 
         first, second, *_ = models
 
@@ -68,24 +77,17 @@ class TestPerturbationModels:
         assert second.sigma.upper_limit == su
 
     def test__perturb_models__prior_overwrite_via_perturb_model_prior_func(
-            self,
-            sensitivity,
+        self,
+        sensitivity,
     ):
-        def perturb_model_prior_func(
-                perturb_instance,
-                perturb_model
-        ):
-
+        def perturb_model_prior_func(perturb_instance, perturb_model):
             perturb_model.centre = af.UniformPrior(lower_limit=-7.0, upper_limit=4.0)
 
             return perturb_model
 
         sensitivity.perturb_model_prior_func = perturb_model_prior_func
         jobs = sensitivity._make_jobs()
-        models = [
-            job.perturb_model
-            for job in jobs
-        ]
+        models = [job.perturb_model for job in jobs]
 
         first, second, *_ = models
 
@@ -94,26 +96,21 @@ class TestPerturbationModels:
         assert first.centre.lower_limit == -7.0
         assert first.centre.upper_limit == 4.0
 
-    def test_physical(
-            self,
-            sensitivity
-    ):
-        sensitivity.perturb_model.sigma = af.UniformPrior(
-            upper_limit=10
-        )
-        model = list(
-            sensitivity._make_jobs()
-        )[0].perturb_model
+    def test_physical(self, sensitivity):
+        sensitivity.perturb_model.sigma = af.UniformPrior(upper_limit=10)
+        model = list(sensitivity._make_jobs())[0].perturb_model
         assert model.sigma.upper_limit == 5
 
     def test_model_with_limits(self):
         model = af.Model(af.Gaussian)
 
-        with_limits = model.with_limits([
-            (0.3, 0.5),
-            (0.3, 0.5),
-            (0.3, 0.5),
-        ])
+        with_limits = model.with_limits(
+            [
+                (0.3, 0.5),
+                (0.3, 0.5),
+                (0.3, 0.5),
+            ]
+        )
         assert with_limits.centre.lower_limit == 0.3
         assert with_limits.centre.upper_limit == 0.5
 
@@ -126,8 +123,21 @@ class TestPerturbationModels:
         assert prior.upper_limit == 5
 
     def test_existing_limits(self):
-        prior = af.UniformPrior(
-            2, 4
-        ).with_limits(3, 5)
+        prior = af.UniformPrior(2, 4).with_limits(3, 5)
         assert prior.lower_limit == 3
         assert prior.upper_limit == 4
+
+
+@pytest.fixture(name="tuple_sensitivity")
+def make_tuple_sensitivity(sensitivity):
+    sensitivity.number_of_steps = (2, 2, 4)
+    return sensitivity
+
+
+def test_models_with_tuple_steps(tuple_sensitivity):
+    tuple_sensitivity.number_of_steps = (2, 2, 4)
+    assert len(list(tuple_sensitivity._perturb_models)) == 16
+
+
+def test_jobs_with_tuple_steps(tuple_sensitivity):
+    assert len(list(tuple_sensitivity._make_jobs())) == 16