import numpy as np
from albums.robinson import RobinsonModes
from tqdm import tqdm
from mpi4py import MPI
from albums.saveload import save_out, save_hdf5, load_out, load_hdf5, load_out_opti, load_out_opti
from albums.plot_func import __plot_modes, __plot_opti, __plot
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
# Helper function to initialize arrays
[docs]
def initialize_arrays(var1, var2, mode_coupling):
N = 2 if mode_coupling else 4
shape_2d = (len(var1), len(var2))
shape_3d = (len(var1), len(var2), N)
return {
"zero_freq_coup": np.zeros(shape_2d, dtype=bool),
"robinson_coup": np.zeros((*shape_2d, 4), dtype=bool),
"modes_coup": np.zeros(shape_3d, dtype=float),
"HOM_coup": np.zeros(shape_2d, dtype=bool),
"xi": np.zeros(shape_2d, dtype=float),
"converged_coup": np.zeros(shape_3d, dtype=bool),
"PTBL_coup": np.zeros(shape_2d, dtype=bool),
"bl": np.zeros(shape_2d, dtype=float),
"R": np.zeros(shape_2d, dtype=float),
}
# Generalized scan function
[docs]
def __scan(func, var1, var2, **kwargs):
if size not in (1, len(var1)):
raise ValueError("len(var1) != size")
mode_coupling = kwargs['other_kwargs'].get("mode_coupling", True)
skip = kwargs.get("skip", False)
var1_local = np.array_split(var1, size)[rank]
arrays_local = initialize_arrays(var1_local, var2, mode_coupling)
for i, v1 in enumerate(var1_local):
var_skip = False
for j, v2 in enumerate(tqdm(var2, desc=f'rank={rank}', position=rank)):
results = func(v1, v2, **kwargs)
bunch_length, zero_frequency, robinson, HOM, Omega, PTBL, converged, xi, R = results
if not np.asarray(converged).any():
if skip:
if var_skip:
break
var_skip = True
continue
var_skip = False
arrays_local["zero_freq_coup"][i, j] = zero_frequency
arrays_local["robinson_coup"][i, j, :] = robinson
arrays_local["modes_coup"][i, j, :] = Omega
arrays_local["HOM_coup"][i, j] = HOM
arrays_local["converged_coup"][i, j, :] = converged
arrays_local["PTBL_coup"][i, j] = PTBL
arrays_local["bl"][i, j] = bunch_length * 1e12
arrays_local["xi"][i, j] = xi
arrays_local["R"][i, j] = R
results = {key: None for key in arrays_local} if rank != 0 else initialize_arrays(var1, var2, mode_coupling)
for key, local_array in arrays_local.items():
comm.Gather(local_array, results[key], root=0)
return (results['zero_freq_coup'], results['robinson_coup'],
results['modes_coup'], results['HOM_coup'], results['converged_coup'],
results['PTBL_coup'], results['bl'], results['xi'], results["R"])
# Scan in one dimension
[docs]
def __scan_1D(MC, HC, ring, psi_HC_vals, current, mode_coupling, tau_boundary, method, **kwargs):
arrays = initialize_arrays(psi_HC_vals, [None], mode_coupling)
for j, psi_HC in enumerate(tqdm(psi_HC_vals)):
HC.psi = psi_HC * np.pi / 180
solver = RobinsonModes(ring, [MC, HC], current, tau_boundary=tau_boundary)
results = solver.solve(method=method, mode_coupling=mode_coupling, **kwargs)
(
bunch_length, zero_frequency, robinson, HOM, Omega, PTBL,
converged
) = results
if not converged.any():
continue
arrays["zero_freq_coup"][j] = zero_frequency
arrays["robinson_coup"][j, :] = robinson
arrays["modes_coup"][j, :] = Omega
arrays["HOM_coup"][j] = HOM
arrays["converged_coup"][j, :] = converged
arrays["PTBL_coup"][j] = PTBL
arrays["bl"][j] = bunch_length * 1e12
arrays["xi"][j] = solver.xi
arrays["R"][j] = solver.R_factor(method)
return tuple(arrays.values())
#%% follow modes at given current
[docs]
def scan_modes(MC,
HC,
ring,
psi_HC_vals,
current,
mode_coupling,
tau_boundary,
method,
**kwargs):
mode_coupling_states = [mode_coupling] if mode_coupling is not None else [False, True]
for mc in mode_coupling_states:
results = __scan_1D(MC, HC, ring, psi_HC_vals, current, mc, tau_boundary, method, **kwargs)
__plot_modes(results, psi_HC_vals, mc)
[docs]
def __scan_opti(func,
var1,
var2,
**kwargs):
if size not in (1, len(var1)):
raise ValueError("len(var1) != size")
shape = (len(var1), len(var2))
results = {key: np.zeros(shape, dtype=float) for key in ["psi", "bl", "xi", "R"]}
local_var1 = np.array_split(var1, size)[rank]
local_results = {key: np.zeros((len(local_var1), len(var2)), dtype=float) for key in results}
for i, v1 in enumerate(local_var1):
for j, v2 in enumerate(tqdm(var2, desc=f"rank={rank}", position=rank)):
psi, bunch_length, xi_val, R = func(v1, v2, **kwargs)
local_results["psi"][i, j] = psi
local_results["bl"][i, j] = bunch_length * 1e12
local_results["xi"][i, j] = xi_val
local_results["R"][i, j] = R
for key in results:
comm.Allgather(local_results[key], results[key])
comm.Barrier()
return tuple(results.values())
[docs]
def __scan_after_opti(func,
psi,
var1,
var2,
psi_add,
**kwargs):
if size not in (1, len(var1)):
raise ValueError("len(var1) != size")
mode_coupling = kwargs.get("other_kwargs", {}).get("mode_coupling", True)
N = 2 if mode_coupling else 4
var1_local = np.array_split(var1, size)[rank]
local_arrays = initialize_arrays(var1_local, var2, mode_coupling)
for i, v1 in enumerate(var1_local):
for j, v2 in enumerate(tqdm(var2, desc=f'rank={rank}', position=rank)):
idx = rank if size > 1 else i
results = func(psi[idx, j] + psi_add, v1, v2, **kwargs)
(
bunch_length, zero_frequency, robinson, HOM, Omega, PTBL,
converged, xi, R
) = results
if not converged.any():
continue
local_arrays["zero_freq_coup"][i, j] = zero_frequency
local_arrays["robinson_coup"][i, j, :] = robinson
local_arrays["modes_coup"][i, j, :] = Omega
local_arrays["HOM_coup"][i, j] = HOM
local_arrays["converged_coup"][i, j, :] = converged
local_arrays["PTBL_coup"][i, j] = PTBL
local_arrays["bl"][i, j] = bunch_length * 1e12
local_arrays["xi"][i, j] = xi
local_arrays["R"][i, j] = R
gathered_results = initialize_arrays(var1, var2, mode_coupling) if rank == 0 else {
key: None for key in local_arrays
}
for key in local_arrays:
comm.Gather(local_arrays[key], gathered_results[key], root=0)
return tuple(gathered_results.values())
[docs]
def __psi_I0(psi, I0, **kwargs):
kwargs["HC"].psi = psi * np.pi / 180
solver = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], I0, tau_boundary=kwargs["tau_boundary"])
results = solver.solve(method=kwargs["method"], **kwargs["other_kwargs"])
return (*results, solver.xi, solver.R_factor(kwargs["method"]))
# Scan psi and current
[docs]
def scan_psi_I0(name, MC, HC, ring, psi_HC_vals, currents, method, tau_boundary=None, save=True, **other_kwargs):
kwargs = {
"MC": MC,
"HC": HC,
"tau_boundary": tau_boundary,
"method": method,
"ring": ring,
"name": name,
"other_kwargs": other_kwargs
}
output = __scan(__psi_I0, psi_HC_vals, currents, **kwargs)
if rank == 0:
if save:
save_out(f"{name}_{method}", output)
plot_type = other_kwargs.get("plot_type", False)
if not plot_type:
for param in ["xi", "bunch_length", "R"]:
__plot(output, psi_HC_vals, currents, 1, 1e3, 'Tuning angle [°]', 'Current [mA]', param, save, **kwargs)
else:
__plot(output, psi_HC_vals, currents, 1, 1e3, 'Tuning angle [°]', 'Current [mA]', plot_type, save, **kwargs)
# Function for psi and R/Q scan
[docs]
def __psi_RoQ(psi, RoQ, **kwargs):
kwargs["HC"].Rs = RoQ * kwargs["HC"].Q
kwargs["HC"].psi = psi * np.pi / 180
solver = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
results = solver.solve(method=kwargs["method"], **kwargs["other_kwargs"])
return (*results, solver.xi, solver.R_factor(kwargs["method"]))
# Scan psi and R/Q
[docs]
def scan_psi_RoQ(name, MC, HC, ring, psi_HC_vals, RoQ_vals, I0, method, tau_boundary=None, save=True, **other_kwargs):
kwargs = {
"MC": MC,
"HC": HC,
"tau_boundary": tau_boundary,
"method": method,
"ring": ring,
"name": f'{name}_RoQ_{np.floor(I0*1e3)}mA',
"I0": I0,
"other_kwargs": other_kwargs
}
output = __scan(__psi_RoQ, psi_HC_vals, RoQ_vals, **kwargs)
if rank == 0:
if save:
save_out(name + "_" + method, output)
for param in ["xi", "bunch_length", "R"]:
__plot(output, psi_HC_vals, RoQ_vals, 1, 1, 'Tuning angle [°]', 'R/Q [ohm]', param, save, **kwargs)
#%% opti RoQ/I
[docs]
def __opti_I0_RoQ(I0, RoQ, **kwargs):
kwargs["HC"].Rs = RoQ*kwargs["HC"].Q
(psi, bunch_length, xi_val, R) = __get_vals(kwargs["ring"],
kwargs["MC"],
kwargs["HC"],
I0,
kwargs["psi0_HC"],
kwargs["tau_boundary"],
kwargs["method"],
kwargs["bounds"],
**kwargs["other_kwargs"])
return (psi, bunch_length, xi_val, R)
[docs]
def scan_RoQ_I0(name,
MC,
HC,
ring,
current_vals,
RoQ_vals,
psi0_HC,
bounds,
tau_boundary,
method,
plot_2D=None,
save=True,
scan_after_opti=True,
psi_add_after_opti=-0.1,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_RoQ_I0',
"psi0_HC":psi0_HC,
"bounds":bounds,
"other_kwargs":other_kwargs}
out = __scan_opti(__opti_I0_RoQ,
current_vals,
RoQ_vals,
**kwargs)
if rank == 0:
if save:
save_out_opti(name + "_" + method, out)
if plot_2D is None:
for _plot_2D in ['xi', 'bunch_length', 'psi', 'R']:
__plot_opti(out,
current_vals,
RoQ_vals,
1e3,
1,
'Current [mA]',
'R/Q [ohm]',
_plot_2D,
save,
**kwargs)
else:
__plot_opti(out,
current_vals,
RoQ_vals,
1e3,
1,
'Current [mA]',
'R/Q [ohm]',
plot_2D,
save,
**kwargs)
comm.Barrier()
if scan_after_opti:
out2 = __scan_after_opti(__psi_I0_RoQ,
out[0],
current_vals,
RoQ_vals,
psi_add=psi_add_after_opti,
**kwargs)
if rank == 0:
if save:
name = name +"_" + "after_opti"
kwargs["name"] = name
save_out(name + "_" + method, out2)
if plot_2D is None:
for _plot_2D in ['xi', 'bunch_length', 'psi', 'R']:
__plot(out2,
current_vals,
RoQ_vals,
1e3,
1,
'Current [mA]',
'R/Q [ohm]',
_plot_2D,
save,
**kwargs)
else:
__plot(out2,
current_vals,
RoQ_vals,
1e3,
1,
'Current [mA]',
'R/Q [ohm]',
plot_2D,
save,
**kwargs)
#%% opti RoQ/I0 equilibrium
[docs]
def __opti_I0_RoQ_equilirium(I0, RoQ, **kwargs):
kwargs["HC"].Rs = RoQ*kwargs["HC"].Q
(psi, bunch_length, xi_val, R) = __get_vals_equilibrium(kwargs["ring"],
kwargs["MC"],
kwargs["HC"],
I0,
kwargs["psi0_HC"],
kwargs["tau_boundary"],
kwargs["method"],
kwargs["bounds"],
**kwargs["other_kwargs"])
return (psi, bunch_length, xi_val, R)
[docs]
def scan_RoQ_I0_equilirium(name,
MC,
HC,
ring,
current_vals,
RoQ_vals,
psi0_HC,
bounds,
tau_boundary,
method,
plot_2D=None,
save=True,
# scan_after_opti=True,
# psi_add_after_opti=-0.1,
**other_kwargs):
# for passive HC ! Set Q0=QL
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_RoQ_I0',
"psi0_HC":psi0_HC,
"bounds":bounds,
"other_kwargs":other_kwargs}
out = __scan_opti(__opti_I0_RoQ_equilirium,
current_vals,
RoQ_vals,
**kwargs)
if rank == 0:
if save:
save_out_opti(name + "_" + method, out)
if plot_2D is None:
for _plot_2D in ['xi', 'bunch_length', 'psi', 'R']:
__plot_opti(out,
current_vals,
RoQ_vals,
1e3,
1,
'Current [mA]',
'R/Q [ohm]',
_plot_2D,
save,
**kwargs)
else:
__plot_opti(out,
current_vals,
RoQ_vals,
1e3,
1,
'Current [mA]',
'R/Q [ohm]',
plot_2D,
save,
**kwargs)
#%% opti RoQ/Q0
[docs]
def __opti_Q0_RoQ(Q0, RoQ, **kwargs):
kwargs["HC"].Q = Q0
kwargs["HC"].QL = Q0
kwargs["HC"].Rs = RoQ*kwargs["HC"].Q
(psi, bunch_length, xi_val, R) = __get_vals(kwargs["ring"],
kwargs["MC"],
kwargs["HC"],
kwargs["I0"],
kwargs["psi0_HC"],
kwargs["tau_boundary"],
kwargs["method"],
kwargs["bounds"],
**kwargs["other_kwargs"])
return (psi, bunch_length, xi_val, R)
[docs]
def scan_RoQ_Q0(name,
MC,
HC,
ring,
Q0_vals,
RoQ_vals,
I0,
psi0_HC,
bounds,
tau_boundary,
method,
plot_2D=None,
save=True,
scan_after_opti=True,
psi_add_after_opti=-0.1,
**other_kwargs):
# for passive HC ! Set Q0=QL
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_RoQ_Q0_{int(I0*1e3)}mA',
"psi0_HC":psi0_HC,
"bounds":bounds,
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan_opti(__opti_Q0_RoQ,
Q0_vals,
RoQ_vals,
**kwargs)
if rank == 0:
if save:
save_out_opti(name + "_" + method, out)
if plot_2D is None:
for _plot_2D in ['xi', 'bunch_length', 'psi', 'R']:
__plot_opti(out,
Q0_vals,
RoQ_vals,
1,
1,
'Harmonic cavity Q0',
'R/Q [ohm]',
_plot_2D,
save,
**kwargs)
else:
__plot_opti(out,
Q0_vals,
RoQ_vals,
1,
1,
'Harmonic cavity Q0',
'R/Q [ohm]',
plot_2D,
save,
**kwargs)
comm.Barrier()
if scan_after_opti:
out2 = __scan_after_opti(__psi_Q0_RoQ,
out[0],
Q0_vals,
RoQ_vals,
psi_add=psi_add_after_opti,
**kwargs)
if rank == 0:
if save:
name = name +"_" + "after_opti"
kwargs["name"] = name
save_out(name + "_" + method, out2)
if plot_2D is None:
for _plot_2D in ['xi', 'bunch_length', 'psi', 'R']:
__plot(out2,
Q0_vals,
RoQ_vals,
1,
1,
'Harmonic cavity Q0',
'R/Q [ohm]',
_plot_2D,
save,
**kwargs)
else:
__plot(out2,
Q0_vals,
RoQ_vals,
1,
1,
'Harmonic cavity Q0',
'R/Q [ohm]',
plot_2D,
save,
**kwargs)
#%% opti RoQ/Q0 equilibrium
[docs]
def __opti_Q0_RoQ_equilirium(Q0, RoQ, **kwargs):
kwargs["HC"].Q = Q0
kwargs["HC"].QL = Q0
kwargs["HC"].Rs = RoQ*kwargs["HC"].Q
(psi, bunch_length, xi_val, R) = __get_vals_equilibrium(kwargs["ring"],
kwargs["MC"],
kwargs["HC"],
kwargs["I0"],
kwargs["psi0_HC"],
kwargs["tau_boundary"],
kwargs["method"],
kwargs["bounds"],
**kwargs["other_kwargs"])
return (psi, bunch_length, xi_val, R)
[docs]
def scan_RoQ_Q0_equilirium(name,
MC,
HC,
ring,
Q0_vals,
RoQ_vals,
I0,
psi0_HC,
bounds,
tau_boundary,
method,
plot_2D=None,
save=True,
# scan_after_opti=True,
# psi_add_after_opti=-0.1,
**other_kwargs):
# for passive HC ! Set Q0=QL
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_RoQ_Q0_{int(I0*1e3)}mA',
"psi0_HC":psi0_HC,
"bounds":bounds,
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan_opti(__opti_Q0_RoQ_equilirium,
Q0_vals,
RoQ_vals,
**kwargs)
if rank == 0:
if save:
save_out_opti(name + "_" + method, out)
if plot_2D is None:
for _plot_2D in ['xi', 'bunch_length', 'psi', 'R']:
__plot_opti(out,
Q0_vals,
RoQ_vals,
1,
1,
'Harmonic cavity Q0',
'R/Q [ohm]',
_plot_2D,
save,
**kwargs)
else:
__plot_opti(out,
Q0_vals,
RoQ_vals,
1,
1,
'Harmonic cavity Q0',
'R/Q [ohm]',
plot_2D,
save,
**kwargs)
#%% generetic find intability after opti
[docs]
def __psi_I0_RoQ(psi, I0, RoQ, **kwargs):
kwargs["HC"].Rs = RoQ * kwargs["HC"].Q
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], I0, tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def __psi_Q0_RoQ(psi, Q0, RoQ, **kwargs):
kwargs["HC"].Q = Q0
kwargs["HC"].QL = Q0
kwargs["HC"].Rs = RoQ * Q0
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
#%% scan_psi_MCQL
[docs]
def __psi_QL(psi, QL, **kwargs):
kwargs["MC"].QL = QL
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_psi_QL(name,
MC,
HC,
ring,
psi_HC_vals,
QL_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_QL_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__psi_QL,
psi_HC_vals,
QL_vals,
**kwargs)
if rank == 0:
if save:
save_out(name + "_" + method, out)
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
psi_HC_vals,
QL_vals,
1,
1,
'Tuning angle [°]',
'Main cavity QL',
_plot_2D,
save,
**kwargs)
#%% scan_psi_MCdet
[docs]
def __psi_MCdet(psi, det, **kwargs):
kwargs["MC"].detune = det
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve_passive(method=kwargs["method"], optimal_tunning=False, **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_psi_MCdet(name,
MC,
HC,
ring,
psi_HC_vals,
MCdet_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_MCdet_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__psi_MCdet,
psi_HC_vals,
MCdet_vals,
**kwargs)
if rank == 0:
if save:
save_out(name + "_" + method, out)
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
psi_HC_vals,
MCdet_vals,
1,
1e-3,
'Tuning angle [°]',
'Main cavity detuning [kHz]',
_plot_2D,
save,
**kwargs)
#%% scan_psi_MC_Rs
[docs]
def __psi_MCRs(psi, Rs, **kwargs):
kwargs["MC"].Rs_per_cavity = Rs
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_psi_MC_Rs(name,
MC,
HC,
ring,
psi_HC_vals,
MC_Rs_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_MC_Rs_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__psi_MCRs,
psi_HC_vals,
MC_Rs_vals,
**kwargs)
if rank == 0:
if save:
save_out(name + "_" + method, out)
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
psi_HC_vals,
MC_Rs_vals,
1,
1e-6,
'Tuning angle [°]',
'Main cavity Rs per cavity [MOhm]',
_plot_2D,
save,
**kwargs)
#%% scan_psi_MC_beta
[docs]
def __psi_MC_beta(psi, beta, **kwargs):
kwargs["MC"].beta = beta
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_psi_MC_beta(name,
MC,
HC,
ring,
psi_HC_vals,
MC_beta_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_MC_beta_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__psi_MC_beta,
psi_HC_vals,
MC_beta_vals,
**kwargs)
if rank == 0:
if save:
save_out(name + "_" + method, out)
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
psi_HC_vals,
MC_beta_vals,
1,
1,
'Tuning angle [°]',
'Main cavity beta',
_plot_2D,
save,
**kwargs)
#%% scan_psi_HC_Q0
[docs]
def __psi_HC_Q0(psi, Q0, **kwargs):
kwargs["HC"].Q = Q0
kwargs["HC"].QL = Q0
kwargs["HC"].Rs = Q0 * kwargs["RoQ"]
kwargs["HC"].psi = psi*np.pi/180
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_psi_HC_Q0(name,
MC,
HC,
ring,
psi_HC_vals,
HC_Q0_vals,
RoQ,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
# for passive HC ! Set Q0=QL
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_HC_Q0_{np.floor(I0*1e3)}mA',
"I0":I0,
"RoQ":RoQ,
"other_kwargs":other_kwargs}
out = __scan(__psi_HC_Q0,
psi_HC_vals,
HC_Q0_vals,
**kwargs)
if rank == 0:
if save:
save_out(name + "_" + method, out)
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
psi_HC_vals,
HC_Q0_vals,
1,
1,
'Tuning angle [°]',
'Harmonic cavity Q0',
_plot_2D,
save,
**kwargs)
#%% scan_MC_Vc_HC_Vc_active
[docs]
def __MC_Vc_HC_Vc_active(MC_Vc, HC_Vc, **kwargs):
kwargs["MC"].Vc = MC_Vc
kwargs["MC"].theta = np.arccos(kwargs["ring"].U0/kwargs["MC"].Vc)
kwargs["MC"].set_optimal_detune(kwargs["I0"])
# kwargs["MC"].psi = kwargs["MC"].psi - 30/180*np.pi
kwargs["MC"].set_generator(kwargs["I0"])
kwargs["HC"].Vc = HC_Vc
kwargs["HC"].theta = -np.pi/2
kwargs["HC"].set_optimal_detune(kwargs["I0"])
# kwargs["HC"].psi = kwargs["HC"].psi - 30/180*np.pi
kwargs["HC"].set_generator(kwargs["I0"])
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_MC_Vc_HC_Vc_active(name,
MC,
HC,
ring,
MC_Vc_vals,
HC_Vc_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_MC_Vc_HC_Vc_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__MC_Vc_HC_Vc_active,
MC_Vc_vals,
HC_Vc_vals,
**kwargs)
if rank == 0:
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
MC_Vc_vals,
HC_Vc_vals,
1e-6,
1e-3,
'MC voltage [MV]',
'HC voltage [kV]',
_plot_2D,
save,
**kwargs)
#%% scan_MC_Vc_MC_psi_active
[docs]
def __MC_Vc_MC_psi_active(MC_Vc, MC_psi, **kwargs):
kwargs["MC"].Vc = MC_Vc
kwargs["MC"].theta = np.arccos(kwargs["ring"].U0/kwargs["MC"].Vc)
kwargs["MC"].set_optimal_detune(kwargs["I0"])
kwargs["MC"].psi = kwargs["MC"].psi + MC_psi/180*np.pi
kwargs["MC"].set_generator(kwargs["I0"])
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_MC_Vc_MC_psi_active(name,
MC,
HC,
ring,
MC_Vc_vals,
MC_psi_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_MC_Vc_MC_psi_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__MC_Vc_MC_psi_active,
MC_Vc_vals,
MC_psi_vals,
**kwargs)
if rank == 0:
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
MC_Vc_vals,
MC_psi_vals,
1e-6,
1,
'MC voltage [MV]',
'MC delta psi from optimal [°]',
_plot_2D,
save,
**kwargs)
#%% scan_MC_psi_HC_psi_active
[docs]
def __MC_psi_HC_psi_active(MC_psi, HC_psi, **kwargs):
kwargs["MC"].set_optimal_detune(kwargs["I0"])
kwargs["MC"].psi = kwargs["MC"].psi + MC_psi/180*np.pi
kwargs["MC"].set_generator(kwargs["I0"])
kwargs["HC"].set_optimal_detune(kwargs["I0"])
kwargs["HC"].psi = kwargs["HC"].psi + HC_psi/180*np.pi
kwargs["HC"].set_generator(kwargs["I0"])
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], kwargs["I0"], tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_MC_psi_HC_psi_active(name,
MC,
HC,
ring,
MC_psi_vals,
HC_psi_vals,
I0,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_MC_psi_HC_psi_{np.floor(I0*1e3)}mA',
"I0":I0,
"other_kwargs":other_kwargs}
out = __scan(__MC_psi_HC_psi_active,
MC_psi_vals,
HC_psi_vals,
**kwargs)
if rank == 0:
for _plot_2D in ['xi', 'bunch_length', 'R']:
__plot(out,
MC_psi_vals,
HC_psi_vals,
1,
1,
'MC delta psi from optimal [°]',
'HC delta psi from optimal [°]',
_plot_2D,
save,
**kwargs)
#%% scan_HC_beta_I0_active
[docs]
def __HC_beta_I0_active(HC_beta, I0, **kwargs):
kwargs["MC"].set_optimal_detune(I0)
kwargs["MC"].set_generator(I0)
kwargs["HC"].beta = HC_beta
kwargs["HC"].set_optimal_detune(I0)
kwargs["HC"].set_generator(I0)
B = RobinsonModes(kwargs["ring"], [kwargs["MC"], kwargs["HC"]], I0, tau_boundary=kwargs["tau_boundary"])
out = B.solve(method=kwargs["method"], **kwargs["other_kwargs"])
out = (*out, B.xi, B.R_factor(kwargs["method"]))
return out
[docs]
def scan_HC_beta_I0_active(name,
MC,
HC,
ring,
HC_beta_vals,
I0_vals,
method,
tau_boundary=None,
save=True,
**other_kwargs):
kwargs = {"MC":MC,
"HC":HC,
"tau_boundary":tau_boundary,
"method":method,
"ring":ring,
"name":f'{name}_HC_beta_I0',
"other_kwargs":other_kwargs}
out = __scan(__HC_beta_I0_active,
HC_beta_vals,
I0_vals,
**kwargs)
if rank == 0:
__plot(out,
HC_beta_vals,
I0_vals,
1,
1e3,
'HC beta',
'Current [mA]',
"xi",
save,
**kwargs)