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

Bug fixings

parent f9f1502b
No related branches found
No related tags found
No related merge requests found
import numpy as np
from waveform import *
import os
cf = MEGA(105) # center frequency
df = MEGA(2.5) # delta frequency
n = 2 # number of tones to generate in either direction of cf, total number of tweezer is 2*n+1
df = MEGA(0.5) # delta frequency
n = 15 # number of tones to generate in either direction of cf, total number of tweezer is 2*n+1
nt = 2*n+1 # total number of tweezers
# sampling_rate and sample_len must follow some relationship explained in the wiki page.
sampling_rate = int(614.4e6) # sampling rate
......@@ -19,17 +21,32 @@ amps = ampMax * np.ones(nt)
# amps *= np.sqrt(corr)
# phase adjustments
phase = np.load("data/optm_phase.npz")['phase'][:nt] # phase table from Caltech
# phase = np.load("data/optm_phase.npz")['phase'][:nt] # phase table from Caltech
twzr.amplitude = amps
twzr.phi = phase
# twzr.phi = phase
# sig_mat = create_static_array(twzr, True)
# empty = np.zeros(sample_len)
t_idx = np.array([0,1,2,3,4])
f_idx = np.array([1,2,4])
table, sig = create_path_table_reduced(twzr, t_idx)
np.savez('data/rtable_5.npz', table=table, static_sig=sig, wfm=twzr, target=t_idx)
# table generations
tgt = np.zeros(nt)
tgt[:n] = 1
t_idx = np.nonzero(tgt)[0]
# savepath = 'table/table-half_31_120722'
# create_path_table_reduced_gpu(twzr, t_idx, max_dist=n, save_path=savepath)
# moving waveform testings
data = np.load("table/table-half_31_120722.npz", allow_pickle=True)
table = data['table'].item()
# twzr = data['wfm'].item()
static_sig = data['static_sig']
# target = data['target']
# t_idx = np.nonzero(target)[0]
f_idx = np.array([5,8,20])
create_moving_array_reduced(table, static_sig, f_idx, t_idx)
np.savez('data/test.npz', signal=static_sig, wfm=twzr)
# data = np.load('data/table_5.npz', allow_pickle=True)
# table = data['table']
# sig = data['static_sig']
......
......@@ -4,7 +4,6 @@ import numpy as np
from scipy.interpolate import interp1d
from AWG import *
class Waveform:
def __init__(self, cf: int, df: int, n: int, sample_rate: int, freq_res):
"""
......@@ -167,7 +166,6 @@ def create_path_table_reduced(
moves[i].append(j)
dw = abs(wfm.omega[j] - wfm.omega[i])
if dw_max < dw: dw_max = dw
# setup basic variables
twopi = 2 * np.pi
vmax = KILO(20) * MEGA(1) # convert units, 20 kHz/us -> 20e3 * 1e6 Hz/s
......@@ -191,10 +189,10 @@ def create_path_table_reduced(
if i == j:
path = (
wfm.amplitude[i] * np.sin(omega_i * t + wfm.phi[i])
).astype(np.int16)
)
static_sig += path
# path = omega_i * t + wfm.phi[i]
path_table[(i, i)] = path
static_sig += path
continue # skip diagonal entries
path = np.zeros(sample_len)
......@@ -255,19 +253,178 @@ def create_path_table_reduced(
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)).astype(np.int16)
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(
wfm: Waveform, target_idx, max_dist=np.inf, 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 save_path: file saving path
:param target_idx: indices of target pattern
:param max_dist: maximum move distance in indices
:param wfm: waveform object already initialized with basic parameters.
:return: dictionary containing rearrange paths
"""
import cupy as cp
# interpolate optimal amplitudes
# data = np.load("data/optimal_amps.npz")
w = wfm.omega
a = wfm.amplitude
omega_interp = interp1d(w, a, kind='cubic')
# obtain all move combinations:
n = len(wfm.omega) # total number of tweezers
moves = []
dw_max = 0 # longest move, this sets the size of path_table
for i in range(n):
moves.append([])
for j in target_idx:
if i < j and True: # only allow uni-direction moves
continue
if abs(i - j) <= max_dist:
moves[i].append(j)
dw = abs(wfm.omega[j] - wfm.omega[i])
if dw_max < dw: dw_max = dw
# 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 = cp.zeros(sample_len) # for fast real-time waveform generation purposes
t = cp.arange(sample_len) / wfm.sample_rate # time series
for i in range(n):
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] * cp.sin(omega_i * t + wfm.phi[i])
)
static_sig += path
# path = cp.asnumpy(path)
# path = omega_i * t + wfm.phi[i]
path_table[(i, i)] = path
continue # skip diagonal entries
# iterate!
for i in range(n):
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] * cp.sin(omega_i * t + wfm.phi[i])
# )
# static_sig += path
# path = cp.asnumpy(path)
# # path = omega_i * t + wfm.phi[i]
# path_table[(i, i)] = path
continue # skip diagonal entries
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
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 = 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 -= path_table[j,j]
path = cp.asnumpy(path).astype(np.int16)
path_table[(i, j)] = path
print(i)
for key in path_table:
# if key[0] != key[1]:
# path_table[key] -= path_table[(key[1], key[1])] # for fast real-time generation
if key[0] == key[1]:
path_table[key] = cp.asnumpy(path_table[key]).astype(np.int16)
# path_table[key] = path_table[key].astype(np.int16)
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, target=target_idx)
return path_table, static_sig.astype(np.int16)
return path_table, static_sig
def get_rearrange_paths(
......@@ -278,7 +435,7 @@ def get_rearrange_paths(
Calculate rearranging paths.
:param filled_idx: indices of tweezer positions filled with atoms.
:param target_idx: indices of tweezer positions in target pattern.
:returns: 2d array containing rearranging paths. 1d array containing tweezer positions to be turned off.
:returns: Tuple containing moving paths and turn-off paths.
"""
n = 0
t_size = target_idx.size
......@@ -340,11 +497,12 @@ def create_moving_array_reduced(
paths, off = get_rearrange_paths(filled_idx, target_idx)
for i, j in paths:
if i == j:
continue
continue # Technically this is useless as get_rearrange_paths took care --
# -- of this, but just in case
if (i, j) in path_table:
sig += path_table[(i, j)]
# else:
# print((i, j), "not in path table")
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
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment