def setup_beam_beam_in_line(
line,
bb_df,
bb_coupling=False,
):
import xfields as xf
assert bb_coupling is False # Not implemented
for ii, (ee, eename) in enumerate(zip(line.elements, line.element_names)):
if isinstance(ee, xf.BeamBeamBiGaussian2D):
ee.n_particles=bb_df.loc[eename, 'other_num_particles']
ee.q0 = bb_df.loc[eename, 'other_particle_charge']
ee.sigma_x = np.sqrt(bb_df.loc[eename, 'other_Sigma_11'])
ee.sigma_y = np.sqrt(bb_df.loc[eename, 'other_Sigma_33'])
ee.beta0 = bb_df.loc[eename, 'other_relativistic_beta']
ee.mean_x = bb_df.loc[eename, 'separation_x']
ee.mean_y = bb_df.loc[eename, 'separation_y']
if isinstance(ee, xf.BeamBeamBiGaussian3D):
params = {}
params['phi'] = bb_df.loc[eename, 'phi']
params['alpha'] = bb_df.loc[eename, 'alpha']
params['x_bb_co'] = bb_df.loc[eename, 'separation_x']
params['y_bb_co'] = bb_df.loc[eename, 'separation_y']
# TODO update xtrack interface to separate charge and b. population
params['charge_slices'] = [(bb_df.loc[eename, 'other_num_particles']
* bb_df.loc[eename, 'other_particle_charge'])]
params['zeta_slices'] = [0.0]
params['sigma_11'] = bb_df.loc[eename, 'other_Sigma_11']
params['sigma_12'] = bb_df.loc[eename, 'other_Sigma_12']
params['sigma_13'] = bb_df.loc[eename, 'other_Sigma_13']
params['sigma_14'] = bb_df.loc[eename, 'other_Sigma_14']
params['sigma_22'] = bb_df.loc[eename, 'other_Sigma_22']
params['sigma_23'] = bb_df.loc[eename, 'other_Sigma_23']
params['sigma_24'] = bb_df.loc[eename, 'other_Sigma_24']
params['sigma_33'] = bb_df.loc[eename, 'other_Sigma_33']
params['sigma_34'] = bb_df.loc[eename, 'other_Sigma_34']
params['sigma_44'] = bb_df.loc[eename, 'other_Sigma_44']
if not (bb_coupling):
params['sigma_13'] = 0.0
params['sigma_14'] = 0.0
params['sigma_23'] = 0.0
params['sigma_24'] = 0.0
params["x_co"] = 0
params["px_co"] = 0
params["y_co"] = 0
params["py_co"] = 0
params["zeta_co"] = 0
params["delta_co"] = 0
params["d_x"] = 0
params["d_px"] = 0
params["d_y"] = 0
params["d_py"] = 0
params["d_zeta"] = 0
params["d_delta"] = 0
newee = xf.BeamBeamBiGaussian3D(old_interface=params)
line.element_dict[eename] = newee
bb = xf.BeamBeamBiGaussian3D(
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,
)
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_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)
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)
sigma_44=Sig_44_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)
print('------------------------')