Repository: rhayes777/PyAutoFit
Pull request #868 (2023-11-09 11:59:57)

Sensitivity Refactor
Previously, sensitivity mapping required input of a non linear search which was used for both model-fits.

However, this had a number of limitations:

1) The way the search was performed could not be very customized.
2) Sensitivity mapping may not even use non-linear searches -- you could imagine simply mapping over a goodness-of-fit metric which is computed in a much simpler way.

I have therefore extracted out the search functionality such that it is now defined by the use. Here is an example of how we would set up the base fit for the previous functionality:

```
class BaseFit:

    def __init__(self, analysis_cls):
        """
        Class used to fit every dataset used for sensitivity mapping with the base model (the model without the 
        perturbed feature sensitivity mapping maps out).

        In this example, the base model therefore does not include the extra Gaussian feature, but the simulated 
        dataset includes one.
        
        The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This 
        maps out the sensitivity of every parameter is (e.g. the sensitivity of the normalization of the Gaussian 
        feature).

        The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is
        fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the
        dataset.
        
        Parameters
        ----------
        analysis_cls
            The `Analysis` class used to fit the model to the dataset.
        """
        self.analysis_cls = analysis_cls

    def __call__(self, dataset, model, paths):
        """
        The base fitting function which fits every dataset used for sensitivity mapping with the base model.

        This function receives as input each simulated dataset of the sensitivity map and fits it, in order to
        quantify how sensitive the model is to the perturbed feature.

        In this example, a full non-linear search is performed to determine how well the model fits the dataset.
        The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.

        Parameters
        ----------
        dataset
            The dataset which is simulated with the perturbed model and which is fitted.
        model
            The model instance which is fitted to the dataset, which does not include the perturbed feature.
        paths
            The `Paths` instance which contains the path to the folder where the results of the fit are written to.
        """

        search = af.DynestyStatic(
            paths= paths.for_sub_analysis(analysis_name="[base]"),
            nlive=25,
            unique_tag=dataset_name,
            session=session,
            iterations_per_update=50000,
        )

        analysis = self.analysis_cls(dataset=dataset)

        return search.fit(
            model=model, analysis=analysis
        )
```

There is one issue, that the following code:

```
            paths = self.paths.for_sub_analysis(
                label,
            )
```

Is not extending the paths with the `label` correctly. Previously, this worked on `search.paths` so I guess its to do with that swap. I think this is leading to a broken unit test.

The issue is see clearly because when we run the following integration test:

https://github.com/Jammy2211/autofit_workspace_test/blob/main/scripts/database/scrape/sensitivity.py

The output paths do not include the parameter name `e.g. normalization_75.000` or the  `[base]` and `[perturbed]` tags.

![image](https://github.com/rhayes777/PyAutoFit/assets/23455639/8ea2784a-4acb-4aa4-8413-eb9ddb416b60)





Patch:
diff --git a/autofit/__init__.py b/autofit/__init__.py
index 6b0483eab..69a0bedad 100644
--- a/autofit/__init__.py
+++ b/autofit/__init__.py
@@ -65,6 +65,7 @@ from .non_linear.search.abstract_search import NonLinearSearch
 from .non_linear.analysis.analysis import Analysis
 from .non_linear.analysis.combined import CombinedAnalysis
 from .non_linear.grid.grid_search import GridSearchResult
+from .non_linear.grid.sensitivity import Sensitivity
 from .non_linear.initializer import InitializerBall
 from .non_linear.initializer import InitializerPrior
 from .non_linear.initializer import SpecificRangeInitializer
diff --git a/autofit/non_linear/grid/grid_list.py b/autofit/non_linear/grid/grid_list.py
index 1e231b462..e836383fa 100644
--- a/autofit/non_linear/grid/grid_list.py
+++ b/autofit/non_linear/grid/grid_list.py
@@ -19,7 +19,7 @@ def as_grid_list(func):
     """
 
     @wraps(func)
-    def wrapper(grid_search_result) -> List:
+    def wrapper(grid_search_result, *args, **kwargs) -> List:
         """
         This decorator converts the output of a function which computes a list of grid search results to a `GridList`.
 
@@ -33,7 +33,7 @@ def as_grid_list(func):
             The function output converted to a `GridList`.
         """
 
-        values = func(grid_search_result)
+        values = func(grid_search_result, *args, **kwargs)
 
         return GridList(values=values, shape=grid_search_result.shape)
 
diff --git a/autofit/non_linear/grid/sensitivity.py b/autofit/non_linear/grid/sensitivity.py
index 3996620f9..bdc2deabf 100644
--- a/autofit/non_linear/grid/sensitivity.py
+++ b/autofit/non_linear/grid/sensitivity.py
@@ -4,37 +4,35 @@ import os
 from copy import copy
 from itertools import count
 from pathlib import Path
-from typing import List, Generator, Callable, ClassVar, Type, Union, Tuple
+from typing import List, Generator, Callable, ClassVar, Union, Tuple
 
-from autoconf import cached_property
 from autoconf.dictable import to_dict
 from autofit.mapper.model import ModelInstance
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
-from autofit.non_linear.search.abstract_search import NonLinearSearch
-from autofit.non_linear.analysis import Analysis
 from autofit.non_linear.grid.grid_search import make_lists
 from autofit.non_linear.parallel import AbstractJob, Process, AbstractJobResult
+from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.result import Result
 from autofit.text.text_util import padding
 
 
 class JobResult(AbstractJobResult):
-    def __init__(self, number: int, result: Result, perturbed_result: Result):
+    def __init__(self, number: int, result: Result, perturb_result: Result):
         """
         The result of a single sensitivity comparison
 
         Parameters
         ----------
         result
-        perturbed_result
+        perturb_result
         """
         super().__init__(number)
         self.result = result
-        self.perturbed_result = perturbed_result
+        self.perturb_result = perturb_result
 
     @property
     def log_likelihood_difference(self):
-        return self.perturbed_result.log_likelihood - self.result.log_likelihood
+        return self.perturb_result.log_likelihood - self.result.log_likelihood
 
     @property
     def log_likelihood_base(self):
@@ -42,62 +40,52 @@ class JobResult(AbstractJobResult):
 
     @property
     def log_likelihood_perturbed(self):
-        return self.perturbed_result.log_likelihood
+        return self.perturb_result.log_likelihood
 
 
 class Job(AbstractJob):
     _number = count()
 
-    use_instance = False
-
     def __init__(
         self,
-        analysis_factory: "AnalysisFactory",
         model: AbstractPriorModel,
-        perturbation_model: AbstractPriorModel,
+        simulate_cls: Callable,
+        perturb_model: AbstractPriorModel,
+        simulate_instance: ModelInstance,
         base_instance: ModelInstance,
-        perturbation_instance: ModelInstance,
-        search: NonLinearSearch,
+        base_fit_cls: Callable,
+        perturb_fit_cls: Callable,
+        paths: AbstractPaths,
         number: int,
     ):
         """
-        Job to run non-linear searches comparing how well a model and a model with a perturbation
-        fit the image.
+        Job to run non-linear searches comparing how well a model and a model with a perturbation fit the image.
 
         Parameters
         ----------
         model
             A base model that fits the image without a perturbation
-        perturbation_model
+        perturb_model
             A model of the perturbation which has been added to the underlying image
-        analysis_factory
-            Factory to generate analysis classes which comprise a model and data
-        search
-            A non-linear search
+        base_fit_cls
+            A class which defines the function which fits the base model to each simulated dataset of the sensitivity
+            map.
+        perturb_fit_cls
+            A class which defines the function which fits the perturbed model to each simulated dataset of the
+            sensitivity map.
+        paths
+            The paths defining the output directory structure of the sensitivity mapping.
         """
         super().__init__(number=number)
 
-        self.analysis_factory = analysis_factory
         self.model = model
-
-        self.perturbation_model = perturbation_model
+        self.simulate_cls = simulate_cls
+        self.perturb_model = perturb_model
+        self.simulate_instance = simulate_instance
         self.base_instance = base_instance
-        self.perturbation_instance = perturbation_instance
-
-        self.search = search.copy_with_paths(
-            search.paths.for_sub_analysis(
-                "[base]",
-            )
-        )
-        self.perturbed_search = search.copy_with_paths(
-            search.paths.for_sub_analysis(
-                "[perturbed]",
-            )
-        )
-
-    @cached_property
-    def analysis(self):
-        return self.analysis_factory()
+        self.base_fit_cls = base_fit_cls
+        self.perturb_fit_cls = perturb_fit_cls
+        self.paths = paths
 
     def perform(self) -> JobResult:
         """
@@ -108,26 +96,30 @@ class Job(AbstractJob):
         -------
         An object comprising the results of the two fits
         """
-        result = self.base_model_func()
 
-        perturbed_model = copy(self.model)
-        perturbed_model.perturbation = self.perturbation_model
+        dataset = self.simulate_cls(
+            instance=self.simulate_instance,
+            simulate_path=self.paths.image_path.with_name("simulate"),
+        )
 
-        # TODO : This is what I added so that the Drawer runs use the correct subhalo model.
+        result = self.base_fit_cls(
+            model=self.model,
+            dataset=dataset,
+            paths=self.paths.for_sub_analysis("[base]"),
+        )
 
-        # if self.use_instance:
-        #     perturbed_model.perturbation = self.perturbation_instance
+        perturb_model = copy(self.model)
+        perturb_model.perturbation = self.perturb_model
 
-        perturbed_result = self.perturbation_model_func(perturbed_model=perturbed_model)
-        return JobResult(
-            number=self.number, result=result, perturbed_result=perturbed_result
+        perturb_result = self.perturb_fit_cls(
+            model=perturb_model,
+            dataset=dataset,
+            paths=self.paths.for_sub_analysis("[perturb]"),
         )
 
-    def base_model_func(self):
-        return self.search.fit(model=self.model, analysis=self.analysis)
-
-    def perturbation_model_func(self, perturbed_model):
-        return self.perturbed_search.fit(model=perturbed_model, analysis=self.analysis)
+        return JobResult(
+            number=self.number, result=result, perturb_result=perturb_result
+        )
 
 
 class SensitivityResult:
@@ -177,11 +169,12 @@ class Sensitivity:
     def __init__(
         self,
         base_model: AbstractPriorModel,
-        perturbation_model: AbstractPriorModel,
+        perturb_model: AbstractPriorModel,
         simulation_instance,
-        simulate_function: Callable,
-        analysis_class: Type[Analysis],
-        search: NonLinearSearch,
+        paths,
+        simulate_cls: Callable,
+        base_fit_cls: Callable,
+        perturb_fit_cls: Callable,
         job_cls: ClassVar = Job,
         number_of_steps: Union[Tuple[int], int] = 4,
         number_of_cores: int = 2,
@@ -192,29 +185,29 @@ class Sensitivity:
         can be detected if it occurs in different parts of an image.
 
         For a range from 0 to 1 with step_size, for each dimension of the
-        perturbation_model, a perturbation is created and used in conjunction
+        perturb_model, a perturbation is created and used in conjunction
         with the instance to create an image.
 
         For each of these images, a fit is run with just the model and with both
-        the model and perturbation_model to compare how much better the image
+        the model and perturb_model to compare how much better the image
         can be fit if the perturbation is included.
 
         Parameters
         ----------
         base_model
             A model that fits the instance well
-        perturbation_model
+        perturb_model
             A model which provides a perturbations to be applied to the instance
             before creating images
         simulation_instance
             An instance of a model to which perturbations are applied prior to
             images being generated
-        simulate_function
-            A function that can convert an instance into an image
-        analysis_class
-            A class which can compare an image to an instance and evaluate fitness
-        search
-            A NonLinear search class which is copied and used to evaluate fitness
+        simulate_cls
+            A class which simulates images from each perturb instance that sensitivity mapping is performed on.
+        base_fit_cls
+            The class which fits the base model to each simulated dataset of the sensitivity map.
+        perturb_fit_cls
+            The class which fits the perturb model to each simulated dataset of the sensitivity map.
         number_of_cores
             How many cores does this computer have? Minimum 2.
         limit_scale
@@ -225,18 +218,19 @@ class Sensitivity:
                 A scale of 0.5 means priors have limits smaller than the grid square
                     with width half a grid square.
         """
-        self.logger = logging.getLogger(f"Sensitivity ({search.name})")
+        self.logger = logging.getLogger(f"Sensitivity ({paths.name})")
 
         self.logger.info("Creating")
 
         self.instance = simulation_instance
         self.model = base_model
+        self.perturb_model = perturb_model
 
-        self.search = search
-        self.analysis_class = analysis_class
+        self.paths = paths
 
-        self.perturbation_model = perturbation_model
-        self.simulate_function = simulate_function
+        self.simulate_cls = simulate_cls
+        self.base_fit_cls = base_fit_cls
+        self.perturb_fit_cls = perturb_fit_cls
 
         self.job_cls = job_cls
 
@@ -266,7 +260,7 @@ class Sensitivity:
         """
         self.logger.info("Running")
 
-        self.search.paths.save_unique_tag(is_grid_search=True)
+        self.paths.save_unique_tag(is_grid_search=True)
 
         headers = [
             "index",
@@ -287,7 +281,7 @@ class Sensitivity:
             results.append(result)
             results = sorted(results)
 
-            os.makedirs(self.search.paths.output_path, exist_ok=True)
+            os.makedirs(self.paths.output_path, exist_ok=True)
             with open(self.results_path, "w+") as f:
                 writer = csv.writer(f)
                 writer.writerow(headers)
@@ -298,30 +292,30 @@ class Sensitivity:
                         for item in [
                             result_.number,
                             *values,
-                            result_.log_likelihood_base,
-                            result_.log_likelihood_perturbed,
-                            result_.log_likelihood_difference,
+                            float(result_.log_likelihood_base),
+                            float(result_.log_likelihood_perturbed),
+                            float(result_.log_likelihood_difference),
                         ]
                     )
 
         result = SensitivityResult(results)
 
-        self.search.paths.save_json("result", to_dict(result))
+        self.paths.save_json("result", to_dict(result))
 
         return SensitivityResult(results)
 
     @property
     def results_path(self) -> Path:
-        return self.search.paths.output_path / "results.csv"
+        return self.paths.output_path / "results.csv"
 
     @property
     def _lists(self) -> List[List[float]]:
         """
         A list of hypercube vectors, used to instantiate
-        the perturbation_model and create the individual
+        the perturb_model and create the individual
         perturbations.
         """
-        return make_lists(self.perturbation_model.prior_count, step_size=self.step_size)
+        return make_lists(self.perturb_model.prior_count, step_size=self.step_size)
 
     @property
     def _physical_values(self) -> List[List[float]]:
@@ -332,7 +326,7 @@ class Sensitivity:
             [
                 prior.value_for(unit_value)
                 for prior, unit_value in zip(
-                    self.perturbation_model.priors_ordered_by_id, unit_values
+                    self.perturb_model.priors_ordered_by_id, unit_values
                 )
             ]
             for unit_values in self._lists
@@ -343,7 +337,7 @@ class Sensitivity:
         """
         A name for each of the perturbed priors
         """
-        for path, _ in self.perturbation_model.prior_tuples:
+        for path, _ in self.perturb_model.prior_tuples:
             yield path
 
     @property
@@ -355,23 +349,23 @@ class Sensitivity:
         """
         for list_ in self._lists:
             strings = list()
-            for value, prior_tuple in zip(list_, self.perturbation_model.prior_tuples):
+            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)
 
     @property
-    def _perturbation_instances(self) -> Generator[ModelInstance, None, None]:
+    def _perturb_instances(self) -> Generator[ModelInstance, None, None]:
         """
         A list of instances each of which defines a perturbation to
         be applied to the image.
         """
         for list_ in self._lists:
-            yield self.perturbation_model.instance_from_unit_vector(list_)
+            yield self.perturb_model.instance_from_unit_vector(list_)
 
     @property
-    def _perturbation_models(self) -> Generator[AbstractPriorModel, None, None]:
+    def _perturb_models(self) -> Generator[AbstractPriorModel, None, None]:
         """
         A list of models representing a perturbation at each grid square.
 
@@ -386,42 +380,9 @@ class Sensitivity:
                     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.perturbation_model.priors_ordered_by_id
-                )
+                for centre, prior in zip(list_, self.perturb_model.priors_ordered_by_id)
             ]
-            yield self.perturbation_model.with_limits(limits)
-
-    @property
-    def _searches(self) -> Generator[NonLinearSearch, None, None]:
-        """
-        A list of non-linear searches, each of which is applied to
-        one perturbation.
-        """
-        for label in self._labels:
-            yield self._search_instance(label)
-
-    def _search_instance(self, name_path: str) -> NonLinearSearch:
-        """
-        Create a search instance, distinguished by its name
-
-        Parameters
-        ----------
-        name_path
-            A path to distinguish this search from other searches
-
-        Returns
-        -------
-        A non linear search, copied from the instance search
-        """
-        paths = self.search.paths
-        search_instance = self.search.copy_with_paths(
-            paths.for_sub_analysis(
-                name_path,
-            )
-        )
-
-        return search_instance
+            yield self.perturb_model.with_limits(limits)
 
     def _make_jobs(self) -> Generator[Job, None, None]:
         """
@@ -430,48 +391,24 @@ class Sensitivity:
         Each job fits a perturbed image with the original model
         and a model which includes a perturbation.
         """
-        for number, (perturbation_instance, perturbation_model, search) in enumerate(
-            zip(self._perturbation_instances, self._perturbation_models, self._searches)
+        for number, (perturb_instance, perturb_model, label) in enumerate(
+            zip(self._perturb_instances, self._perturb_models, self._labels)
         ):
-            instance = copy(self.instance)
-            instance.perturbation = perturbation_instance
+            simulate_instance = copy(self.instance)
+            simulate_instance.perturb = perturb_instance
+
+            paths = self.paths.for_sub_analysis(
+                label,
+            )
 
             yield self.job_cls(
-                analysis_factory=AnalysisFactory(
-                    instance=instance,
-                    simulate_function=self.simulate_function,
-                    analysis_class=self.analysis_class,
-                    paths=search.paths,
-                ),
+                simulate_instance=simulate_instance,
                 model=self.model,
-                perturbation_model=perturbation_model,
+                perturb_model=perturb_model,
                 base_instance=self.instance,
-                perturbation_instance=perturbation_instance,
-                search=search,
+                simulate_cls=self.simulate_cls,
+                base_fit_cls=self.base_fit_cls,
+                perturb_fit_cls=self.perturb_fit_cls,
+                paths=paths,
                 number=number,
             )
-
-
-class AnalysisFactory:
-    def __init__(
-        self,
-        instance,
-        simulate_function,
-        analysis_class,
-        paths,
-    ):
-        """
-        Callable to delay simulation such that it is performed
-        on the Job subprocess
-        """
-        self.instance = instance
-        self.simulate_function = simulate_function
-        self.analysis_class = analysis_class
-        self.paths = paths
-
-    def __call__(self):
-        dataset = self.simulate_function(
-            instance=self.instance,
-            simulate_path=self.paths.image_path.with_name("simulate"),
-        )
-        return self.analysis_class(dataset)
diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py
index 01a055b06..e3ac8f15d 100644
--- a/autofit/non_linear/search/abstract_search.py
+++ b/autofit/non_linear/search/abstract_search.py
@@ -81,6 +81,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         iterations_per_update: int = None,
         number_of_cores: int = 1,
         session: Optional[sa.orm.Session] = None,
+        paths: Optional[AbstractPaths] = None,
         **kwargs,
     ):
         """
@@ -133,9 +134,11 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
 
         self.unique_tag = unique_tag
 
-        if session is not None:
+        if paths:
+            self.paths = paths
+        elif session is not None:
             logger.debug("Session found. Using database.")
-            paths = DatabasePaths(
+            self.paths = DatabasePaths(
                 name=name,
                 path_prefix=path_prefix,
                 session=session,
@@ -144,13 +147,11 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
             )
         elif name is not None or path_prefix:
             logger.debug("Session not found. Using directory output.")
-            paths = DirectoryPaths(
+            self.paths = DirectoryPaths(
                 name=name, path_prefix=path_prefix, unique_tag=unique_tag
             )
         else:
-            paths = NullPaths()
-
-        self.paths: AbstractPaths = paths
+            self.paths = NullPaths()
 
         self.force_pickle_overwrite = conf.instance["general"]["output"][
             "force_pickle_overwrite"
@@ -519,12 +520,12 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         """
 
         if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
-
             from autofit.non_linear.search.nest.nautilus.search import Nautilus
 
             if isinstance(self, Nautilus):
-
-                from autofit.non_linear.search.nest.dynesty.search.static import DynestyStatic
+                from autofit.non_linear.search.nest.dynesty.search.static import (
+                    DynestyStatic,
+                )
 
                 from pathlib import Path
 
@@ -638,10 +639,10 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
             timeout_seconds = get_timeout_seconds()
 
             if timeout_seconds is not None:
-
-                logger.info(f"\n\n ***Log Likelihood Function timeout is "
-                            f"turned on and set to {timeout_seconds} seconds.***\n")
-
+                logger.info(
+                    f"\n\n ***Log Likelihood Function timeout is "
+                    f"turned on and set to {timeout_seconds} seconds.***\n"
+                )
 
     def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Result:
         """
@@ -670,7 +671,9 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         and errors on the model parameters.
         """
         if self.is_master:
-            if not isinstance(self.paths, DatabasePaths) and not isinstance(self.paths, NullPaths):
+            if not isinstance(self.paths, DatabasePaths) and not isinstance(
+                self.paths, NullPaths
+            ):
                 self.timer.start()
 
         model.freeze()
@@ -682,7 +685,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
             model=model,
             analysis=analysis,
             search_internal=search_internal,
-            during_analysis=False
+            during_analysis=False,
         )
 
         result = analysis.make_result(
@@ -699,7 +702,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         return result
 
     def result_via_completed_fit(
-        self, analysis: Analysis, model: AbstractPriorModel, search_internal = None
+        self, analysis: Analysis, model: AbstractPriorModel, search_internal=None
     ) -> Result:
         """
         Returns the result of the non-linear search of a completed model-fit.
@@ -742,7 +745,10 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
 
             if self.force_visualize_overwrite:
                 self.perform_visualization(
-                    model=model, analysis=analysis, search_internal=search_internal, during_analysis=False
+                    model=model,
+                    analysis=analysis,
+                    search_internal=search_internal,
+                    during_analysis=False,
                 )
 
             if self.force_pickle_overwrite:
@@ -850,7 +856,11 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         return self._class_config[section][attribute_name]
 
     def perform_update(
-        self, model: AbstractPriorModel, analysis: Analysis, during_analysis: bool, search_internal = None,
+        self,
+        model: AbstractPriorModel,
+        analysis: Analysis,
+        during_analysis: bool,
+        search_internal=None,
     ) -> Samples:
         """
         Perform an update of the non-linear search's model-fitting results.
@@ -890,7 +900,9 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
                 f"output folder for final visualization, samples, etc.)"
             )
 
-        if not isinstance(self.paths, DatabasePaths) and not isinstance(self.paths, NullPaths):
+        if not isinstance(self.paths, DatabasePaths) and not isinstance(
+            self.paths, NullPaths
+        ):
             self.timer.update()
 
         samples = self.samples_from(model=model, search_internal=search_internal)
@@ -947,7 +959,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         model: AbstractPriorModel,
         analysis: AbstractPriorModel,
         during_analysis: bool,
-        search_internal = None,
+        search_internal=None,
     ):
         """
         Perform visualization of the non-linear search's model-fitting results.
@@ -1003,7 +1015,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
     def remove_state_files(self):
         pass
 
-    def samples_from(self, model: AbstractPriorModel, search_internal = None) -> Samples:
+    def samples_from(self, model: AbstractPriorModel, search_internal=None) -> Samples:
         """
         Loads the samples of a non-linear search from its output files.
 
@@ -1021,11 +1033,15 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
             The model which generates instances for different points in parameter space.
         """
         try:
-            return self.samples_via_internal_from(model=model, search_internal=search_internal)
+            return self.samples_via_internal_from(
+                model=model, search_internal=search_internal
+            )
         except (FileNotFoundError, NotImplementedError, AttributeError):
             return self.samples_via_csv_from(model=model)
 
-    def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=None):
+    def samples_via_internal_from(
+        self, model: AbstractPriorModel, search_internal=None
+    ):
         raise NotImplementedError
 
     def samples_via_csv_from(self, model: AbstractPriorModel) -> Samples:
diff --git a/autofit/non_linear/search/nest/dynesty/search/static.py b/autofit/non_linear/search/nest/dynesty/search/static.py
index 25bb92d00..0e0402d16 100644
--- a/autofit/non_linear/search/nest/dynesty/search/static.py
+++ b/autofit/non_linear/search/nest/dynesty/search/static.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
-from typing import Optional
+
+from pathlib import Path
+from typing import Optional, Union
 
 from dynesty import NestedSampler as StaticSampler
 
@@ -26,7 +28,7 @@ class DynestyStatic(AbstractDynesty):
     def __init__(
         self,
         name: Optional[str] = None,
-        path_prefix: Optional[str] = None,
+        path_prefix: Optional[Union[str, Path]] = None,
         unique_tag: Optional[str] = None,
         iterations_per_update: int = None,
         number_of_cores: int = None,
@@ -106,7 +108,9 @@ class DynestyStatic(AbstractDynesty):
         """
 
         if checkpoint_exists:
-            search_internal = StaticSampler.restore(fname=self.checkpoint_file, pool=pool)
+            search_internal = StaticSampler.restore(
+                fname=self.checkpoint_file, pool=pool
+            )
 
             uses_pool = self.read_uses_pool()
 
@@ -115,9 +119,7 @@ class DynestyStatic(AbstractDynesty):
             return search_internal
 
         else:
-            live_points = self.live_points_init_from(
-                model=model, fitness=fitness
-            )
+            live_points = self.live_points_init_from(model=model, fitness=fitness)
 
             if pool is not None:
                 self.write_uses_pool(uses_pool=True)
@@ -146,4 +148,4 @@ class DynestyStatic(AbstractDynesty):
 
     @property
     def number_live_points(self):
-        return self.config_dict_search["nlive"]
\ No newline at end of file
+        return self.config_dict_search["nlive"]
diff --git a/docs/features/sensitivity_mapping.rst b/docs/features/sensitivity_mapping.rst
index 9742d3c06..344d21cc5 100644
--- a/docs/features/sensitivity_mapping.rst
+++ b/docs/features/sensitivity_mapping.rst
@@ -98,12 +98,12 @@ The ``base_model`` corresponds to the ``gaussian_main`` above.
 Perturbation Model
 ------------------
 
-We now define the ``perturbation_model``, which is the model component whose parameters we iterate over to perform
-sensitivity mapping. Many instances of the ``perturbation_model`` are created and used to simulate the many datasets
+We now define the ``perturb_model``, which is the model component whose parameters we iterate over to perform
+sensitivity mapping. Many instances of the ``perturb_model`` are created and used to simulate the many datasets
 that we fit. However, it is only included in half of the model-fits corresponding to the more complex models whose
 Bayesian evidence we compare to the simpler model-fits consisting of just the ``base_model``.
 
-The ``perturbation_model`` is therefore another ``Gaussian`` but now corresponds to the ``gaussian_feature`` above.
+The ``perturb_model`` is therefore another ``Gaussian`` but now corresponds to the ``gaussian_feature`` above.
 
 By fitting both of these models to every simulated dataset, we will therefore infer the Bayesian evidence of every
 model to every dataset. Sensitivity mapping therefore maps out for what values of ``normalization`` in the ``gaussian_feature``
@@ -112,10 +112,10 @@ values ot the ``centre`` and ``sigma`` of the ``Gaussian`` so we only map over i
 
 .. code-block:: bash
 
-    perturbation_model = af.Model(m.Gaussian)
-    perturbation_model.centre = 70.0
-    perturbation_model.sigma = 0.5
-    perturbation_model.normalization = af.UniformPrior(lower_limit=0.01, upper_limit=100.0)
+    perturb_model = af.Model(m.Gaussian)
+    perturb_model.centre = 70.0
+    perturb_model.sigma = 0.5
+    perturb_model.normalization = af.UniformPrior(lower_limit=0.01, upper_limit=100.0)
 
 Simulation
 ----------
@@ -132,19 +132,19 @@ performed above.
 
     simulation_instance = result_single.instance
 
-We now write the ``simulate_function``, which takes the ``instance`` of our model (defined above) and uses it to
+We now write the ``simulate_cls``, which takes the ``instance`` of our model (defined above) and uses it to
 simulate a dataset which is subsequently fitted.
 
-Note that when this dataset is simulated, the quantity ``instance.perturbation`` is used in the ``simulate_function``.
-This is an instance of the ``gaussian_feature``, and it is different every time the ``simulate_function`` is called.
+Note that when this dataset is simulated, the quantity ``instance.perturb`` is used in the ``simulate_cls``.
+This is an instance of the ``gaussian_feature``, and it is different every time the ``simulate_cls`` is called.
 
-In this example, this ``instance.perturbation`` corresponds to different ``gaussian_feature``'s with values of
+In this example, this ``instance.perturb`` corresponds to different ``gaussian_feature``'s with values of
 ``normalization`` ranging over 0.01 -> 100.0, such that our simulated datasets correspond to a very faint and very bright
 gaussian features.
 
 .. code-block:: bash
 
-    def simulate_function(instance, simulate_path):
+    def __call__(instance, simulate_path):
 
         """
         Specify the number of pixels used to create the xvalues on which the 1D line of the profile is generated using and
@@ -161,11 +161,11 @@ gaussian features.
         values of ``centre=70`` and ``sigma=0.5``, whereas the normalization varies over the ``step_size`` based on its prior.
         """
 
-        print(instance.perturbation.centre)
-        print(instance.perturbation.normalization)
-        print(instance.perturbation.sigma)
+        print(instance.perturb.centre)
+        print(instance.perturb.normalization)
+        print(instance.perturb.sigma)
 
-        model_line = instance.gaussian_main.model_data_1d_via_xvalues_from(xvalues=xvalues) + instance.perturbation.model_data_1d_via_xvalues_from(xvalues=xvalues)
+        model_line = instance.gaussian_main.model_data_1d_via_xvalues_from(xvalues=xvalues) + instance.perturb.model_data_1d_via_xvalues_from(xvalues=xvalues)
 
         """Determine the noise (at a specified signal to noise level) in every pixel of our model profile."""
         signal_to_noise_ratio = 25.0
@@ -200,11 +200,11 @@ object below are:
 
 - ``base_model``: This is the simpler model that is fitted to every simulated dataset, which in this example is composed of a single ``Gaussian`` called the ``gaussian_main``.
 
-- ``perturbation_model``: This is the extra model component that alongside the ``base_model`` is fitted to every simulated dataset, which in this example  is composed of two ``Gaussians`` called the ``gaussian_main`` and ``gaussian_feature``.
+- ``perturb_model``: This is the extra model component that alongside the ``base_model`` is fitted to every simulated dataset, which in this example  is composed of two ``Gaussians`` called the ``gaussian_main`` and ``gaussian_feature``.
 
-- ``simulate_function``: This is the function that uses the ``instance`` and many instances of the ``perturbation_model`` to simulate many datasets that are fitted with the ``base_model`` and ``base_model`` + ``perturbation_model``.
+- ``simulate_cls``: This is the function that uses the ``instance`` and many instances of the ``perturb_model`` to simulate many datasets that are fitted with the ``base_model`` and ``base_model`` + ``perturb_model``.
 
-- ``step_size``: The size of steps over which the parameters in the ``perturbation_model`` are iterated. In this example, normalization has a ``LogUniformPrior`` with lower limit 1e-4 and upper limit 1e2, therefore the ``step_size`` of 0.5 will simulate and fit just 2 datasets where the normalization is 1e-4 and 1e2.
+- ``step_size``: The size of steps over which the parameters in the ``perturb_model`` are iterated. In this example, normalization has a ``LogUniformPrior`` with lower limit 1e-4 and upper limit 1e2, therefore the ``step_size`` of 0.5 will simulate and fit just 2 datasets where the normalization is 1e-4 and 1e2.
 
 - ``number_of_cores``: The number of cores over which the sensitivity mapping is performed, enabling parallel processing.
 
@@ -217,8 +217,8 @@ full example script on the ``autofit_workspace``).
         search=search,
         simulation_instance=simulation_instance,
         base_model=base_model,
-        perturbation_model=perturbation_model,
-        simulate_function=simulate_function,
+        perturb_model=perturb_model,
+        simulate_cls=simulate_cls,
         analysis_class=Analysis,
         step_size=0.5,
         number_of_cores=2,


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 f2bb8cff8..212bcb95d 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/conftest.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/conftest.py
@@ -8,89 +8,96 @@ from autofit.non_linear.grid import sensitivity as s
 x = np.array(range(10))
 
 
-def simulate_function(instance: af.ModelInstance, simulate_path:Optional[str]):
-    image = instance.gaussian(x)
-    if hasattr(instance, "perturbation"):
-        image += instance.perturbation(x)
-    return image
+class Simulate:
+    def __init__(self):
+        pass
 
+    def __call__(self, instance: af.ModelInstance, simulate_path: Optional[str]):
+        image = instance.gaussian(x)
 
-class Analysis(af.Analysis):
+        if hasattr(instance, "perturbation"):
+            image += instance.perturbation(x)
+
+        return image
 
-    def __init__(self, image: np.array):
-        self.image = image
+
+class Analysis(af.Analysis):
+    def __init__(self, dataset: np.array):
+        self.dataset = dataset
 
     def log_likelihood_function(self, instance):
-        image = simulate_function(instance, simulate_path=None)
-        return np.mean(np.multiply(-0.5, np.square(np.subtract(self.image, image))))
+        simulate = Simulate()
 
+        dataset = simulate(instance, simulate_path=None)
 
-@pytest.fixture(
-    name="perturbation_model"
-)
-def make_perturbation_model():
-    return af.Model(af.Gaussian)
+        return np.mean(np.multiply(-0.5, np.square(np.subtract(self.dataset, dataset))))
+
+
+class BaseFit:
+    def __init__(self, analysis_cls):
+        self.analysis_cls = analysis_cls
+
+    def __call__(self, dataset, model, paths):
+        search = af.m.MockSearch(return_sensitivity_results=True)
+
+        analysis = self.analysis_cls(dataset=dataset)
+
+        return search.fit(model=model, analysis=analysis)
 
 
-@pytest.fixture(
-    name="search"
-)
-def make_search():
-    return af.m.MockSearch(return_sensitivity_results=True)
+class PerturbFit:
+    def __init__(self, analysis_cls):
+        self.analysis_cls = analysis_cls
 
+    def __call__(self, dataset, model, paths):
+        search = af.m.MockSearch(return_sensitivity_results=True)
 
-@pytest.fixture(
-    name="sensitivity"
-)
+        analysis = self.analysis_cls(dataset=dataset)
+
+        return search.fit(model=model, analysis=analysis)
+
+
+@pytest.fixture(name="perturb_model")
+def make_perturb_model():
+    return af.Model(af.Gaussian)
+
+
+@pytest.fixture(name="sensitivity")
 def make_sensitivity(
-        perturbation_model,
-        search
+    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)
-        ),
-        perturbation_model=perturbation_model,
-        simulate_function=simulate_function,
-        analysis_class=Analysis,
-        search=search,
+        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,
     )
 
 
-class MockAnalysisFactory:
-    def __init__(self, analysis):
-        self.analysis = analysis
-
-    def __call__(self):
-        return self.analysis
-
-
-@pytest.fixture(
-    name="job"
-)
+@pytest.fixture(name="job")
 def make_job(
-        perturbation_model,
-        search
+    perturb_model,
 ):
     instance = af.ModelInstance()
     instance.gaussian = af.Gaussian()
     base_instance = instance
-    instance.perturbation = af.Gaussian()
-    image = simulate_function(instance, "")
+    instance.perturb = af.Gaussian()
     # noinspection PyTypeChecker
     return s.Job(
-        model=af.Collection(
-            gaussian=af.Model(af.Gaussian)
-        ),
-        perturbation_model=af.Model(af.Gaussian),
+        model=af.Collection(gaussian=af.Model(af.Gaussian)),
+        perturb_model=af.Model(af.Gaussian),
+        simulate_instance=instance,
         base_instance=base_instance,
-        perturbation_instance=instance,
-        analysis_factory=MockAnalysisFactory(Analysis(image)),
-        search=search,
-        number=1
+        simulate_cls=Simulate(),
+        base_fit_cls=BaseFit(Analysis),
+        perturb_fit_cls=PerturbFit(Analysis),
+        paths=af.DirectoryPaths(),
+        number=1,
     )
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 5e3599b2d..e9b5fb2b0 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/test_functionality.py
@@ -5,7 +5,7 @@ from autofit.non_linear.grid import sensitivity as s
 
 
 def test_lists(sensitivity):
-    assert len(list(sensitivity._perturbation_instances)) == 8
+    assert len(list(sensitivity._perturb_instances)) == 8
 
 
 def test_tuple_step_size(sensitivity):
@@ -28,28 +28,14 @@ def test_labels(sensitivity):
     ]
 
 
-def test_searches(sensitivity):
-    assert len(list(sensitivity._searches)) == 8
-
 
 def test_perform_job(job):
     result = job.perform()
     assert isinstance(result, s.JobResult)
-    assert isinstance(result.perturbed_result, af.Result)
+    assert isinstance(result.perturb_result, af.Result)
     assert isinstance(result.result, af.Result)
 
 
-def test_job_paths(
-        job,
-        search
-):
-
-    output_path = search.paths.output_path
-
-    assert job.perturbed_search.paths.output_path == output_path / "[perturbed]"
-    assert job.search.paths.output_path == output_path / "[base]"
-
-
 class TestPerturbationModels:
     @pytest.mark.parametrize(
         "limit_scale, fl, fu, sl, su",
@@ -59,7 +45,7 @@ class TestPerturbationModels:
             (4.0, 0.0, 1.0, 0.0, 1.0,),
         ]
     )
-    def test_perturbation_models(
+    def test_perturb_models(
             self,
             sensitivity,
             limit_scale,
@@ -68,7 +54,7 @@ class TestPerturbationModels:
         sensitivity.limit_scale = limit_scale
         jobs = sensitivity._make_jobs()
         models = [
-            job.perturbation_model
+            job.perturb_model
             for job in jobs
         ]
 
@@ -85,12 +71,12 @@ class TestPerturbationModels:
             self,
             sensitivity
     ):
-        sensitivity.perturbation_model.sigma = af.UniformPrior(
+        sensitivity.perturb_model.sigma = af.UniformPrior(
             upper_limit=10
         )
         model = list(
             sensitivity._make_jobs()
-        )[0].perturbation_model
+        )[0].perturb_model
         assert model.sigma.upper_limit == 5
 
     def test_model_with_limits(self):
diff --git a/test_autofit/non_linear/grid/test_sensitivity/test_results.py b/test_autofit/non_linear/grid/test_sensitivity/test_results.py
index 7a9a79bb4..01c9fb473 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/test_results.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/test_results.py
@@ -12,7 +12,7 @@ def make_result():
     return JobResult(
         number=0,
         result=Result(log_likelihood=1.0),
-        perturbed_result=Result(log_likelihood=2.0),
+        perturb_result=Result(log_likelihood=2.0),
     )
 
 
diff --git a/test_autofit/non_linear/grid/test_sensitivity/test_run.py b/test_autofit/non_linear/grid/test_sensitivity/test_run.py
index 98f025f1d..aa0dc1f51 100644
--- a/test_autofit/non_linear/grid/test_sensitivity/test_run.py
+++ b/test_autofit/non_linear/grid/test_sensitivity/test_run.py
@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from autoconf.conf import with_config
 
 
@@ -8,7 +6,7 @@ def test_sensitivity(sensitivity):
     results = sensitivity.run()
     assert len(results) == 8
 
-    output_path = sensitivity.search.paths.output_path
+    output_path = sensitivity.paths.output_path
 
     assert (output_path / ".is_grid_search").exists()
     path = output_path / "results.csv"