Repository: rhayes777/PyAutoFit (python)

Issue:
Customizing Priors After analysis.with_free_parameter
We often use the following API to extend a model over summed analysis classes:

```
...

analysis = sum(analysis_list)

dataset_model = af.Model(al.DatasetModel)

model = af.Collection(
    dataset_model=dataset_model,
    galaxies=af.Collection(lens=lens, source=source)
)

analysis = analysis.with_free_parameters(
    model.dataset_model.grid_offset
)
```

I now have a use case where I want to customize the prior on the `grid_offset` of the `DatasetModel` corresponding to the second `analysis` only. 

I don't think we currently have an API that allows us to do this?

I cannot achieve the desired behaviour by editing prior configs, because I want the prior on one `grid_offset` to differ from the other.

Let me know if we have an API hidden, else lets think about how we extend the API.

Pull request description:
feature/free parameter override
Resolves #1015

Specific model or parameter overrides can now be defined by calling get-set on the free parameter analysis.

```python
analysis = sum(analysis_list)

dataset_model = af.Model(al.DatasetModel)

model = af.Collection(
    dataset_model=dataset_model,
    galaxies=af.Collection(lens=lens, source=source)
)

analysis = analysis.with_free_parameters(
    model.dataset_model.grid_offset
)

analysis[0][model.dataset_model.background_sky_level] = 1.0


Patch:
diff --git a/autofit/mapper/prior_model/abstract.py b/autofit/mapper/prior_model/abstract.py
index 70acc8548..b4540910e 100644
--- a/autofit/mapper/prior_model/abstract.py
+++ b/autofit/mapper/prior_model/abstract.py
@@ -206,6 +206,23 @@ class AbstractPriorModel(AbstractModel):
                 f"{number_of_failed_assertions} assertions failed!\n{name_string}"
             )
 
+    def set_item_at_path(self, path: Tuple[str, ...], value):
+        """
+        Set an item at a path in the model.
+
+        Parameters
+        ----------
+        path
+            A tuple of strings representing a path to an attribute
+        value
+            The value to be set at the path
+        """
+        obj = self
+        for attribute in path[:-1]:
+            obj = getattr(obj, attribute)
+
+        setattr(obj, path[-1], value)
+
     def cast(
         self,
         value_dict: Dict["AbstractModel", dict],
@@ -1506,6 +1523,26 @@ class AbstractPriorModel(AbstractModel):
         """
         return sorted(priors, key=lambda prior: self.path_for_prior(prior))
 
+    def path_for_object(self, obj) -> Optional[Path]:
+        """
+        Find a path that points at the given object.
+
+        Parameters
+        ----------
+        obj
+            An object in the model.
+
+        Returns
+        -------
+        A path, a series of attributes that point to one location of the object.
+        """
+        for path, instance in self.path_instance_tuples_for_class(
+            object, ignore_children=False
+        ):
+            if instance is obj:
+                return path
+        return None
+
     def path_for_prior(self, prior: Prior) -> Optional[Path]:
         """
         Find a path that points at the given tuple.
diff --git a/autofit/non_linear/analysis/free_parameter.py b/autofit/non_linear/analysis/free_parameter.py
index 8497f0997..c3e42b8a4 100644
--- a/autofit/non_linear/analysis/free_parameter.py
+++ b/autofit/non_linear/analysis/free_parameter.py
@@ -1,4 +1,5 @@
 import logging
+from collections import defaultdict
 from typing import Tuple
 
 from autofit.mapper.prior.abstract import Prior
@@ -13,6 +14,62 @@ from ..paths.abstract import AbstractPaths
 logger = logging.getLogger(__name__)
 
 
+def _unpack(free_parameters: Tuple[Prior, ...]):
+    """
+    Unpack free parameters from a tuple of free parameters and priors.
+
+    Parameters
+    ----------
+    free_parameters
+        A tuple of free parameters and priors.
+
+    Returns
+    -------
+    A list of free parameters.
+    """
+    return [
+        parameter for parameter in free_parameters if isinstance(parameter, Prior)
+    ] + [
+        prior
+        for parameter in free_parameters
+        if isinstance(parameter, (AbstractPriorModel, TuplePrior))
+        for prior in parameter.priors
+    ]
+
+
+class PositionalParameters:
+    def __init__(
+        self,
+        analysis: "FreeParameterAnalysis",
+        position: int,
+    ):
+        """
+        Manage overriding positional parameters for a given analysis.
+
+        Parameters
+        ----------
+        analysis
+            The analysis
+        position
+            The position of the model in the collection
+        """
+        self.analysis = analysis
+        self.position = position
+
+    def __setitem__(self, key, value):
+        """
+        Override some component of the model at the index.
+
+        Parameters
+        ----------
+        key
+            The key of the component (e.g. a prior)
+        value
+            The new value
+        """
+        self.analysis.positional_parameters[self.position][key] = value
+
+
 class FreeParameterAnalysis(IndexCollectionAnalysis):
     def __init__(self, *analyses: Analysis, free_parameters: Tuple[Prior, ...]):
         """
@@ -30,16 +87,23 @@ class FreeParameterAnalysis(IndexCollectionAnalysis):
             A list of priors which are independent for each analysis
         """
         super().__init__(*analyses)
-        self.free_parameters = [
-            parameter for parameter in free_parameters if isinstance(parameter, Prior)
-        ]
-        # noinspection PyUnresolvedReferences
-        self.free_parameters += [
-            prior
-            for parameter in free_parameters
-            if isinstance(parameter, (AbstractPriorModel, TuplePrior))
-            for prior in parameter.priors
-        ]
+        self.free_parameters = _unpack(free_parameters)
+        self.positional_parameters = defaultdict(dict)
+
+    def __getitem__(self, item: int):
+        """
+        Used to override some model component for the model at a given index.
+
+        Parameters
+        ----------
+        item
+            The index of the model in the collection
+
+        Returns
+        -------
+        A manager for overriding the model components
+        """
+        return PositionalParameters(self, item)
 
     def modify_model(self, model: AbstractPriorModel) -> AbstractPriorModel:
         """
@@ -60,19 +124,30 @@ class FreeParameterAnalysis(IndexCollectionAnalysis):
         A new model with all the same priors except for those associated
         with free parameters.
         """
-        return Collection(
+        collection = Collection(
             [
                 analysis.modify_model(
-                    model.mapper_from_partial_prior_arguments(
-                        {
-                            free_parameter: free_parameter.new()
-                            for free_parameter in self.free_parameters
-                        }
-                    )
+                    model.mapper_from_partial_prior_arguments(self._arguments)
                 )
                 for analysis in self.analyses
             ]
         )
+        for i, positional_parameters in self.positional_parameters.items():
+            for key, value in positional_parameters.items():
+                path = model.path_for_object(key)
+                if path == ():
+                    collection[i] = value
+                else:
+                    collection[i].set_item_at_path(path, value)
+
+        return collection
+
+    @property
+    def _arguments(self):
+        return {
+            free_parameter: free_parameter.new()
+            for free_parameter in self.free_parameters
+        }
 
     def modify_before_fit(self, paths: AbstractPaths, model: Collection):
         """


Test patch:
diff --git a/test_autofit/analysis/test_free_parameter.py b/test_autofit/analysis/test_free_parameter.py
index 5acb0fbb9..951828321 100644
--- a/test_autofit/analysis/test_free_parameter.py
+++ b/test_autofit/analysis/test_free_parameter.py
@@ -41,6 +41,92 @@ def make_combined_analysis(model, Analysis):
     return (Analysis() + Analysis()).with_free_parameters(model.centre)
 
 
+def test_override_specific_free_parameter(model, combined_analysis):
+    combined_analysis[0][model.centre] = 2
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].centre == 2
+    assert new_model[1].centre != 2
+
+
+def test_override_multiple(model, combined_analysis):
+    combined_analysis[0][model.centre] = 2
+    combined_analysis[1][model.centre] = 3
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].centre == 2
+    assert new_model[1].centre == 3
+
+
+def test_override_multiple_one_analysis(model, combined_analysis):
+    combined_analysis[0][model.centre] = 2
+    combined_analysis[0][model.sigma] = 3
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].centre == 2
+    assert new_model[0].sigma == 3
+
+
+def test_complex_path(Analysis):
+    model = af.Collection(
+        collection=af.Collection(
+            gaussian=af.Model(af.Gaussian),
+        )
+    )
+    combined_analysis = (Analysis() + Analysis()).with_free_parameters(
+        model.collection.gaussian
+    )
+
+    combined_analysis[0][model.collection.gaussian.centre] = 2
+    combined_analysis[0][model.collection.gaussian.sigma] = 3
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].collection.gaussian.centre == 2
+    assert new_model[0].collection.gaussian.sigma == 3
+
+
+def test_tuple_prior_override(Analysis):
+    model = af.Collection(
+        model=af.Model(
+            af.mock.MockChildTuple,
+        )
+    )
+
+    combined_analysis = (Analysis() + Analysis()).with_free_parameters(
+        model.model.tup,
+    )
+
+    first = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)
+    second = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)
+
+    combined_analysis[0][model.model.tup.tup_0] = first
+    combined_analysis[0][model.model.tup.tup_1] = second
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].model.tup.tup_0 is first
+    assert new_model[0].model.tup.tup_1 is second
+
+
+def test_override_model(model, combined_analysis):
+    new_model = af.Model(af.Gaussian, centre=2)
+    combined_analysis[0][model] = new_model
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].centre == 2
+    assert new_model[1].centre != 2
+
+
+def test_override_child_model(Analysis):
+    model = af.Collection(gaussian=af.Gaussian)
+    combined_analysis = (Analysis() + Analysis()).with_free_parameters(
+        model.gaussian.centre
+    )
+    combined_analysis[0][model.gaussian] = af.Model(af.Gaussian, centre=2)
+
+    new_model = combined_analysis.modify_model(model)
+    assert new_model[0].gaussian.centre == 2
+
+
 def test_multiple_free_parameters(model, Analysis):
     combined_analysis = (Analysis() + Analysis()).with_free_parameters(
         model.centre, model.sigma