diff --git a/Python/waveform.py b/Python/waveform.py
index 9b4305850e99e722e7f74ca7a6b2a047eb660ba6..6d962943bf6620f3b9337b367241a85a232465b7 100644
--- a/Python/waveform.py
+++ b/Python/waveform.py
@@ -58,7 +58,7 @@ def create_static_array(wfm: Waveform, full=False) -> np.ndarray:
return sig.astype(np.int16)
-def create_path_table(wfm: Waveform) -> any:
+def create_path_table_old(wfm: Waveform) -> any:
"""
create a dim-3 look up table where the table[i,j] contains a sine wave to move tweezer i to tweezer j
:param wfm: waveform object already initialized with basic parameters.
@@ -175,171 +175,7 @@ def stack_right(i_start, i_end, offset, stack_size=0):
return moves, max_dist
-def get_paths_dist(nt, target_idx, max_dist):
- moves = []
- for i in range(nt):
- for j in target_idx:
- if i < j and True: # only allow uni-direction moves
- continue
- if abs(i - j) <= max_dist:
- moves.append(i,j)
- return moves
-
-
-def create_path_table_reduced(
- wfm: Waveform, target_idx, dist_offset=np.inf, save_path=None, partition=False
-) -> Tuple[Dict[Tuple[int, int], np.ndarray], np.ndarray]:
- """
- create a dim-3 look up table where the table[i,j] contains a sine wave to move tweezer i to tweezer j
- :param save_path: file saving path
- :param target_idx: indices of target pattern
- :param dist_offset: maximum move distance in indices
- :param wfm: waveform object already initialized with basic parameters.
- :return: dictionary containing rearrange paths
- """
- # interpolate optimal amplitudes
- # data = np.load("data/optimal_amps.npz")
- w = wfm.omega
- a = wfm.amplitude
- omega_interp = interp1d(w, a, kind='cubic')
-
- nt = len(wfm.omega) # total number of tweezers
- moves = []
- target = np.zeros(nt)
- target[target_idx] = 1
- dw_max = 0 # longest move, this sets the size of path_table
-
- if not partition:
- # obtain all move combinations, target based, non-partitioned:
- for i in range(nt):
- moves.append([])
- for j in target_idx:
- if i < j and True: # only allow uni-direction moves
- continue
- if abs(i - j) <= dist_offset:
- moves[i].append(j)
- dw = abs(wfm.omega[j] - wfm.omega[i])
- if dw_max < dw: dw_max = dw
- if partition:
- offset = dist_offset
- divide_idx = int(np.floor(np.median(target_idx)))
- left_size = np.sum(target[:divide_idx], dtype=int)
- right_size = np.sum(target[divide_idx:], dtype=int)
- moves_l, dw_max_l = stack_right(0, divide_idx, offset, left_size) # stack left side to right
- moves_r, dw_max_r = stack_left(divide_idx, nt, offset, right_size)
- # print("stack size, half size, middle:", len(t_idx), left_size, right_size)
- moves_l.extend(moves_r)
- moves = moves_l
- dw_max = dw_max_l if dw_max_l > dw_max_r else dw_max_r
- dw_max = abs(wfm.omega[dw_max] - wfm.omega[0])
-
- print("max dw:", dw_max / 2 / np.pi)
-
- # setup basic variables
- twopi = 2 * np.pi
- vmax = KILO(20) * MEGA(1) # convert units, 20 kHz/us -> 20e3 * 1e6 Hz/s
- t_max = 2 * dw_max / vmax # Longest move sets the maximum moving time
- a_max = vmax * 2 / t_max # maximum acceleration, negative sign because of magic
- # get number of samples required for longest move,this sets the size of lookup table
- sample_len = int(np.ceil(t_max * wfm.sample_rate))
- # sample_len += (512 - sample_len % 512) # make overall length a multiple of 512 so AWG doesn't freak out
- sample_len += wfm.sample_len_min - sample_len % wfm.sample_len_min
- sample_len = int(sample_len)
- # now we calculate all possible trajectories, go to Group Notes/Projects/Rearrangement for detail
- path_table = {} # lookup table to store all moves
- static_sig = np.zeros(sample_len) # for fast real-time waveform generation purposes
- t = np.arange(sample_len) / wfm.sample_rate # time series
-
- # iterate!
- for i in range(nt):
- omega_i = wfm.omega[i]
- for j in moves[i]: # j is the target position, i is starting position
- omega_j = wfm.omega[j]
- if i == j:
- path = (
- wfm.amplitude[i] * np.sin(omega_i * t + wfm.phi[i])
- )
- static_sig += path
- # path = omega_i * t + wfm.phi[i]
- path_table[(i, i)] = path
- continue # skip diagonal entries
-
- path = np.zeros(sample_len)
-
- # I advise reading through the notes page first before going further
- dw = omega_j - omega_i # delta omega in the equation
- adw = abs(dw)
- t_tot = np.sqrt(abs(4 * dw / a_max)) # calculate minimum time to complete move
-
- phi_j = wfm.phi[j] % twopi # wrap around two pi
- phi_i = wfm.phi[i] % twopi
- dphi = (phi_j - phi_i) % twopi # delta phi in the equation
- if dphi < 0: dphi = abs(dphi) + twopi - phi_i # warp around for negative phase shift
- t_tot += 12 * np.pi / adw - (
- (t_tot - 6 * dphi / adw) %
- (12 * np.pi / adw)) # extend move time to arrive at the correct phase
- a = 4 * (omega_i - omega_j) / (t_tot ** 2) # adjust acceleration accordingly to ensure we still get to omega_j
-
- end = int(np.round(t_tot * wfm.sample_rate)) # convert to an index in samples
- # print(f'({i},{j}), {end}')
- half = int(end / 2) + 1 # index of sample half-way through the move where equation changes
- # t_tot = t[end]
- t1 = t[:half] # first half of the move, slicing to make life easier
- t2 = t[half:end] - t_tot / 2 # time series for second half of the move
- # a = 4 * adw / (t_tot ** 2) # adjust acceleration accordingly to ensure we still get to omega_j
-
- # interpolate amplitudes during the move
- amps = np.zeros(sample_len)
- inst_w = np.zeros(end)
- inst_w[0] = omega_i
- inst_w[-1] = omega_j
- inst_w[1:half] = omega_i - 0.5 * a * t1[1:] ** 2
- inst_w[half:end - 1] = omega_i - \
- a / 2 * (t_tot / 2) ** 2 - \
- a * t_tot / 2 * t2[:-1] + \
- a / 2 * t2[:-1] ** 2
- sw = omega_i
- bw = omega_j
- if omega_i > omega_j:
- sw = omega_j
- bw = omega_i
- inst_w[inst_w < sw] = sw
- inst_w[inst_w > bw] = bw
- amps[:end] = omega_interp(inst_w)
- amps[end:] = wfm.amplitude[j]
- # frequency/phase diagnostic
- # print(i,j)
- # print(inst_w[-2] - omega_j)
- # print(a*t_tot**3/24 % np.pi - dphi)
- # print(end)
- # print()
-
- # calculate sine wave
- path[:half] = wfm.phi[i] + 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
- path[half:end] = path[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
- path[end:] = path[end-1] + omega_j * (t[end:] - t[end-1])
- path = (amps * np.sin(path))
- path_table[(i, j)] = path
-
- for key in path_table:
- if key[0] != key[1]:
- path_table[key] -= path_table[(key[1], key[1])] # for fast real-time generation
- path_table[key] = path_table[key].astype(np.int16)
-
- static_sig = static_sig.astype(np.int16)
-
- # save stuff if prompted
- if save_path is not None:
- np.savez(save_path, table=path_table, static_sig=static_sig, wfm=wfm, target=target_idx)
-
- return path_table, static_sig
-
-
-def create_path_table_reduced_gpu(
+def create_path_table_gpu(
wfm: Waveform, t_idx, pre_paths, save_path=None,
) -> Tuple[Dict[Tuple[int, int], np.ndarray], np.ndarray]:
"""
@@ -366,7 +202,7 @@ def create_path_table_reduced_gpu(
# setup basic variables
twopi = 2 * np.pi
- vmax = KILO(40) * MEGA(1) # convert units, 20 kHz/us -> 20e3 * 1e6 Hz/s
+ vmax = KILO(60) * MEGA(1) # convert units, 20 kHz/us -> 20e3 * 1e6 Hz/s
t_max = 2 * dw_max / vmax # Longest move sets the maximum moving time
a_max = vmax * 2 / t_max # maximum acceleration, negative sign because of magic
# get number of samples required for longest move,this sets the size of lookup table
@@ -472,44 +308,38 @@ def create_path_table_reduced_gpu(
def get_rearrange_paths(
- filled_idx: np.ndarray,
- target_idx: np.ndarray,
-) -> List[Tuple[Any, Any]]:
+ f_idx: np.ndarray,
+ t_idx: np.ndarray,
+) -> np.ndarray:
"""
- Calculate rearranging paths.
- :param filled_idx: indices of tweezer positions filled with atoms.
- :param target_idx: indices of tweezer positions in target pattern.
- :returns: list containing tuples of moving paths
+ Finds the minimum weight perfect matching between f_idx and t_idx
+ :param f_idx: indices of tweezer positions filled with atoms.
+ :param t_idx: indices of tweezer positions in target pattern.
+ :returns: 2d numpy array containing moving path trajectories
"""
- t_size = target_idx.size
- f_size = filled_idx.size
- reserve = f_size - t_size
- paths = []
- i = 0
- j = 0
- while i < f_size:
- if j == t_size: break
- if filled_idx[i] == target_idx[j]:
- # paths.append((filled_idx[i], filled_idx[i]))
- j += 1
- 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] - target_idx[j])):
- i += 1
- reserve -= 1
+ if len(f_idx) < len(t_idx):
+ return np.array([])
+ l_ptr = np.searchsorted(f_idx, t_idx[0])
+ r_ptr = np.searchsorted(f_idx, t_idx[-1], side='right') - 1
+ n_unpaired = len(t_idx) - len(f_idx[l_ptr:r_ptr+1])
+ while n_unpaired > 0:
+ if l_ptr == 0:
+ r_ptr += n_unpaired
+ break
+ if r_ptr == len(f_idx) - 1:
+ l_ptr -= n_unpaired
+ break
+ l_dist = abs(t_idx[0] - f_idx[l_ptr - 1])
+ r_dist = abs(f_idx[r_ptr + 1] - t_idx[-1])
+ if l_dist < r_dist:
+ l_ptr -= 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
+ r_ptr += 1
+ n_unpaired -= 1
+ return np.vstack((f_idx[l_ptr:r_ptr+1], t_idx)).T
-def create_moving_array(path_table: np.ndarray, paths: np.ndarray) -> np.ndarray:
+def create_moving_array_old(path_table: np.ndarray, paths: np.ndarray) -> np.ndarray:
"""
create a rearranging signal that moves tweezers as specified by paths
:param path_table: lookup table returned from create_path_table().
@@ -518,66 +348,51 @@ def create_moving_array(path_table: np.ndarray, paths: np.ndarray) -> np.ndarray
return np.sum(path_table[paths[:, 0], paths[:, 1]], axis=0)
-def create_moving_array_reduced(
+def create_moving_array(
path_table: Dict,
sig: np.ndarray,
filled_idx: np.ndarray,
target_idx: np.ndarray,
- # paths: np.ndarray,
- # off: np.ndarray
):
"""
create a rearranging signal that moves tweezers as specified by paths.
- :param sig: initially a static-array-generating waveform.
+ :param sig: initially a static-array-generating waveform, this function
+ modifies sig directly
:param path_table: lookup table returned from create_path_table_reduced().
:param filled_idx: see get_rearrange_paths for detail.
:param target_idx: see get_rearrange_paths for detail.
- :param paths: 2d array with moving trajectories, [:,0] stores start pos, [:,1] stores end pos.
- :param off: 1d array with tweezer indices that need to be set to 0.
+
"""
- paths, off = get_rearrange_paths(filled_idx, target_idx)
- if len(off) != 0:
+ paths = get_rearrange_paths(filled_idx, target_idx)
+ if len(paths) == 0:
return
for i, j in paths:
if i == j:
- continue # Technically this is useless as get_rearrange_paths took care --
- # -- of this, but just in case
+ continue # skip stationary paths
if (i, j) in path_table:
sig += path_table[(i, j)]
- # else:
- # sig -= path_table[(j, j)] # (j,j) does not appear in off, must manually do this
- # for i in off:
- # sig -= path_table[(i, i)]
- pass
+ return
-def create_moving_array_reduced_GPUOTF(
+def create_moving_array_GPUOTF(
path_table: Dict,
sig: np.ndarray,
filled_idx: np.ndarray,
target_idx: np.ndarray,
- # paths: np.ndarray,
- # off: np.ndarray
):
"""
same function as above, with running gpu arrays on the fly
"""
- paths, off = get_rearrange_paths(filled_idx, target_idx)
- if len(off) != 0:
+ paths = get_rearrange_paths(filled_idx, target_idx)
+ if len(paths) == 0:
return
n_moves = len(paths)
- all_paths = cp.zeros((n_moves, sig.size))
for k in range(n_moves):
(i,j) = paths[k]
if i == j:
- continue # Technically this is useless as get_rearrange_paths took care --
- # -- of this, but just in case
+ continue
if (i, j) in path_table:
- all_paths[k] = cp.array(path_table[(i, j)], dtype=cp.int16)
- # for k in range(len(off)):
- # i = off[k]
- # all_paths[paths.shape[0] + k] = -cp.array(path_table[(i, i)], dtype=cp.int16)
- sig += cp.sum(all_paths, axis=0, dtype=cp.int16)
+ sig += cp.array(path_table[(i, j)], dtype=cp.int16)
return