Add callback to SGD and CMA-ES optimizer

src/qibo/optimizers.py

@@ -90,7 +90,7 @@ def myloss(parameters, circuit):
 
         backend = _check_backend(backend)
 
-        return sgd(loss, initial_parameters, args, options, compile, backend)
+        return sgd(loss, initial_parameters, args, callback, options, compile, backend)
     else:
         from qibo.backends import _check_backend
 
@@ -213,7 +213,15 @@ def newtonian(
     return m.fun, m.x, m
 
 
-def sgd(loss, initial_parameters, args=(), options=None, compile=False, backend=None):
+def sgd(
+    loss,
+    initial_parameters,
+    args=(),
+    callback=None,
+    options=None,
+    compile=False,
+    backend=None,
+):
     """Stochastic Gradient Descent (SGD) optimizer using Tensorflow backpropagation.
 
     See `tf.keras.Optimizers <https://www.tensorflow.org/api_docs/python/tf/keras/optimizers>`_
@@ -225,6 +233,7 @@ def sgd(loss, initial_parameters, args=(), options=None, compile=False, backend=
         initial_parameters (np.ndarray): Initial guess for the variational
             parameters.
         args (tuple): optional arguments for the loss function.
+        callback (callable): Called after each iteration.
         options (dict): Dictionary with options for the SGD optimizer. Supports
             the following keys:
 
@@ -234,6 +243,7 @@ def sgd(loss, initial_parameters, args=(), options=None, compile=False, backend=
             - ``'nmessage'`` (int, default: ``1e3``): Every how many epochs to print
               a message of the loss function.
     """
+
     if not backend.name == "tensorflow":
         raise_error(RuntimeError, "SGD optimizer requires Tensorflow backend.")
 
@@ -265,6 +275,8 @@ def opt_step():
 
     for e in range(sgd_options["nepochs"]):
         l = opt_step()
+        if callback is not None:
+            callback(vparams)
         if e % sgd_options["nmessage"] == 1:
             log.info("ite %d : loss %f", e, l.numpy())
 

tests/test_models_variational.py

@@ -128,8 +128,16 @@ def test_vqe(backend, method, options, compile, filename):
     np.random.seed(0)
     initial_parameters = np.random.uniform(0, 2 * np.pi, 2 * nqubits * layers + nqubits)
     v = models.VQE(circuit, hamiltonian)
+
+    def callback(parameters):
+        pass
+
     best, params, _ = v.minimize(
-        initial_parameters, method=method, options=options, compile=compile
+        initial_parameters,
+        method=method,
+        options=options,
+        compile=compile,
+        callback=callback,
     )
     if method == "cma":
         # remove `outcmaes` folder