Repository: rhayes777/PyAutoFit
Pull request #1062 (2024-10-11 11:41:18)

feature/sensitivity skip
Partially resolves #1059

Essentially if a job has been performed we skip simulation and pass the data as None. It's a bit tricky to make the overall result as responsibility for creating base and peturb results is delegated to the base_fit_cls and peturb_fit_cls respectively. It is assumed that if these are called again with dataset = None that: 1) it won't crash. 2) they will just load the previously output results



Patch:
diff --git a/autofit/example/analysis.py b/autofit/example/analysis.py
index 83db61645..581415228 100644
--- a/autofit/example/analysis.py
+++ b/autofit/example/analysis.py
@@ -132,7 +132,7 @@ class Analysis(af.Analysis):
         Parameters
         ----------
         paths
-            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         """
         paths.save_json(name="data", object_dict=self.data.tolist(), prefix="dataset")
diff --git a/autofit/example/visualize.py b/autofit/example/visualize.py
index a66f49824..7493021d6 100644
--- a/autofit/example/visualize.py
+++ b/autofit/example/visualize.py
@@ -33,7 +33,7 @@ class VisualizerExample(af.Visualizer):
         analysis
             The analysis class used to perform the model-fit whose quantities are being visualized.
         paths
-            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         model
             The model which is fitted to the data, which may be used to customize the visualization.
@@ -84,7 +84,7 @@ class VisualizerExample(af.Visualizer):
         analysis
             The analysis class used to perform the model-fit whose quantities are being visualized.
         paths
-            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
@@ -163,7 +163,7 @@ class VisualizerExample(af.Visualizer):
         analyses
             A list of the analysis classes used to perform the model-fit whose quantities are being visualized.
         paths
-            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         model
             The model which is fitted to the data, which may be used to customize the visualization.
@@ -202,7 +202,7 @@ class VisualizerExample(af.Visualizer):
         analyses
             A list of the analysis classes used to perform the model-fit whose quantities are being visualized.
         paths
-            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization, and the pickled objects used by the aggregator output by this function.
         model
             The model which is fitted to the data, which may be used to customize the visualization.
diff --git a/autofit/non_linear/analysis/visualize.py b/autofit/non_linear/analysis/visualize.py
index 7b0e74313..d0000c28a 100644
--- a/autofit/non_linear/analysis/visualize.py
+++ b/autofit/non_linear/analysis/visualize.py
@@ -32,7 +32,7 @@ class Visualizer:
         Parameters
         ----------
         paths
-            The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+            The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
             visualization and the pickled objects used by the aggregator output by this function.
 
         Returns
diff --git a/autofit/non_linear/grid/sensitivity/__init__.py b/autofit/non_linear/grid/sensitivity/__init__.py
index 4b8e672e5..889da6e55 100644
--- a/autofit/non_linear/grid/sensitivity/__init__.py
+++ b/autofit/non_linear/grid/sensitivity/__init__.py
@@ -1,7 +1,7 @@
-import csv
 import logging
 import os
 from copy import copy
+import numpy as np
 from pathlib import Path
 from typing import List, Generator, Callable, ClassVar, Optional, Union, Tuple
 
@@ -9,12 +9,11 @@ from autoconf.dictable import to_dict
 from autofit.mapper.model import ModelInstance
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
 from autofit.non_linear.grid.grid_search import make_lists, Sequential
-from autofit.non_linear.grid.sensitivity.job import Job
+from autofit.non_linear.grid.sensitivity.job import Job, MaskedJobResult
 from autofit.non_linear.grid.sensitivity.job import JobResult
 from autofit.non_linear.grid.sensitivity.result import SensitivityResult
 from autofit.non_linear.parallel import Process
 from autofit.text.formatter import write_table
-from autofit.text.text_util import padding
 
 
 class Sensitivity:
@@ -30,6 +29,7 @@ class Sensitivity:
         job_cls: ClassVar = Job,
         perturb_model_prior_func: Optional[Callable] = None,
         number_of_steps: Union[Tuple[int, ...], int] = 4,
+        mask: Optional[List[bool]] = None,
         number_of_cores: int = 2,
         limit_scale: int = 1,
     ):
@@ -65,6 +65,10 @@ class Sensitivity:
             The number of steps for each dimension of the sensitivity grid. If input as a float the dimensions are
             all that value. If input as a tuple of length the number of dimensions, each tuple value is the number of
             steps in that dimension.
+        mask
+            A mask to apply to the sensitivity grid, such that all `True` values are not included in the sensitivity
+            mapping. This is useful for removing regions of the sensitivity grid that are expected to have no
+            sensitivity, for example because they have no signal.
         number_of_cores
             How many cores does this computer have?
         limit_scale
@@ -94,6 +98,22 @@ class Sensitivity:
         self.job_cls = job_cls
 
         self.number_of_steps = number_of_steps
+        self.mask = None
+
+        if mask is not None:
+            self.mask = np.asarray(mask)
+            if self.shape != self.mask.shape:
+                raise ValueError(
+                    f"""
+                    The mask of the Sensitivity object must have the same shape as the sensitivity grid.
+                    
+                    For your inputs, the shape of each are as follows:
+                    
+                    Sensitivity Grid: {self.shape}
+                    Mask: {self.mask.shape}
+                    """
+                )
+
         self.number_of_cores = number_of_cores
 
         self.limit_scale = limit_scale
@@ -117,9 +137,17 @@ class Sensitivity:
 
         process_class = Process if self.number_of_cores > 1 else Sequential
 
-        results = list()
+        results = []
+        jobs = []
+
+        for number in range(len(self._perturb_instances)):
+            if self._should_bypass(number=number):
+                results.append(MaskedJobResult(number=number))
+            else:
+                jobs.append(self._make_job(number))
+
         for result in process_class.run_jobs(
-            self._make_jobs(), number_of_cores=self.number_of_cores
+            jobs, number_of_cores=self.number_of_cores
         ):
             if isinstance(result, Exception):
                 raise result
@@ -177,6 +205,21 @@ class Sensitivity:
             self.number_of_steps for _ in range(self.perturb_model.prior_count)
         )
 
+    def shape_index_from_number(self, number: int) -> Tuple[int, ...]:
+        """
+        Returns the index of the sensitivity grid from a number.
+
+        Parameters
+        ----------
+        number
+            The number of the sensitivity grid.
+
+        Returns
+        -------
+        The index of the sensitivity grid.
+        """
+        return np.unravel_index(number, self.shape)
+
     @property
     def step_size(self) -> Union[float, Tuple]:
         """
@@ -228,31 +271,36 @@ class Sensitivity:
             yield path
 
     @property
-    def _labels(self) -> Generator[str, None, None]:
+    def _labels(self) -> List[str]:
         """
         One label for each perturbation, used to distinguish
         fits for each perturbation by placing them in separate
         directories.
         """
+        labels = []
         for list_ in self._lists:
             strings = list()
             for value, prior_tuple in zip(list_, self.perturb_model.prior_tuples):
                 path, prior = prior_tuple
                 value = prior.value_for(value)
                 strings.append(f"{path}_{value}")
-            yield "_".join(strings)
+            labels.append("_".join(strings))
+
+        return labels
 
     @property
-    def _perturb_instances(self) -> Generator[ModelInstance, None, None]:
+    def _perturb_instances(self) -> List[ModelInstance]:
         """
         A list of instances each of which defines a perturbation to
         be applied to the image.
         """
-        for list_ in self._lists:
-            yield self.perturb_model.instance_from_unit_vector(list_)
+
+        return [
+            self.perturb_model.instance_from_unit_vector(list_) for list_ in self._lists
+        ]
 
     @property
-    def _perturb_models(self) -> Generator[AbstractPriorModel, None, None]:
+    def _perturb_models(self) -> List[AbstractPriorModel]:
         """
         A list of models representing a perturbation at each grid square.
 
@@ -266,6 +314,8 @@ class Sensitivity:
             step_sizes = (self.step_size,) * self.perturb_model.prior_count
 
         half_steps = [self.limit_scale * step_size / 2 for step_size in step_sizes]
+
+        perturb_models = []
         for list_ in self._lists:
             limits = [
                 (
@@ -278,38 +328,48 @@ class Sensitivity:
                     half_steps,
                 )
             ]
-            yield self.perturb_model.with_limits(limits)
+            perturb_models.append(self.perturb_model.with_limits(limits))
+        return perturb_models
+
+    def _should_bypass(self, number: int) -> bool:
+        shape_index = self.shape_index_from_number(number=number)
+        return self.mask is not None and np.asarray(self.mask)[shape_index]
 
     def _make_jobs(self) -> Generator[Job, None, None]:
+        for number, _ in enumerate(self._perturb_instances):
+            yield self._make_job(number)
+
+    def _make_job(self, number) -> 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)
-        ):
-            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
+        perturb_instance = self._perturb_instances[number]
+        perturb_model = self._perturb_models[number]
+        label = self._labels[number]
 
-            paths = self.paths.for_sub_analysis(
-                label,
+        if self.perturb_model_prior_func is not None:
+            perturb_model = self.perturb_model_prior_func(
+                perturb_instance=perturb_instance, perturb_model=perturb_model
             )
 
-            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,
-            )
+        simulate_instance = copy(self.instance)
+        simulate_instance.perturb = perturb_instance
+
+        paths = self.paths.for_sub_analysis(
+            label,
+        )
+
+        return 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,
+        )
diff --git a/autofit/non_linear/grid/sensitivity/job.py b/autofit/non_linear/grid/sensitivity/job.py
index c5089eb0f..03c2b2db6 100644
--- a/autofit/non_linear/grid/sensitivity/job.py
+++ b/autofit/non_linear/grid/sensitivity/job.py
@@ -35,7 +35,10 @@ class JobResult(AbstractJobResult):
 
         if hasattr(self.result.samples, "log_evidence"):
             if self.result.samples.log_evidence is not None:
-                return float(self.perturb_result.samples.log_evidence - self.result.samples.log_evidence)
+                return float(
+                    self.perturb_result.samples.log_evidence
+                    - self.result.samples.log_evidence
+                )
 
     @property
     def log_likelihood_increase(self) -> Optional[float]:
@@ -48,6 +51,23 @@ class JobResult(AbstractJobResult):
         return float(self.perturb_result.log_likelihood - self.result.log_likelihood)
 
 
+class MaskedJobResult(AbstractJobResult):
+    """
+    A placeholder result for a job that has been masked out.
+    """
+
+    @property
+    def result(self):
+        return self
+
+    @property
+    def perturb_result(self):
+        return self
+
+    def __getattr__(self, item):
+        return None
+
+
 class Job(AbstractJob):
     _number = count()
 
@@ -92,6 +112,24 @@ class Job(AbstractJob):
         self.perturb_fit_cls = perturb_fit_cls
         self.paths = paths
 
+    @property
+    def base_paths(self):
+        return self.paths.for_sub_analysis("[base]")
+
+    @property
+    def perturb_paths(self):
+        return self.paths.for_sub_analysis("[perturb]")
+
+    @property
+    def is_complete(self) -> bool:
+        """
+        Returns True if the job has been completed, False otherwise.
+        """
+        return (self.base_paths.is_complete and self.perturb_paths.is_complete) or (
+            (self.paths.output_path / "[base].zip").exists()
+            and (self.paths.output_path / "[perturb].zip").exists()
+        )
+
     def perform(self) -> JobResult:
         """
         - Create one model with a perturbation and another without
@@ -102,15 +140,19 @@ class Job(AbstractJob):
         An object comprising the results of the two fits
         """
 
-        dataset = self.simulate_cls(
-            instance=self.simulate_instance,
-            simulate_path=self.paths.image_path.with_name("simulate"),
-        )
+        if self.is_complete:
+            dataset = None
+        else:
+            dataset = self.simulate_cls(
+                instance=self.simulate_instance,
+                simulate_path=self.paths.image_path.with_name("simulate"),
+            )
 
         result = self.base_fit_cls(
             model=self.model,
             dataset=dataset,
-            paths=self.paths.for_sub_analysis("[base]"),
+            paths=self.base_paths,
+            instance=self.simulate_instance,
         )
 
         perturb_model = copy(self.model)
@@ -119,9 +161,12 @@ class Job(AbstractJob):
         perturb_result = self.perturb_fit_cls(
             model=perturb_model,
             dataset=dataset,
-            paths=self.paths.for_sub_analysis("[perturb]"),
+            paths=self.perturb_paths,
+            instance=self.simulate_instance,
         )
 
         return JobResult(
-            number=self.number, result=result, perturb_result=perturb_result
+            number=self.number,
+            result=result,
+            perturb_result=perturb_result,
         )
diff --git a/autofit/non_linear/mock/mock_search.py b/autofit/non_linear/mock/mock_search.py
index 4dbd53c12..c68b5d9da 100644
--- a/autofit/non_linear/mock/mock_search.py
+++ b/autofit/non_linear/mock/mock_search.py
@@ -136,6 +136,7 @@ class MockSearch(NonLinearSearch):
         )
 
         self.paths.save_samples_summary(self.samples_summary)
+        self.paths.completed()
 
         return analysis.make_result(
             samples_summary=samples_summary,
diff --git a/autofit/non_linear/plot/nest_plotters.py b/autofit/non_linear/plot/nest_plotters.py
index 64d4cd43a..2d4c470cb 100644
--- a/autofit/non_linear/plot/nest_plotters.py
+++ b/autofit/non_linear/plot/nest_plotters.py
@@ -1,5 +1,3 @@
-from anesthetic.samples import NestedSamples
-from anesthetic import make_2d_axes
 from functools import wraps
 import numpy as np
 import warnings
@@ -59,6 +57,8 @@ class NestPlotter(SamplesPlotter):
 
         config_dict = conf.instance["visualize"]["plots_settings"]["corner_anesthetic"]
 
+        from anesthetic.samples import NestedSamples
+        from anesthetic import make_2d_axes
         import matplotlib.pylab as pylab
 
         params = {'font.size' : int(config_dict["fontsize"])}
diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py
index 289f1a010..603236dc2 100644
--- a/autofit/non_linear/search/nest/nautilus/search.py
+++ b/autofit/non_linear/search/nest/nautilus/search.py
@@ -290,23 +290,6 @@ class Nautilus(abstract_nest.AbstractNest):
 
         finished = False
 
-        minimum_iterations_per_updates = 3 * self.config_dict_search["n_live"]
-        if self.iterations_per_update < minimum_iterations_per_updates:
-
-            self.iterations_per_update = minimum_iterations_per_updates
-
-            logger.info(
-                f"""
-                The number of iterations_per_update is less than 3 times the number of live points, which can cause
-                issues where Nautilus loses sampling information due to stopping to output results. The number of
-                iterations per update has been increased to 3 times the number of live points, therefore a value
-                of {minimum_iterations_per_updates}.
-
-                To remove this warning, increase the number of iterations_per_update to three or more times the
-                number of live points.
-                """
-            )
-
         while not finished:
 
             iterations, total_iterations = self.iterations_from(
diff --git a/docs/cookbooks/analysis.rst b/docs/cookbooks/analysis.rst
index 8b766c3d6..de6ce6d0d 100644
--- a/docs/cookbooks/analysis.rst
+++ b/docs/cookbooks/analysis.rst
@@ -641,7 +641,7 @@ These files can then also be loaded via the database, as described in the databa
             Parameters
             ----------
             paths
-                The PyAutoFit paths object which manages all paths, e.g. where
+                The paths object which manages all paths, e.g. where
                 the non-linear search outputs are stored, visualization, and the
                 pickled objects used by the aggregator output by this function.
             """
@@ -674,7 +674,7 @@ These files can then also be loaded via the database, as described in the databa
             Parameters
             ----------
             paths
-                The PyAutoFit paths object which manages all paths, e.g. where the
+                The paths object which manages all paths, e.g. where the
                 non-linear search outputs are stored, visualization and the pickled
                 objects used by the aggregator output by this function.
             result
diff --git a/docs/cookbooks/result.rst b/docs/cookbooks/result.rst
index ccbb3f56e..7102e759f 100644
--- a/docs/cookbooks/result.rst
+++ b/docs/cookbooks/result.rst
@@ -553,7 +553,7 @@ as 1D numpy arrays, are converted to a suitable dictionary output format. This u
             Parameters
             ----------
             paths
-                The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+                The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
                 visualization, and the pickled objects used by the aggregator output by this function.
             """
             from autoconf.dictable import to_dict
@@ -575,7 +575,7 @@ as 1D numpy arrays, are converted to a suitable dictionary output format. This u
             Parameters
             ----------
             paths
-                The PyAutoFit paths object which manages all paths, e.g. where the non-linear search outputs are stored,
+                The paths object which manages all paths, e.g. where the non-linear search outputs are stored,
                 visualization and the pickled objects used by the aggregator output by this function.
             result
                 The result of a model fit, including the non-linear search, samples and maximum likelihood model.


Test patch:
diff --git a/test_autofit/non_linear/grid/test_sensitivity/conftest.py b/test_autofit/non_linear/grid/test_sensitivity/conftest.py
index f808ffc39..e2a71d2eb 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/conftest.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/conftest.py
@@ -38,10 +38,11 @@ class BaseFit:
     def __init__(self, analysis_cls):
         self.analysis_cls = analysis_cls
 
-    def __call__(self, dataset, model, paths):
+    def __call__(self, dataset, model, paths, instance):
         search = af.m.MockSearch(
             return_sensitivity_results=True,
             samples_summary=MockSamplesSummary(model=model),
+            paths=paths,
         )
 
         analysis = self.analysis_cls(dataset=dataset)
@@ -53,10 +54,11 @@ class PerturbFit:
     def __init__(self, analysis_cls):
         self.analysis_cls = analysis_cls
 
-    def __call__(self, dataset, model, paths):
+    def __call__(self, dataset, model, paths, instance):
         search = af.m.MockSearch(
             return_sensitivity_results=True,
             samples_summary=MockSamplesSummary(model=model),
+            paths=paths,
         )
 
         analysis = self.analysis_cls(dataset=dataset)
@@ -88,6 +90,37 @@ def make_sensitivity(
     )
 
 
+@pytest.fixture(name="masked_sensitivity")
+def make_masked_sensitivity(
+    perturb_model,
+):
+    # noinspection PyTypeChecker
+    instance = af.ModelInstance()
+    instance.gaussian = af.Gaussian()
+    return s.Sensitivity(
+        simulation_instance=instance,
+        base_model=af.Collection(gaussian=af.Model(af.Gaussian)),
+        perturb_model=perturb_model,
+        simulate_cls=Simulate(),
+        base_fit_cls=BaseFit(Analysis),
+        perturb_fit_cls=PerturbFit(Analysis),
+        paths=af.DirectoryPaths(),
+        number_of_steps=2,
+        mask=np.array(
+            [
+                [
+                    [True, True],
+                    [True, True],
+                ],
+                [
+                    [True, True],
+                    [True, True],
+                ],
+            ]
+        ),
+    )
+
+
 @pytest.fixture(name="job")
 def make_job(
     perturb_model,
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 4d1330d1c..c6d348067 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
@@ -29,11 +29,27 @@ def test_labels(sensitivity):
 
 
 def test_perform_job(job):
+    assert not job.is_complete
+
+    result = job.perform()
+    assert isinstance(result, s.JobResult)
+    assert isinstance(result.perturb_result, af.Result)
+    assert isinstance(result.result, af.Result)
+
+    assert job.is_complete
+
+
+def test_perform_twice(job):
+    job.perform()
+    assert job.is_complete
+
     result = job.perform()
     assert isinstance(result, s.JobResult)
     assert isinstance(result.perturb_result, af.Result)
     assert isinstance(result.result, af.Result)
 
+    assert job.is_complete
+
 
 class TestPerturbationModels:
     @pytest.mark.parametrize(
diff --git a/test_autofit/non_linear/grid/test_sensitivity/test_masked_sensitivity.py b/test_autofit/non_linear/grid/test_sensitivity/test_masked_sensitivity.py
new file mode 100644
index 000000000..4ea67f2ca
--- /dev/null
+++ b/test_autofit/non_linear/grid/test_sensitivity/test_masked_sensitivity.py
@@ -0,0 +1,7 @@
+from math import prod
+
+
+def test_run(masked_sensitivity):
+    result = masked_sensitivity.run()
+    number_elements = prod(masked_sensitivity.shape)
+    assert len(result.samples) == number_elements