Repository: CUQI-DTU/CUQIpy

Issue:
Better documentation/interface on RestorationPrior
## Issue
- With the current implementation of `RestorationPrior`, we expect the user-supplied `restorator` to return a tuple of two variables, i.e., the "restored" solution and extra information.
- It's quite easy for the user (and us...) to make a mistake here as they might provide a function that only returns the "restored" solution when there is no meaningful extra information to provide. Particularly in a 2D case, we won't even notice this error as the code runs with such an incorrectly defined restorator (but gives wrong results, of course).

## DoD
- [x] Propose a solution and discuss it with team
  - Option 1: CUQIpy internally detects and handles no matter if the user-supplied restorator returns a solution or a tuple
  - Option 2: CUQIpy accepts only a function that returns a tuple of length 2 and raises error if it doesn't
  - Stand-up suggests that option 1 might add confusion to the user and we go with option 2
- [x] Investigate the use of `inspect` module to infer the return type of a user-supplied restorator
  - `inspect` seems not what we want here and we shouldn't expect our user to annotate their functions 
- [x] Check return type at run time and raise value error if the supplied restorator doesn't return a tuple of length 2.
- [x] Add unit test to check of the value error is correctly catpured
- [x] Update the docstring of `RestorationPrior` class, the `restore` method and ValueError message to guide the user.

Hints from issue comments:
@chaozg: @jakobsj Could you please review the DoD here? It's intended to make the interface of `restorator` more robust to avoid mistakes (that are very easy to make).@jakobsj: @chaozg Thank you, could you explain why one may want to return such extra information as a second output and ideally point me to one or more specific examples of "extra information"?   Just considering whether it could be removed altogether such that always a single output is returned.@chaozg: Hi @jakobsj , as I understand, @amal-ghamdi and @remilaumont reserved this "extra information" for potential information in computing the proximal operator, for example number of loss function evaluations. Maybe @amal-ghamdi could give more comments on the design here?@chaozg: @jakobsj I just updated the DoD here a bit after suggestions from our stand-up@jakobsj: @chaozg would it make sense to add a unit test checking that returning different number of outputs than 2 is caught? If so, please add a tickbox for unit test in the DoD.@chaozg: Thank you @jakobsj just update the DoD and the PR as you suggested.@jakobsj: THanks @chaozg DoD approved!

Solution patch:
diff --git a/cuqi/implicitprior/_restorator.py b/cuqi/implicitprior/_restorator.py
index 1e64a4f8d..20d03a2bd 100644
--- a/cuqi/implicitprior/_restorator.py
+++ b/cuqi/implicitprior/_restorator.py
@@ -15,8 +15,10 @@ class RestorationPrior(Distribution):
     Parameters
     ---------- 
     restorator : callable f(x, restoration_strength)
-        Function f that accepts input x to be restored and returns the
-        restored version of x and information about the restoration operation.
+        Function f that accepts input x to be restored and returns a two-element 
+        tuple of the restored version of x and extra information about the 
+        restoration operation. The second element can be of any type, including
+        `None` in case there is no information.
             
     restorator_kwargs : dictionary
         Dictionary containing information about the restorator.
@@ -41,8 +43,9 @@ def __init__(self, restorator, restorator_kwargs
         super().__init__(**kwargs)
 
     def restore(self, x, restoration_strength):
-        """This function allows us to restore the input x and returns the
-        restored version of x.
+        """This function allows us to restore the input x with the user-supplied
+        restorator. Extra information about the restoration operation is stored
+        in the `RestorationPrior` info attribute.
         
         Parameters
         ---------- 
@@ -54,9 +57,18 @@ def restore(self, x, restoration_strength):
             restorator is a denoiser, this parameter might correspond to the
             noise level.
         """
-        solution, info = self.restorator(x, restoration_strength=restoration_strength,
+        restorator_return = self.restorator(x, restoration_strength=restoration_strength,
                                          **self.restorator_kwargs)
-        self.info = info
+
+        if type(restorator_return) == tuple and len(restorator_return) == 2:
+            solution, self.info = restorator_return
+        else:
+            raise ValueError("Unsupported return type from the user-supplied restorator function. "+ 
+                             "Please ensure that the restorator function returns a two-element tuple with the "+
+                             "restored solution as the first element and additional information about the "+
+                             "restoration as the second element. The second element can be of any type, "+
+                             "including `None` in case there is no particular information.")
+
         return solution
     
     def logpdf(self, x):


Test patch:
diff --git a/tests/test_implicit_priors.py b/tests/test_implicit_priors.py
index ad617a3c1..312943f87 100644
--- a/tests/test_implicit_priors.py
+++ b/tests/test_implicit_priors.py
@@ -28,7 +28,7 @@ def test_RegularizedGaussian_guarding_statements():
 
     with pytest.raises(ValueError, match="Projector should take 1 argument"):
         cuqi.implicitprior.RegularizedGaussian(np.zeros(5), 1, projector=lambda s,z: s)
-        
+
 def test_creating_restorator():
     """ Test creating the object from restorator class."""
 
@@ -38,6 +38,32 @@ def func(x, restoration_strength=0.1):
     assert np.allclose(restorator.restore(np.ones(4), 0.1), np.ones(4))
     assert restorator.info == True
 
+def test_handling_invalid_restorator():
+    """ Test handling invalid restorator."""
+    # Invalid return type 1: None
+    def func_1(x, restoration_strength=0.1):
+        return
+    restore_prior_1 = cuqi.implicitprior.RestorationPrior(func_1)
+    with pytest.raises(ValueError, match=r"Unsupported return type .*"):
+        restore_prior_1.restore(np.ones(4), 0.1)
+    # Invalid return type 2: one parameter
+    def func_2(x, restoration_strength=0.1):
+        return x
+    restore_prior_2 = cuqi.implicitprior.RestorationPrior(func_2)
+    with pytest.raises(ValueError, match=r"Unsupported return type .*"):
+        restore_prior_2.restore(np.ones(4), 0.1)
+    # Invalid return type 3: tuple with 3 elements
+    def func_3(x, restoration_strength=0.1):
+        return x, None, False
+    restore_prior_3 = cuqi.implicitprior.RestorationPrior(func_3)
+    with pytest.raises(ValueError, match=r"Unsupported return type .*"):
+        restore_prior_3.restore(np.ones(4), 0.1)
+    # Invalid return type 4: list with 2 elements
+    def func_4(x, restoration_strength=0.1):
+        return [x, None]
+    restore_prior_4 = cuqi.implicitprior.RestorationPrior(func_4)
+    with pytest.raises(ValueError, match=r"Unsupported return type .*"):
+        restore_prior_4.restore(np.ones(4), 0.1)
 
 def test_creating_restorator_with_potential():
     """ Test creating the object from restorator class with a potential."""