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)