Skip to content
Snippets Groups Projects
Commit 2495e5d3 authored by xiyehu2's avatar xiyehu2
Browse files

Minimum weight matching path finding.

parent 06feb429
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 43 additions and 228 deletions
Python/waveform.py
+ 43
− 228
View file @ 2495e5d3
......@@ -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
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment