Merge pull request #71 from akataba/cf/2-qubit-main

Cf/2 qubit main

scripts/two-qubits/CZ.py

@@ -4,13 +4,11 @@
 import ray
 from ray.rllib.algorithms.ddpg import DDPGConfig
 from ray.tune.registry import register_env
-from relaqs.environments.gate_synth_env_rllib_Haar import TwoQubitGateSynth
+from relaqs.environments import NoisyTwoQubitEnv
 from relaqs.save_results import SaveResults
 from relaqs.plot_data import plot_data
 
-from relaqs.quantum_noise_data.get_data import get_month_of_single_qubit_data, get_month_of_all_qubit_data
-from relaqs import quantum_noise_data
-from relaqs import QUANTUM_NOISE_DATA_DIR
+
 from relaqs import RESULTS_DIR
 
 from qutip.operators import *
@@ -21,17 +19,18 @@
 
 
 def env_creator(config):
-    return TwoQubitGateSynth(config)
-
-def save_grad_to_file(resultdict):
-    try:
-        policydict = resultdict["default_policy"]
-        stats = policydict["learner_stats"]
-        grad_gnorm = stats["grad_gnorm"]
-        with open("gradfile", "a") as f:
-            f.write(f"{grad_gnorm}\n")
-    except KeyError:
-        print(f"Failed to extract grad_gnorm from: {resultdict}")
+    return NoisyTwoQubitEnv(config)
+
+# def save_grad_to_file(resultdict):
+#     try:
+#         policydict = resultdict["default_policy"]
+#         stats = policydict["learner_stats"]
+#         grad_gnorm = stats["grad_gnorm"]
+#         with open("gradfile", "a") as f:
+#             f.write(f"{grad_gnorm}\n")
+#     except KeyError:
+#         pass
+        # print(f"Failed to extract grad_gnorm from: {resultdict}")
 
 def inject_logging(alg, logging_func):
     og_ts = alg.training_step
@@ -43,45 +42,52 @@ def new_training_step():
     alg.training_step = new_training_step
 
 def run(n_training_iterations=1, save=True, plot=True):
-    ray.init()
+    ray.init(num_gpus=1)
+    # ray.init()
     try:
         register_env("my_env", env_creator)
 
-
         # ---------------------> Configure algorithm and Environment <-------------------------
         alg_config = DDPGConfig()
+        # alg_config = DDPGConfig()
         alg_config.framework("torch")
 
-        env_config = TwoQubitGateSynth.get_default_env_config()
+        env_config = NoisyTwoQubitEnv.get_default_env_config()
         CZ = cphase(np.pi).data.toarray()
         env_config["U_target"] = CZ
 
         alg_config.environment("my_env", env_config=env_config)
 
         alg_config.rollouts(batch_mode="complete_episodes")
-        alg_config.train_batch_size = TwoQubitGateSynth.get_default_env_config()["steps_per_Haar"]
+        alg_config.train_batch_size = NoisyTwoQubitEnv.get_default_env_config()["steps_per_Haar"]
 
         ### working 1-3 sets
-        alg_config.actor_lr = 4e-5
-        alg_config.critic_lr = 5e-4
+        alg_config.actor_lr = 1e-4
+        alg_config.critic_lr = 1e-4
 
         alg_config.actor_hidden_activation = "relu"
         alg_config.critic_hidden_activation = "relu"
-        alg_config.num_steps_sampled_before_learning_starts = 10000
-        alg_config.actor_hiddens = [500,500,500,500]
-        alg_config.critic_hiddens = [500,500,500,500]
-        alg_config.exploration_config["scale_timesteps"] = 1E5
+        alg_config.num_steps_sampled_before_learning_starts = 5000
+        # alg_config.actor_hiddens = [500,20000,500]
+        # alg_config.critic_hiddens = [500,20000,500]
+        alg_config.actor_hiddens = [1000, 1000, 1000]
+        alg_config.critic_hiddens = [1000, 1000, 1000]
+        # alg_config.exploration_config["scale_timesteps"] = 200000
+        alg_config.exploration_config["scale_timesteps"] = 10000
+
         print(alg_config.algo_class)
         print(alg_config["framework"])
 
         alg = alg_config.build()
-        inject_logging(alg, save_grad_to_file)
+        # inject_logging(alg, save_grad_to_file)
         # ---------------------------------------------------------------------
         list_of_results = []
         # ---------------------> Train Agent <-------------------------
-        for _ in range(n_training_iterations):
+        for ii in range(n_training_iterations):
             result = alg.train()
             list_of_results.append(result['hist_stats'])
+            if np.mod(ii,5)==0:
+                print("currently",ii,"/",n_training_iterations)
         # -------------------------------------------------------------
 
         # ---------------------> Save Results <-------------------------
@@ -102,7 +108,7 @@ def run(n_training_iterations=1, save=True, plot=True):
         ray.shutdown()
 
 if __name__ == "__main__":
-    n_training_iterations = 100
+    n_training_iterations = 50
     save = True
     plot = True
     run(n_training_iterations, save, plot)

src/relaqs/api/utils.py

@@ -16,6 +16,10 @@
                                    int(np.sqrt(X.shape[0])),
                                    order="F") # inverse vectorization operation, column-order. X is a numpy array.
 
+def normalize(quantity, list_of_values):
+    """ normalize quantity to [0, 1] range based on list of values """
+    return (quantity - min(list_of_values) + 1E-15) / (max(list_of_values) - min(list_of_values) + 1E-15)
+
 def polar_vec_to_complex_matrix(vec, return_flat=False):
     """ 
     The intended use of this function is to convert from the representation of the unitary

src/relaqs/environments/__init__.py

@@ -1,2 +1,3 @@
 from .single_qubit_env import SingleQubitEnv
 from .noisy_single_qubit_env import NoisySingleQubitEnv
+from .noisy_two_qubit_env import NoisyTwoQubitEnv

src/relaqs/environments/noisy_single_qubit_env.py

@@ -7,7 +7,7 @@
 from qutip.operators import sigmam, sigmaz
 from relaqs.environments.single_qubit_env import SingleQubitEnv
 from relaqs.api import gates
-from relaqs.api.utils import sample_noise_parameters
+from relaqs.api.utils import sample_noise_parameters, normalize
 
 I = gates.I().get_matrix()
 X = gates.X().get_matrix()
@@ -56,15 +56,15 @@ def get_relaxation_rate(self):
             sampled_rate_list.append(random.sample(self.relaxation_rates_list[ii],k=1)[0])
 
         return sampled_rate_list
-            
+
     def get_observation(self):
-        normalized_detuning = [(self.detuning - min(self.detuning_list) + 1E-15) / (max(self.detuning_list) - min(self.detuning_list) + 1E-15)]
-        normalized_t1 = (self.relaxation_rate[0] - min(self.relaxation_rates_list[0]) + 1E-15) / (max(self.relaxation_rates_list[0]) - min(self.relaxation_rates_list[0]) + 1E-15)
-        normalized_t2 = (self.relaxation_rate[1] - min(self.relaxation_rates_list[1]) + 1E-15) / (max(self.relaxation_rates_list[1]) - min(self.relaxation_rates_list[1]) + 1E-15)
+        normalized_detuning = [normalize(self.detuning, self.detuning_list)]
+        normalized_relaxation_rates = [normalize(self.relaxation_rate[0], self.relaxation_rates_list[0]),
+                                       normalize(self.relaxation_rate[1], self.relaxation_rates_list[1])] # could do list comprehension
         return np.append([self.compute_fidelity()] +
-                         [normalized_t1, normalized_t2] +
+                         normalized_relaxation_rates +
                          normalized_detuning,
-                         self.unitary_to_observation(self.U)) 
+                         self.unitary_to_observation(self.U))
 
     def hamiltonian(self, detuning, alpha, gamma_magnitude, gamma_phase):
         return (detuning + alpha)*Z + gamma_magnitude*(np.cos(gamma_phase)*X + np.sin(gamma_phase)*Y)