From 6a62f2e9d8258805f7745c8a83ffdeb69c324cfb Mon Sep 17 00:00:00 2001
From: xiyehu2 <35537592+xiyehu2@users.noreply.github.com>
Date: Thu, 15 Dec 2022 11:53:32 -0600
Subject: [PATCH] bug fixings
---
Python/waveform.py | 75 ++++++++++++++++++++++++++++++++++------------
1 file changed, 56 insertions(+), 19 deletions(-)
diff --git a/Python/waveform.py b/Python/waveform.py
index e936327..7e24798 100644
--- a/Python/waveform.py
+++ b/Python/waveform.py
@@ -1,8 +1,10 @@
-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
--
GitLab