def track(self, particles):
if self._update_flag:
mean_x, sigma_x = mean_and_std(particles.x,
weights=particles.weight)
mean_y, sigma_y = mean_and_std(particles.y,
weights=particles.weight)
if self.update_mean_x_on_track:
self.mean_x = mean_x
if self.update_mean_y_on_track:
self.mean_y = mean_y
if self.update_sigma_x_on_track:
self.sigma_x = sigma_x
if self.update_sigma_y_on_track:
self.sigma_y = sigma_y
super().track(particles)
def test_mean_and_std():
for ctx in xo.context.get_test_contexts():
print(f"Test {ctx.__class__}")
n_x=100
a_host = np.array(np.random.rand(n_x))
a_dev = ctx.nparray_to_context_array(a_host)
mm, ss = xf.mean_and_std(a_dev)
assert np.isclose(mm, np.mean(a_host))
assert np.isclose(ss, np.std(a_host))
weights_host = np.zeros_like(a_host)+.2
weights_dev = ctx.nparray_to_context_array(weights_host)
mm, ss = xf.mean_and_std(a_dev, weights=weights_dev)
assert np.isclose(mm, np.mean(a_host))
assert np.isclose(ss, np.std(a_host))
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)))
mean_y=None, # needs to be specified only for weak-strong
min_sigma_diff=1e-10)
bbeam_b2 = BeamBeamBiGaussian2D(
_context=context,
n_particles=bunch_intensity_b1,
q0=particles_b1.q0,
beta0=particles_b1_gen.beta0[0],
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
print('Track...')
bbeam_b1.track(particles_b1)
#############################
# Compare against ducktrack #
#############################
print('Check against ducktrack...')
alpha_y_s1 = 0.0, beta_y_s1 = betastar_y, D_y_s1 = 0.0,
dQ_x = Qx/2, dQ_y=Qy/2)
#################################################################
# Tracking #
#################################################################
print('Track...')
nTurn = 1024
positions_x_b1 = np.zeros(nTurn,dtype=float)
positions_y_b1 = np.zeros(nTurn,dtype=float)
positions_x_b2 = np.zeros(nTurn,dtype=float)
positions_y_b2 = np.zeros(nTurn,dtype=float)
for turn in range(nTurn):
time0 = time.time()
# Measure beam properties at IP1
mean_x_meas_b1, sigma_x_meas_b1 = xf.mean_and_std(particles_b1.x)
mean_y_meas_b1, sigma_y_meas_b1 = xf.mean_and_std(particles_b1.y)
mean_x_meas_b2, sigma_x_meas_b2 = xf.mean_and_std(particles_b2.x)
mean_y_meas_b2, sigma_y_meas_b2 = xf.mean_and_std(particles_b2.y)
#Record positions for post-processing
positions_x_b1[turn] = mean_x_meas_b1
positions_y_b1[turn] = mean_y_meas_b1
positions_x_b2[turn] = mean_x_meas_b2
positions_y_b2[turn] = mean_y_meas_b2
# Update bb lens with measured properties of the other beam
bbeamIP1_b1.sigma_x = sigma_x_meas_b2
bbeamIP1_b1.mean_x = mean_x_meas_b2
bbeamIP1_b1.sigma_y = sigma_y_meas_b2
bbeamIP1_b1.mean_y = mean_y_meas_b2
bbeamIP1_b2.sigma_x = sigma_x_meas_b1
bbeamIP1_b2.mean_x = mean_x_meas_b1
