Spaces:
Running
Running
| from __future__ import annotations | |
| import math | |
| from typing import Dict | |
| import numpy as np | |
| import pandas as pd | |
| from .calibration import EPS, RFCalibrationParamsV2 | |
| from .utils import sigmoid, stable_softmax | |
| def calibrated_surrogate_v4_core( | |
| df: pd.DataFrame, | |
| params: RFCalibrationParamsV2, | |
| baseline_eplus_per_pulse=5.0e6, | |
| combined_jitter_fs=math.sqrt(50.0**2 + 1.2**2), | |
| steps=45, | |
| ) -> Dict[str, np.ndarray]: | |
| x0 = np.column_stack([ | |
| np.clip(df["Vc_MV"].values / 12.0, 0, 4), | |
| np.clip((combined_jitter_fs / (1e15 / df["freq_Hz"].values)), 0, 0.45), | |
| np.clip(0.12 + 0.015 * np.maximum(df["freq_GHz"].values - 1.3, 0), 0, 0.5), | |
| np.clip(0.08 + 0.02 * np.maximum(df["freq_GHz"].values - 1.0, 0), 0, 0.8), | |
| np.clip((df["Vc_MV"].values / np.maximum(df["R_over_Q_ohm"].values, 1.0)) * 0.06, 0, 2), | |
| np.clip( | |
| np.log10(np.maximum(df.get("P_wall_peak_kW", pd.Series(1.0, index=df.index)).values, 1e-6) + 1.0) / 6.0, | |
| 0, | |
| 1.2, | |
| ), | |
| ]) | |
| u_scalar = np.clip(df["Vc_MV"].values / 10.0, 0, 3) | |
| u_vec = np.column_stack([u_scalar, 0 * u_scalar, 0 * u_scalar, 0.4 * u_scalar, 0.2 * u_scalar, 0.2 * u_scalar]) | |
| wdiag = np.array([1.6, 1.2, 1.0, 0.9, 1.4, 1.1], dtype=float) | |
| x = x0.copy() | |
| z = np.clip(0.1 + 0.25 * x[:, 5] + 0.1 * x[:, 4], 0, 1.5) | |
| p_damp = np.zeros_like(x) | |
| p_drive = np.zeros_like(x) | |
| p_coup = np.zeros_like(x) | |
| p_growth = np.zeros_like(x) | |
| e_prev = np.sum((x * wdiag) * x, axis=1) + 0.8 * z * z + EPS | |
| pph_acc = np.zeros(len(df)) | |
| h_last = np.zeros(len(df)) | |
| zdot_last = np.zeros(len(df)) | |
| eta_fill = np.where( | |
| df["type"].eq("SC_elliptical_1p3GHz"), 0.96, | |
| np.where(df["type"].eq("NC_pillbox_Lband"), 0.90, | |
| np.where(df["type"].eq("NC_Cband_TW"), 0.88, 0.84)) | |
| ) | |
| eta_ap = np.clip(0.96 - 0.015 * np.maximum(df["freq_GHz"].values - 1.3, 0), 0.72, 0.97) | |
| for _ in range(steps): | |
| e_eff = np.sum((x * wdiag) * x, axis=1) + 0.8 * z * z + EPS | |
| omega_d = 0.8 * z * z / (e_eff + EPS) | |
| alpha_b = params.sigma * omega_d | |
| h = 0.5 * (np.log(e_eff + EPS) - np.log(e_prev + EPS)) | |
| h_last = h | |
| f_damp = -0.14 * x | |
| f_drive = np.tanh(0.8 * x + 0.45 * u_vec) | |
| coup_vec = np.column_stack([0.7 * u_scalar, 0.25 * u_scalar, 0.3 * u_scalar, 0.25 * u_scalar, 0.5 * u_scalar, 0.2 * u_scalar]) | |
| f_coup = (alpha_b * zdot_last)[:, None] * coup_vec | |
| f_growth = (params.beta_h * h)[:, None] * x | |
| p_damp += (-p_damp + f_damp) / params.tau | |
| p_drive += (-p_drive + f_drive) / params.tau | |
| p_coup += (-p_coup + f_coup) / params.tau | |
| p_growth += (-p_growth + f_growth) / params.tau | |
| xdot = ( | |
| (f_damp - params.lambda_ * p_damp) | |
| + (f_drive - params.lambda_ * p_drive) | |
| + (f_coup - params.lambda_ * p_coup) | |
| + (f_growth - params.lambda_ * p_growth) | |
| ) | |
| x = np.clip(x + 0.18 * xdot + 0.02 * u_vec, -3.0, 3.0) | |
| zdot = -0.20 * z + 0.28 * x[:, 0] + 0.16 * x[:, 4] - 0.10 * x[:, 1] + 0.08 * x[:, 5] | |
| z = np.clip(z + 0.2 * zdot, -2.0, 2.0) | |
| zdot_last = zdot | |
| pph = np.sum(x * params.lambda_ * (p_damp + p_drive + p_coup + p_growth), axis=1) + 0.16 * params.lambda_ * z * z | |
| pph_acc += np.maximum(pph, 0.0) | |
| e_prev = e_eff | |
| ttf = np.clip(np.sinc(0.35 * df["freq_GHz"].values / 10.0), 0.78, 1.0) | |
| eta_phase = np.exp(-params.c_phase * (np.abs(x[:, 1]) + x[:, 2]) ** 2) | |
| eta_loss = np.exp(-params.c_loss * pph_acc / (steps * (e_prev + EPS))) | |
| eta_capture = sigmoid(params.c_cap * (x[:, 0] - 0.65 * x[:, 2] - 0.55 * x[:, 3] - 0.45 * np.abs(x[:, 5]) + 0.25)) | |
| eta_overlap = np.exp(-0.14 * np.maximum(df["freq_GHz"].values - 2.0, 0)) | |
| omega_eff = 2.0 * np.pi * df["freq_Hz"].values * (1.0 + params.k_det * z + params.k_bl * x[:, 4] + params.k_th * h_last) | |
| sigma_t_phys = (combined_jitter_fs * 1e-15) * (1.0 + 0.55 * x[:, 2]) | |
| delta_e = np.abs(x[:, 3]) | |
| v_eff = df["Vc_MV"].values * ttf * eta_fill * eta_phase * eta_ap * eta_loss | |
| rq_eff = df["R_over_Q_ohm"].values * eta_overlap | |
| ql_eff = df["Q_loaded"].values * np.exp(-0.08 * np.abs(z)) | |
| g_bunch_eff = ( | |
| 1.0 | |
| + params.c_g | |
| * np.power(np.maximum(v_eff, 1e-6), params.a_exp) | |
| * np.exp(-(omega_eff * sigma_t_phys) ** 2) | |
| * np.exp(-(delta_e ** 2) / (params.alphaE ** 2)) | |
| * eta_capture | |
| ) | |
| g_bunch_eff = np.clip(g_bunch_eff, 0.5, 8.0) | |
| vc_v = v_eff * 1e6 | |
| p_cav_eff = (vc_v ** 2) / (np.maximum(rq_eff, 1e-6) * np.maximum(ql_eff, 1e-6)) | |
| cryo_overhead = np.where(df["type"].eq("SC_elliptical_1p3GHz"), 700.0, 1.0) | |
| p_wall_peak_kW = 1.2 * p_cav_eff * cryo_overhead / 1000.0 | |
| p_floor = np.where( | |
| df["type"].eq("SC_elliptical_1p3GHz"), 8.0, | |
| np.where(df["type"].eq("NC_pillbox_Lband"), 4.0, | |
| np.where(df["type"].eq("NC_Cband_TW"), 9.0, 12.0)) | |
| ) | |
| p_wall_peak_kW = np.maximum(p_wall_peak_kW, p_floor) | |
| p_limit_kw = np.where(df["type"].eq("SC_elliptical_1p3GHz"), 24.0, 50.0) | |
| f_rf = np.where(p_wall_peak_kW <= p_limit_kw, 0.88, 0.88 * np.exp(-(p_wall_peak_kW / p_limit_kw - 1.0))) | |
| f_rf = np.clip(f_rf, 0.04, 0.92) | |
| delivered_base = params.y_scale * 0.35 * baseline_eplus_per_pulse * g_bunch_eff * eta_capture * eta_loss * f_rf | |
| brightness = delivered_base * np.sqrt(df["freq_GHz"].values) | |
| sdot_i = 0.14 * np.sum(x * x, axis=1) + 0.20 * z * z | |
| log_capture = 2.8 * eta_capture + 0.22 * v_eff - 1.2 * x[:, 2] - 1.0 * x[:, 3] | |
| log_detune = 2.6 * np.abs(params.k_det * z) + 1.4 * np.abs(x[:, 1]) + 0.25 * df["freq_GHz"].values | |
| log_beam = 2.2 * x[:, 4] + 0.8 * np.maximum(df["freq_GHz"].values - 1.3, 0) + 0.2 * v_eff | |
| log_thermal = 1.7 * x[:, 5] + 0.9 * np.log1p(p_wall_peak_kW) + 0.6 * np.abs(h_last) | |
| log_loss = 2.3 * (1.0 - eta_loss) + 1.0 * (1.0 - eta_phase) + 0.7 * (1.0 - f_rf) | |
| coeffs = stable_softmax(np.column_stack([log_capture, log_detune, log_beam, log_thermal, log_loss]), params.mode_temp) | |
| c_capture, c_detune, c_beamload, c_thermal, c_loss = coeffs.T | |
| regime = np.array(["capture", "detune", "beamload", "thermal", "loss"])[np.argmax(coeffs, axis=1)] | |
| delivered = delivered_base * ( | |
| 1.10 * c_capture + 0.93 * c_beamload + 0.86 * c_detune + 0.82 * c_thermal + 0.74 * c_loss | |
| ) | |
| return { | |
| "delivered": delivered, | |
| "brightness": brightness, | |
| "V_eff_MV": v_eff, | |
| "omega_eff_Hz": omega_eff / (2.0 * np.pi), | |
| "eta_phase": eta_phase, | |
| "eta_capture": eta_capture, | |
| "eta_loss": eta_loss, | |
| "f_RF": f_rf, | |
| "P_wall_peak_kW": p_wall_peak_kW, | |
| "Sdot_i": sdot_i, | |
| "z_latent": z, | |
| "c_capture": c_capture, | |
| "c_detune": c_detune, | |
| "c_beamload": c_beamload, | |
| "c_thermal": c_thermal, | |
| "c_loss": c_loss, | |
| "regime": regime, | |
| } | |