from typing import Dict, Tuple, Any, List
import numpy as np
from scipy.interpolate import interp1d
from AWG import *
# import cupy as cp
import scipy.optimize as spopt
class Waveform:
def __init__(self, cf: int, df: int, n: int, sample_rate: int, freq_res):
"""
helper class to store basic waveform information.
:param cf: center frequency tone of tweezer array.
:param df: differential frequency between neighboring tweezers.
:param n: number of tweezers to the left/right of center frequency tone, total number of tweezer is 2n+1.
:param sample_rate: sampling rate of the AWG to generate correct number of samples.
"""
# define some useful numbers
scale = 2 ** 12 # sets the amplitude scale, max must not exceed 2**15-1
num_tz = 2 * n + 2 # total number of tweezers to be generated
max_amp = scale / np.sqrt(num_tz) # scale down by number of tweezers
self.amplitude = max_amp * np.ones(num_tz)
# self.amplitude: np.ndarray = max_amp # uniform amplitude
self.omega = 2 * np.pi * np.linspace(cf - n * df, cf + (n+1) * df, num_tz) # frequency tones
self.phi = 2 * np.pi * np.random.rand(num_tz) # random initial phases from 0-2pi
self.sample_rate: int = sample_rate
self.freq_res = freq_res
# formula for minimum sample length
# sample_len_min = 512 * m
# m * 512 * freq_resolution / sampling_rate = k, k % 2 == 0
self.sample_len_min = 2 * sample_rate / freq_res / 100 # !!! this only works for sample_rate = 614.4e6 !!!
def create_static_array(wfm: Waveform, full=False) -> np.ndarray:
"""
create a static-array-generating waveform with user set number of samples
:param wfm: waveform object already initialized with basic parameters.
:return: returns a 1D array with static-array-generating waveform.
"""
# construct time axis, t_total(s) = sample_len / sample_rate, dt = t_total / sample_len
t = np.arange(wfm.sample_len_min) / wfm.sample_rate
# calculate individual sin waves, sig_mat[i] corresponds to data for ith tweezer
# sin_mat = wfm.amplitude * np.sin(np.outer(wfm.omega,t) + np.expand_dims(wfm.phi, axis=1)) # shape=(number of tweezers x sample_len)
sin_mat = np.sin(
np.outer(wfm.omega, t) + np.expand_dims(wfm.phi, axis=1)
# shape=(number of tweezers x sample_len)
)
sin_mat = (wfm.amplitude * sin_mat.T).T # this works, trust me
# sum up all rows to get final signal
sig = np.sum(sin_mat, axis=0)
if np.max(sig) >= 2 ** 15 - 1:
print("Signal amp exceeds maximum")
if full:
return sin_mat.astype(np.int16)
return sig.astype(np.int16)
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.
:return: dim-3 ndarray
"""
# interpolate optimal amplitudes
# data = np.load("data/optimal_amps.npz")
w = wfm.omega
a = wfm.amplitude
omega_interp = interp1d(w, a, kind='cubic')
# setup basic variables
twopi = 2 * np.pi
vmax = KILO(20) * MEGA(1) # convert units, 20 kHz/us -> 20e3 * 1e6 Hz/s
dw_max = wfm.omega[-1] - wfm.omega[0] # Longest move in frequency
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
# now we calculate all possible trajectories, go to Group Notes/Projects/Rearrangement for detail
n = len(wfm.omega) # total number of tweezers
path_table = np.zeros((n, n, sample_len)) # lookup table to store all moves
t = np.arange(sample_len) / wfm.sample_rate # time series
# iterate! I think this part can be vectorized as well... but unnecessary.
for i, omega_i in enumerate(wfm.omega):
for j, omega_j in enumerate(wfm.omega): # j is the target position, i is starting position
if i == j:
path_table[i, i] = wfm.amplitude[i] * np.sin(omega_i * t + wfm.phi[i])
continue # skip diagonal entries
dw = omega_j - omega_i # delta omega in the equation
adw = abs(dw)
# I advise reading through the notes page first before going further
phi_j = wfm.phi[j] % twopi # wrap around two pi
phi_i = wfm.phi[i] % twopi
dphi = phi_j - phi_i # delta phi in the equation
if dphi < 0: dphi = abs(dphi) + twopi - phi_i # warp around for negative phase shift
t_tot = np.sqrt(abs(4 * dw / a_max)) # calculate minimum time to complete move
t_tot = ((t_tot - 6 * dphi / adw) // (
12 * np.pi / adw) + 1) * 12 * np.pi / adw # extend move time to arrive at the correct phase
a = 4 * dw / (t_tot ** 2) # adjust acceleration accordingly to ensure we still get to omega_j
end = int(np.ceil(t_tot * wfm.sample_rate)) # convert to an index in samples
half = int(end / 2) # index of sample half-way through the move where equation changes
# if end % 2 == 0: half += 1
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
# 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
amps[:end] = omega_interp(inst_w)
amps[end:] = wfm.amplitude[j]
# calculate sine wave
path_table[i, j, :half] = wfm.phi[i] + omega_i * t1 + a / 6 * t1 ** 3 # t<=T/2
path_table[i, j, half:end] = path_table[i, j, 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_table[i, j, end:] = omega_j * t[end:] + (phi_j - omega_j * t_tot) % twopi
path_table[i, j] = amps * np.sin(path_table[i, j])
# now compile everything into sine wave
path_table = path_table.astype(np.int16)
return path_table
# return path_table.astype(int), np.sum(path_table.diagonal().T, axis=0, dtype=int)
def stack_left(i_start, i_end, offset, stack_size=0):
# calculate first index where the reduced path algorithm is applied
# threshold = 0.01
# cutoff = np.ceil(np.log(threshold) / np.log(1-load_p))
# cutoff = int(cutoff)
# print(cutoff)
if stack_size == 0:
stack_size = np.floor((i_end - i_start) / 2)
stack_last = int(stack_size + i_start) - 1
dist_mod = (i_end - i_start - stack_size) / (i_end - i_start) # max_distance ratio
dist_add = offset
# get a list of moves to pre-generate
moves = []
max_dist = 0
for i in range(i_start, i_end):
moves.append([])
j_max = i if i < stack_last else stack_last
dist = np.ceil((i - i_start) * dist_mod + dist_add)
j_min = int(i - dist) if i - dist >= i_start else i_start
for j in range(j_min, j_max + 1):
moves[i - i_start].append(j) # add all paths between j_min and j_max
if max_dist < abs(j-i):
max_dist = abs(j-i)
return moves, max_dist
def stack_right(i_start, i_end, offset, stack_size=0):
moves, max_dist = stack_left(i_start, i_end, offset=offset, stack_size=stack_size)
moves.reverse()
for i in range(len(moves)):
moves[i].reverse()
for j in range(len(moves[i])):
moves[i][j] = i_end - 1 - moves[i][j] + i_start
return moves, max_dist
def create_path_table_gpu(
wfm: Waveform, t_idx, pre_paths, save_path=None,
) -> 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 pre_paths: list of pre-defined paths to generate signals for
:param save_path: file saving path
:param t_idx: indices of target pattern
: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')
dw_max = 0
for i,j in pre_paths:
dw = abs(wfm.omega[j] - wfm.omega[i])
if dw_max < dw: dw_max = dw
# setup basic variables
twopi = 2 * np.pi
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
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 = cp.zeros(sample_len) # for fast real-time waveform generation purposes
t = cp.arange(sample_len) / wfm.sample_rate # time series
dt = t[1]
t += dt
nt = len(wfm.omega)
diagonal_mat = cp.sin(
cp.outer(cp.array(wfm.omega), t) + cp.expand_dims(cp.array(wfm.phi), axis=1)
# shape=(number of tweezers x sample_len)
)
diagonal_mat = (cp.array(wfm.amplitude) * diagonal_mat.T).T # this works, trust me
# diagonal_mat = cp.array(diagonal_mat)
static_sig = cp.sum(diagonal_mat[t_idx], axis=0)
# iterate!
time_counter = 0
for i, j in pre_paths:
time_counter += 1
if time_counter % 100 == 0:
print(time_counter)
if i == j:
continue
omega_i = wfm.omega[i]
omega_j = wfm.omega[j]
path = cp.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
end = int(np.round(t_tot * wfm.sample_rate)) # convert to an index in samples
t_tot = end / wfm.sample_rate + t[1]
'''
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
half = int(end / 2) + 1 # index of sample half-way through the move where equation changes
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
# interpolate amplitudes during the move
amps = cp.zeros(sample_len)
inst_w = cp.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] = cp.array(omega_interp(inst_w.get()))
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 * cp.sin(path)
if i != j:
path -= diagonal_mat[j]
path = cp.asnumpy(path).astype(np.int16)
path_table[(i, j)] = path
static_sig = cp.asnumpy(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, t_idx=t_idx)
return path_table, static_sig
def get_rearrange_paths(
f_idx: np.ndarray,
t_idx: np.ndarray,
) -> np.ndarray:
"""
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
"""
if len(f_idx) < len(t_idx):
return np.array([])
cm = abs(np.subtract.outer(f_idx, t_idx))
row, col = spopt.linear_sum_assignment(cm)
return np.stack([f_idx[row], t_idx]).T
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().
:param paths: 2d array with moving trajectories, [:,0] stores start pos, [:,1] stores end pos.
"""
return np.sum(path_table[paths[:, 0], paths[:, 1]], axis=0)
def create_moving_array(
path_table: Dict,
sig: np.ndarray,
filled_idx: np.ndarray,
target_idx: np.ndarray,
):
"""
create a rearranging signal that moves tweezers as specified by paths.
: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.
"""
paths = get_rearrange_paths(filled_idx, target_idx)
if len(paths) == 0:
return
for i, j in paths:
if i == j:
continue # skip stationary paths
if (i, j) in path_table:
sig += path_table[(i, j)]
return
def create_moving_array_GPUOTF(
path_table: Dict,
sig: np.ndarray,
filled_idx: np.ndarray,
target_idx: np.ndarray,
):
"""
same function as above, with running gpu arrays on the fly
"""
paths = get_rearrange_paths(filled_idx, target_idx)
if len(paths) == 0:
return
n_moves = len(paths)
for k in range(n_moves):
(i,j) = paths[k]
if i == j:
continue
if (i, j) in path_table:
sig += cp.array(path_table[(i, j)], dtype=cp.int16)
return
def create_moving_signal_single(
omega_i, omega_f, sample_rate, signal_time, amp=2**12, phi=0
):
min_len = 2 * sample_rate / (1e3)
sample_len = sample_rate * signal_time
sample_len += min_len - sample_len % min_len
sample_len = int(sample_len)
t = np.arange(sample_len) / sample_rate
t += t[1]
t_tot = sample_len / sample_rate + t[1]
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
signal = np.zeros(sample_len)
signal[:half] = phi + 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])
phi_end = signal[-1]
signal = amp * np.sin(signal)
return signal.astype(np.int16), phi_end
def create_static_signal_single(
omega, sample_rate, signal_time, amp=2 ** 12, phi=0
):
min_len = 2 * sample_rate / (1e3)
sample_len = sample_rate * signal_time
sample_len += min_len - sample_len % min_len
sample_len = int(sample_len)
t = np.arange(sample_len) / sample_rate
t += t[1]
signal = phi + omega * t
phi_end = signal[-1]
signal = amp * np.sin(signal)
return signal.astype(np.int16), phi_end
def create_move_then_back(omega_i, omega_f, sample_rate, move_time, stay_time):
move0, phi0 = create_moving_signal_single(
omega_i, omega_f, sample_rate, move_time
)
stay, phi1 = create_static_signal_single(
omega_f, sample_rate, stay_time, phi=phi0
)
move1, phi2 = create_moving_signal_single(
omega_f, omega_i, sample_rate, move_time, phi=phi1
)
signal = np.concatenate((move0, stay, move1))
return signal