feat: add callback support to VQA optimization functions

Author: JulienCalistoTD

mpqp/execution/vqa/vqa.py

@@ -23,6 +23,15 @@
     [OptimizableFunc, Optional[OptimizerInput], Optional[OptimizerOptions]],
     tuple[float, OptimizerInput],
 ]
+OptimizerCallback = Optional[
+    Union[
+        Callable[[OptimizeResult], None],
+        Callable[
+            [Union[list[float], npt.NDArray[np.float32], tuple[float, ...]]], None
+        ],
+    ]
+]
+
 
 # TODO: all those functions with almost or exactly the same signature look like
 #  a code smell to me.
@@ -47,6 +56,7 @@ def minimize(
     init_params: Optional[OptimizerInput] = None,
     nb_params: Optional[int] = None,
     optimizer_options: Optional[dict[str, Any]] = None,
+    callback: Optional[OptimizerCallback] = None,
 ) -> tuple[float, OptimizerInput]:
     """This function runs an optimization on the parameters of the circuit, in order to
     minimize the measured expectation value of observables associated with the given circuit.
@@ -71,6 +81,7 @@ def minimize(
         optimizer_options: Options used to configure the VQA optimizer (maximum
             iterations, convergence threshold, etc...). These options are passed
             as is to the minimizer.
+        callback:  A callable called after each iteration.
 
     Returns:
         The optimal value reached and the parameters corresponding to this value.
@@ -117,11 +128,25 @@ def minimize(
         if device is None:
             raise ValueError("A device is needed to optimize a circuit")
         optimizer = _minimize_remote if device.is_remote() else _minimize_local
-        return optimizer(optimizable, method, device, init_params, nb_params)
+        return optimizer(
+            optimizable,
+            method,
+            device,
+            init_params,
+            nb_params,
+            optimizer_options,
+            callback,
+        )
     else:
         # TODO: find a way to know if the job is remote or local from the function
         return _minimize_local(
-            optimizable, method, device, init_params, nb_params, optimizer_options
+            optimizable,
+            method,
+            device,
+            init_params,
+            nb_params,
+            optimizer_options,
+            callback,
         )
 
 
@@ -133,6 +158,7 @@ def _minimize_remote(
     init_params: Optional[OptimizerInput] = None,
     nb_params: Optional[int] = None,
     optimizer_options: Optional[dict[str, Any]] = None,
+    callback: Optional[OptimizerCallback] = None,
 ) -> tuple[float, OptimizerInput]:
     """This function runs an optimization on the parameters of the circuit, to
     minimize the expectation value of the measure of the circuit by it's
@@ -158,6 +184,7 @@ def _minimize_remote(
         optimizer_options: Options used to configure the VQA optimizer (maximum
             iterations, convergence threshold, etc...). These options are passed
             as is to the minimizer.
+        callback:  A callable called after each iteration.
 
     Returns:
         The optimal value reached and the parameters used to reach this value.
@@ -175,6 +202,7 @@ def _minimize_local(
     init_params: Optional[OptimizerInput] = None,
     nb_params: Optional[int] = None,
     optimizer_options: Optional[dict[str, Any]] = None,
+    callback: Optional[OptimizerCallback] = None,
 ) -> tuple[float, OptimizerInput]:
     """This function runs an optimization on the parameters of the circuit, to
     minimize the expectation value of the measure of the circuit by it's
@@ -200,6 +228,7 @@ def _minimize_local(
         optimizer_options: Options used to configure the VQA optimizer (maximum
             iterations, convergence threshold, etc...). These options are passed
             as is to the minimizer.
+        callback:  A callable called after each iteration.
 
     Returns:
         the optimal value reached and the parameters used to reach this value.
@@ -208,11 +237,11 @@ def _minimize_local(
         if device is None:
             raise ValueError("A device is needed to optimize a circuit")
         return _minimize_local_circ(
-            optimizable, device, method, init_params, optimizer_options
+            optimizable, device, method, init_params, optimizer_options, callback
         )
     else:
         return _minimize_local_func(
-            optimizable, method, init_params, nb_params, optimizer_options
+            optimizable, method, init_params, nb_params, optimizer_options, callback
         )
 
 
@@ -223,6 +252,7 @@ def _minimize_local_circ(
     method: Optimizer | OptimizerCallable,
     init_params: Optional[OptimizerInput] = None,
     optimizer_options: Optional[dict[str, Any]] = None,
+    callback: Optional[OptimizerCallback] = None,
 ) -> tuple[float, OptimizerInput]:
     """This function runs an optimization on the parameters of the circuit, to
     minimize the expectation value of the measure of the circuit by it's
@@ -244,6 +274,7 @@ def _minimize_local_circ(
         optimizer_options: Options used to configure the VQA optimizer (maximum
             iterations, convergence threshold, etc...). These options are passed
             as is to the minimizer.
+        callback:  A callable called after each iteration.
 
     Returns:
         The optimal value reached and the parameters used to reach this value.
@@ -264,7 +295,7 @@ def eval_circ(params: OptimizerInput):
         ).expectation_value
 
     return _minimize_local_func(
-        eval_circ, method, init_params, len(variables), optimizer_options
+        eval_circ, method, init_params, len(variables), optimizer_options, callback
     )
 
 
@@ -275,6 +306,7 @@ def _minimize_local_func(
     init_params: Optional[OptimizerInput] = None,
     nb_params: Optional[int] = None,
     optimizer_options: Optional[OptimizerOptions] = None,
+    callback: Optional[OptimizerCallback] = None,
 ) -> tuple[float, OptimizerInput]:
     """This function runs an optimization on the parameters of the circuit, to
     minimize the expectation value of the measure of the circuit by it's
@@ -298,6 +330,8 @@ def _minimize_local_func(
         optimizer_options: Options used to configure the VQA optimizer (maximum
             iterations, convergence threshold, etc...). These options are passed
             as is to the minimizer.
+        callback:  A callable called after each iteration.
+
 
     Returns:
         The optimal value reached and the parameters used to reach this value.
@@ -317,6 +351,7 @@ def _minimize_local_func(
             x0=np.array(init_params),
             method=method.name.lower(),
             options=optimizer_options,
+            callback=callback,
         )
         return res.fun, res.x
     else: