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