def get_optics_and_orbit_at_start_ring(mad, seq_name, with_bb_forces=False,
skip_mad_use=False):
initial_bb_state = None
try:
initial_bb_state = mad.globals.on_bb_switch
mad.globals.on_bb_switch = {True: 1, False: 0}[with_bb_forces]
except AttributeError:
print('Warning! on_bb_switch not present')
# Twiss and get closed-orbit
if not skip_mad_use:
mad.use(sequence=seq_name)
twiss_table = mad.twiss(rmatrix=True)
if initial_bb_state is not None:
mad.globals.on_bb_switch = initial_bb_state
mad_beam = mad.sequence[seq_name].beam
assert mad_beam.deltap == 0, "Not implemented."
particle_on_madx_co = xp.Particles(
p0c = mad_beam.pc*1e9,
q0 = mad_beam.charge,
mass0 = mad_beam.mass*1e9,
s = 0,
x = twiss_table.x[0],
px = twiss_table.px[0],
y = twiss_table.y[0],
py = twiss_table.py[0],
psigma = twiss_table.pt[0]/mad_beam.beta,
)
particle_on_madx_co.zeta = (twiss_table.t[0]
*particle_on_madx_co.beta0
*particle_on_madx_co.rvv)
RR_madx = np.zeros([6,6])
for ii in range(6):
for jj in range(6):
RR_madx[ii, jj] = getattr(twiss_table, f're{ii+1}{jj+1}')[0]
optics_and_co_at_start_ring_from_madx = {
'betx': twiss_table.betx[0],
'bety': twiss_table.bety[0],
'alfx': twiss_table.alfx[0],
'alfy': twiss_table.alfy[0],
'dx': twiss_table.dx[0],
'dy': twiss_table.dy[0],
'dpx': twiss_table.dpx[0],
'dpy': twiss_table.dpy[0],
'RR_madx': RR_madx,
'particle_on_madx_co': particle_on_madx_co.to_dict()
}
return optics_and_co_at_start_ring_from_madx
def electroncloud_dipolar_kicks_of_fieldmap(fieldmap=None, p0c=None):
assert p0c is not None
assert fieldmap is not None
part = xp.Particles(p0c=p0c)
ecloud = xf.ElectronCloud(
length=1,
fieldmap=fieldmap,
_buffer=fieldmap._buffer)
ecloud.track(part)
px = part.px[0]
py = part.py[0]
ptau = part.ptau[0]
return [px, py, ptau]
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
def test_spacecharge_gauss_qgauss():
for frozen in [True, False]:
for context in xo.context.get_test_contexts():
print(f"Test {context.__class__}")
#################################
# Generate particles and probes #
#################################
n_macroparticles = int(1e6)
bunch_intensity = 2.5e11
sigma_x = 3e-3
sigma_y = 2e-3
sigma_z = 30e-2
x0 = 1e-3
y0 = -4e-3
p0c = 25.92e9
mass = pmass,
theta_probes = 30 * np.pi/180
r_max_probes = 2e-2
z_probes = 1.2*sigma_z
n_probes = 1000
from xfields.test_support.temp_makepart import generate_particles_object
(particles_dtk, r_probes, _, _, _) = generate_particles_object(
n_macroparticles,
bunch_intensity,
sigma_x,
sigma_y,
sigma_z,
p0c,
mass,
n_probes,
r_max_probes,
z_probes,
theta_probes)
particles = xp.Particles(
_context=context, **particles_dtk.to_dict())
particles.x += x0
particles.y += y0
################
# Space charge #
################
from xfields import LongitudinalProfileQGaussian
lprofile = LongitudinalProfileQGaussian(
_context=context,
number_of_particles=bunch_intensity,
sigma_z=sigma_z,
z0=0.,
q_parameter=1. # there is a bug in ducktrack,
# only q=1 can be tested
)
from xfields import SpaceChargeBiGaussian
# Just not to fool myself in the test
if frozen:
x0_init = x0
y0_init = y0
sx_init = sigma_x
sy_init = sigma_y
else:
x0_init = None
y0_init = None
sx_init = None
sy_init = None
scgauss = SpaceChargeBiGaussian(
_context=context,
update_on_track=not(frozen),
length=1.,
apply_z_kick=False,
longitudinal_profile=lprofile,
mean_x=x0_init,
mean_y=y0_init,
sigma_x=sx_init,
sigma_y=sy_init,
min_sigma_diff=1e-10)
scgauss.track(particles)
#############################
# Compare against ducktrack #
#############################
p2np = context.nparray_from_context_array
x_probes = p2np(particles.x[:n_probes])
y_probes = p2np(particles.y[:n_probes])
z_probes = p2np(particles.zeta[:n_probes])
scdtk = dtk.SCQGaussProfile(
number_of_particles = bunch_intensity,
bunchlength_rms=sigma_z,
sigma_x=sigma_x,
sigma_y=sigma_y,
length=scgauss.length,
q_parameter=scgauss.longitudinal_profile.q_parameter,
x_co=x0,
y_co=y0)
p_dtk = dtk.TestParticles(p0c=p0c,
mass=mass,
x=x_probes.copy(),
y=y_probes.copy(),
zeta=z_probes.copy())
scdtk.track(p_dtk)
assert np.allclose(
p2np(particles.px[:n_probes]),
p_dtk.px,
atol={True:1e-7, False:1e2}[frozen]
* np.max(np.abs(p_dtk.px)))
assert np.allclose(
p2np(particles.py[:n_probes]),
p_dtk.py,
atol={True:1e-7, False:1e2}[frozen]
* np.max(np.abs(p_dtk.py)))
sigma_y = 2e-3
sigma_z = 30e-2
p0c = 25.92e9
mass = xp.pmass,
theta_probes = 30 * np.pi / 180
r_max_probes = 2e-2
z_probes = 1.2 * sigma_z
n_probes = 1000
from xfields.test_support.temp_makepart import generate_particles_object
(particles_dtk, r_probes, x_probes, y_probes,
z_probes) = generate_particles_object(n_macroparticles, bunch_intensity,
sigma_x, sigma_y, sigma_z, p0c, mass,
n_probes, r_max_probes, z_probes,
theta_probes)
particles = xp.Particles(_context=context, **particles_dtk.to_dict())
######################
# Space charge (PIC) #
######################
x_lim = 5. * sigma_x
y_lim = 5. * sigma_y
z_lim = 5. * sigma_z
from xfields import SpaceCharge3D
spcharge = SpaceCharge3D(_context=context,
length=1,
update_on_track=True,
apply_z_kick=False,
x_range=(-x_lim, x_lim),
sigma_z = 30e-2
p0c = 25.92e9
mass = xp.pmass,
theta_probes = 30 * np.pi / 180
r_max_probes = 2e-2
z_probes = 1.2 * sigma_z
n_probes = 1000
from xfields.test_support.temp_makepart import generate_particles_object
print('Generate particles b1...')
(particles_b1_gen, r_probes, _, _,
_) = generate_particles_object(n_macroparticles_b1, bunch_intensity_b1,
sigma_x_b1, sigma_y_b1, sigma_z, p0c, mass,
n_probes, r_max_probes, z_probes, theta_probes)
# Move to context
particles_b1 = xp.Particles(_context=context, **particles_b1_gen.to_dict())
particles_b1.x += mean_x_b1
particles_b1.y += mean_y_b1
print('Generate particles b2...')
(particles_b2_gen, r_probes, _, _,
_) = generate_particles_object(n_macroparticles_b2, bunch_intensity_b2,
sigma_x_b2, sigma_y_b2, sigma_z, p0c, mass,
n_probes, r_max_probes, z_probes, theta_probes)
# Move to context
particles_b2 = xp.Particles(_context=context, **particles_b2_gen.to_dict())
particles_b2.x += mean_x_b2
particles_b2.y += mean_y_b2
d_x=d_x,
d_px=d_px,
d_y=d_y,
d_py=d_py,
d_zeta=d_zeta,
d_delta=d_delta)
dtk_part = dtk.TestParticles(p0c=6500e9,
x=-1.23e-3,
px=50e-3,
y=2e-3,
py=27e-3,
sigma=3.,
delta=2e-4)
part = xp.Particles(_context=context, **dtk_part.to_dict())
part.name = 'beam1_bunch1'
ret = bb.track(part, _force_suspend=True)
assert ret.on_hold
ret = bb.track(part)
assert ret is None
print('------------------------')
bb_dtk.track(dtk_part)
for cc in 'x px y py zeta delta'.split():
val_test = getattr(part, cc)[0]
val_ref = getattr(dtk_part, cc)
print('')
Qy = 0.32
#Offsets in sigma
mean_x_b1 = 1E-2
mean_y_b1 = 0.0
mean_x_b2 = -1E-2
mean_y_b2 = 0.0
p0c = 7000e9
print('Initialising particles')
particles_b1 = xp.Particles(_context=context,
p0c=p0c,
x=np.sqrt(physemit_x*betastar_x)*(np.random.randn(n_macroparticles)+mean_x_b1),
px=np.sqrt(physemit_x/betastar_x)*np.random.randn(n_macroparticles),
y=np.sqrt(physemit_y*betastar_y)*(np.random.randn(n_macroparticles)+mean_y_b1),
py=np.sqrt(physemit_y/betastar_y)*np.random.randn(n_macroparticles),
zeta=sigma_z*np.random.randn(n_macroparticles),
delta=sigma_delta*np.random.randn(n_macroparticles),
)
particles_b2 = xp.Particles(_context=context,
p0c=p0c,
x=np.sqrt(physemit_x*betastar_x)*(np.random.randn(n_macroparticles)+mean_x_b2),
px=np.sqrt(physemit_x/betastar_x)*np.random.randn(n_macroparticles),
y=np.sqrt(physemit_y*betastar_y)*(np.random.randn(n_macroparticles)+mean_y_b2),
py=np.sqrt(physemit_y/betastar_y)*np.random.randn(n_macroparticles),
zeta=sigma_z*np.random.randn(n_macroparticles),
delta=sigma_delta*np.random.randn(n_macroparticles),
)
def test_electronlens_interpolated():
for context in xo.context.get_test_contexts():
print(f"Test {context.__class__}")
outer_radius = 3.e-3
inner_radius = 1.5e-3
x_center = 0.
y_center = 0.
x_range = (-0.5e-2, 0.5e-2)
y_range = (-0.55e-2, 0.55e-2)
nx = 101
ny = 151
x_grid = np.linspace(x_range[0], x_range[1], nx)
y_grid = np.linspace(y_range[0], y_range[1], ny)
dx = x_grid[1] - x_grid[0]
dy = y_grid[1] - y_grid[0]
X, Y = np.meshgrid(x_grid, y_grid, indexing="ij")
rho = np.zeros_like(X)
rho[:] = 0
R = np.sqrt((X - x_center)**2 + (Y - y_center)**2)
rho[(R > inner_radius) & (R < outer_radius)] = 1.
#rho[ X < 0 ] = 0
norm_rho = np.sum(rho[:, :]) * dx * dy
rho[:] /= norm_rho
elens = xf.ElectronLensInterpolated(current=1,
length=1,
voltage=15e3,
x_grid=x_grid,
y_grid=y_grid,
rho=rho)
elens_ideal = xt.Elens(current=1,
elens_length=1,
voltage=15e3,
inner_radius=inner_radius,
outer_radius=outer_radius)
npart = 1000
x_init = np.linspace(-0.42e-2, 0.42e-2, npart)
y_init = np.linspace(-0.42e-2, 0.42e-2, npart)
X_init, Y_init = np.meshgrid(x_init, y_init)
X_init = X_init.flatten()
Y_init = Y_init.flatten()
part = xp.Particles(x=X_init[:], y=Y_init[:], zeta=[0], p0c=450e9)
part_ideal = xp.Particles(x=X_init[:],
y=Y_init[:],
zeta=[0],
p0c=450e9)
elens.track(part)
elens_ideal.track(part_ideal)
sort_mask = np.argsort(part.particle_id)
sort_ideal_mask = np.argsort(part_ideal.particle_id)
print(
np.max(np.abs(part.px[sort_mask] -
part_ideal.px[sort_ideal_mask])))
print(
np.max(np.abs(part.py[sort_mask] -
part_ideal.py[sort_ideal_mask])))
assert np.allclose(part.px[sort_mask],
part_ideal.px[sort_ideal_mask],
atol=1.e-8,
rtol=1.e-15)
assert np.allclose(part.py[sort_mask],
part_ideal.py[sort_ideal_mask],
atol=1.e-8,
rtol=1.e-15)
assert np.all(part.delta == 0.)
assert np.all(part.ptau == 0.)
def test_beambeam3d_collective():
for context in xo.context.get_test_contexts():
if not isinstance(context, xo.ContextCpu):
print(f'skipping test_beambeam3d_collective for context {context}')
continue
print(repr(context))
# crossing plane
alpha = 0.7
# crossing angle
phi = 0.8
# separations
x_bb_co=5e-3
y_bb_co=-4e-3
charge_slices=np.array([1e16, 2e16, 5e16])
z_slices=np.array([-6., 0.2, 5.5])
x_co = 2e-3
px_co= 1e-6
y_co=-3e-3
py_co=-2e-6
zeta_co=0.01
delta_co=1.2e-3
d_x=1.5e-3
d_px=1.6e-6
d_y=-1.7e-3
d_py=-1.8e-6
d_zeta=0.019
d_delta=3e-4
for ss in sigma_configurations():
(Sig_11_0, Sig_12_0, Sig_13_0, Sig_14_0, Sig_22_0, Sig_23_0, Sig_24_0,
Sig_33_0, Sig_34_0, Sig_44_0) = ss
Sig_11_0 = Sig_11_0 + np.zeros_like(charge_slices)
Sig_12_0 = Sig_12_0 + np.zeros_like(charge_slices)
Sig_13_0 = Sig_13_0 + np.zeros_like(charge_slices)
Sig_14_0 = Sig_14_0 + np.zeros_like(charge_slices)
Sig_22_0 = Sig_22_0 + np.zeros_like(charge_slices)
Sig_23_0 = Sig_23_0 + np.zeros_like(charge_slices)
Sig_24_0 = Sig_24_0 + np.zeros_like(charge_slices)
Sig_33_0 = Sig_33_0 + np.zeros_like(charge_slices)
Sig_34_0 = Sig_34_0 + np.zeros_like(charge_slices)
Sig_44_0 = Sig_44_0 + np.zeros_like(charge_slices)
print('------------------------')
print(ss)
bb_dtk = dtk.elements.BeamBeam6D(
phi=phi, alpha=alpha,
x_bb_co=x_bb_co,
y_bb_co=y_bb_co,
charge_slices=charge_slices,
zeta_slices=z_slices,
sigma_11=Sig_11_0[0],
sigma_12=Sig_12_0[0],
sigma_13=Sig_13_0[0],
sigma_14=Sig_14_0[0],
sigma_22=Sig_22_0[0],
sigma_23=Sig_23_0[0],
sigma_24=Sig_24_0[0],
sigma_33=Sig_33_0[0],
sigma_34=Sig_34_0[0],
sigma_44=Sig_44_0[0],
x_co=x_co,
px_co=px_co,
y_co=y_co,
py_co=py_co,
zeta_co=zeta_co,
delta_co=delta_co,
d_x=d_x,
d_px=d_px,
d_y=d_y,
d_py=d_py,
d_zeta=d_zeta,
d_delta=d_delta
)
slicer = xf.TempSlicer(bin_edges = [-10, -5, 0, 5, 10])
config_for_update=xf.ConfigForUpdateBeamBeamBiGaussian3D(
pipeline_manager=None,
element_name=None,
partner_element_name=None,
slicer=slicer,
collision_schedule={'beam1_bunch1': 'beam2_bunch1',},
update_every=None # Never updates (test in weakstrong mode)
)
bb = xf.BeamBeamBiGaussian3D(
_context=context,
config_for_update=config_for_update,
phi=phi, alpha=alpha, other_beam_q0=1,
slices_other_beam_num_particles=charge_slices[::-1],
slices_other_beam_zeta_center=z_slices[::-1],
slices_other_beam_Sigma_11=Sig_11_0,
slices_other_beam_Sigma_12=Sig_12_0,
slices_other_beam_Sigma_13=Sig_13_0,
slices_other_beam_Sigma_14=Sig_14_0,
slices_other_beam_Sigma_22=Sig_22_0,
slices_other_beam_Sigma_23=Sig_23_0,
slices_other_beam_Sigma_24=Sig_24_0,
slices_other_beam_Sigma_33=Sig_33_0,
slices_other_beam_Sigma_34=Sig_34_0,
slices_other_beam_Sigma_44=Sig_44_0,
ref_shift_x=x_co,
ref_shift_px=px_co,
ref_shift_y=y_co,
ref_shift_py=py_co,
ref_shift_zeta=zeta_co,
ref_shift_pzeta=delta_co,
other_beam_shift_x=x_bb_co,
other_beam_shift_y=y_bb_co,
post_subtract_x=d_x,
post_subtract_px=d_px,
post_subtract_y=d_y,
post_subtract_py=d_py,
post_subtract_zeta=d_zeta,
post_subtract_pzeta=d_delta,
)
dtk_part = dtk.TestParticles(
p0c=6500e9,
x=-1.23e-3,
px = 50e-3,
y = 2e-3,
py = 27e-3,
sigma = 3.,
delta = 2e-4)
part= xp.Particles(_context=context, **dtk_part.to_dict())
part.name = 'beam1_bunch1'
ret = bb.track(part, _force_suspend=True)
assert ret.on_hold
ret = bb.track(part)
assert ret is None
bb_dtk.track(dtk_part)
for cc in 'x px y py zeta delta'.split():
val_test = getattr(part, cc)[0]
val_ref = getattr(dtk_part, cc)
print('')
print(f'ducktrack: {cc} = {val_ref:.12e}')
print(f'xsuite: {cc} = {val_test:.12e}')
assert np.isclose(val_test, val_ref, rtol=0, atol=5e-12)
def track_particle_sixtrack(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,
input_folder='./'):
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)
n_part = len(Dx_wrt_CO_m)
wfold = 'temp_trackfun'
if not os.path.exists(wfold):
os.mkdir(wfold)
os.system(f'cp {input_folder}/fort.* {wfold}')
with open(f'{wfold}/fort.3', 'r') as fid:
lines_f3 = fid.readlines()
# Set initial coordinates
i_start_ini = None
for ii, ll in enumerate(lines_f3):
if ll.startswith('INITIAL COO'):
i_start_ini = ii
break
lines_f3[i_start_ini + 2] = ' 0.\n'
lines_f3[i_start_ini + 3] = ' 0.\n'
lines_f3[i_start_ini + 4] = ' 0.\n'
lines_f3[i_start_ini + 5] = ' 0.\n'
lines_f3[i_start_ini + 6] = ' 0.\n'
lines_f3[i_start_ini + 7] = ' 0.\n'
lines_f3[i_start_ini + 2 + 6] = ' 0.\n'
lines_f3[i_start_ini + 3 + 6] = ' 0.\n'
lines_f3[i_start_ini + 4 + 6] = ' 0.\n'
lines_f3[i_start_ini + 5 + 6] = ' 0.\n'
lines_f3[i_start_ini + 6 + 6] = ' 0.\n'
lines_f3[i_start_ini + 7 + 6] = ' 0.\n'
lines_f13 = []
for i_part in range(n_part):
if type(partCO) == xp.Particles:
temp_part = partCO.copy()
else:
temp_part = xp.Particles(**partCO)
temp_part.x += Dx_wrt_CO_m[i_part]
temp_part.px += Dpx_wrt_CO_rad[i_part]
temp_part.y += Dy_wrt_CO_m[i_part]
temp_part.py += Dpy_wrt_CO_rad[i_part]
temp_part.zeta += Dzeta_wrt_CO_m[i_part]
temp_part.update_delta(temp_part.delta + Ddelta_wrt_CO[i_part])
lines_f13.append('%.10e\n' % ((temp_part.x) * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.px) * temp_part.rpp * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.y) * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.py) * temp_part.rpp * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.zeta) * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.delta)))
if i_part % 2 == 1:
lines_f13.append('%.10e\n' % (temp_part.energy0 * 1e-6))
lines_f13.append('%.10e\n' % (prev_part.energy * 1e-6))
lines_f13.append('%.10e\n' % (temp_part.energy * 1e-6))
prev_part = temp_part
with open(wfold + '/fort.13', 'w') as fid:
fid.writelines(lines_f13)
if np.mod(n_part, 2) != 0:
raise ValueError('SixTrack does not like this!')
i_start_tk = None
for ii, ll in enumerate(lines_f3):
if ll.startswith('TRACKING PAR'):
i_start_tk = ii
break
# Set number of turns and number of particles
temp_list = lines_f3[i_start_tk + 1].split(' ')
temp_list[0] = '%d' % n_turns
temp_list[2] = '%d' % (n_part / 2)
lines_f3[i_start_tk + 1] = ' '.join(temp_list)
# Set number of idfor = 2
temp_list = lines_f3[i_start_tk + 2].split(' ')
temp_list[2] = '2'
lines_f3[i_start_tk + 2] = ' '.join(temp_list)
# Setup turn-by-turn dump
i_start_dp = None
for ii, ll in enumerate(lines_f3):
if ll.startswith('DUMP'):
i_start_dp = ii
break
lines_f3[i_start_dp + 1] = 'StartDUMP 1 664 101 dumtemp.dat\n'
with open(wfold + '/fort.3', 'w') as fid:
fid.writelines(lines_f3)
os.system('(cd temp_trackfun; sixtrack >fort.6)')
# Load sixtrack tracking data
sixdump_all = sixtracktools.SixDump101('%s/dumtemp.dat' % wfold)
x_tbt = np.zeros((n_turns, n_part))
px_tbt = np.zeros((n_turns, n_part))
y_tbt = np.zeros((n_turns, n_part))
py_tbt = np.zeros((n_turns, n_part))
zeta_tbt = np.zeros((n_turns, n_part))
delta_tbt = np.zeros((n_turns, n_part))
for i_part in range(n_part):
sixdump_part = sixdump_all[i_part::n_part]
x_tbt[:, i_part] = sixdump_part.x
px_tbt[:, i_part] = sixdump_part.px
y_tbt[:, i_part] = sixdump_part.y
py_tbt[:, i_part] = sixdump_part.py
zeta_tbt[:, i_part] = sixdump_part.zeta
delta_tbt[:, i_part] = sixdump_part.delta
extra = {}
return x_tbt, px_tbt, y_tbt, py_tbt, zeta_tbt, delta_tbt, extra
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 test_tricubic_interpolation():
for context in xo.context.get_test_contexts():
print(f"Test {context.__class__}")
scale=0.05
ff = lambda x, y, z: sum([ scale * x**i * y**j * z**k
for i in range(4) for j in range(4) for k in range(4)])
dfdx = lambda x, y, z: sum([ i * scale * x**(i-1) * y**j * z**k
for i in range(1,4) for j in range(4) for k in range(4)])
dfdy = lambda x, y, z: sum([ j * scale * x**i * y**(j-1) * z**k
for i in range(4) for j in range(1,4) for k in range(4)])
dfdz = lambda x, y, z: sum([ k * scale * x**i * y**j * z**(k-1)
for i in range(4) for j in range(4) for k in range(1,4)])
dfdxy = lambda x, y, z: sum([ i * j * scale * x**(i-1) * y**(j-1) * z**k
for i in range(1,4) for j in range(1,4) for k in range(4)])
dfdxz = lambda x, y, z: sum([ i * k * scale * x**(i-1) * y**j * z**(k-1)
for i in range(1,4) for j in range(4) for k in range(1,4)])
dfdyz = lambda x, y, z: sum([ j * k * scale * x**i * y**(j-1) * z**(k-1)
for i in range(4) for j in range(1,4) for k in range(1,4)])
dfdxyz = lambda x, y, z: sum([ i * j * k * scale * x**(i-1) * y**(j-1) * z**(k-1)
for i in range(1,4) for j in range(1,4) for k in range(1,4)])
NN=21
x_grid = np.linspace(-0.5, 0.5, NN)
y_grid = np.linspace(-0.5, 0.5, NN)
z_grid = np.linspace(-0.5, 0.5, NN)
fieldmap = xf.TriCubicInterpolatedFieldMap(_context=context,
x_grid=x_grid, y_grid=y_grid, z_grid=z_grid)
ecloud = xf.ElectronCloud(length=1, fieldmap=fieldmap,
_buffer=fieldmap._buffer)
x0 = fieldmap.x_grid[0]
y0 = fieldmap.y_grid[0]
z0 = fieldmap.z_grid[0]
dx = fieldmap.dx
dy = fieldmap.dy
dz = fieldmap.dz
nx = fieldmap.nx
ny = fieldmap.ny
for ix in range(NN):
for iy in range(NN):
for iz in range(NN):
index = 0 + 8 * ix + 8 * nx * iy + 8 * nx * ny * iz
fieldmap._phi_taylor[index + 0] = ff(x_grid[ix], y_grid[iy], z_grid[iz])
fieldmap._phi_taylor[index + 1] = dfdx(x_grid[ix], y_grid[iy], z_grid[iz]) * dx
fieldmap._phi_taylor[index + 2] = dfdy(x_grid[ix], y_grid[iy], z_grid[iz]) * dy
fieldmap._phi_taylor[index + 3] = dfdz(x_grid[ix], y_grid[iy], z_grid[iz]) * dz
fieldmap._phi_taylor[index + 4] = dfdxy(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dy
fieldmap._phi_taylor[index + 5] = dfdxz(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dz
fieldmap._phi_taylor[index + 6] = dfdyz(x_grid[ix], y_grid[iy], z_grid[iz]) * dy * dz
fieldmap._phi_taylor[index + 7] = dfdxyz(x_grid[ix], y_grid[iy], z_grid[iz]) * dx * dy * dz
n_parts = 1000
rng = default_rng(12345)
x_test = rng.random(n_parts) * 1.2 - 0.6
y_test = rng.random(n_parts) * 1.2 - 0.6
tau_test = rng.random(n_parts) * 1.2 - 0.6
p0c = 450e9
testp0 = xp.Particles(p0c=p0c)
beta0 = testp0.beta0
part = xp.Particles(_context=context, x=x_test, y=y_test,
zeta=beta0*tau_test, p0c=p0c)
ecloud.track(part)
part.move(_context=xo.ContextCpu())
mask_p = part.state != -11
true_px = np.array([-dfdx(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p],
part.zeta[mask_p] / part.beta0[mask_p])])
true_py = np.array([-dfdy(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p],
part.zeta[mask_p] / part.beta0[mask_p])])
true_ptau = np.array([-dfdz(xx, yy, zz) for xx, yy, zz in zip(part.x[mask_p], part.y[mask_p],
part.zeta[mask_p] / part.beta0[mask_p])])
# print(true_px[:5])
# print(part.ptau[:5])
# print(part.state[:5])
# print(f"px kick diff.: {np.mean(part.px[mask_p]-true_px):.2e} +- {np.std(part.px[mask_p] - true_px):.2e}")
# print(f"py kick diff.: {np.mean(part.py[mask_p]-true_py):.2e} +- {np.std(part.py[mask_p] - true_py):.2e}")
# print(f"ptau kick diff.: {np.mean(part.ptau[mask_p]-true_ptau):.2e} +- {np.std(part.ptau[mask_p] - true_ptau):.2e}")
# print(np.allclose(part.px[mask_p], true_px, atol=1.e-13, rtol=1.e-13))
# print(np.allclose(part.py[mask_p], true_py, atol=1.e-13, rtol=1.e-13))
# print(np.allclose(part.ptau[mask_p], true_ptau, atol=1.e-13, rtol=1.e-13))
# print(np.max(np.abs(part.px[mask_p]- true_px)))
# print(np.max(np.abs(part.py[mask_p]- true_py)))
# print(np.max(np.abs(part.ptau[mask_p]- true_ptau)))
assert np.allclose(part.px[mask_p], true_px, atol=1.e-13, rtol=1.e-13)
assert np.allclose(part.py[mask_p], true_py, atol=1.e-13, rtol=1.e-13)
assert np.allclose(part.ptau[mask_p], true_ptau, atol=1.e-13, rtol=1.e-13)
def test_spacecharge_pic():
for solver in ['FFTSolver2p5D', 'FFTSolver3D']:
for context in xo.context.get_test_contexts():
print(f"Test {context.__class__}")
print(repr(context))
#################################
# Generate particles and probes #
#################################
n_macroparticles = int(5e6)
bunch_intensity = 2.5e11
sigma_x = 3e-3
sigma_y = 2e-3
sigma_z = 30e-2
p0c = 25.92e9
mass = pmass,
theta_probes = 30 * np.pi/180
r_max_probes = 2e-2
z_probes = 1.2*sigma_z
n_probes = 1000
from xfields.test_support.temp_makepart import generate_particles_object
(particles_gen, r_probes, x_probes,
y_probes, z_probes) = generate_particles_object(
n_macroparticles,
bunch_intensity,
sigma_x,
sigma_y,
sigma_z,
p0c,
mass,
n_probes,
r_max_probes,
z_probes,
theta_probes)
# Transfer particles to context
particles = xp.Particles(
_context=context, **particles_gen.to_dict())
######################
# Space charge (PIC) #
######################
x_lim = 4.*sigma_x
y_lim = 4.*sigma_y
z_lim = 4.*sigma_z
from xfields import SpaceCharge3D
spcharge = SpaceCharge3D(
_context=context,
length=1, update_on_track=True, apply_z_kick=False,
x_range=(-x_lim, x_lim),
y_range=(-y_lim, y_lim),
z_range=(-z_lim, z_lim),
nx=128, ny=128, nz=25,
solver=solver,
gamma0=particles_gen.gamma0[0],
)
spcharge.track(particles)
#############################
# Compare against ducktrack #
#############################
p2np = context.nparray_from_context_array
scdtk = dtk.SCQGaussProfile(
number_of_particles = bunch_intensity,
bunchlength_rms=sigma_z,
sigma_x=sigma_x,
sigma_y=sigma_y,
length=spcharge.length,
x_co=0.,
y_co=0.)
p_dtk = dtk.TestParticles(p0c=p0c,
mass=mass,
x=x_probes.copy(),
y=y_probes.copy(),
zeta=z_probes.copy())
scdtk.track(p_dtk)
mask_inside_grid = ((np.abs(x_probes)<0.9*x_lim) &
(np.abs(y_probes)<0.9*y_lim))
assert np.allclose(
p2np(particles.px[:n_probes])[mask_inside_grid],
p_dtk.px[mask_inside_grid],
atol=3e-2*np.max(np.abs(p_dtk.px[mask_inside_grid])))
assert np.allclose(
p2np(particles.py[:n_probes])[mask_inside_grid],
p_dtk.py[mask_inside_grid],
atol=3e-2*np.max(np.abs(p_dtk.py[mask_inside_grid])))
def test_beambeam3d_old_interface():
for context in xo.context.get_test_contexts():
print(repr(context))
# crossing plane
alpha = 0.7
# crossing angle
phi = 0.8
# separations
x_bb_co=5e-3
y_bb_co=-4e-3
charge_slices=np.array([1e16, 2e16, 5e16])
z_slices=np.array([-6., 0.2, 5.5])
x_co = 2e-3
px_co= 1e-6
y_co=-3e-3
py_co=-2e-6
zeta_co=0.01
delta_co=1.2e-3
d_x=1.5e-3
d_px=1.6e-6
d_y=-1.7e-3
d_py=-1.8e-6
d_zeta=0.019
d_delta=3e-4
for ss in sigma_configurations():
(Sig_11_0, Sig_12_0, Sig_13_0, Sig_14_0, Sig_22_0, Sig_23_0, Sig_24_0,
Sig_33_0, Sig_34_0, Sig_44_0) = ss
Sig_11_0 = Sig_11_0 + np.zeros_like(charge_slices)
Sig_12_0 = Sig_12_0 + np.zeros_like(charge_slices)
Sig_13_0 = Sig_13_0 + np.zeros_like(charge_slices)
Sig_14_0 = Sig_14_0 + np.zeros_like(charge_slices)
Sig_22_0 = Sig_22_0 + np.zeros_like(charge_slices)
Sig_23_0 = Sig_23_0 + np.zeros_like(charge_slices)
Sig_24_0 = Sig_24_0 + np.zeros_like(charge_slices)
Sig_33_0 = Sig_33_0 + np.zeros_like(charge_slices)
Sig_34_0 = Sig_34_0 + np.zeros_like(charge_slices)
Sig_44_0 = Sig_44_0 + np.zeros_like(charge_slices)
print('------------------------')
print(ss)
bb_dtk = dtk.elements.BeamBeam6D(
phi=phi, alpha=alpha,
x_bb_co=x_bb_co,
y_bb_co=y_bb_co,
charge_slices=charge_slices,
zeta_slices=z_slices,
sigma_11=Sig_11_0[0],
sigma_12=Sig_12_0[0],
sigma_13=Sig_13_0[0],
sigma_14=Sig_14_0[0],
sigma_22=Sig_22_0[0],
sigma_23=Sig_23_0[0],
sigma_24=Sig_24_0[0],
sigma_33=Sig_33_0[0],
sigma_34=Sig_34_0[0],
sigma_44=Sig_44_0[0],
x_co=x_co,
px_co=px_co,
y_co=y_co,
py_co=py_co,
zeta_co=zeta_co,
delta_co=delta_co,
d_x=d_x,
d_px=d_px,
d_y=d_y,
d_py=d_py,
d_zeta=d_zeta,
d_delta=d_delta
)
bb = xf.BeamBeamBiGaussian3D(old_interface=bb_dtk.to_dict(), _context=context)
dtk_part = dtk.TestParticles(
p0c=6500e9,
x=-1.23e-3,
px = 50e-3,
y = 2e-3,
py = 27e-3,
sigma = 3.,
delta = 2e-4)
part=xp.Particles(_context=context, **dtk_part.to_dict())
bb.track(part)
bb_dtk.track(dtk_part)
part.move(_context=xo.ContextCpu())
for cc in 'x px y py zeta delta'.split():
val_test = getattr(part, cc)[0]
val_ref = getattr(dtk_part, cc)
print('')
print(f'ducktrack: {cc} = {val_ref:.12e}')
print(f'xsuite: {cc} = {val_test:.12e}')
assert np.isclose(val_test, val_ref, rtol=0, atol=5e-12)
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
def test_beambeam3d():
for context in xo.context.get_test_contexts():
if not isinstance(context, xo.ContextCpu):
print(f'skipping test_beambeam3d_collective for context {context}')
continue
print(repr(context))
# crossing plane
alpha = 0.7
# crossing angle
phi = 0.8
# separations
x_bb_co=5e-3
y_bb_co=-4e-3
charge_slices=np.array([1e16, 2e16, 5e16])
z_slices=np.array([-6., 0.2, 5.5])
x_co = 2e-3
px_co= 1e-6
y_co=-3e-3
py_co=-2e-6
zeta_co=0.01
delta_co=1.2e-3
d_x=1.5e-3
d_px=1.6e-6
d_y=-1.7e-3
d_py=-1.8e-6
d_zeta=0.019
d_delta=3e-4
for ss in sigma_configurations():
(Sig_11_0, Sig_12_0, Sig_13_0, Sig_14_0, Sig_22_0, Sig_23_0, Sig_24_0,
Sig_33_0, Sig_34_0, Sig_44_0) = ss
Sig_11_0 = Sig_11_0 + np.zeros_like(charge_slices)
Sig_12_0 = Sig_12_0 + np.zeros_like(charge_slices)
Sig_13_0 = Sig_13_0 + np.zeros_like(charge_slices)
Sig_14_0 = Sig_14_0 + np.zeros_like(charge_slices)
Sig_22_0 = Sig_22_0 + np.zeros_like(charge_slices)
Sig_23_0 = Sig_23_0 + np.zeros_like(charge_slices)
Sig_24_0 = Sig_24_0 + np.zeros_like(charge_slices)
Sig_33_0 = Sig_33_0 + np.zeros_like(charge_slices)
Sig_34_0 = Sig_34_0 + np.zeros_like(charge_slices)
Sig_44_0 = Sig_44_0 + np.zeros_like(charge_slices)
# I modify one slice to check that properties are working correctly
Sig_11_0[1] *= 1000
Sig_12_0[1] *= 1000
Sig_13_0[1] *= 1000
Sig_14_0[1] *= 1000
Sig_22_0[1] *= 1000
Sig_23_0[1] *= 1000
Sig_24_0[1] *= 1000
Sig_33_0[1] *= 1000
Sig_34_0[1] *= 1000
Sig_44_0[1] *= 1000
print('------------------------')
print(ss)
bb_dtk = dtk.elements.BeamBeam6D(
phi=phi, alpha=alpha,
x_bb_co=x_bb_co,
y_bb_co=y_bb_co,
charge_slices=charge_slices,
zeta_slices=z_slices,
sigma_11=Sig_11_0[0],
sigma_12=Sig_12_0[0],
sigma_13=Sig_13_0[0],
sigma_14=Sig_14_0[0],
sigma_22=Sig_22_0[0],
sigma_23=Sig_23_0[0],
sigma_24=Sig_24_0[0],
sigma_33=Sig_33_0[0],
sigma_34=Sig_34_0[0],
sigma_44=Sig_44_0[0],
x_co=x_co,
px_co=px_co,
y_co=y_co,
py_co=py_co,
zeta_co=zeta_co,
delta_co=delta_co,
d_x=d_x,
d_px=d_px,
d_y=d_y,
d_py=d_py,
d_zeta=d_zeta,
d_delta=d_delta
)
bb = xf.BeamBeamBiGaussian3D(
_context=context,
phi=phi, alpha=alpha, other_beam_q0=1,
slices_other_beam_num_particles=charge_slices[::-1],
slices_other_beam_zeta_center=z_slices[::-1],
slices_other_beam_Sigma_11=Sig_11_0,
slices_other_beam_Sigma_12=Sig_12_0,
slices_other_beam_Sigma_13=Sig_13_0,
slices_other_beam_Sigma_14=Sig_14_0,
slices_other_beam_Sigma_22=Sig_22_0,
slices_other_beam_Sigma_23=Sig_23_0,
slices_other_beam_Sigma_24=Sig_24_0,
slices_other_beam_Sigma_33=Sig_33_0,
slices_other_beam_Sigma_34=Sig_34_0,
slices_other_beam_Sigma_44=Sig_44_0,
ref_shift_x=x_co,
ref_shift_px=px_co,
ref_shift_y=y_co,
ref_shift_py=py_co,
ref_shift_zeta=zeta_co,
ref_shift_pzeta=delta_co,
other_beam_shift_x=x_bb_co,
other_beam_shift_y=y_bb_co,
post_subtract_x=d_x,
post_subtract_px=d_px,
post_subtract_y=d_y,
post_subtract_py=d_py,
post_subtract_zeta=d_zeta,
post_subtract_pzeta=d_delta,
)
bb.slices_other_beam_Sigma_11[1] = bb.slices_other_beam_Sigma_11[0]
bb.slices_other_beam_Sigma_12[1] = bb.slices_other_beam_Sigma_12[0]
bb.slices_other_beam_Sigma_13[1] = bb.slices_other_beam_Sigma_13[0]
bb.slices_other_beam_Sigma_14[1] = bb.slices_other_beam_Sigma_14[0]
bb.slices_other_beam_Sigma_22[1] = bb.slices_other_beam_Sigma_22[0]
bb.slices_other_beam_Sigma_23[1] = bb.slices_other_beam_Sigma_23[0]
bb.slices_other_beam_Sigma_24[1] = bb.slices_other_beam_Sigma_24[0]
bb.slices_other_beam_Sigma_33[1] = bb.slices_other_beam_Sigma_33[0]
bb.slices_other_beam_Sigma_34[1] = bb.slices_other_beam_Sigma_34[0]
bb.slices_other_beam_Sigma_44[1] = bb.slices_other_beam_Sigma_44[0]
dtk_part = dtk.TestParticles(
p0c=6500e9,
x=-1.23e-3,
px = 50e-3,
y = 2e-3,
py = 27e-3,
sigma = 3.,
delta = 2e-4)
part= xp.Particles(_context=context, **dtk_part.to_dict())
part.name = 'beam1_bunch1'
ret = bb.track(part)
bb_dtk.track(dtk_part)
for cc in 'x px y py zeta delta'.split():
val_test = getattr(part, cc)[0]
val_ref = getattr(dtk_part, cc)
print('')
print(f'ducktrack: {cc} = {val_ref:.12e}')
print(f'xsuite: {cc} = {val_test:.12e}')
assert np.isclose(val_test, val_ref, rtol=0, atol=5e-12)
def test_beambeam():
for context in xo.context.get_test_contexts():
print(repr(context))
#################################
# Generate particles and probes #
#################################
n_macroparticles_b1 = int(1e6)
bunch_intensity_b1 = 2.5e11
sigma_x_b1 = 3e-3
sigma_y_b1 = 2e-3
mean_x_b1 = 1.3e-3
mean_y_b1 = -1.2e-3
n_macroparticles_b2 = int(1e6)
bunch_intensity_b2 = 3e11
sigma_x_b2 = 1.7e-3
sigma_y_b2 = 2.1e-3
mean_x_b2 = -1e-3
mean_y_b2 = 1.4e-3
sigma_z = 30e-2
p0c = 25.92e9
mass = pmass,
theta_probes = 30 * np.pi / 180
r_max_probes = 2e-2
z_probes = 1.2 * sigma_z
n_probes = 1000
from xfields.test_support.temp_makepart import generate_particles_object
(particles_b1, r_probes, _, _, _) = generate_particles_object(
n_macroparticles_b1, bunch_intensity_b1, sigma_x_b1, sigma_y_b1,
sigma_z, p0c, mass, n_probes, r_max_probes, z_probes, theta_probes)
# Move to right context
particles_b1 = xp.Particles(_context=context, **particles_b1.to_dict())
particles_b1.x += mean_x_b1
particles_b1.y += mean_y_b1
(particles_b2, r_probes, _, _, _) = generate_particles_object(
n_macroparticles_b2, bunch_intensity_b2, sigma_x_b2, sigma_y_b2,
sigma_z, p0c, mass, n_probes, r_max_probes, z_probes, theta_probes)
particles_b2 = xp.Particles(_context=context, **particles_b2.to_dict())
particles_b2.x += mean_x_b2
particles_b2.y += mean_y_b2
#############
# Beam-beam #
#############
from xfields import BeamBeamBiGaussian2D, mean_and_std
# if beta0 is array I just take the first
beta0_b2 = context.nparray_from_context_array(particles_b2.beta0)[0]
beta0_b1 = context.nparray_from_context_array(particles_b1.beta0)[0]
bbeam_b1 = BeamBeamBiGaussian2D(
_context=context,
n_particles=bunch_intensity_b2,
q0=particles_b2.q0,
beta0=beta0_b2,
sigma_x=None, # needs to be specified only for weak-strong
sigma_y=None, # needs to be specified only for weak-strong
mean_x=None, # needs to be specified only for weak-strong
mean_y=None, # needs to be specified only for weak-strong
min_sigma_diff=1e-10)
# Measure beam properties
mean_x_meas, sigma_x_meas = mean_and_std(particles_b2.x)
mean_y_meas, sigma_y_meas = mean_and_std(particles_b2.y)
# Update bb lens
bbeam_b1.update(sigma_x=sigma_x_meas,
mean_x=mean_x_meas,
sigma_y=sigma_y_meas,
mean_y=mean_y_meas)
#Track
bbeam_b1.track(particles_b1)
#############################
# Compare against ducktrack #
#############################
p2np = context.nparray_from_context_array
x_probes = p2np(particles_b1.x[:n_probes])
y_probes = p2np(particles_b1.y[:n_probes])
z_probes = p2np(particles_b1.zeta[:n_probes])
from ducktrack.elements import BeamBeam4D
bb_b1_dtk = BeamBeam4D(charge=bunch_intensity_b2,
sigma_x=sigma_x_b2,
sigma_y=sigma_y_b2,
x_bb=mean_x_b2,
y_bb=mean_y_b2,
beta_r=np.float64(beta0_b2))
p_dtk = dtk.TestParticles(p0c=p0c,
mass=mass,
x=x_probes.copy(),
y=y_probes.copy(),
zeta=z_probes.copy())
bb_b1_dtk.track(p_dtk)
assert np.allclose(p_dtk.px,
p2np(particles_b1.px[:n_probes]),
atol=2e-2 * np.max(np.abs(p_dtk.px)))
assert np.allclose(p_dtk.py,
p2np(particles_b1.py[:n_probes]),
atol=2e-2 * np.max(np.abs(p_dtk.py)))
elens = xf.ElectronLensInterpolated(current=1,
length=1,
voltage=15e3,
x_grid=x_grid,
y_grid=y_grid,
rho=rho)
elens_ideal = xt.Elens(current=1,
elens_length=1,
voltage=15e3,
inner_radius=inner_radius,
outer_radius=outer_radius)
npart = 10000
part = xp.Particles(x=np.linspace(-1.e-2, 1.e-2, npart),
y=[y_center],
zeta=[0],
p0c=450e9)
part_ideal = xp.Particles(x=np.linspace(-1.e-2, 1.e-2, npart),
y=[0],
zeta=[0],
p0c=450e9)
elens.track(part)
elens_ideal.track(part_ideal)
plt.figure(1)
plt.plot(part.x * 1000., part.px, 'r.')
plt.plot(part_ideal.x * 1000., part_ideal.px, 'k.')
plt.figure(2)
