change GENETIC to Genetic

Author: rochisha0

doc/guide/guide-control.rst

@@ -351,8 +351,8 @@ have passed without improvement (stagnation criterion).
 Each generation represents a full evaluation of the population, making the method inherently parallelizable 
 and effective in high-dimensional control landscapes.
 
-In QuTiP, the GA optimization is implemented via the ``_GENETIC`` class, and can be invoked using the 
-standard ``optimize_pulses`` interface by setting the algorithm to ``"GENETIC"``.
+In QuTiP, the GA optimization is implemented via the ``_Genetic`` class, and can be invoked using the 
+standard ``optimize_pulses`` interface by setting the algorithm to ``"Genetic"``.
 
 Optimal Quantum Control in QuTiP (Genetic Algorithm)
 ====================================================
@@ -388,7 +388,7 @@ Defining a control problem in QuTiP using the Genetic Algorithm follows the same
 
     # set genetic algorithm hyperparameters
     algorithm_kwargs = {
-        "alg": "GENETIC",
+        "alg": "Genetic",
         "population_size": 50,
         "generations": 100,
         "mutation_rate": 0.2,

src/qutip_qoc/_genetic.py

@@ -3,13 +3,16 @@
 import time
 from qutip_qoc.result import Result
 
-class _GENETIC:
+
+class _Genetic:
     """
     Genetic Algorithm-based optimizer for quantum control problems.
 
     This class implements a global optimization routine using a Genetic Algorithm
-    to optimize control pulses that drive a quantum system from an initial state 
-    to a target state (or unitary). 
+    to optimize control pulses that drive a quantum system from an initial state
+    to a target state (or unitary).
+
+    Based on the code from Jonathan Brown
     """
 
     def __init__(
@@ -26,7 +29,9 @@ def __init__(
     ):
         self._objective = objectives[0]
         self._Hd = self._objective.H[0]
-        self._Hc_lst = [H[0] if isinstance(H, list) else H for H in self._objective.H[1:]]
+        self._Hc_lst = [
+            H[0] if isinstance(H, list) else H for H in self._objective.H[1:]
+        ]
         self._initial = self._objective.initial
         self._target = self._objective.target
         self._norm_fac = 1 / self._target.norm()
@@ -41,14 +46,17 @@ def __init__(
         self.generations = alg_kwargs.get("generations", 100)
         self.mutation_rate = alg_kwargs.get("mutation_rate", 0.3)
         self.fid_err_targ = alg_kwargs.get("fid_err_targ", 1e-4)
-        self._stagnation_patience = 20  # Internally fixed
+        self._stagnation_patience = 50  # Internally fixed
 
         self._integrator_kwargs = integrator_kwargs
         self._fid_type = alg_kwargs.get("fid_type", "PSU")
 
         self._generator = self._prepare_generator()
-        self._solver = qt.MESolver(H=self._generator, options=self._integrator_kwargs) \
-            if self._Hd.issuper else qt.SESolver(H=self._generator, options=self._integrator_kwargs)
+        self._solver = (
+            qt.MESolver(H=self._generator, options=self._integrator_kwargs)
+            if self._Hd.issuper
+            else qt.SESolver(H=self._generator, options=self._integrator_kwargs)
+        )
 
         self._result = Result(
             objectives=[self._objective],
@@ -64,10 +72,18 @@ def __init__(
         self._result._final_states = []
 
     def _prepare_generator(self):
-        args = {f"p{i+1}_{j}": 0.0 for i in range(self.N_controls) for j in range(self.N_steps)}
+        args = {
+            f"p{i+1}_{j}": 0.0
+            for i in range(self.N_controls)
+            for j in range(self.N_steps)
+        }
 
         def make_coeff(i, j):
-            return lambda t, args: args[f"p{i+1}_{j}"] if int(t / (self._evo_time / self.N_steps)) == j else 0
+            return lambda t, args: (
+                args[f"p{i+1}_{j}"]
+                if int(t / (self._evo_time / self.N_steps)) == j
+                else 0
+            )
 
         H_qev = [self._Hd]
         for i, Hc in enumerate(self._Hc_lst):
@@ -77,7 +93,11 @@ def make_coeff(i, j):
         return qt.QobjEvo(H_qev, args=args)
 
     def _infid(self, params):
-        args = {f"p{i+1}_{j}": params[i * self.N_steps + j] for i in range(self.N_controls) for j in range(self.N_steps)}
+        args = {
+            f"p{i+1}_{j}": params[i * self.N_steps + j]
+            for i in range(self.N_controls)
+            for j in range(self.N_steps)
+        }
         result = self._solver.run(self._initial, [0.0, self._evo_time], args=args)
         final_state = result.final_state
         self._result._final_states.append(final_state)
@@ -87,7 +107,9 @@ def _infid(self, params):
             fid = 0.5 * np.real((diff.dag() * diff).tr())
         else:
             overlap = self._norm_fac * self._target.overlap(final_state)
-            fid = 1 - np.abs(overlap) if self._fid_type == "PSU" else 1 - np.real(overlap)
+            fid = (
+                1 - np.abs(overlap) if self._fid_type == "PSU" else 1 - np.real(overlap)
+            )
 
         return fid
 
@@ -101,7 +123,7 @@ def initial_population(self):
         return np.random.uniform(-1, 1, (self.N_pop, self.N_var))
 
     def darwin(self, population, fitness):
-        indices = np.argsort(-fitness)[:self.N_pop // 2]
+        indices = np.argsort(-fitness)[: self.N_pop // 2]
         return population[indices], fitness[indices]
 
     def pairing(self, survivors, survivor_fitness):
@@ -130,7 +152,9 @@ def build_next_gen(self, survivors, offspring):
         return np.vstack((survivors, offspring))
 
     def mutate(self, population):
-        n_mut = int((population.shape[0] - 1) * population.shape[1] * self.mutation_rate)
+        n_mut = int(
+            (population.shape[0] - 1) * population.shape[1] * self.mutation_rate
+        )
         row = np.random.randint(1, population.shape[0], size=n_mut)
         col = np.random.randint(0, population.shape[1], size=n_mut)
         population[row, col] += np.random.normal(0, 0.3, size=n_mut)
@@ -183,13 +207,15 @@ def optimize(self):
 
         self._result.message = (
             f"Stopped early: reached infidelity target {self.fid_err_targ}"
-            if -best_fit <= self.fid_err_targ else
-            f"Stopped due to stagnation after {self._stagnation_patience} generations"
-            if no_improvement_counter >= self._stagnation_patience else
-            "Optimization completed successfully"
+            if -best_fit <= self.fid_err_targ
+            else (
+                f"Stopped due to stagnation after {self._stagnation_patience} generations"
+                if no_improvement_counter >= self._stagnation_patience
+                else "Optimization completed successfully"
+            )
         )
         return self._result
-    
+
     def result(self):
         self._result.start_local_time = time.strftime(
             "%Y-%m-%d %H:%M:%S", time.localtime(self._result.start_local_time)

tests/test_result.py

@@ -200,8 +200,6 @@ def sin_z_jax(t, r, **kwargs):
 
 # --------------------------- Genetic ---------------------------
 
-# TODO: this is the input for optimiz_pulses() function
-# you can use this routine to test your implementation
 
 # state to state transfer
 init = qt.basis(2, 0)
@@ -219,16 +217,15 @@ def sin_z_jax(t, r, **kwargs):
     control_parameters={
         "p": {"bounds": [(-13, 13)]},
     },
-    tlist=np.linspace(0, 10, 100), # TODO: derive single step duration and max evo time / max num steps from this
+    tlist=np.linspace(0, 10, 100), 
     algorithm_kwargs={
         "fid_err_targ": 0.01,
-        "alg": "GENETIC",
+        "alg": "Genetic",
         "max_iter": 100,
     },
     optimizer_kwargs={},
 )
 
-# TODO: no big difference for unitary evolution
 
 initial = qt.qeye(2) # Identity
 target  = qt.gates.hadamard_transform()
@@ -237,8 +234,6 @@ def sin_z_jax(t, r, **kwargs):
     objectives=[Objective(initial, H, target)],
 )
 
-
-
 @pytest.fixture(
     params=[
         pytest.param(state2state_grape, id="State to state (GRAPE)"),