194 lines
7.8 KiB
Python
194 lines
7.8 KiB
Python
|
import numpy as np
|
||
|
|
import os
|
||
|
|
import logging
|
||
|
|
import copy
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from typing import Any, Dict, List, Optional, Union, Callable
|
||
|
|
from enum import Enum
|
||
|
|
from env_gen.flattop import flattop
|
||
|
|
from env_gen.acz import aczwave
|
||
|
|
from env_gen.accz_gen import accz_wave
|
||
|
|
from matplotlib import gridspec
|
||
|
|
from FourChZreg_define import *
|
||
|
|
|
||
|
|
|
||
|
|
class EnvelopeGenerator:
|
||
|
|
_axes_list: List[tuple] = []
|
||
|
|
|
||
|
|
PLOT_STYLES = [
|
||
|
|
'ggplot', 'bmh', 'fivethirtyeight', 'Solarize_Light2',
|
||
|
|
'fast', 'tableau-colorblind10', 'seaborn-poster', 'seaborn-bright'
|
||
|
|
]
|
||
|
|
PLOT_COLORS = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7']
|
||
|
|
|
||
|
|
SINGLE_PLOT_CONFIG = {
|
||
|
|
'figsize': (8, 6),
|
||
|
|
'dpi': 220,
|
||
|
|
'style': 'seaborn-v0_8-poster',
|
||
|
|
'linewidth': 3,
|
||
|
|
'fontsize': {'xlabel': 22, 'ylabel': 22, 'title': 20, 'legend': 18, 'tick': 18}
|
||
|
|
}
|
||
|
|
|
||
|
|
MULTI_PLOT_CONFIG = {
|
||
|
|
'dpi': 200,
|
||
|
|
'fontsize': {'xlabel': 18, 'ylabel': 18, 'title': 16, 'tick': 16},
|
||
|
|
'grid_alpha': 0.3
|
||
|
|
}
|
||
|
|
|
||
|
|
def __init__(self, mk_instance, **kwargs):
|
||
|
|
self.mk = mk_instance
|
||
|
|
self.config_file = kwargs.get('config_file')
|
||
|
|
|
||
|
|
def write_register(self, address, value):
|
||
|
|
self.mk.rw_once('w', address, value, self.config_file)
|
||
|
|
|
||
|
|
def _generate_rect_envelope(self, **kwargs):
|
||
|
|
amp = kwargs.pop('amp')
|
||
|
|
wave_time = kwargs.pop('wave_time')
|
||
|
|
rect_wave = [amp] * wave_time
|
||
|
|
return rect_wave
|
||
|
|
|
||
|
|
def _generate_rect_hold_envelope(self, **kwargs):
|
||
|
|
amp = kwargs.pop('amp')
|
||
|
|
rect_rising_edge = [amp] * 4
|
||
|
|
rect_falling_edge = [0] * 4
|
||
|
|
return rect_rising_edge, rect_falling_edge
|
||
|
|
|
||
|
|
def _generate_flattop_envelope(self, **kwargs):
|
||
|
|
amp = kwargs.pop('amp')
|
||
|
|
edge_time = kwargs.pop('edge_time')
|
||
|
|
wave_time = kwargs.pop('wave_time')
|
||
|
|
flattop_wave = flattop(float(amp), float(edge_time), float(wave_time), 1.0)
|
||
|
|
return flattop_wave
|
||
|
|
|
||
|
|
def _generate_flattop_hold_envelope(self, **kwargs):
|
||
|
|
amp = kwargs.pop('amp')
|
||
|
|
edge_time = kwargs.pop('edge_time')
|
||
|
|
wave_time = kwargs.pop('wave_time')
|
||
|
|
flattop_wave = flattop(float(amp), float(edge_time), float(wave_time), 1.0)
|
||
|
|
hold_value = flattop_wave.max()
|
||
|
|
hold_idx = np.where(flattop_wave == hold_value)[0]
|
||
|
|
rising_edge_end_idx = hold_idx[0]
|
||
|
|
falling_edge_start_idx = hold_idx[-1]
|
||
|
|
flattop_rising_edge = flattop_wave[:rising_edge_end_idx + 1]
|
||
|
|
flattop_falling_edge = flattop_wave[falling_edge_start_idx:-2]
|
||
|
|
return flattop_rising_edge, flattop_falling_edge
|
||
|
|
|
||
|
|
def _generate_acz_envelope(self, **kwargs):
|
||
|
|
amp = kwargs.pop('amp')
|
||
|
|
wave_time = kwargs.pop('wave_time')
|
||
|
|
acz_wave = aczwave(amp, wave_time, 0.0, 0.0, 0.0, 0.864, 0.05, -0.18, 0.04)
|
||
|
|
acz_wave_real = [val.real for val in acz_wave]
|
||
|
|
return acz_wave_real
|
||
|
|
|
||
|
|
def _generate_accz_envelope(self, **kwargs):
|
||
|
|
wave_time = kwargs.pop('wave_time')
|
||
|
|
env_accz = accz_wave(T=wave_time, A=0.8, plot=False)
|
||
|
|
return env_accz
|
||
|
|
|
||
|
|
def _cosine_envelope(self, **kwargs):
|
||
|
|
alpha = kwargs.pop('alpha', 1)
|
||
|
|
phi = kwargs.pop('phi', np.pi)
|
||
|
|
amp = kwargs.pop('amp')
|
||
|
|
wave_time = kwargs.pop('wave_time')
|
||
|
|
t_norm = np.arange(wave_time) / wave_time
|
||
|
|
cosine_wave = amp * (1 + alpha * np.cos(2 * np.pi * t_norm + phi)) / 2
|
||
|
|
return cosine_wave
|
||
|
|
|
||
|
|
def _generate_envelope_data(self, envelope_type, **kwargs):
|
||
|
|
if envelope_type == 'rect':
|
||
|
|
env_data = self._generate_rect_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'rect_hold':
|
||
|
|
env_data = self._generate_rect_hold_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'flattop':
|
||
|
|
env_data = self._generate_flattop_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'flattop_hold':
|
||
|
|
env_data = self._generate_flattop_hold_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'acz':
|
||
|
|
env_data = self._generate_acz_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'accz':
|
||
|
|
env_data = self._generate_accz_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'cosine':
|
||
|
|
env_data = self._cosine_envelope(**kwargs)
|
||
|
|
elif envelope_type == 'file_read_direct':
|
||
|
|
env_data = kwargs.get('external_envelope_data', [])
|
||
|
|
elif envelope_type == 'file_read_txt':
|
||
|
|
file_path = kwargs.get('file_path')
|
||
|
|
if file_path is None:
|
||
|
|
raise ValueError("Missing 'file_path'")
|
||
|
|
|
||
|
|
txt_data = np.loadtxt(file_path)
|
||
|
|
env_data = np.asarray(txt_data, dtype=float).reshape(-1).tolist()
|
||
|
|
return env_data
|
||
|
|
|
||
|
|
def _next_env_idx(self, idx_num, env_idx_mem, envelope_length):
|
||
|
|
if idx_num == 0:
|
||
|
|
return envelope_length
|
||
|
|
else:
|
||
|
|
last_env_idx = env_idx_mem[-1]
|
||
|
|
env_base_addr = (last_env_idx >> 16) + ((last_env_idx & 0xFFFF) << 1)
|
||
|
|
return env_base_addr << 16 | envelope_length
|
||
|
|
|
||
|
|
def _env_data_pack(self, float_data_array):
|
||
|
|
data_int = np.round(float_data_array).astype(int)
|
||
|
|
data_int[data_int < 0] += 65536
|
||
|
|
hex_pairs = []
|
||
|
|
for data0, data1 in zip(data_int[::2], data_int[1::2]):
|
||
|
|
hex_pairs.append((data1 << 16) | data0)
|
||
|
|
|
||
|
|
return hex_pairs
|
||
|
|
|
||
|
|
def _generate_envelope_batch(self, **kwargs):
|
||
|
|
kwargs_copy = copy.deepcopy(kwargs)
|
||
|
|
channel = kwargs_copy.pop('channel')
|
||
|
|
envelope_configs = kwargs_copy.pop('envelope_configs', [])
|
||
|
|
env_data_mem = []
|
||
|
|
env_idx_mem = []
|
||
|
|
idx_num = 0
|
||
|
|
for envelope_config in envelope_configs:
|
||
|
|
envelope_type = envelope_config.pop('envelope_type')
|
||
|
|
env_data = self._generate_envelope_data(envelope_type, **envelope_config)
|
||
|
|
if isinstance(env_data, tuple):
|
||
|
|
retun_param_count = len(env_data)
|
||
|
|
else:
|
||
|
|
retun_param_count = 1
|
||
|
|
if retun_param_count == 1:
|
||
|
|
envelope = env_data
|
||
|
|
envelope_arr = np.asarray(envelope, dtype=float).reshape(-1)
|
||
|
|
if envelope_arr.size % 4 != 0:
|
||
|
|
raise ValueError("Envelope length must be multiple of 4")
|
||
|
|
env_data_mem += envelope_arr.astype(int).tolist()
|
||
|
|
envelope_length = int(envelope_arr.size)
|
||
|
|
current_env_idx = self._next_env_idx(idx_num, env_idx_mem, envelope_length)
|
||
|
|
env_idx_mem.append(current_env_idx)
|
||
|
|
idx_num += 1
|
||
|
|
elif retun_param_count == 2:
|
||
|
|
rising_edge, falling_edge = env_data
|
||
|
|
rising_edge_arr = np.asarray(rising_edge, dtype=float).reshape(-1)
|
||
|
|
falling_edge_arr = np.asarray(falling_edge, dtype=float).reshape(-1)
|
||
|
|
if rising_edge_arr.size % 4 != 0:
|
||
|
|
raise ValueError("Envelope length must be multiple of 4")
|
||
|
|
env_data_mem += rising_edge_arr.astype(int).tolist()
|
||
|
|
rising_edge_length = int(rising_edge_arr.size)
|
||
|
|
rising_edge_idx = self._next_env_idx(idx_num, env_idx_mem, rising_edge_length)
|
||
|
|
env_idx_mem.append(rising_edge_idx)
|
||
|
|
idx_num += 1
|
||
|
|
env_data_mem += falling_edge_arr.astype(int).tolist()
|
||
|
|
falling_edge_length = int(falling_edge_arr.size)
|
||
|
|
falling_edge_idx = self._next_env_idx(idx_num, env_idx_mem, falling_edge_length)
|
||
|
|
env_idx_mem.append(falling_edge_idx)
|
||
|
|
idx_num += 1
|
||
|
|
|
||
|
|
env2mem_format = self._env_data_pack(env_data_mem)
|
||
|
|
|
||
|
|
ENVI_BASE = addr_base['ENVI0_BASE'] + channel*0x600000
|
||
|
|
ENVM_BASE = addr_base['ENVM0_BASE'] + channel*0x600000
|
||
|
|
self.write_register(ENVI_BASE, env_idx_mem)
|
||
|
|
self.write_register(ENVM_BASE, env2mem_format)
|
||
|
|
|
||
|
|
|
||
|
|
def env_config(mk_instance, **kwargs):
|
||
|
|
env_gen = EnvelopeGenerator(mk_instance, **kwargs)
|
||
|
|
env_gen._generate_envelope_batch(**kwargs)
|