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

bug fixings

parent 16582067
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 56 additions and 19 deletions
Python/waveform.py
+ 56
− 19
View file @ 6a62f2e9
from typing import Dict, Tuple, Any
from typing import Dict, Tuple, Any, List
import numpy as np
from scipy.interpolate import interp1d
from AWG import *
import cupy as cp
class Waveform:
def __init__(self, cf: int, df: int, n: int, sample_rate: int, freq_res):
......@@ -15,12 +17,12 @@ class Waveform:
"""
# define some useful numbers
scale = 2 ** 12 # sets the amplitude scale, max must not exceed 2**15-1
num_tz = 2 * n + 1 # total number of tweezers to be generated
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 * df, num_tz) # frequency tones
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
......@@ -338,7 +340,7 @@ def create_path_table_reduced(
def create_path_table_reduced_gpu(
wfm: Waveform, pre_paths, save_path=None,
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
......@@ -349,7 +351,6 @@ def create_path_table_reduced_gpu(
: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")
......@@ -365,7 +366,7 @@ def create_path_table_reduced_gpu(
# setup basic variables
twopi = 2 * np.pi
vmax = KILO(20) * MEGA(1) # convert units, 20 kHz/us -> 20e3 * 1e6 Hz/s
vmax = KILO(40) * 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
......@@ -380,16 +381,21 @@ def create_path_table_reduced_gpu(
nt = len(wfm.omega)
diagonal_mat = np.sin(
np.outer(wfm.omega, t) + np.expand_dims(wfm.phi, axis=1)
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 = (wfm.amplitude * diagonal_mat.T).T # this works, trust me
diagonal_mat = cp.array(diagonal_mat)
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, axis=0)
# iterate!
time_counter = 0
for i, j in pre_paths:
time_counter += 1
if time_counter % 100 == 0:
print(time_counter)
omega_i = wfm.omega[i]
omega_j = wfm.omega[j]
path = cp.zeros(sample_len)
......@@ -452,16 +458,15 @@ def create_path_table_reduced_gpu(
path[end:] = path[end-1] + omega_j * (t[end:] - t[end-1])
path = (amps * cp.sin(path))
if i != j:
path -= diagonal_mat[i]
path -= diagonal_mat[j]
path = cp.asnumpy(path).astype(np.int16)
path_table[(i, j)] = path
print(i)
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)
np.savez(save_path, table=path_table, static_sig=static_sig, wfm=wfm, t_idx=t_idx)
return path_table, static_sig
......@@ -469,7 +474,7 @@ def create_path_table_reduced_gpu(
def get_rearrange_paths(
filled_idx: np.ndarray,
target_idx: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray]:
) -> Tuple[List[Tuple[Any, Any]], List[Any]]:
"""
Calculate rearranging paths.
:param filled_idx: indices of tweezer positions filled with atoms.
......@@ -504,7 +509,7 @@ def get_rearrange_paths(
if reserve < 0:
for i in range(abs(reserve)):
off.append(target_idx[-1 - i])
return np.array(paths), np.array(off)
return paths, off
def create_moving_array(path_table: np.ndarray, paths: np.ndarray) -> np.ndarray:
......@@ -534,19 +539,51 @@ def create_moving_array_reduced(
: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:
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
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)]
# 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
def create_moving_array_reduced_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:
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
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)
return
def create_moving_signal_single(omega_i, omega_f, sample_rate, signal_time):
min_len = 2 * sample_rate / (10e3)
sample_len = sample_rate * signal_time
......
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