Source code for albums.optimiser

from scipy.optimize import Bounds, minimize
import numpy as np
# from robinsonmodes.robinson import RobinsonModes
from albums.robinson import RobinsonModes
# Global counter for evaluations
eval_num = 0

[docs] def evaluate_R(B, method, is_equilibrium=False): """ Evaluate the R-factor or equilibrium solution. Parameters: - B: RobinsonModes instance - method: The method used for solving - is_equilibrium: Boolean indicating if equilibrium-only evaluation is performed Returns: - Evaluation score (negative R-factor if successful, penalized score otherwise) """ try: if is_equilibrium: _, R, _, converged = B.solve_equilibrium_only(method=method) if not converged: return 10 return -R if not np.isnan(R) else 10 else: out = B.solve(method=method) _, _, _, _, _, PTBL, converged = out if not converged.any() or PTBL: return 10 return -B.R_factor(method) except: return 10
[docs] def optimize_R(ring, MC, HC, I0, psi0_HC, tau_boundary, bounds, is_equilibrium=False, **kwargs): """ General optimization for the R-factor or equilibrium. Parameters: - ring: Ring configuration - MC: Main cavity parameters - HC: Harmonic cavity parameters - I0: Beam current - psi0_HC: Initial guess for psi_HC - tau_boundary: Boundary for tau - bounds: Bounds for psi_HC - is_equilibrium: Boolean flag for equilibrium-only optimization - kwargs: Additional configuration for the optimizer Returns: - Optimal value of psi_HC or a fallback value (90 on failure) """ method_opti = kwargs.get("method_opti", "COBYLA") tol_opti = kwargs.get("tol_opti", 0.01) maxiter_opti = kwargs.get("maxiter_opti", 1000) rhobeg_opti = kwargs.get("rhobeg_opti", 0.1) method = kwargs.get("method", "default_method") def to_eval(psi_HC): global eval_num eval_num += 1 HC.psi = psi_HC * np.pi / 180 B = RobinsonModes(ring, [MC, HC], I0, tau_boundary=tau_boundary) return evaluate_R(B, method, is_equilibrium) bd = Bounds([bounds[0]], [bounds[1]]) res = minimize(fun=to_eval, x0=[psi0_HC], bounds=bd, method=method_opti, tol=tol_opti, options={"maxiter": maxiter_opti, "rhobeg": rhobeg_opti}) return res.x[0] if res.success else 90
[docs] def minimize_psi(ring, MC, HC, I0, psi0_HC, tau_boundary, bounds, **kwargs): """ Minimize the psi parameter for stability. Parameters: - ring: Ring configuration - MC: Main cavity parameters - HC: Harmonic cavity parameters - I0: Beam current - psi0_HC: Initial guess for psi_HC - tau_boundary: Boundary for tau - bounds: Bounds for psi_HC - kwargs: Additional configuration for the optimizer Returns: - Optimal value of psi_HC or np.nan on failure """ method = kwargs.get("method", "Nelder-Mead") tol = kwargs.get("tol", 0.05) def to_eval(psi_HC): HC.psi = psi_HC * np.pi / 180 B = RobinsonModes(ring, [MC, HC], I0, tau_boundary=tau_boundary) try: out = B.solve(method=method, **kwargs) _, zero_frequency, robinson, _, _, PTBL, converged = out if (not converged.any() or robinson.any() or PTBL or zero_frequency): return psi_HC + 100 except: return psi_HC + 100 return psi_HC bd = Bounds([bounds[0]], [bounds[1]]) res = minimize(fun=to_eval, x0=[psi0_HC], bounds=bd, method=method, tol=tol) return res.x[0] if res.success else np.nan
[docs] def __get_vals(ring, MC, HC, I0, psi0_HC, tau_boundary, method, bounds, **kwargs): debug = kwargs.get("debug", False) loop_option = kwargs.get("loop_option", False) add_psi_loop = kwargs.get("add_psi_loop", 0.02) auto_psi_input = kwargs.get("auto_psi_input", False) xi_init_input = kwargs.get("xi_init_input", 0.8) if auto_psi_input: xi = xi_init_input HC_det = I0*HC.Rs/HC.Q*ring.f1/MC.Vc*HC.m**2/xi HC_fr = HC_det + HC.m * ring.f1 HC_psi = np.arctan(HC.QL * (HC_fr / (HC.m * ring.f1) - (HC.m * ring.f1) / HC_fr)) * 180 / np.pi psi0_HC = HC_psi if debug: print(f"psi0 input = {psi0_HC}") psi = maximize_R(ring, MC, HC, I0, psi0_HC, tau_boundary, method, bounds, **kwargs) if debug: print(f"psi out of maximize_R = {psi}") if loop_option: count = 0 add_psi = add_psi_loop while(True): psi = (psi + count*add_psi) HC.psi = psi*np.pi/180 B = RobinsonModes(ring, [MC, HC], I0, tau_boundary=tau_boundary) out = B.solve(method=method, **kwargs) (bunch_length, zero_frequency, robinson, _, _, PTBL, converged) = out try: cond = not converged.any() is_stable = not (cond or robinson.any() or PTBL or zero_frequency) except AttributeError: cond = not converged is_stable = False if is_stable: if debug: print(f"psi sent to result = {psi}") print(out) return (psi, bunch_length, B.xi, B.R_factor(method)) else: if debug: print(f"unstable: {count}") print(f"psi unstable = {psi}") print(out) if count == 200: raise ValueError() count += 1 else: HC.psi = psi*np.pi/180 B = RobinsonModes(ring, [MC, HC], I0, tau_boundary=tau_boundary) out = B.solve(method=method, **kwargs) (bunch_length, zero_frequency, robinson, _, _, PTBL, converged) = out if debug: print(out) return (psi, bunch_length, B.xi, B.R_factor(method))
[docs] def __get_vals_equilibrium(ring, MC, HC, I0, psi0_HC, tau_boundary, method, bounds, **kwargs): debug = kwargs.get("debug", False) loop_option = kwargs.get("loop_option", False) add_psi_loop = kwargs.get("add_psi_loop", 0.02) auto_psi_input = kwargs.get("auto_psi_input", False) xi_init_input = kwargs.get("xi_init_input", 0.8) if auto_psi_input: xi = xi_init_input HC_det = I0*HC.Rs/HC.Q*ring.f1/MC.Vc*HC.m**2/xi HC_fr = HC_det + HC.m * ring.f1 HC_psi = np.arctan(HC.QL * (HC_fr / (HC.m * ring.f1) - (HC.m * ring.f1) / HC_fr)) * 180 / np.pi psi0_HC = HC_psi if debug: print(f"psi0 input = {psi0_HC}") psi = maximize_R_equilibrium(ring, MC, HC, I0, psi0_HC, tau_boundary, method, bounds, **kwargs) if debug: print(f"psi out of maximize_R = {psi}") if loop_option: count = 0 add_psi = add_psi_loop while(True): psi = (psi + count*add_psi) HC.psi = psi*np.pi/180 B = RobinsonModes(ring, [MC, HC], I0, tau_boundary=tau_boundary) out = B.solve_equilibrium_only(method=method, **kwargs) (bunch_length, _, xi, converged) = out is_stable = converged if is_stable: if debug: print(f"psi sent to result = {psi}") print(out) return (psi, bunch_length, B.xi, B.R_factor(method)) else: if debug: print(f"unstable: {count}") print(f"psi unstable = {psi}") print(out) if count == 200: raise ValueError() count += 1 else: HC.psi = psi*np.pi/180 B = RobinsonModes(ring, [MC, HC], I0, tau_boundary=tau_boundary) out = B.solve_equilibrium_only(method=method, **kwargs) (bunch_length, _, xi, converged) = out if debug: print(out) return (psi, bunch_length, B.xi, B.R_factor(method))