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, }