def full_electroncloud_setup(line=None, ecloud_info=None, filenames=None, context=None,
tau_max=None, subtract_dipolar_kicks=True, shift_to_closed_orbit=True):
buffer = context.new_buffer()
fieldmaps = {
ecloud_type: get_electroncloud_fieldmap_from_h5(
filename=filename,
buffer=buffer,
tau_max=tau_max,
ecloud_name=ecloud_type) for (
ecloud_type,
filename) in filenames.items()}
for ecloud_type, fieldmap in fieldmaps.items():
print(f"Inserting \"{ecloud_type}\" electron clouds...")
insert_electronclouds(
ecloud_info[ecloud_type],
fieldmap=fieldmap,
line=line)
tracker = xt.Tracker(_context=context, line=line, _buffer=buffer)
twiss_without_ecloud = tracker.twiss()
config_electronclouds(
line,
twiss=twiss_without_ecloud,
ecloud_info=ecloud_info,
subtract_dipolar_kicks=subtract_dipolar_kicks,
shift_to_closed_orbit=shift_to_closed_orbit,
fieldmaps=fieldmaps,
ecloud_strength=1)
twiss_with_ecloud = tracker.twiss()
return tracker, twiss_without_ecloud, twiss_with_ecloud
def track_particle_xtrack(line,
partCO,
Dx_wrt_CO_m,
Dpx_wrt_CO_rad,
Dy_wrt_CO_m,
Dpy_wrt_CO_rad,
Dzeta_wrt_CO_m,
Ddelta_wrt_CO,
n_turns,
_context=None):
Dx_wrt_CO_m, Dpx_wrt_CO_rad,\
Dy_wrt_CO_m, Dpy_wrt_CO_rad,\
Dzeta_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords(
Dx_wrt_CO_m, Dpx_wrt_CO_rad,
Dy_wrt_CO_m, Dpy_wrt_CO_rad,
Dzeta_wrt_CO_m, Ddelta_wrt_CO)
tracker = xt.Tracker(_context=_context, line=line)
particles = xp.Particles(_context=_context,
p0c=partCO.p0c,
x=partCO.x + Dx_wrt_CO_m,
px=partCO.px + Dpx_wrt_CO_rad,
y=partCO.y + Dy_wrt_CO_m,
py=partCO.py + Dpy_wrt_CO_rad,
zeta=partCO.zeta + Dzeta_wrt_CO_m,
delta=partCO.delta + Ddelta_wrt_CO)
print('Start track')
tracker.track(particles, num_turns=n_turns, turn_by_turn_monitor=True)
print('Done track')
#print(res.particles[0])
x_tbt = tracker.record_last_track.x.copy().T
px_tbt = tracker.record_last_track.px.copy().T
y_tbt = tracker.record_last_track.y.copy().T
py_tbt = tracker.record_last_track.py.copy().T
zeta_tbt = tracker.record_last_track.zeta.copy().T
delta_tbt = tracker.record_last_track.delta.copy().T
print('Done loading!')
extra = {'particles': particles, 'tracker': tracker}
return x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra
import json
import numpy as np
import xtrack as xt
import xpart as xp
with open('../xsuite_lines/line_bb_for_tracking.json', 'r') as fid:
line_dict = json.load(fid)
line = xt.Line.from_dict(line_dict)
partCO = xp.Particles.from_dict(line_dict['particle_on_tracker_co'])
tracker = xt.Tracker(line=line)
particles = xp.Particles(**partCO.to_dict())
for _ in range(10):
print(particles.at_turn[0], particles.x[0], particles.y[0],
particles.zeta[0])
tracker.track(particles)
for nn in 'x px y py zeta delta'.split():
assert np.abs(getattr(particles, nn) - getattr(partCO, nn)) < 3e-10
WW = np.array(line_dict['WW_finite_diffs'])
WWinv = np.array(line_dict['WWInv_finite_diffs'])
assert np.max(np.abs(np.dot(WW, WWinv) - np.eye(6))) < 1e-10
ampl_sigmas = 0.2
norm_emit_x = 2.5e-6
partCO = xl.Particles.from_dict(dict_line_xtrack['particle_on_tracker_co'])
(x_tbt_sixtrack, px_tbt_sixtrack, y_tbt_sixtrack, py_tbt_sixtrack,
sigma_tbt_sixtrack, delta_tbt_sixtrack,
extra) = hp.track_particle_sixtrack(partCO=partCO,
Dx_wrt_CO_m=np.array(displace_x),
Dpx_wrt_CO_rad=0,
Dy_wrt_CO_m=np.array(displace_y),
Dpy_wrt_CO_rad=0.,
Dsigma_wrt_CO_m=0.,
Ddelta_wrt_CO=0.,
n_turns=num_turns,
input_folder='../')
tracker = xt.Tracker(sequence=line)
part_track = partCO.copy()
part_track.x += np.array(displace_x)
part_track.y += np.array(displace_y)
particles = xt.Particles(**part_track.to_dict())
particles.particle_id = np.arange(particles.num_particles)
tracker.track(particles, turn_by_turn_monitor=True, num_turns=num_turns)
print('Xtrack')
print(tracker.record_last_track.x)
print('Sixtrack')
print(x_tbt_sixtrack.T)
def install_spacecharge_frozen(line, particle_ref, longitudinal_profile,
nemitt_x, nemitt_y, sigma_z,
num_spacecharge_interactions,
tol_spacecharge_position):
line_no_sc = line
tracker_no_sc = xt.Tracker(line=line_no_sc)
# Make a matched bunch just to get the matched momentum spread
bunch = xp.generate_matched_gaussian_bunch(num_particles=int(2e6),
total_intensity_particles=1.,
nemitt_x=nemitt_x,
nemitt_y=nemitt_y,
sigma_z=sigma_z,
particle_ref=particle_ref,
tracker=tracker_no_sc)
delta_rms = np.std(bunch.delta)
# Remove all drifts
s_no_drifts = []
e_no_drifts = []
n_no_drifts = []
for ss, ee, nn in zip(line_no_sc.get_s_elements(), line_no_sc.elements,
line_no_sc.element_names):
if not _is_drift(ee):
assert not _is_thick(ee)
s_no_drifts.append(ss)
e_no_drifts.append(ee)
n_no_drifts.append(nn)
s_no_drifts = np.array(s_no_drifts)
# Generate spacecharge positions
s_spacecharge = np.linspace(0, line_no_sc.get_length(),
num_spacecharge_interactions + 1)[:-1]
# Adjust spacecharge positions where possible
for ii, ss in enumerate(s_spacecharge):
s_closest = np.argmin(np.abs(ss - s_no_drifts))
if np.abs(ss - s_closest) < tol_spacecharge_position:
s_spacecharge[ii] = s_closest
sc_lengths = 0 * s_spacecharge
sc_lengths[:-1] = np.diff(s_spacecharge)
sc_lengths[-1] = line_no_sc.get_length() - s_spacecharge[-1]
# Create spacecharge elements (dummy)
sc_elements = []
sc_names = []
for ii, ll in enumerate(sc_lengths):
sc_elements.append(
SpaceChargeBiGaussian(length=ll,
apply_z_kick=False,
longitudinal_profile=longitudinal_profile,
mean_x=0.,
mean_y=0.,
sigma_x=1.,
sigma_y=1.))
sc_names.append(f'spacecharge_{ii}')
# Merge lattice and spacecharge elements
df_lattice = pd.DataFrame({
's': s_no_drifts,
'elements': e_no_drifts,
'element_names': n_no_drifts
})
df_spacecharge = pd.DataFrame({
's': s_spacecharge,
'elements': sc_elements,
'element_names': sc_names
})
df_elements = pd.concat([df_lattice, df_spacecharge]).sort_values('s')
# Build new line with drifts
new_elements = []
new_names = []
s_curr = 0
i_drift = 0
for ss, ee, nn, in zip(df_elements['s'].values,
df_elements['elements'].values,
df_elements['element_names'].values):
if ss > s_curr + 1e-10:
new_elements.append(xt.Drift(length=(ss - s_curr)))
new_names.append(f'drift_{i_drift}')
s_curr = ss
i_drift += 1
new_elements.append(ee)
new_names.append(nn)
if s_curr < line_no_sc.get_length():
new_elements.append(xt.Drift(length=line_no_sc.get_length() - s_curr))
new_names.append(f'drift_{i_drift}')
line = xt.Line(elements=new_elements, element_names=new_names)
assert np.isclose(line.get_length(),
line_no_sc.get_length(),
rtol=0,
atol=1e-10)
# Twiss at spacecharge
line_sc_off = line.filter_elements(exclude_types_starting_with='SpaceCh')
tracker_sc_off = xt.Tracker(line=line_sc_off,
element_classes=tracker_no_sc.element_classes,
track_kernel=tracker_no_sc.track_kernel)
tw_at_sc = tracker_sc_off.twiss(particle_ref=particle_ref,
at_elements=sc_names)
# Configure lenses
for ii, sc in enumerate(sc_elements):
sc.mean_x = tw_at_sc['x'][ii]
sc.mean_y = tw_at_sc['y'][ii]
sc.sigma_x = np.sqrt(tw_at_sc['betx'][ii] * nemitt_x /
particle_ref.beta0 / particle_ref.gamma0 +
(tw_at_sc['dx'][ii] * delta_rms)**2)
sc.sigma_y = np.sqrt(tw_at_sc['bety'][ii] * nemitt_y /
particle_ref.beta0 / particle_ref.gamma0 +
(tw_at_sc['dy'][ii] * delta_rms)**2)
return line
def generate_xsuite_line(mad, seq_name, bb_df,
optics_and_co_at_start_ring_from_madx,
folder_name=None, skip_mad_use=False,
prepare_line_for_xtrack=True,
steps_for_finite_diffs={'dx': 1e-8, 'dpx': 1e-11,
'dy': 1e-8, 'dpy': 1e-11, 'dzeta': 1e-7, 'ddelta': 1e-8},
deferred_expressions=True):
# Build xsuite model
print('Start building xtrack line...')
line = xt.Line.from_madx_sequence(
mad.sequence[seq_name], apply_madx_errors=True,
deferred_expressions=deferred_expressions)
print('Done building xtrack.')
if bb_df is not None:
bb.setup_beam_beam_in_line(line, bb_df, bb_coupling=False)
# Temporary fix due to bug in mad loader
cavities, cav_names = line.get_elements_of_type(xt.Cavity)
for cc, nn in zip(cavities, cav_names):
if cc.frequency ==0.:
ii_mad = mad.sequence[seq_name].element_names().index(nn)
cc_mad = mad.sequence[seq_name].elements[ii_mad]
f0_mad = mad.sequence[seq_name].beam.freq0 * 1e6 # mad has it in MHz
cc.frequency = f0_mad*cc_mad.parent.harmon
line_bb_dipole_not_cancelled_dict = line.to_dict()
line_bb_dipole_not_cancelled_dict['particle_on_madx_co'] = (
optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])
line_bb_dipole_not_cancelled_dict['RR_madx'] = (
optics_and_co_at_start_ring_from_madx['RR_madx'])
if folder_name is not None:
os.makedirs(folder_name, exist_ok=True)
# Note that full separation and not strong beam position is present
# in bb lenses (for comparison with sixtrack input)
with open(folder_name + '/line_bb_dipole_not_cancelled.json', 'w') as fid:
json.dump(line_bb_dipole_not_cancelled_dict, fid, cls=JEncoder)
if prepare_line_for_xtrack:
tracker = xt.Tracker(line=line)
_disable_beam_beam(tracker.line)
particle_on_tracker_co = tracker.find_closed_orbit(
particle_co_guess=xp.Particles(
**optics_and_co_at_start_ring_from_madx['particle_on_madx_co']))
_restore_beam_beam(tracker.line)
xf.configure_orbit_dependent_parameters_for_bb(tracker,
particle_on_co=particle_on_tracker_co)
_disable_beam_beam(tracker.line)
RR_finite_diffs = tracker.compute_one_turn_matrix_finite_differences(
particle_on_tracker_co,
steps_r_matrix=steps_for_finite_diffs)
_restore_beam_beam(tracker.line)
(WW_finite_diffs, WWInv_finite_diffs, RotMat_finite_diffs
) = xp.compute_linear_normal_form(RR_finite_diffs)
line_bb_for_tracking_dict = line.to_dict()
line_bb_for_tracking_dict['particle_on_tracker_co'] = (
particle_on_tracker_co.to_dict())
line_bb_for_tracking_dict['RR_finite_diffs'] = RR_finite_diffs
line_bb_for_tracking_dict['WW_finite_diffs'] = WW_finite_diffs
line_bb_for_tracking_dict['WWInv_finite_diffs'] = WWInv_finite_diffs
line_bb_for_tracking_dict['RotMat_finite_diffs'] = RotMat_finite_diffs
if folder_name is not None:
os.makedirs(folder_name, exist_ok=True)
with open(folder_name +
'/line_bb_for_tracking.json', 'w') as fid:
json.dump(line_bb_for_tracking_dict, fid, cls=JEncoder)
return tracker, line_bb_for_tracking_dict
def generate_xline(
mad,
seq_name,
bb_df,
optics_and_co_at_start_ring_from_madx,
folder_name=None,
skip_mad_use=False,
prepare_line_for_xtrack=True,
steps_for_finite_diffs={
'dx': 1e-9,
'dpx': 1e-12,
'dy': 1e-9,
'dpy': 1e-12,
'dzeta': 1e-9,
'ddelta': 1e-9
}):
# Build xline model
print('Start building xline...')
line = xl.Line.from_madx_sequence(mad.sequence[seq_name],
apply_madx_errors=True)
print('Done building xline.')
if bb_df is not None:
bb.setup_beam_beam_in_line(line, bb_df, bb_coupling=False)
# Temporary fix due to bug in mad loader
cavities, cav_names = line.get_elements_of_type(xl.elements.Cavity)
for cc, nn in zip(cavities, cav_names):
if cc.frequency == 0.:
ii_mad = mad.sequence[seq_name].element_names().index(nn)
cc_mad = mad.sequence[seq_name].elements[ii_mad]
f0_mad = mad.sequence[
seq_name].beam.freq0 * 1e6 # mad has it in MHz
cc.frequency = f0_mad * cc_mad.parent.harmon
line_bb_dipole_not_cancelled_dict = line.to_dict(keepextra=True)
line_bb_dipole_not_cancelled_dict['particle_on_madx_co'] = (
optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])
line_bb_dipole_not_cancelled_dict['RR_madx'] = (
optics_and_co_at_start_ring_from_madx['RR_madx'])
if folder_name is not None:
os.makedirs(folder_name, exist_ok=True)
# Note that full separation and not strong beam position is present
# in bb lenses (for comparison with sixtrack input)
with open(folder_name + '/line_bb_dipole_not_cancelled.json',
'w') as fid:
json.dump(line_bb_dipole_not_cancelled_dict, fid, cls=JEncoder)
if prepare_line_for_xtrack:
tracker = xt.Tracker(sequence=line)
# Disable beam-beam
for ee in tracker.line.elements:
if ee.__class__.__name__.startswith('BeamBeam'):
ee._temp_q0 = ee.q0
ee.q0 = 0
particle_on_tracker_co = find_closed_orbit_from_tracker(
tracker,
optics_and_co_at_start_ring_from_madx['particle_on_madx_co'])
RR_finite_diffs = compute_R_matrix_finite_differences(
particle_on_tracker_co,
tracker,
symplectify=True,
**steps_for_finite_diffs)
(WW_finite_diffs, WWInv_finite_diffs,
RotMat_finite_diffs) = compute_linear_normal_form(RR_finite_diffs)
# (Re-activates bb in line and tracker)
_set_orbit_dependent_parameters_for_bb(line, tracker,
particle_on_tracker_co)
line_bb_for_tracking_dict = line.to_dict(keepextra=True)
line_bb_for_tracking_dict['particle_on_tracker_co'] = (
particle_on_tracker_co.to_dict())
line_bb_for_tracking_dict['RR_finite_diffs'] = RR_finite_diffs
line_bb_for_tracking_dict['WW_finite_diffs'] = WW_finite_diffs
line_bb_for_tracking_dict['WWInv_finite_diffs'] = WWInv_finite_diffs
line_bb_for_tracking_dict['RotMat_finite_diffs'] = RotMat_finite_diffs
if folder_name is not None:
os.makedirs(folder_name, exist_ok=True)
with open(folder_name + '/line_bb_for_tracking.json', 'w') as fid:
json.dump(line_bb_for_tracking_dict, fid, cls=JEncoder)
# # Switch off all beam-beam lenses
crabs, crab_names = ltest.get_elements_of_type([xt.RFMultipole])
#for cc in crabs:
# cc.pn = [-90]
# cc.ps = [-90]
# for cc in crabs:
# cc.knl[0] *= -1
with open('../optics_orbit_at_start_ring_from_madx.json', 'r') as fid:
ddd = json.load(fid)
partco = xp.Particles.from_dict(ddd['particle_on_madx_co'])
z_slices = s_rel * 2.0
partco = xp.build_particles(particle_on_co=partco, mode='shift', zeta=z_slices)
tracker = xt.Tracker(line=ltest)
tracker.track(partco, turn_by_turn_monitor='ONE_TURN_EBE')
first_turn = tracker.record_last_track
tracker.track(partco, turn_by_turn_monitor='ONE_TURN_EBE')
second_turn = tracker.record_last_track
plt.figure(2)
axcox = plt.subplot(2, 1, 1)
axcoy = plt.subplot(2, 1, 2, sharex=axcox)
plt.suptitle('Check closed orbit 2 turns')
axcox.plot(first_turn.s[iho, :], first_turn.x[iho, :])
axcox.plot(second_turn.s[iho, :], second_turn.x[iho, :])
axcoy.plot(first_turn.s[iho, :], first_turn.y[iho, :])
axcoy.plot(second_turn.s[iho, :], second_turn.y[iho, :])
def install_spacecharge_frozen(line,
particle_ref,
longitudinal_profile,
nemitt_x,
nemitt_y,
sigma_z,
num_spacecharge_interactions,
tol_spacecharge_position,
s_spacecharge=None):
tracker_no_sc = xt.Tracker(line=line.copy())
# Make a matched bunch just to get the matched momentum spread
bunch = xp.generate_matched_gaussian_bunch(num_particles=int(2e6),
total_intensity_particles=1.,
nemitt_x=nemitt_x,
nemitt_y=nemitt_y,
sigma_z=sigma_z,
particle_ref=particle_ref,
tracker=tracker_no_sc)
delta_rms = np.std(bunch.delta)
# Generate spacecharge positions
if s_spacecharge is None:
s_spacecharge = np.linspace(0, line.get_length(),
num_spacecharge_interactions + 1)[:-1]
# Create spacecharge elements (dummy)
sc_elements = []
sc_names = []
for ii, ss in enumerate(s_spacecharge):
sc_elements.append(
SpaceChargeBiGaussian(length=-9999,
apply_z_kick=False,
longitudinal_profile=longitudinal_profile,
mean_x=0.,
mean_y=0.,
sigma_x=1.,
sigma_y=1.))
sc_names.append(f'spacecharge_{ii}')
#TODO Replace loop with single insert_element when available in xtrack
line.insert_element(name=sc_names[-1],
element=sc_elements[-1],
at_s=ss,
s_tol=tol_spacecharge_position)
actual_s_spch = line.get_s_position(sc_names)
sc_lengths = 0 * s_spacecharge
sc_lengths[:-1] = np.diff(actual_s_spch)
sc_lengths[-1] = line.get_length() - np.sum(sc_lengths[:-1])
# Twiss at spacecharge
line_sc_off = line.filter_elements(exclude_types_starting_with='SpaceCh')
tracker_sc_off = xt.Tracker(line=line_sc_off,
element_classes=tracker_no_sc.element_classes,
track_kernel=tracker_no_sc.track_kernel)
tw_at_sc = tracker_sc_off.twiss(particle_ref=particle_ref,
at_elements=sc_names)
# Configure lenses
for ii, sc in enumerate(sc_elements):
sc.mean_x = tw_at_sc['x'][ii]
sc.mean_y = tw_at_sc['y'][ii]
sc.sigma_x = np.sqrt(tw_at_sc['betx'][ii] * nemitt_x /
particle_ref.beta0 / particle_ref.gamma0 +
(tw_at_sc['dx'][ii] * delta_rms)**2)
sc.sigma_y = np.sqrt(tw_at_sc['bety'][ii] * nemitt_y /
particle_ref.beta0 / particle_ref.gamma0 +
(tw_at_sc['dy'][ii] * delta_rms)**2)
sc.length = sc_lengths[ii]