Cleaned up and added tests

Author: Ashutosh-Mishra2

c3/experiment.py

@@ -631,7 +631,7 @@ def expect_oper(self, state, lindbladian, oper):
         trace = np.trace(np.matmul(rho, oper))
         return [[np.real(trace)]]  # ,[np.imag(trace)]]
 
-    def compute_states(self, solver="rk4", step_function="schroedinger"):
+    def compute_states(self, solver="rk4", step_function="schrodinger"):
         """
         Use a state solver to compute the trajectory of the system.
 
@@ -681,7 +681,7 @@ def compute_states(self, solver="rk4", step_function="schroedinger"):
 
         return {"states": state_list, "ts": ts_list}
 
-    def compute_final_state(self, solver="rk4", step_function="schroedinger"):
+    def compute_final_state(self, solver="rk4", step_function="schrodinger"):
         """
         Solve the Lindblad master equation by integrating the differential
         equation of the density matrix

c3/libraries/fidelities.py

@@ -791,10 +791,15 @@ def orbit_infid(
 
 
 @fid_reg_deco
-def state_transfer_from_states(states: List[tf.Tensor], index, dims, params, n_eval=-1):
+def state_transfer_from_states(states: tf.Tensor, index, dims, params, n_eval=-1):
     infids = []
     psi_0 = params["target"]
-    overlap = calculateStateOverlap(states[-1], psi_0)
+
+    if len(states.shape) > 2:
+        overlap = calculateStateOverlap(states[-1], psi_0)
+    else:
+        overlap = calculateStateOverlap(states, psi_0)
+
     infid = 1 - overlap
     infids.append(infid)
     return tf.reduce_max(tf.math.real(infids))

c3/libraries/propagation.py

@@ -689,7 +689,6 @@ def ode_solver(
     model: Model, gen: Generator, instr: Instruction, init_state, solver, step_function
 ) -> Dict:
 
-    init_state = model.get_init_state()
     signal = gen.generate_signals(instr)
 
     if model.lindbladian:
@@ -728,7 +727,6 @@ def ode_solver_final_state(
     model: Model, gen: Generator, instr: Instruction, init_state, solver, step_function
 ) -> Dict:
 
-    init_state = model.get_init_state()
     signal = gen.generate_signals(instr)
 
     if model.lindbladian:
@@ -899,7 +897,7 @@ def lindblad(rho, h, dt, col):
 
 
 @step_deco
-def schroedinger(psi, h, dt, col=None):
+def schrodinger(psi, h, dt, col=None):
     return -1j * tf.matmul(h, psi) * dt
 
 

c3/optimizers/optimalcontrol.py

@@ -12,17 +12,6 @@
 from c3.libraries.fidelities import fidelities
 
 
-goal_functions_dict = dict()
-
-
-def goal_functions_deco(func):
-    """
-    Decorator for making registry of functions
-    """
-    goal_functions_dict[str(func.__name__)] = func
-    return func
-
-
 class OptimalControl(Optimizer):
     """
     Object that deals with the open loop optimal control.
@@ -69,7 +58,7 @@ def __init__(
         logger=None,
         fid_func_kwargs={},
         ode_solver=None,
-        ode_step_function="schroedinger",
+        ode_step_function="schrodinger",
         only_final_state=False,
     ) -> None:
         if type(algorithm) is str:
@@ -96,17 +85,29 @@ def __init__(
             self.optimize_controls
         )  # Alias the legacy name for the method running the
         # optimization
+        self.set_goal_function(
+            ode_solver=ode_solver,
+            ode_step_function=ode_step_function,
+            only_final_state=only_final_state,
+        )
+
+    def set_goal_function(
+        self,
+        ode_solver=None,
+        ode_step_function="schrodinger",
+        only_final_state=False,
+    ):
         self.ode_solver = ode_solver
         self.ode_step_function = ode_step_function
         self.only_final_state = only_final_state
 
-        self.goal_function = goal_functions_dict["goal_run"]
-
         if self.ode_solver is not None:
             if self.only_final_state:
-                self.goal_function = goal_functions_dict["goal_run_ode_only_final"]
+                self.goal_function = self.goal_run_ode_only_final
             else:
-                self.goal_function = goal_functions_dict["goal_run_ode"]
+                self.goal_function = self.goal_run_ode
+        else:
+            self.goal_function = self.goal_run
 
     def set_fid_func(self, fid_func) -> None:
         if type(fid_func) is str:
@@ -182,7 +183,6 @@ def optimize_controls(self, setup_log: bool = True) -> None:
         self.end_log()
 
     @tf.function
-    @goal_functions_deco
     def goal_run(self, current_params: tf.Tensor) -> tf.float64:
         """
         Evaluate the goal function for current parameters.
@@ -213,7 +213,6 @@ def goal_run(self, current_params: tf.Tensor) -> tf.float64:
         return goal
 
     @tf.function
-    @goal_functions_deco
     def goal_run_ode(self, current_params: tf.Tensor) -> tf.float64:
         """
         Evaluate the goal function using ode solver for current parameters.
@@ -246,7 +245,6 @@ def goal_run_ode(self, current_params: tf.Tensor) -> tf.float64:
         return goal
 
     @tf.function
-    @goal_functions_deco
     def goal_run_ode_only_final(self, current_params: tf.Tensor) -> tf.float64:
         """
         Evaluate the goal function using ode solver for current parameters.

c3/optimizers/optimizer.py

@@ -208,7 +208,7 @@ def goal_run(
     def goal_run_with_grad(self, current_params):
         with tf.GradientTape(watch_accessed_variables=False) as t:
             t.watch(current_params)
-            goal = self.goal_function(self, current_params)
+            goal = self.goal_function(current_params)
         grad = t.gradient(goal, current_params)
         return goal, grad
 
@@ -262,12 +262,12 @@ def fct_to_min(
         self.optim_status["params"] = pars
         if isinstance(input_parameters, np.ndarray):
             current_params = tf.constant(input_parameters)
-            goal = self.goal_run(current_params)
+            goal = self.goal_function(current_params)
             self.optim_status["goal"] = float(goal)
             goal = float(goal)
         else:
             current_params = input_parameters
-            goal = self.goal_run(current_params)
+            goal = self.goal_function(current_params)
             self.optim_status["goal"] = float(goal)
         self.log_parameters(input_parameters)
         return goal

test/test_two_qubits.py

@@ -7,9 +7,11 @@
 import pytest
 from numpy.testing import assert_array_almost_equal as almost_equal
 import c3.libraries.algorithms as algorithms
+from c3.libraries.fidelities import state_transfer_from_states
 
 # Libs and helpers
 from c3.libraries.propagation import rk4_unitary
+import tensorflow as tf
 
 with open("test/two_qubit_data.pickle", "rb") as filename:
     test_data = pickle.load(filename)
@@ -82,6 +84,64 @@ def test_optim(get_OC_optimizer) -> None:
     assert opt.evaluation == opt.options[maxiterKey] - 1
 
 
+@pytest.mark.slow
+@pytest.mark.tensorflow
+@pytest.mark.optimizers
+@pytest.mark.integration
+def test_optim_ode_solver(get_OC_optimizer) -> None:
+    """
+    check if optimization result is below 1e-1
+    """
+    opt = get_OC_optimizer
+    opt.fid_func = state_transfer_from_states
+    opt.set_goal_function(ode_solver="rk4", ode_step_function="schrodinger")
+
+    model = opt.exp.pmap.model
+    psi_init = [[0] * model.tot_dim]
+    psi_init[0][0] = 1 / np.sqrt(2)
+    index = model.get_state_indeces([(1, 0)])[0]
+    psi_init[0][index] = 1 / np.sqrt(2)
+    target_state = tf.transpose(tf.constant(psi_init, tf.complex128))
+    params = {"target": target_state}
+    opt.fid_func_kwargs = {"params": params}
+
+    assert opt.evaluation == 0
+    opt.optimize_controls()
+    assert opt.current_best_goal < 0.5  # TODO - Make it 0.1
+    maxiterKey = "maxiters" if opt.algorithm == algorithms.tf_sgd else "maxiter"
+    assert opt.evaluation == opt.options[maxiterKey] - 1
+
+
+@pytest.mark.slow
+@pytest.mark.tensorflow
+@pytest.mark.optimizers
+@pytest.mark.integration
+def test_optim_ode_solver_final(get_OC_optimizer) -> None:
+    """
+    check if optimization result is below 1e-1
+    """
+    opt = get_OC_optimizer
+    opt.fid_func = state_transfer_from_states
+    opt.set_goal_function(
+        ode_solver="rk4", ode_step_function="schrodinger", only_final_state=True
+    )
+
+    model = opt.exp.pmap.model
+    psi_init = [[0] * model.tot_dim]
+    psi_init[0][0] = 1 / np.sqrt(2)
+    index = model.get_state_indeces([(1, 0)])[0]
+    psi_init[0][index] = 1 / np.sqrt(2)
+    target_state = tf.transpose(tf.constant(psi_init, tf.complex128))
+    params = {"target": target_state}
+    opt.fid_func_kwargs = {"params": params}
+
+    assert opt.evaluation == 0
+    opt.optimize_controls()
+    assert opt.current_best_goal < 0.5  # TODO - Make it 0.1
+    maxiterKey = "maxiters" if opt.algorithm == algorithms.tf_sgd else "maxiter"
+    assert opt.evaluation == opt.options[maxiterKey] - 1
+
+
 @pytest.mark.tensorflow
 @pytest.mark.integration
 def test_rk4_unitary(get_two_qubit_chip) -> None:
@@ -101,3 +161,6 @@ def test_ode_solver(get_two_qubit_chip) -> None:
     exp.compute_states(solver="rk4")
     exp.compute_states(solver="rk5")
     exp.compute_states(solver="Tsit5")
+    exp.compute_states(solver="rk4", step_function="vonNeumann")
+    exp.compute_states(solver="rk5", step_function="vonNeumann")
+    exp.compute_states(solver="Tsit5", step_function="vonNeumann")