waveform.py

from typing import Dict, Tuple, Any

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):
        """
        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 + 1  # 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.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  # !!! 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(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 create_path_table_reduced(
        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
    """
    # 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 = 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(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] * np.sin(omega_i * t + wfm.phi[i])
                ).astype(np.int16)
                # path = omega_i * t + wfm.phi[i]
                path_table[(i, i)] = path
                static_sig += 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)).astype(np.int16)
            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)

    # 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)


def get_rearrange_paths(
        filled_idx: np.ndarray,
        target_idx: np.ndarray,
) -> Tuple[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: 2d array containing rearranging paths. 1d array containing tweezer positions to be turned off.
    """
    n = 0
    t_size = target_idx.size
    f_size = filled_idx.size
    reserve = f_size - t_size
    # move_paths = []
    # static_paths = []
    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 = j
        elif (reserve > 0
              and filled_idx[i] < target_idx[j]
              and abs(filled_idx[i + 1] - target_idx[j]) < abs(filled_idx[i + 1] - target_idx[j])):
            i += 1
            reserve -= 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 np.array(paths), np.array(off)


def create_moving_array(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_reduced(
        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 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)
    for i, j in paths:
        if i == j:
            continue
        if (i, j) in path_table:
            sig += path_table[(i, j)]
        # else:
        #     print((i, j), "not in path table")
    for i in off:
        sig -= path_table[(i, i)]
    pass


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
    sample_len += min_len - sample_len % min_len
    sample_len = int(sample_len)

    t = np.arange(sample_len) / sample_rate
    t_tot = sample_len / sample_rate
    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

    amps = 2**12
    signal = np.zeros(sample_len)

    signal[:half] = 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])
    signal = (amps * np.sin(signal)).astype(np.int16)
    return signal