xy_norm = temp_data['xy_norm']
DpxDpy_wrt_CO = temp_data['DpxDpy_wrt_CO']
if track_with == 'PySixtrack':
part = pysixtrack.Particles(**partCO)
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_pysixtrack(
line, part=part, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0, Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns, verbose=True)
elif track_with == 'Sixtrack':
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_sixtrack(
partCO=partCO, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0, Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns)
else:
raise ValueError('What?!')
n_part = x_tbt.shape[1]
Qx = np.zeros(n_part)
Qy = np.zeros(n_part)
for i_part in range(n_part):
Qx[i_part] = NAFFlib.get_tune(x_tbt[:, i_part])
Qy[i_part] = NAFFlib.get_tune(y_tbt[:, i_part])
Qxy_fp = np.zeros_like(xy_norm)
if track_with == 'PySixtrack':
part = pysixtrack.Particles(**partCO)
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_pysixtrack(
line, part=part, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0, Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns, verbose=True)
info = track_with
elif track_with == 'Sixtrack':
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_sixtrack(
partCO=partCO, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0, Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns)
info = track_with
elif track_with == 'Sixtracklib':
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_sixtracklib(
line=line, partCO=partCO, Dx_wrt_CO_m=0., Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0., Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dsigma_wrt_CO_m=0., Ddelta_wrt_CO=0., n_turns=n_turns, device_opencl=device_opencl)
info = track_with
if device_opencl is None:
info += ' (CPU)'
else:
info += ' (GPU %s)'%device_opencl
else:
raise ValueError('What?!')
Ddelta = 0.
part = pysixtrack.Particles(**partCO)
dict_tbt = {}
print('Tracking PyST')
x_tbt_pyST, px_tbt_pyST, y_tbt_pyST, py_tbt_pyST, sigma_tbt_pyST, delta_tbt_pyST = hp.track_particle_pysixtrack(
line, part=part, Dx_wrt_CO_m=Dx_m, Dpx_wrt_CO_rad=Dpx_rad,
Dy_wrt_CO_m=Dy_m, Dpy_wrt_CO_rad=Dpy_rad,
Dsigma_wrt_CO_m=Dsigma_m, Ddelta_wrt_CO=Ddelta, n_turns=n_turns, verbose=True)
dict_tbt['PyST'] = {kk:vv for kk,vv in zip('x y px py sigma delta'.split(),
[x_tbt_pyST, px_tbt_pyST, y_tbt_pyST, py_tbt_pyST, sigma_tbt_pyST, delta_tbt_pyST])}
print('Tracking ST')
x_tbt_ST, px_tbt_ST, y_tbt_ST, py_tbt_ST, sigma_tbt_ST, delta_tbt_ST = hp.track_particle_sixtrack(
partCO=partCO, Dx_wrt_CO_m=np.array([Dx_m, Dx_m]), Dpx_wrt_CO_rad=Dpx_rad,
Dy_wrt_CO_m=Dy_m, Dpy_wrt_CO_rad=Dpy_rad,
Dsigma_wrt_CO_m=Dsigma_m, Ddelta_wrt_CO=Ddelta, n_turns=n_turns)
dict_tbt['ST'] = {kk:vv for kk,vv in zip('x y px py sigma delta'.split(),
[x_tbt_ST, px_tbt_ST, y_tbt_ST, py_tbt_ST, sigma_tbt_ST, delta_tbt_ST])}
print('Tracking STlib')
x_tbt_STl, px_tbt_STl, y_tbt_STl, py_tbt_STl, sigma_tbt_STl, delta_tbt_STl = hp.track_particle_sixtracklib(
line, partCO=partCO, Dx_wrt_CO_m=np.array([Dx_m, Dx_m]), Dpx_wrt_CO_rad=Dpx_rad,
Dy_wrt_CO_m=Dy_m, Dpy_wrt_CO_rad=Dpy_rad,
Dsigma_wrt_CO_m=Dsigma_m, Ddelta_wrt_CO=Ddelta, n_turns=n_turns)
dict_tbt['STlib'] = {kk:vv for kk,vv in zip('x y px py sigma delta'.split(),
[x_tbt_STl, px_tbt_STl, y_tbt_STl, py_tbt_STl, sigma_tbt_STl, delta_tbt_STl])}
plt.close('all')
fig1 = plt.figure(1, figsize=(8 * 1.5, 6 * 1.2))
axx = fig1.add_subplot(3, 2, 1)
axx.plot(x_tbt_pyST[:, 0], label='PyST')
my_sigma_xx_mm2=1.
particleMomentum_eV=7e12
x_test= np.linspace(-2e-2, 2e-2, 100)
y_test = 0*np.linspace(-2e-2, 2e-2, 100)
write_fort3(filename='./fort.3',partnum=my_partnum, xang_rad=my_xang_rad, xplane_rad=my_xplane_rad, h_sep_mm=my_h_sep_mm, v_sep_mm=my_v_sep_mm, sigma_xx_mm2=my_sigma_xx_mm2)
part_on_CO = pysixtrack.Particles(p0c=particleMomentum_eV, x=0.)
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_sixtrack(partCO=part_on_CO.to_dict(),
Dx_wrt_CO_m=x_test,
Dpx_wrt_CO_rad=0, Dy_wrt_CO_m=y_test, Dpy_wrt_CO_rad=0,
Dsigma_wrt_CO_m=0, Ddelta_wrt_CO=0, n_turns=2,
mode = 'ebe')
plt.figure(figsize=myFigSize)
plt.subplot(221)
plt.plot(x_test, px_tbt[1, :],label='SixTrack results')
plt.grid()
plt.xlabel('x [m]')
plt.ylabel('px')
plt.title('The 6D kick')
plt.plot(x_test,BBKickVector(x_test+my_h_sep_mm, my_partnum, part_on_CO.gamma0, np.sqrt(my_sigma_xx_mm2)*1e-3)\
-BBKickVector(0+my_h_sep_mm, my_partnum, part_on_CO.gamma0, np.sqrt(my_sigma_xx_mm2)*1e-3),'.r',label='analytical reference')
plt.legend(loc='best')
Dpx_rad = 1e-6
Dy_m = 2e-4
Dpy_rad = 3e-6
Dsigma_m = 2e-3
Ddelta = 2e-5
part = pysixtrack.Particles(**partCO)
dict_tbt = {}
print('Tracking ST')
x_tbt_ST, px_tbt_ST, y_tbt_ST, py_tbt_ST, sigma_tbt_ST, delta_tbt_ST = hp.track_particle_sixtrack(
partCO=partCO,
Dx_wrt_CO_m=np.array([Dx_m, Dx_m]),
Dpx_wrt_CO_rad=Dpx_rad,
Dy_wrt_CO_m=Dy_m,
Dpy_wrt_CO_rad=Dpy_rad,
Dsigma_wrt_CO_m=Dsigma_m,
Ddelta_wrt_CO=Ddelta,
n_turns=n_turns)
dict_tbt['ST'] = {
kk: vv
for kk, vv in zip('x y px py sigma delta'.split(
), [x_tbt_ST, px_tbt_ST, y_tbt_ST, py_tbt_ST, sigma_tbt_ST, delta_tbt_ST])
}
print('Tracking STlib')
x_tbt_STl, px_tbt_STl, y_tbt_STl, py_tbt_STl, sigma_tbt_STl, delta_tbt_STl = hp.track_particle_sixtracklib(
line,
partCO=partCO,
Dx_wrt_CO_m=np.array([Dx_m, Dx_m]),
Dpx_wrt_CO_rad=Dpx_rad,
DpxDpy_wrt_CO = np.zeros_like(xy_norm)
for ii in range(xy_norm.shape[0]):
for jj in range(xy_norm.shape[1]):
DpxDpy_wrt_CO[ii, jj, 0] = xy_norm[ii, jj, 0] * np.sqrt(
epsn_x / part.beta0 / part.gamma0 / beta_x)
DpxDpy_wrt_CO[ii, jj, 1] = xy_norm[ii, jj, 1] * np.sqrt(
epsn_y / part.beta0 / part.gamma0 / beta_y)
if track_with == 'sixtrack':
x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra = hp.track_particle_sixtrack(
partCO=partCO,
Dx_wrt_CO_m=0.,
Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0,
Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dzeta_wrt_CO_m=0.,
Ddelta_wrt_CO=0.,
n_turns=n_turns,
input_folder='../')
info = track_with
elif track_with == 'xtrack':
x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra = hp.track_particle_xtrack(
line=line,
partCO=partCO,
Dx_wrt_CO_m=0.,
Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0.,
Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dzeta_wrt_CO_m=0.,
Ddelta_wrt_CO=0.,
#with open('../xline/line_bb_dipole_cancelled.json', 'r') as fid:
# dict_line_old = json.load(fid)
line = xl.Line.from_dict(dict_line_xtrack)
#line = xl.Line.from_dict(dict_line_old)
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)