From d19d6637ef408d797f56b662258f2fde88f4a2b3 Mon Sep 17 00:00:00 2001 From: Mahmood Sabooni Date: Mon, 1 Jun 2026 12:11:31 -0400 Subject: [PATCH 1/2] Add BSI.py for BSI integration --- src/oqd_trical/BSI_004.py | 422 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 422 insertions(+) create mode 100644 src/oqd_trical/BSI_004.py diff --git a/src/oqd_trical/BSI_004.py b/src/oqd_trical/BSI_004.py new file mode 100644 index 0000000..a9680e4 --- /dev/null +++ b/src/oqd_trical/BSI_004.py @@ -0,0 +1,422 @@ +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.animation as animation +from matplotlib.widgets import RadioButtons, TextBox +import scipy.constants as scipy_cst +from scipy.optimize import minimize +from matplotlib.ticker import MultipleLocator +import math + +# Import the exact solver from your notebook +try: + import oqd_trical +except ImportError: + oqd_trical = None + +# ========================================== +# 1. Initial Parameters & State Caching +# ========================================== +init_N = "20" +init_mode = "1" +init_wx = "3.00" # MHz (Radial X frequency) +init_wy = "3.00" # MHz (Radial Y frequency) +init_wz = "0.200" # MHz (Axial Z frequency) + +# Physical Constants for Lamb-Dicke Calculation +MASS_ION = 171 * scipy_cst.m_u # Ytterbium-171 mass +LAMBDA_LASER = 355e-9 # 355 nm Raman beams +DELTA_K = 4 * np.pi / LAMBDA_LASER + +trical_cache = { + 'N': None, 'wx': None, 'wy': None, 'wz': None, + 'freqs_x': None, 'freqs_y': None, 'freqs_z': None, + 'evecs_x': None, 'evecs_y': None, 'evecs_z': None, + 'x_eq': None, 'raw_z_eq': None, + 'trical_stable': True, 'steane_stable': True, 'w_crit': 0.0, + 'table_drawn_for': None +} + +# ========================================== +# 2. Physics Engine (OQD_TRICAL + Bare-Metal Fallback) +# ========================================== +def solve_trical(N, wx_mhz, wy_mhz, wz_mhz): + if (trical_cache['N'] == N and + trical_cache['wx'] == wx_mhz and + trical_cache['wy'] == wy_mhz and + trical_cache['wz'] == wz_mhz): + return + + w_rad_min = min(wx_mhz, wy_mhz) + beta_c = 0.73 * (N ** 0.86) + omega_r_critical = beta_c * wz_mhz + steane_stable = w_rad_min > omega_r_critical + + trical_cache['steane_stable'] = steane_stable + trical_cache['w_crit'] = omega_r_critical + + try: + if oqd_trical is None: raise Exception("oqd_trical not installed") + + omega_x = 2 * np.pi * wx_mhz * 1e6 + omega_y = 2 * np.pi * wy_mhz * 1e6 + omega_z = 2 * np.pi * wz_mhz * 1e6 + + alpha = np.zeros((3, 3, 3)) + alpha[2, 0, 0] = MASS_ION * (omega_x) ** 2 / 2 + alpha[0, 2, 0] = MASS_ION * (omega_y) ** 2 / 2 + alpha[0, 0, 2] = MASS_ION * (omega_z) ** 2 / 2 + + pp = oqd_trical.mechanical.PolynomialPotential(alpha, N=N) + ti = oqd_trical.mechanical.TrappedIons(N, pp, m=MASS_ION) + + ti.equilibrium_position() + ti.normal_modes() + + if not hasattr(ti, 'mode_vectors') and not hasattr(ti, 'eigenvectors'): + raise Exception("Trical missing eigenvector attributes.") + + except Exception as e: + trical_cache['trical_stable'] = False + + # --- CUSTOM BARE-METAL EXACT NUMERICAL SOLVER --- + def coulomb_potential(u): + axial_confinement = 0.5 * np.sum(u**2) + coulomb_repulsion = 0.0 + for i in range(N): + for j in range(i + 1, N): + coulomb_repulsion += 1.0 / np.abs(u[i] - u[j]) + return axial_confinement + coulomb_repulsion + + u_guess = np.linspace(-N/2, N/2, N) + res = minimize(coulomb_potential, u_guess, method='BFGS') + u_eq = np.sort(res.x) + + # Build Hessian Matrices (CORRECTED SIGNS) + H_z, H_x, H_y = np.zeros((N, N)), np.zeros((N, N)), np.zeros((N, N)) + alpha_x = (wx_mhz / wz_mhz)**2 + alpha_y = (wy_mhz / wz_mhz)**2 + + for i in range(N): + for j in range(N): + if i != j: + inv_dist_cubed = 1.0 / np.abs(u_eq[i] - u_eq[j])**3 + # Axial off-diagonals are negative (-2/d^3) + H_z[i, j] = -2.0 * inv_dist_cubed + # Radial off-diagonals are positive (+1/d^3) + H_x[i, j] = 1.0 * inv_dist_cubed + H_y[i, j] = 1.0 * inv_dist_cubed + + for i in range(N): + # Diagonal elements derived from the COM eigenvalue conditions + H_z[i, i] = 1.0 - np.sum(H_z[i, :]) + H_z[i, i] + H_x[i, i] = alpha_x - np.sum(H_x[i, :]) + H_x[i, i] + H_y[i, i] = alpha_y - np.sum(H_y[i, :]) + H_y[i, i] + + def solve_modes(H, omega_base): + evals, evecs = np.linalg.eigh(H) + freqs = np.sqrt(np.maximum(0, evals)) * omega_base + return freqs, evecs + + f_z, e_z = solve_modes(H_z, wz_mhz) + f_x, e_x = solve_modes(H_x, wz_mhz) + f_y, e_y = solve_modes(H_y, wz_mhz) + + idx_z = np.argsort(f_z) + trical_cache['freqs_z'] = f_z[idx_z] + trical_cache['evecs_z'] = e_z[:, idx_z] + + idx_x = np.argsort(f_x)[::-1] + trical_cache['freqs_x'] = f_x[idx_x] + trical_cache['evecs_x'] = e_x[:, idx_x] + + idx_y = np.argsort(f_y)[::-1] + trical_cache['freqs_y'] = f_y[idx_y] + trical_cache['evecs_y'] = e_y[:, idx_y] + + omega_z_si = 2 * np.pi * wz_mhz * 1e6 + l_0 = ((scipy_cst.e**2 / (4 * np.pi * scipy_cst.epsilon_0)) / (MASS_ION * omega_z_si**2))**(1/3) + + trical_cache['raw_z_eq'] = u_eq * l_0 + trical_cache['x_eq'] = u_eq + + trical_cache['N'] = N + trical_cache['wx'] = wx_mhz + trical_cache['wy'] = wy_mhz + trical_cache['wz'] = wz_mhz + +# ========================================== +# 3. Figure and UI Setup (OPTIMIZED LAYOUT) +# ========================================== +fig, (ax_x, ax_y, ax_z) = plt.subplots(3, 1, figsize=(16, 9)) + +# Gives maximum vertical height to the three spectrum plots on the left +plt.subplots_adjust(top=0.94, bottom=0.22, left=0.06, right=0.52, hspace=0.45) + +# --- RIGHT SIDE PLOTS --- +# Main modes table (Top Right) +ax_table = plt.axes([0.56, 0.52, 0.42, 0.42]) +ax_table.axis('off') + +# Eigenvector Plot (Bottom Right-Middle) +ax_evec = plt.axes([0.56, 0.05, 0.24, 0.40]) + +# Lamb-Dicke Table (Bottom Right Edge - Maximized for Height) +ax_eta = plt.axes([0.82, 0.05, 0.16, 0.40]) +ax_eta.axis('off') + +# --- Spectrum Plots Setup --- +axes = [ax_x, ax_y, ax_z] +labels = ['Radial X', 'Radial Y', 'Axial Z'] +colors = ['tab:red', 'tab:blue', 'tab:green'] + +lines = {} +highlights = {} + +for ax, label, color in zip(axes, labels, colors): + ax.grid(axis='x', linestyle='--', alpha=0.6) + ax.set_yticks([]) + ax.set_ylim(-1, 1) + ax.set_ylabel(label, fontsize=11, fontweight='bold') + ax.xaxis.set_minor_locator(MultipleLocator(0.1)) + ax.grid(axis='x', which='minor', linestyle=':', alpha=0.4) + + lines[label] = ax.plot([], [], '|', color=color, markersize=30, markeredgewidth=1.5)[0] + highlights[label] = ax.plot([], [], 'o', color='gold', markeredgecolor='black', markersize=12, zorder=10)[0] + +ax_z.set_xlabel(r"Frequency $\omega/2\pi$ (MHz)", fontsize=12) + +# --- TEXT BOX UI ELEMENTS (GRID LAYOUT) --- +# Column 1 +ax_N = plt.axes([0.08, 0.13, 0.05, 0.04]) +ax_mode = plt.axes([0.08, 0.06, 0.05, 0.04]) + +# Column 2 +ax_wx = plt.axes([0.22, 0.13, 0.05, 0.04]) +ax_wy = plt.axes([0.22, 0.06, 0.05, 0.04]) + +# Column 3 +ax_wz = plt.axes([0.35, 0.13, 0.05, 0.04]) + +# Column 4 (Radio Buttons) +ax_radio = plt.axes([0.43, 0.06, 0.08, 0.11], facecolor='lightgray') + +# Status Bar (Bottom spanning) +ax_status = plt.axes([0.06, 0.01, 0.45, 0.03]) +ax_status.axis('off') +status_text = ax_status.text(0.0, 0.5, "STATUS: INITIALIZING...", ha='left', va='center', fontsize=11, fontweight='bold', color='gray') + +# Text Boxes with Compact Labels +text_N = TextBox(ax_N, 'N: ', initial=init_N) +text_mode = TextBox(ax_mode, 'Mode: ', initial=init_mode) +text_wx = TextBox(ax_wx, r'$\omega_x$ (MHz): ', initial=init_wx) +text_wy = TextBox(ax_wy, r'$\omega_y$ (MHz): ', initial=init_wy) +text_wz = TextBox(ax_wz, r'$\omega_z$ (MHz): ', initial=init_wz) +radio_dir = RadioButtons(ax_radio, ('Radial X', 'Radial Y', 'Axial Z')) + +def get_val(text_box, default, dtype=float): + try: + return dtype(text_box.text) + except ValueError: + return default + +# ========================================== +# 4. Unified Animation Engine +# ========================================== +def init(): + return list(lines.values()) + list(highlights.values()) + [status_text] + +def animate(frame): + try: + N = get_val(text_N, int(init_N), int) + mode = get_val(text_mode, int(init_mode), int) + wx = get_val(text_wx, float(init_wx)) + wy = get_val(text_wy, float(init_wy)) + wz = get_val(text_wz, float(init_wz)) + direction = radio_dir.value_selected + + if N < 2: N = 2 + if mode < 1: mode = 1 + if mode > N: mode = N + + solve_trical(N, wx, wy, wz) + + is_stable = trical_cache['steane_stable'] + w_crit = trical_cache['w_crit'] + raw_z = trical_cache['raw_z_eq'] + + freqs_z, freqs_x, freqs_y = trical_cache['freqs_z'], trical_cache['freqs_x'], trical_cache['freqs_y'] + + # --- UPDATE MAIN TABLE --- + current_state = (N, wx, wy, wz, mode, is_stable) + if trical_cache['table_drawn_for'] != current_state: + ax_table.clear() + ax_table.axis('off') + ax_table.set_title("Normal Modes & Z-Positions", fontweight='bold', pad=10) + + if N <= 15: + col_labels = ['m', 'Axial Z', 'Radial X', 'Radial Y', 'Z Pos (μm)'] + cell_text = [] + for i in range(N): + pos_str = f"{raw_z[i] * 1e6:.3f}" if not trical_cache['trical_stable'] else f"{raw_z[i] * 1e6:.3f}" + cell_text.append([f"{i+1}", f"{freqs_z[i]:.4f}", f"{freqs_x[i]:.4f}", f"{freqs_y[i]:.4f}", pos_str]) + + table = ax_table.table(cellText=cell_text, colLabels=col_labels, loc='center', cellLoc='center') + table.auto_set_font_size(False) + table.set_fontsize(9) + table.scale(1, 1.3) + + active_row = mode + for j in range(5): + table[(active_row, j)].set_facecolor('gold') + + else: + col_labels = ['m', 'Z', 'X', 'Y', 'Z(μm)', 'm', 'Z', 'X', 'Y', 'Z(μm)'] + cell_text = [] + half = math.ceil(N / 2) + for i in range(half): + pos_str1 = f"{raw_z[i] * 1e6:.2f}" + row = [f"{i+1}", f"{freqs_z[i]:.2f}", f"{freqs_x[i]:.2f}", f"{freqs_y[i]:.2f}", pos_str1] + + idx2 = i + half + if idx2 < N: + pos_str2 = f"{raw_z[idx2] * 1e6:.2f}" + row.extend([f"{idx2+1}", f"{freqs_z[idx2]:.2f}", f"{freqs_x[idx2]:.2f}", f"{freqs_y[idx2]:.2f}", pos_str2]) + else: + row.extend(["", "", "", "", ""]) + cell_text.append(row) + + table = ax_table.table(cellText=cell_text, colLabels=col_labels, loc='center', cellLoc='center') + table.auto_set_font_size(False) + table.set_fontsize(7) + table.scale(1, 1.1) + + if mode <= half: + active_row = mode + for j in range(5): table[(active_row, j)].set_facecolor('gold') + else: + active_row = mode - half + for j in range(5, 10): table[(active_row, j)].set_facecolor('gold') + + trical_cache['table_drawn_for'] = current_state + + # --- UPDATE STATUS BAR --- + if is_stable: + status_text.set_text(f"STATUS: STABLE (1D Linear Chain) ✓ | Critical Radial w_r > {w_crit:.3f} MHz") + status_text.set_color("darkgreen") + else: + status_text.set_text(f"STATUS: UNSTABLE (Buckles into 2D Zigzag) ✗ | Critical Radial w_r > {w_crit:.3f} MHz") + status_text.set_color("red") + + # --- UPDATE STACKED SPECTRUM PLOTS --- + lines['Radial X'].set_data(freqs_x, np.zeros_like(freqs_x)) + lines['Radial Y'].set_data(freqs_y, np.zeros_like(freqs_y)) + lines['Axial Z'].set_data(freqs_z, np.zeros_like(freqs_z)) + + for hl in highlights.values(): hl.set_data([], []) + + if direction == 'Radial X': + active_freq = freqs_x[mode - 1] if mode <= len(freqs_x) else 0.0 + highlights['Radial X'].set_data([active_freq], [0]) + active_ax = ax_x + evecs = trical_cache['evecs_x'] + elif direction == 'Radial Y': + active_freq = freqs_y[mode - 1] if mode <= len(freqs_y) else 0.0 + highlights['Radial Y'].set_data([active_freq], [0]) + active_ax = ax_y + evecs = trical_cache['evecs_y'] + else: + active_freq = freqs_z[mode - 1] if mode <= len(freqs_z) else 0.0 + highlights['Axial Z'].set_data([active_freq], [0]) + active_ax = ax_z + evecs = trical_cache['evecs_z'] + + for ax, freqs in zip(axes, [freqs_x, freqs_y, freqs_z]): + if len(freqs) > 0 and np.max(freqs) > 0: + ax.set_xlim(np.min(freqs)*0.9 - 0.05, np.max(freqs)*1.1 + 0.05) + else: + ax.set_xlim(0, 1) + + for ax in axes: ax.set_title("") + + if not is_stable: + fig.suptitle(f"2D Zigzag Threshold Crossed (Modes Invalid)", color='red', fontsize=14, y=0.98) + else: + title_str = f"Yb: Bare-Metal {direction} Spectrum (Active: {active_freq:.6f} MHz)" + active_ax.set_title(title_str, color='black', fontsize=11) + fig.suptitle("", y=0.98) + + # --- UPDATE EIGENVECTOR & LAMBD-DICKE PLOT --- + ax_evec.clear() + ax_eta.clear() + ax_eta.axis('off') + + if evecs is not None and mode <= evecs.shape[1]: + vec = evecs[:, mode - 1].copy() + + # Phase correction + if vec[np.argmax(np.abs(vec))] < 0: + vec *= -1 + + x_pos = np.arange(1, N + 1) + ax_evec.stem(x_pos, vec, basefmt="C3-", linefmt="C0-", markerfmt="C0o") + ax_evec.set_ylim(-1.0, 1.0) + ax_evec.set_title(f"Eigenvector $b_m^{{(j)}}$", fontweight='bold', fontsize=11) + + # Add numerical values beside each mode marker + if N <= 15: + for i, val in enumerate(vec): + y_offset = 0.12 if val >= 0 else -0.15 + ax_evec.text(x_pos[i], val + y_offset, f"{val:.3f}", + ha='center', va='center', fontsize=9, color='black', fontweight='bold') + + if N <= 10: + ax_evec.set_xticks(x_pos) + ax_evec.set_xticklabels([f'{i}' for i in x_pos]) + else: + ax_evec.set_xticks([1, N//2, N]) + ax_evec.set_xticklabels(['1', f'{N//2}', f'{N}']) + + ax_evec.set_xlabel("Ion Index (j)") + ax_evec.grid(True, linestyle=':', alpha=0.6) + + # --- Lamb-Dicke Calculation --- + if active_freq > 1e-6: + omega_m = 2 * np.pi * active_freq * 1e6 + zero_point_spread = np.sqrt(scipy_cst.hbar / (2 * MASS_ION * omega_m)) + eta_vec = vec * DELTA_K * zero_point_spread + else: + eta_vec = np.zeros_like(vec) + + # --- Rendering Lamb-Dicke Table --- + ax_eta.set_title(r"Lamb-Dicke $\eta_m^{(j)}$", fontweight='bold', fontsize=11) + col_labels = ['Ion', r'$b_m^{(j)}$', r'$\eta_m^{(j)}$'] + cell_text = [] + + display_N = min(N, 30) + for j in range(display_N): + cell_text.append([f"{j+1}", f"{vec[j]:.3f}", f"{eta_vec[j]:.4f}"]) + + if N > 30: + cell_text.append(["...", "...", "..."]) + + table_eta = ax_eta.table(cellText=cell_text, colLabels=col_labels, loc='center', cellLoc='center') + table_eta.auto_set_font_size(False) + + # Dynamically shrink table row height slightly as N grows to prevent overflow + scale_factor = 1.0 if N <= 20 else 0.85 + table_eta.scale(1, scale_factor) + table_eta.set_fontsize(7) + + return list(lines.values()) + list(highlights.values()) + [status_text] + + except Exception as e: + print(f"Update Error: {e}") + status_text.set_text(f"STATUS: ERROR ({e})") + status_text.set_color("red") + return list(lines.values()) + list(highlights.values()) + [status_text] + +# ========================================== +# 5. Execute +# ========================================== +ani = animation.FuncAnimation(fig, animate, init_func=init, frames=1000, interval=100, blit=False) +plt.show() \ No newline at end of file From 672ee2667d4be385eb77a8c696c6537e38c690a0 Mon Sep 17 00:00:00 2001 From: yhteoh Date: Mon, 1 Jun 2026 13:19:45 -0400 Subject: [PATCH 2/2] [refactor] moved bloodstone trapped ion visualization to oqd_trical.visualization.bloodstone and added it as a script to the project with command trical_bloodstone. --- pyproject.toml | 4 + src/oqd_trical/BSI_004.py | 422 ---------------- src/oqd_trical/visualization/bloodstone.py | 558 +++++++++++++++++++++ 3 files changed, 562 insertions(+), 422 deletions(-) delete mode 100644 src/oqd_trical/BSI_004.py create mode 100644 src/oqd_trical/visualization/bloodstone.py diff --git a/pyproject.toml b/pyproject.toml index 4c0349e..013b0b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,10 @@ docs = [ ] tests = ["pytest"] + +[project.scripts] +trical_bloodstone = "oqd_trical.visualization.bloodstone:main" + [tool.setuptools.packages.find] where = ["src"] include = ["oqd_trical*"] diff --git a/src/oqd_trical/BSI_004.py b/src/oqd_trical/BSI_004.py deleted file mode 100644 index a9680e4..0000000 --- a/src/oqd_trical/BSI_004.py +++ /dev/null @@ -1,422 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt -import matplotlib.animation as animation -from matplotlib.widgets import RadioButtons, TextBox -import scipy.constants as scipy_cst -from scipy.optimize import minimize -from matplotlib.ticker import MultipleLocator -import math - -# Import the exact solver from your notebook -try: - import oqd_trical -except ImportError: - oqd_trical = None - -# ========================================== -# 1. Initial Parameters & State Caching -# ========================================== -init_N = "20" -init_mode = "1" -init_wx = "3.00" # MHz (Radial X frequency) -init_wy = "3.00" # MHz (Radial Y frequency) -init_wz = "0.200" # MHz (Axial Z frequency) - -# Physical Constants for Lamb-Dicke Calculation -MASS_ION = 171 * scipy_cst.m_u # Ytterbium-171 mass -LAMBDA_LASER = 355e-9 # 355 nm Raman beams -DELTA_K = 4 * np.pi / LAMBDA_LASER - -trical_cache = { - 'N': None, 'wx': None, 'wy': None, 'wz': None, - 'freqs_x': None, 'freqs_y': None, 'freqs_z': None, - 'evecs_x': None, 'evecs_y': None, 'evecs_z': None, - 'x_eq': None, 'raw_z_eq': None, - 'trical_stable': True, 'steane_stable': True, 'w_crit': 0.0, - 'table_drawn_for': None -} - -# ========================================== -# 2. Physics Engine (OQD_TRICAL + Bare-Metal Fallback) -# ========================================== -def solve_trical(N, wx_mhz, wy_mhz, wz_mhz): - if (trical_cache['N'] == N and - trical_cache['wx'] == wx_mhz and - trical_cache['wy'] == wy_mhz and - trical_cache['wz'] == wz_mhz): - return - - w_rad_min = min(wx_mhz, wy_mhz) - beta_c = 0.73 * (N ** 0.86) - omega_r_critical = beta_c * wz_mhz - steane_stable = w_rad_min > omega_r_critical - - trical_cache['steane_stable'] = steane_stable - trical_cache['w_crit'] = omega_r_critical - - try: - if oqd_trical is None: raise Exception("oqd_trical not installed") - - omega_x = 2 * np.pi * wx_mhz * 1e6 - omega_y = 2 * np.pi * wy_mhz * 1e6 - omega_z = 2 * np.pi * wz_mhz * 1e6 - - alpha = np.zeros((3, 3, 3)) - alpha[2, 0, 0] = MASS_ION * (omega_x) ** 2 / 2 - alpha[0, 2, 0] = MASS_ION * (omega_y) ** 2 / 2 - alpha[0, 0, 2] = MASS_ION * (omega_z) ** 2 / 2 - - pp = oqd_trical.mechanical.PolynomialPotential(alpha, N=N) - ti = oqd_trical.mechanical.TrappedIons(N, pp, m=MASS_ION) - - ti.equilibrium_position() - ti.normal_modes() - - if not hasattr(ti, 'mode_vectors') and not hasattr(ti, 'eigenvectors'): - raise Exception("Trical missing eigenvector attributes.") - - except Exception as e: - trical_cache['trical_stable'] = False - - # --- CUSTOM BARE-METAL EXACT NUMERICAL SOLVER --- - def coulomb_potential(u): - axial_confinement = 0.5 * np.sum(u**2) - coulomb_repulsion = 0.0 - for i in range(N): - for j in range(i + 1, N): - coulomb_repulsion += 1.0 / np.abs(u[i] - u[j]) - return axial_confinement + coulomb_repulsion - - u_guess = np.linspace(-N/2, N/2, N) - res = minimize(coulomb_potential, u_guess, method='BFGS') - u_eq = np.sort(res.x) - - # Build Hessian Matrices (CORRECTED SIGNS) - H_z, H_x, H_y = np.zeros((N, N)), np.zeros((N, N)), np.zeros((N, N)) - alpha_x = (wx_mhz / wz_mhz)**2 - alpha_y = (wy_mhz / wz_mhz)**2 - - for i in range(N): - for j in range(N): - if i != j: - inv_dist_cubed = 1.0 / np.abs(u_eq[i] - u_eq[j])**3 - # Axial off-diagonals are negative (-2/d^3) - H_z[i, j] = -2.0 * inv_dist_cubed - # Radial off-diagonals are positive (+1/d^3) - H_x[i, j] = 1.0 * inv_dist_cubed - H_y[i, j] = 1.0 * inv_dist_cubed - - for i in range(N): - # Diagonal elements derived from the COM eigenvalue conditions - H_z[i, i] = 1.0 - np.sum(H_z[i, :]) + H_z[i, i] - H_x[i, i] = alpha_x - np.sum(H_x[i, :]) + H_x[i, i] - H_y[i, i] = alpha_y - np.sum(H_y[i, :]) + H_y[i, i] - - def solve_modes(H, omega_base): - evals, evecs = np.linalg.eigh(H) - freqs = np.sqrt(np.maximum(0, evals)) * omega_base - return freqs, evecs - - f_z, e_z = solve_modes(H_z, wz_mhz) - f_x, e_x = solve_modes(H_x, wz_mhz) - f_y, e_y = solve_modes(H_y, wz_mhz) - - idx_z = np.argsort(f_z) - trical_cache['freqs_z'] = f_z[idx_z] - trical_cache['evecs_z'] = e_z[:, idx_z] - - idx_x = np.argsort(f_x)[::-1] - trical_cache['freqs_x'] = f_x[idx_x] - trical_cache['evecs_x'] = e_x[:, idx_x] - - idx_y = np.argsort(f_y)[::-1] - trical_cache['freqs_y'] = f_y[idx_y] - trical_cache['evecs_y'] = e_y[:, idx_y] - - omega_z_si = 2 * np.pi * wz_mhz * 1e6 - l_0 = ((scipy_cst.e**2 / (4 * np.pi * scipy_cst.epsilon_0)) / (MASS_ION * omega_z_si**2))**(1/3) - - trical_cache['raw_z_eq'] = u_eq * l_0 - trical_cache['x_eq'] = u_eq - - trical_cache['N'] = N - trical_cache['wx'] = wx_mhz - trical_cache['wy'] = wy_mhz - trical_cache['wz'] = wz_mhz - -# ========================================== -# 3. Figure and UI Setup (OPTIMIZED LAYOUT) -# ========================================== -fig, (ax_x, ax_y, ax_z) = plt.subplots(3, 1, figsize=(16, 9)) - -# Gives maximum vertical height to the three spectrum plots on the left -plt.subplots_adjust(top=0.94, bottom=0.22, left=0.06, right=0.52, hspace=0.45) - -# --- RIGHT SIDE PLOTS --- -# Main modes table (Top Right) -ax_table = plt.axes([0.56, 0.52, 0.42, 0.42]) -ax_table.axis('off') - -# Eigenvector Plot (Bottom Right-Middle) -ax_evec = plt.axes([0.56, 0.05, 0.24, 0.40]) - -# Lamb-Dicke Table (Bottom Right Edge - Maximized for Height) -ax_eta = plt.axes([0.82, 0.05, 0.16, 0.40]) -ax_eta.axis('off') - -# --- Spectrum Plots Setup --- -axes = [ax_x, ax_y, ax_z] -labels = ['Radial X', 'Radial Y', 'Axial Z'] -colors = ['tab:red', 'tab:blue', 'tab:green'] - -lines = {} -highlights = {} - -for ax, label, color in zip(axes, labels, colors): - ax.grid(axis='x', linestyle='--', alpha=0.6) - ax.set_yticks([]) - ax.set_ylim(-1, 1) - ax.set_ylabel(label, fontsize=11, fontweight='bold') - ax.xaxis.set_minor_locator(MultipleLocator(0.1)) - ax.grid(axis='x', which='minor', linestyle=':', alpha=0.4) - - lines[label] = ax.plot([], [], '|', color=color, markersize=30, markeredgewidth=1.5)[0] - highlights[label] = ax.plot([], [], 'o', color='gold', markeredgecolor='black', markersize=12, zorder=10)[0] - -ax_z.set_xlabel(r"Frequency $\omega/2\pi$ (MHz)", fontsize=12) - -# --- TEXT BOX UI ELEMENTS (GRID LAYOUT) --- -# Column 1 -ax_N = plt.axes([0.08, 0.13, 0.05, 0.04]) -ax_mode = plt.axes([0.08, 0.06, 0.05, 0.04]) - -# Column 2 -ax_wx = plt.axes([0.22, 0.13, 0.05, 0.04]) -ax_wy = plt.axes([0.22, 0.06, 0.05, 0.04]) - -# Column 3 -ax_wz = plt.axes([0.35, 0.13, 0.05, 0.04]) - -# Column 4 (Radio Buttons) -ax_radio = plt.axes([0.43, 0.06, 0.08, 0.11], facecolor='lightgray') - -# Status Bar (Bottom spanning) -ax_status = plt.axes([0.06, 0.01, 0.45, 0.03]) -ax_status.axis('off') -status_text = ax_status.text(0.0, 0.5, "STATUS: INITIALIZING...", ha='left', va='center', fontsize=11, fontweight='bold', color='gray') - -# Text Boxes with Compact Labels -text_N = TextBox(ax_N, 'N: ', initial=init_N) -text_mode = TextBox(ax_mode, 'Mode: ', initial=init_mode) -text_wx = TextBox(ax_wx, r'$\omega_x$ (MHz): ', initial=init_wx) -text_wy = TextBox(ax_wy, r'$\omega_y$ (MHz): ', initial=init_wy) -text_wz = TextBox(ax_wz, r'$\omega_z$ (MHz): ', initial=init_wz) -radio_dir = RadioButtons(ax_radio, ('Radial X', 'Radial Y', 'Axial Z')) - -def get_val(text_box, default, dtype=float): - try: - return dtype(text_box.text) - except ValueError: - return default - -# ========================================== -# 4. Unified Animation Engine -# ========================================== -def init(): - return list(lines.values()) + list(highlights.values()) + [status_text] - -def animate(frame): - try: - N = get_val(text_N, int(init_N), int) - mode = get_val(text_mode, int(init_mode), int) - wx = get_val(text_wx, float(init_wx)) - wy = get_val(text_wy, float(init_wy)) - wz = get_val(text_wz, float(init_wz)) - direction = radio_dir.value_selected - - if N < 2: N = 2 - if mode < 1: mode = 1 - if mode > N: mode = N - - solve_trical(N, wx, wy, wz) - - is_stable = trical_cache['steane_stable'] - w_crit = trical_cache['w_crit'] - raw_z = trical_cache['raw_z_eq'] - - freqs_z, freqs_x, freqs_y = trical_cache['freqs_z'], trical_cache['freqs_x'], trical_cache['freqs_y'] - - # --- UPDATE MAIN TABLE --- - current_state = (N, wx, wy, wz, mode, is_stable) - if trical_cache['table_drawn_for'] != current_state: - ax_table.clear() - ax_table.axis('off') - ax_table.set_title("Normal Modes & Z-Positions", fontweight='bold', pad=10) - - if N <= 15: - col_labels = ['m', 'Axial Z', 'Radial X', 'Radial Y', 'Z Pos (μm)'] - cell_text = [] - for i in range(N): - pos_str = f"{raw_z[i] * 1e6:.3f}" if not trical_cache['trical_stable'] else f"{raw_z[i] * 1e6:.3f}" - cell_text.append([f"{i+1}", f"{freqs_z[i]:.4f}", f"{freqs_x[i]:.4f}", f"{freqs_y[i]:.4f}", pos_str]) - - table = ax_table.table(cellText=cell_text, colLabels=col_labels, loc='center', cellLoc='center') - table.auto_set_font_size(False) - table.set_fontsize(9) - table.scale(1, 1.3) - - active_row = mode - for j in range(5): - table[(active_row, j)].set_facecolor('gold') - - else: - col_labels = ['m', 'Z', 'X', 'Y', 'Z(μm)', 'm', 'Z', 'X', 'Y', 'Z(μm)'] - cell_text = [] - half = math.ceil(N / 2) - for i in range(half): - pos_str1 = f"{raw_z[i] * 1e6:.2f}" - row = [f"{i+1}", f"{freqs_z[i]:.2f}", f"{freqs_x[i]:.2f}", f"{freqs_y[i]:.2f}", pos_str1] - - idx2 = i + half - if idx2 < N: - pos_str2 = f"{raw_z[idx2] * 1e6:.2f}" - row.extend([f"{idx2+1}", f"{freqs_z[idx2]:.2f}", f"{freqs_x[idx2]:.2f}", f"{freqs_y[idx2]:.2f}", pos_str2]) - else: - row.extend(["", "", "", "", ""]) - cell_text.append(row) - - table = ax_table.table(cellText=cell_text, colLabels=col_labels, loc='center', cellLoc='center') - table.auto_set_font_size(False) - table.set_fontsize(7) - table.scale(1, 1.1) - - if mode <= half: - active_row = mode - for j in range(5): table[(active_row, j)].set_facecolor('gold') - else: - active_row = mode - half - for j in range(5, 10): table[(active_row, j)].set_facecolor('gold') - - trical_cache['table_drawn_for'] = current_state - - # --- UPDATE STATUS BAR --- - if is_stable: - status_text.set_text(f"STATUS: STABLE (1D Linear Chain) ✓ | Critical Radial w_r > {w_crit:.3f} MHz") - status_text.set_color("darkgreen") - else: - status_text.set_text(f"STATUS: UNSTABLE (Buckles into 2D Zigzag) ✗ | Critical Radial w_r > {w_crit:.3f} MHz") - status_text.set_color("red") - - # --- UPDATE STACKED SPECTRUM PLOTS --- - lines['Radial X'].set_data(freqs_x, np.zeros_like(freqs_x)) - lines['Radial Y'].set_data(freqs_y, np.zeros_like(freqs_y)) - lines['Axial Z'].set_data(freqs_z, np.zeros_like(freqs_z)) - - for hl in highlights.values(): hl.set_data([], []) - - if direction == 'Radial X': - active_freq = freqs_x[mode - 1] if mode <= len(freqs_x) else 0.0 - highlights['Radial X'].set_data([active_freq], [0]) - active_ax = ax_x - evecs = trical_cache['evecs_x'] - elif direction == 'Radial Y': - active_freq = freqs_y[mode - 1] if mode <= len(freqs_y) else 0.0 - highlights['Radial Y'].set_data([active_freq], [0]) - active_ax = ax_y - evecs = trical_cache['evecs_y'] - else: - active_freq = freqs_z[mode - 1] if mode <= len(freqs_z) else 0.0 - highlights['Axial Z'].set_data([active_freq], [0]) - active_ax = ax_z - evecs = trical_cache['evecs_z'] - - for ax, freqs in zip(axes, [freqs_x, freqs_y, freqs_z]): - if len(freqs) > 0 and np.max(freqs) > 0: - ax.set_xlim(np.min(freqs)*0.9 - 0.05, np.max(freqs)*1.1 + 0.05) - else: - ax.set_xlim(0, 1) - - for ax in axes: ax.set_title("") - - if not is_stable: - fig.suptitle(f"2D Zigzag Threshold Crossed (Modes Invalid)", color='red', fontsize=14, y=0.98) - else: - title_str = f"Yb: Bare-Metal {direction} Spectrum (Active: {active_freq:.6f} MHz)" - active_ax.set_title(title_str, color='black', fontsize=11) - fig.suptitle("", y=0.98) - - # --- UPDATE EIGENVECTOR & LAMBD-DICKE PLOT --- - ax_evec.clear() - ax_eta.clear() - ax_eta.axis('off') - - if evecs is not None and mode <= evecs.shape[1]: - vec = evecs[:, mode - 1].copy() - - # Phase correction - if vec[np.argmax(np.abs(vec))] < 0: - vec *= -1 - - x_pos = np.arange(1, N + 1) - ax_evec.stem(x_pos, vec, basefmt="C3-", linefmt="C0-", markerfmt="C0o") - ax_evec.set_ylim(-1.0, 1.0) - ax_evec.set_title(f"Eigenvector $b_m^{{(j)}}$", fontweight='bold', fontsize=11) - - # Add numerical values beside each mode marker - if N <= 15: - for i, val in enumerate(vec): - y_offset = 0.12 if val >= 0 else -0.15 - ax_evec.text(x_pos[i], val + y_offset, f"{val:.3f}", - ha='center', va='center', fontsize=9, color='black', fontweight='bold') - - if N <= 10: - ax_evec.set_xticks(x_pos) - ax_evec.set_xticklabels([f'{i}' for i in x_pos]) - else: - ax_evec.set_xticks([1, N//2, N]) - ax_evec.set_xticklabels(['1', f'{N//2}', f'{N}']) - - ax_evec.set_xlabel("Ion Index (j)") - ax_evec.grid(True, linestyle=':', alpha=0.6) - - # --- Lamb-Dicke Calculation --- - if active_freq > 1e-6: - omega_m = 2 * np.pi * active_freq * 1e6 - zero_point_spread = np.sqrt(scipy_cst.hbar / (2 * MASS_ION * omega_m)) - eta_vec = vec * DELTA_K * zero_point_spread - else: - eta_vec = np.zeros_like(vec) - - # --- Rendering Lamb-Dicke Table --- - ax_eta.set_title(r"Lamb-Dicke $\eta_m^{(j)}$", fontweight='bold', fontsize=11) - col_labels = ['Ion', r'$b_m^{(j)}$', r'$\eta_m^{(j)}$'] - cell_text = [] - - display_N = min(N, 30) - for j in range(display_N): - cell_text.append([f"{j+1}", f"{vec[j]:.3f}", f"{eta_vec[j]:.4f}"]) - - if N > 30: - cell_text.append(["...", "...", "..."]) - - table_eta = ax_eta.table(cellText=cell_text, colLabels=col_labels, loc='center', cellLoc='center') - table_eta.auto_set_font_size(False) - - # Dynamically shrink table row height slightly as N grows to prevent overflow - scale_factor = 1.0 if N <= 20 else 0.85 - table_eta.scale(1, scale_factor) - table_eta.set_fontsize(7) - - return list(lines.values()) + list(highlights.values()) + [status_text] - - except Exception as e: - print(f"Update Error: {e}") - status_text.set_text(f"STATUS: ERROR ({e})") - status_text.set_color("red") - return list(lines.values()) + list(highlights.values()) + [status_text] - -# ========================================== -# 5. Execute -# ========================================== -ani = animation.FuncAnimation(fig, animate, init_func=init, frames=1000, interval=100, blit=False) -plt.show() \ No newline at end of file diff --git a/src/oqd_trical/visualization/bloodstone.py b/src/oqd_trical/visualization/bloodstone.py new file mode 100644 index 0000000..129139d --- /dev/null +++ b/src/oqd_trical/visualization/bloodstone.py @@ -0,0 +1,558 @@ +# Copyright 2024-2025 Open Quantum Design + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import matplotlib.animation as animation +import matplotlib.pyplot as plt +import numpy as np +import scipy.constants as scipy_cst +from matplotlib.ticker import MultipleLocator +from matplotlib.widgets import RadioButtons, TextBox +from scipy.optimize import minimize + +# Import the exact solver from your notebook +try: + import oqd_trical +except ImportError: + oqd_trical = None + + +def main(): + # ========================================== + # 1. Initial Parameters & State Caching + # ========================================== + init_N = "20" + init_mode = "1" + init_wx = "3.00" # MHz (Radial X frequency) + init_wy = "3.00" # MHz (Radial Y frequency) + init_wz = "0.200" # MHz (Axial Z frequency) + + # Physical Constants for Lamb-Dicke Calculation + MASS_ION = 171 * scipy_cst.m_u # Ytterbium-171 mass + LAMBDA_LASER = 355e-9 # 355 nm Raman beams + DELTA_K = 4 * np.pi / LAMBDA_LASER + + trical_cache = { + "N": None, + "wx": None, + "wy": None, + "wz": None, + "freqs_x": None, + "freqs_y": None, + "freqs_z": None, + "evecs_x": None, + "evecs_y": None, + "evecs_z": None, + "x_eq": None, + "raw_z_eq": None, + "trical_stable": True, + "steane_stable": True, + "w_crit": 0.0, + "table_drawn_for": None, + } + + # ========================================== + # 2. Physics Engine (OQD_TRICAL + Bare-Metal Fallback) + # ========================================== + def solve_trical(N, wx_mhz, wy_mhz, wz_mhz): + if ( + trical_cache["N"] == N + and trical_cache["wx"] == wx_mhz + and trical_cache["wy"] == wy_mhz + and trical_cache["wz"] == wz_mhz + ): + return + + w_rad_min = min(wx_mhz, wy_mhz) + beta_c = 0.73 * (N**0.86) + omega_r_critical = beta_c * wz_mhz + steane_stable = w_rad_min > omega_r_critical + + trical_cache["steane_stable"] = steane_stable + trical_cache["w_crit"] = omega_r_critical + + try: + if oqd_trical is None: + raise Exception("oqd_trical not installed") + + omega_x = 2 * np.pi * wx_mhz * 1e6 + omega_y = 2 * np.pi * wy_mhz * 1e6 + omega_z = 2 * np.pi * wz_mhz * 1e6 + + alpha = np.zeros((3, 3, 3)) + alpha[2, 0, 0] = MASS_ION * (omega_x) ** 2 / 2 + alpha[0, 2, 0] = MASS_ION * (omega_y) ** 2 / 2 + alpha[0, 0, 2] = MASS_ION * (omega_z) ** 2 / 2 + + pp = oqd_trical.mechanical.PolynomialPotential(alpha, N=N) + ti = oqd_trical.mechanical.TrappedIons(N, pp, m=MASS_ION) + + ti.equilibrium_position() + ti.normal_modes() + + if not hasattr(ti, "mode_vectors") and not hasattr(ti, "eigenvectors"): + raise Exception("Trical missing eigenvector attributes.") + + except Exception as e: + trical_cache["trical_stable"] = False + + # --- CUSTOM BARE-METAL EXACT NUMERICAL SOLVER --- + def coulomb_potential(u): + axial_confinement = 0.5 * np.sum(u**2) + coulomb_repulsion = 0.0 + for i in range(N): + for j in range(i + 1, N): + coulomb_repulsion += 1.0 / np.abs(u[i] - u[j]) + return axial_confinement + coulomb_repulsion + + u_guess = np.linspace(-N / 2, N / 2, N) + res = minimize(coulomb_potential, u_guess, method="BFGS") + u_eq = np.sort(res.x) + + # Build Hessian Matrices (CORRECTED SIGNS) + H_z, H_x, H_y = np.zeros((N, N)), np.zeros((N, N)), np.zeros((N, N)) + alpha_x = (wx_mhz / wz_mhz) ** 2 + alpha_y = (wy_mhz / wz_mhz) ** 2 + + for i in range(N): + for j in range(N): + if i != j: + inv_dist_cubed = 1.0 / np.abs(u_eq[i] - u_eq[j]) ** 3 + # Axial off-diagonals are negative (-2/d^3) + H_z[i, j] = -2.0 * inv_dist_cubed + # Radial off-diagonals are positive (+1/d^3) + H_x[i, j] = 1.0 * inv_dist_cubed + H_y[i, j] = 1.0 * inv_dist_cubed + + for i in range(N): + # Diagonal elements derived from the COM eigenvalue conditions + H_z[i, i] = 1.0 - np.sum(H_z[i, :]) + H_z[i, i] + H_x[i, i] = alpha_x - np.sum(H_x[i, :]) + H_x[i, i] + H_y[i, i] = alpha_y - np.sum(H_y[i, :]) + H_y[i, i] + + def solve_modes(H, omega_base): + evals, evecs = np.linalg.eigh(H) + freqs = np.sqrt(np.maximum(0, evals)) * omega_base + return freqs, evecs + + f_z, e_z = solve_modes(H_z, wz_mhz) + f_x, e_x = solve_modes(H_x, wz_mhz) + f_y, e_y = solve_modes(H_y, wz_mhz) + + idx_z = np.argsort(f_z) + trical_cache["freqs_z"] = f_z[idx_z] + trical_cache["evecs_z"] = e_z[:, idx_z] + + idx_x = np.argsort(f_x)[::-1] + trical_cache["freqs_x"] = f_x[idx_x] + trical_cache["evecs_x"] = e_x[:, idx_x] + + idx_y = np.argsort(f_y)[::-1] + trical_cache["freqs_y"] = f_y[idx_y] + trical_cache["evecs_y"] = e_y[:, idx_y] + + omega_z_si = 2 * np.pi * wz_mhz * 1e6 + l_0 = ( + (scipy_cst.e**2 / (4 * np.pi * scipy_cst.epsilon_0)) + / (MASS_ION * omega_z_si**2) + ) ** (1 / 3) + + trical_cache["raw_z_eq"] = u_eq * l_0 + trical_cache["x_eq"] = u_eq + + trical_cache["N"] = N + trical_cache["wx"] = wx_mhz + trical_cache["wy"] = wy_mhz + trical_cache["wz"] = wz_mhz + + # ========================================== + # 3. Figure and UI Setup (OPTIMIZED LAYOUT) + # ========================================== + fig, (ax_x, ax_y, ax_z) = plt.subplots(3, 1, figsize=(16, 9)) + + # Gives maximum vertical height to the three spectrum plots on the left + plt.subplots_adjust(top=0.94, bottom=0.22, left=0.06, right=0.52, hspace=0.45) + + # --- RIGHT SIDE PLOTS --- + # Main modes table (Top Right) + ax_table = plt.axes([0.56, 0.52, 0.42, 0.42]) + ax_table.axis("off") + + # Eigenvector Plot (Bottom Right-Middle) + ax_evec = plt.axes([0.56, 0.05, 0.24, 0.40]) + + # Lamb-Dicke Table (Bottom Right Edge - Maximized for Height) + ax_eta = plt.axes([0.82, 0.05, 0.16, 0.40]) + ax_eta.axis("off") + + # --- Spectrum Plots Setup --- + axes = [ax_x, ax_y, ax_z] + labels = ["Radial X", "Radial Y", "Axial Z"] + colors = ["tab:red", "tab:blue", "tab:green"] + + lines = {} + highlights = {} + + for ax, label, color in zip(axes, labels, colors): + ax.grid(axis="x", linestyle="--", alpha=0.6) + ax.set_yticks([]) + ax.set_ylim(-1, 1) + ax.set_ylabel(label, fontsize=11, fontweight="bold") + ax.xaxis.set_minor_locator(MultipleLocator(0.1)) + ax.grid(axis="x", which="minor", linestyle=":", alpha=0.4) + + lines[label] = ax.plot( + [], [], "|", color=color, markersize=30, markeredgewidth=1.5 + )[0] + highlights[label] = ax.plot( + [], [], "o", color="gold", markeredgecolor="black", markersize=12, zorder=10 + )[0] + + ax_z.set_xlabel(r"Frequency $\omega/2\pi$ (MHz)", fontsize=12) + + # --- TEXT BOX UI ELEMENTS (GRID LAYOUT) --- + # Column 1 + ax_N = plt.axes([0.08, 0.13, 0.05, 0.04]) + ax_mode = plt.axes([0.08, 0.06, 0.05, 0.04]) + + # Column 2 + ax_wx = plt.axes([0.22, 0.13, 0.05, 0.04]) + ax_wy = plt.axes([0.22, 0.06, 0.05, 0.04]) + + # Column 3 + ax_wz = plt.axes([0.35, 0.13, 0.05, 0.04]) + + # Column 4 (Radio Buttons) + ax_radio = plt.axes([0.43, 0.06, 0.08, 0.11], facecolor="lightgray") + + # Status Bar (Bottom spanning) + ax_status = plt.axes([0.06, 0.01, 0.45, 0.03]) + ax_status.axis("off") + status_text = ax_status.text( + 0.0, + 0.5, + "STATUS: INITIALIZING...", + ha="left", + va="center", + fontsize=11, + fontweight="bold", + color="gray", + ) + + # Text Boxes with Compact Labels + text_N = TextBox(ax_N, "N: ", initial=init_N) + text_mode = TextBox(ax_mode, "Mode: ", initial=init_mode) + text_wx = TextBox(ax_wx, r"$\omega_x$ (MHz): ", initial=init_wx) + text_wy = TextBox(ax_wy, r"$\omega_y$ (MHz): ", initial=init_wy) + text_wz = TextBox(ax_wz, r"$\omega_z$ (MHz): ", initial=init_wz) + radio_dir = RadioButtons(ax_radio, ("Radial X", "Radial Y", "Axial Z")) + + def get_val(text_box, default, dtype=float): + try: + return dtype(text_box.text) + except ValueError: + return default + + # ========================================== + # 4. Unified Animation Engine + # ========================================== + def init(): + return list(lines.values()) + list(highlights.values()) + [status_text] + + def animate(frame): + try: + N = get_val(text_N, int(init_N), int) + mode = get_val(text_mode, int(init_mode), int) + wx = get_val(text_wx, float(init_wx)) + wy = get_val(text_wy, float(init_wy)) + wz = get_val(text_wz, float(init_wz)) + direction = radio_dir.value_selected + + if N < 2: + N = 2 + if mode < 1: + mode = 1 + if mode > N: + mode = N + + solve_trical(N, wx, wy, wz) + + is_stable = trical_cache["steane_stable"] + w_crit = trical_cache["w_crit"] + raw_z = trical_cache["raw_z_eq"] + + freqs_z, freqs_x, freqs_y = ( + trical_cache["freqs_z"], + trical_cache["freqs_x"], + trical_cache["freqs_y"], + ) + + # --- UPDATE MAIN TABLE --- + current_state = (N, wx, wy, wz, mode, is_stable) + if trical_cache["table_drawn_for"] != current_state: + ax_table.clear() + ax_table.axis("off") + ax_table.set_title( + "Normal Modes & Z-Positions", fontweight="bold", pad=10 + ) + + if N <= 15: + col_labels = ["m", "Axial Z", "Radial X", "Radial Y", "Z Pos (μm)"] + cell_text = [] + for i in range(N): + pos_str = ( + f"{raw_z[i] * 1e6:.3f}" + if not trical_cache["trical_stable"] + else f"{raw_z[i] * 1e6:.3f}" + ) + cell_text.append( + [ + f"{i + 1}", + f"{freqs_z[i]:.4f}", + f"{freqs_x[i]:.4f}", + f"{freqs_y[i]:.4f}", + pos_str, + ] + ) + + table = ax_table.table( + cellText=cell_text, + colLabels=col_labels, + loc="center", + cellLoc="center", + ) + table.auto_set_font_size(False) + table.set_fontsize(9) + table.scale(1, 1.3) + + active_row = mode + for j in range(5): + table[(active_row, j)].set_facecolor("gold") + + else: + col_labels = [ + "m", + "Z", + "X", + "Y", + "Z(μm)", + "m", + "Z", + "X", + "Y", + "Z(μm)", + ] + cell_text = [] + half = math.ceil(N / 2) + for i in range(half): + pos_str1 = f"{raw_z[i] * 1e6:.2f}" + row = [ + f"{i + 1}", + f"{freqs_z[i]:.2f}", + f"{freqs_x[i]:.2f}", + f"{freqs_y[i]:.2f}", + pos_str1, + ] + + idx2 = i + half + if idx2 < N: + pos_str2 = f"{raw_z[idx2] * 1e6:.2f}" + row.extend( + [ + f"{idx2 + 1}", + f"{freqs_z[idx2]:.2f}", + f"{freqs_x[idx2]:.2f}", + f"{freqs_y[idx2]:.2f}", + pos_str2, + ] + ) + else: + row.extend(["", "", "", "", ""]) + cell_text.append(row) + + table = ax_table.table( + cellText=cell_text, + colLabels=col_labels, + loc="center", + cellLoc="center", + ) + table.auto_set_font_size(False) + table.set_fontsize(7) + table.scale(1, 1.1) + + if mode <= half: + active_row = mode + for j in range(5): + table[(active_row, j)].set_facecolor("gold") + else: + active_row = mode - half + for j in range(5, 10): + table[(active_row, j)].set_facecolor("gold") + + trical_cache["table_drawn_for"] = current_state + + # --- UPDATE STATUS BAR --- + if is_stable: + status_text.set_text( + f"STATUS: STABLE (1D Linear Chain) ✓ | Critical Radial w_r > {w_crit:.3f} MHz" + ) + status_text.set_color("darkgreen") + else: + status_text.set_text( + f"STATUS: UNSTABLE (Buckles into 2D Zigzag) ✗ | Critical Radial w_r > {w_crit:.3f} MHz" + ) + status_text.set_color("red") + + # --- UPDATE STACKED SPECTRUM PLOTS --- + lines["Radial X"].set_data(freqs_x, np.zeros_like(freqs_x)) + lines["Radial Y"].set_data(freqs_y, np.zeros_like(freqs_y)) + lines["Axial Z"].set_data(freqs_z, np.zeros_like(freqs_z)) + + for hl in highlights.values(): + hl.set_data([], []) + + if direction == "Radial X": + active_freq = freqs_x[mode - 1] if mode <= len(freqs_x) else 0.0 + highlights["Radial X"].set_data([active_freq], [0]) + active_ax = ax_x + evecs = trical_cache["evecs_x"] + elif direction == "Radial Y": + active_freq = freqs_y[mode - 1] if mode <= len(freqs_y) else 0.0 + highlights["Radial Y"].set_data([active_freq], [0]) + active_ax = ax_y + evecs = trical_cache["evecs_y"] + else: + active_freq = freqs_z[mode - 1] if mode <= len(freqs_z) else 0.0 + highlights["Axial Z"].set_data([active_freq], [0]) + active_ax = ax_z + evecs = trical_cache["evecs_z"] + + for ax, freqs in zip(axes, [freqs_x, freqs_y, freqs_z]): + if len(freqs) > 0 and np.max(freqs) > 0: + ax.set_xlim(np.min(freqs) * 0.9 - 0.05, np.max(freqs) * 1.1 + 0.05) + else: + ax.set_xlim(0, 1) + + for ax in axes: + ax.set_title("") + + if not is_stable: + fig.suptitle( + f"2D Zigzag Threshold Crossed (Modes Invalid)", + color="red", + fontsize=14, + y=0.98, + ) + else: + title_str = f"Yb: Bare-Metal {direction} Spectrum (Active: {active_freq:.6f} MHz)" + active_ax.set_title(title_str, color="black", fontsize=11) + fig.suptitle("", y=0.98) + + # --- UPDATE EIGENVECTOR & LAMBD-DICKE PLOT --- + ax_evec.clear() + ax_eta.clear() + ax_eta.axis("off") + + if evecs is not None and mode <= evecs.shape[1]: + vec = evecs[:, mode - 1].copy() + + # Phase correction + if vec[np.argmax(np.abs(vec))] < 0: + vec *= -1 + + x_pos = np.arange(1, N + 1) + ax_evec.stem(x_pos, vec, basefmt="C3-", linefmt="C0-", markerfmt="C0o") + ax_evec.set_ylim(-1.0, 1.0) + ax_evec.set_title( + f"Eigenvector $b_m^{{(j)}}$", fontweight="bold", fontsize=11 + ) + + # Add numerical values beside each mode marker + if N <= 15: + for i, val in enumerate(vec): + y_offset = 0.12 if val >= 0 else -0.15 + ax_evec.text( + x_pos[i], + val + y_offset, + f"{val:.3f}", + ha="center", + va="center", + fontsize=9, + color="black", + fontweight="bold", + ) + + if N <= 10: + ax_evec.set_xticks(x_pos) + ax_evec.set_xticklabels([f"{i}" for i in x_pos]) + else: + ax_evec.set_xticks([1, N // 2, N]) + ax_evec.set_xticklabels(["1", f"{N // 2}", f"{N}"]) + + ax_evec.set_xlabel("Ion Index (j)") + ax_evec.grid(True, linestyle=":", alpha=0.6) + + # --- Lamb-Dicke Calculation --- + if active_freq > 1e-6: + omega_m = 2 * np.pi * active_freq * 1e6 + zero_point_spread = np.sqrt( + scipy_cst.hbar / (2 * MASS_ION * omega_m) + ) + eta_vec = vec * DELTA_K * zero_point_spread + else: + eta_vec = np.zeros_like(vec) + + # --- Rendering Lamb-Dicke Table --- + ax_eta.set_title( + r"Lamb-Dicke $\eta_m^{(j)}$", fontweight="bold", fontsize=11 + ) + col_labels = ["Ion", r"$b_m^{(j)}$", r"$\eta_m^{(j)}$"] + cell_text = [] + + display_N = min(N, 30) + for j in range(display_N): + cell_text.append([f"{j + 1}", f"{vec[j]:.3f}", f"{eta_vec[j]:.4f}"]) + + if N > 30: + cell_text.append(["...", "...", "..."]) + + table_eta = ax_eta.table( + cellText=cell_text, + colLabels=col_labels, + loc="center", + cellLoc="center", + ) + table_eta.auto_set_font_size(False) + + # Dynamically shrink table row height slightly as N grows to prevent overflow + scale_factor = 1.0 if N <= 20 else 0.85 + table_eta.scale(1, scale_factor) + table_eta.set_fontsize(7) + + return list(lines.values()) + list(highlights.values()) + [status_text] + + except Exception as e: + print(f"Update Error: {e}") + status_text.set_text(f"STATUS: ERROR ({e})") + status_text.set_color("red") + return list(lines.values()) + list(highlights.values()) + [status_text] + + # ========================================== + # 5. Execute + # ========================================== + ani = animation.FuncAnimation( + fig, animate, init_func=init, frames=1000, interval=100, blit=False + ) + plt.show()