get_rearrange_paths bug fixing

Python/waveform.py

@@ -474,19 +474,16 @@ def create_path_table_reduced_gpu(
 def get_rearrange_paths(
         filled_idx: np.ndarray,
         target_idx: np.ndarray,
-) -> Tuple[List[Tuple[Any, Any]], List[Any]]:
+) -> List[Tuple[Any, Any]]:
     """
     Calculate rearranging paths.
     :param filled_idx: indices of tweezer positions filled with atoms.
     :param target_idx: indices of tweezer positions in target pattern.
-    :returns: Tuple containing moving paths and turn-off paths.
+    :returns: list containing tuples of moving paths
     """
-    n = 0
     t_size = target_idx.size
     f_size = filled_idx.size
     reserve = f_size - t_size
-    # move_paths = []
-    # static_paths = []
     paths = []
     i = 0
     j = 0
@@ -495,21 +492,21 @@ def get_rearrange_paths(
         if filled_idx[i] == target_idx[j]:
             # paths.append((filled_idx[i], filled_idx[i]))
             j += 1
-            i = j
+            i += 1
         elif (reserve > 0
               and filled_idx[i] < target_idx[j]
-              and abs(filled_idx[i + 1] - target_idx[j]) < abs(filled_idx[i + 1] - target_idx[j])):
+              and abs(filled_idx[i + 1] - target_idx[j]) < abs(filled_idx[i] - target_idx[j])):
             i += 1
             reserve -= 1
         else:
             paths.append((filled_idx[i], target_idx[j]))
             i += 1
             j += 1
-    off = []
-    if reserve < 0:
-        for i in range(abs(reserve)):
-            off.append(target_idx[-1 - i])
-    return paths, off
+    # off = []
+    # if reserve < 0:
+    #     for i in range(abs(reserve)):
+    #         off.append(target_idx[-1 - i])
+    return paths
 
 
 def create_moving_array(path_table: np.ndarray, paths: np.ndarray) -> np.ndarray:
@@ -611,3 +608,30 @@ def create_moving_signal_single(omega_i, omega_f, sample_rate, signal_time):
     signal = (amps * np.sin(signal)).astype(np.int16)
     return signal
 
+
+def create_move_then_back(omega_i, omega_f, sample_rate, move_time, stay_time):
+    min_len = 2 * sample_rate / 0.1e3
+    sample_len = sample_rate * (move_time * 2 + stay_time)
+    sample_len += min_len - sample_len % min_len
+    sample_len = int(sample_len)
+
+    t = np.arange(sample_len) / sample_rate
+    t_tot = sample_len / sample_rate
+    a = 4 * (omega_i - omega_f) / (t_tot ** 2)
+    end = sample_len
+    half = int(end / 2) + 1
+    t1 = t[:half]
+    t2 = t[half:end] - t_tot / 2
+
+    amps = 2**12
+    signal = np.zeros(sample_len)
+
+    signal[:half] = omega_i * t1 - a / 6 * t1 ** 3  # t<=T/2
+    # ph = wfm.phi[i] + omega_i * t_tot / 2 + a / 6 * (t_tot / 2) ** 3
+    signal[half:end] = signal[half - 1] + \
+                     (omega_i - a / 2 * (t_tot / 2) ** 2) * t2 - \
+                     a / 2 * t_tot / 2 * t2 ** 2 + \
+                     a / 6 * t2 ** 3  # t>=T/2
+    signal[end:] = signal[end - 1] + omega_f * (t[end:] - t[end - 1])
+    signal = (amps * np.sin(signal)).astype(np.int16)
+    return signal