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))