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)))
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])))
n_macroparticles = int(1e6)
bunch_intensity = 2.5e11
sigma_x = 3e-3
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,
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)))
mean_x_b2 = -1e-3
mean_y_b2 = 1.4e-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
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())