def create_transverse_map(self, chromaticity_on=True, amplitude_detuning_on=True): detuners = [] if chromaticity_on: detuners.append(Chromaticity(self.Qp_x, self.Qp_y)) if amplitude_detuning_on: detuners.append( AmplitudeDetuning(self.app_x, self.app_y, self.app_xy) ) self.transverse_map = TransverseMap( self.s, self.alpha_x, self.beta_x, self.D_x, self.alpha_y, self.beta_y, self.D_y, self.Q_x, self.Q_y, detuners, printer=self._printer)
def test_transverse_track_with_Cdetuning(self): ''' Track the GPU bunch through a TransverseMap with chromaticity (detuning) Check if results on CPU and GPU are the same ''' bunch_cpu = self.create_all1_bunch() bunch_gpu = self.create_all1_bunch() detuner = Chromaticity(Qp_x=[5, 1], Qp_y=[7, 2]) transverse_map = tt.TransverseMap(self.s, self.alpha_x, self.beta_x, self.Dx, self.alpha_y, self.beta_y, self.Dy, self.Qx, self.Qy, detuner, printer=SilentPrinter()) self.assertTrue( self._track_cpu_gpu(transverse_map, bunch_cpu, bunch_gpu), 'Transverse tracking with chromaticity CPU/GPU differs')
yoffset = 0e-4 bunch.x += xoffset bunch.y += yoffset afile = open('bunch', 'wb') pickle.dump(bunch, afile) afile.close() # =================================== # CREATE TRANVERSE AND LONGITUDINAL MAPS # ================================== scale_factor = 2*bunch.p0 # for detuning coefficients transverse_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, [Chromaticity(Qp_x, Qp_y), AmplitudeDetuning(app_x*scale_factor, app_y*scale_factor, app_xy*scale_factor)]) longitudinal_map = LinearMap([alpha], circumference, Q_s) # ====================================================================== # SET UP ACCELERATOR MAP AND START TRACKING # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ t0 = time.clock() #reload object from file file2 = open('bunch', 'rb') bunch = pickle.load(file2)
import pickle import csv import NAFFlib as pnf from scipy.constants import c from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning from PyHEADTAIL.trackers.transverse_tracking import TransverseMap from PyHEADTAIL.trackers.simple_long_tracking import RFSystems, LinearMap import matplotlib.pyplot as plt n_turns = 1000 # 1000 turns are enough for the NAFF algorithm to compute the tune # 1. CREATE ONE TURN MAP transverse_map = TransverseMap( pp.s, pp.alpha_x, pp.beta_x, pp.D_x, pp.alpha_y, pp.beta_y, pp.D_y, pp.Q_x, pp.Q_y, [ Chromaticity(pp.Qp_x, pp.Qp_y), AmplitudeDetuning(pp.app_x, pp.app_y, pp.app_xy) ]) longitudinal_map = LinearMap([pp.alpha], pp.circumference, pp.Q_s) one_turn_map = [transverse_map[0]] + [longitudinal_map] # 2. LOAD OBJECTS FROM FILES, BUNCH AND NOISE KICKS bfile = open('input/bunch', 'rb') bunch = pickle.load(bfile) bfile.close() # calculate initial actions Jx = 1 / 2 * ((1 + pp.alpha_x**2) / pp.beta_x * bunch.x**2 + 2 * pp.alpha_x * bunch.x * bunch.xp + pp.beta_x * bunch.xp**2)
def _construct_transverse_map( self, optics_mode=None, circumference=None, n_segments=None, s=None, name=None, alpha_x=None, beta_x=None, D_x=None, alpha_y=None, beta_y=None, D_y=None, accQ_x=None, accQ_y=None, Qp_x=None, Qp_y=None, app_x=None, app_y=None, app_xy=None, other_detuners=None, use_cython=None): if optics_mode == 'smooth': if circumference is None: raise ValueError('circumference has to be specified ' 'if optics_mode = "smooth"') if n_segments is None: raise ValueError('n_segments has to be specified ' 'if optics_mode = "smooth"') if s is not None: raise ValueError('s vector should not be provided ' 'if optics_mode = "smooth"') s = (np.arange(0, n_segments + 1) * circumference / n_segments) alpha_x = 0.*s beta_x = 0.*s+beta_x D_x = 0.*s+D_x alpha_y = 0.*s beta_y = 0.*s+beta_y D_y = 0.*s+D_y elif optics_mode == 'non-smooth': if circumference is not None: raise ValueError('circumference should not be provided ' 'if optics_mode = "non-smooth"') if n_segments is not None: raise ValueError('n_segments should not be provided ' 'if optics_mode = "non-smooth"') if s is None: raise ValueError('s has to be specified ' 'if optics_mode = "smooth"') else: raise ValueError('optics_mode not recognized') detuners = [] if Qp_x != 0 or Qp_y != 0: detuners.append(Chromaticity(Qp_x, Qp_y)) if app_x != 0 or app_y != 0 or app_xy != 0: detuners.append(AmplitudeDetuning(app_x, app_y, app_xy)) detuners += other_detuners self.transverse_map = TransverseMap( s=s, alpha_x=alpha_x, beta_x=beta_x, D_x=D_x, alpha_y=alpha_y, beta_y=beta_y, D_y=D_y, accQ_x=accQ_x, accQ_y=accQ_y, detuners=detuners) self.circumference = s[-1] self.transverse_map.n_segments = len(s)-1 if name is None: self.transverse_map.name = ['P_%d' % ip for ip in range(len(s)-1)] self.transverse_map.name.append('end_ring') else: self.transverse_map.name = name for i_seg, m in enumerate(self.transverse_map): m.i0 = i_seg m.i1 = i_seg+1 m.s0 = self.transverse_map.s[i_seg] m.s1 = self.transverse_map.s[i_seg+1] m.name0 = self.transverse_map.name[i_seg] m.name1 = self.transverse_map.name[i_seg+1] m.beta_x0 = self.transverse_map.beta_x[i_seg] m.beta_x1 = self.transverse_map.beta_x[i_seg+1] m.beta_y0 = self.transverse_map.beta_y[i_seg] m.beta_y1 = self.transverse_map.beta_y[i_seg+1] # insert transverse map in the ring for m in self.transverse_map: self.one_turn_map.append(m)
def run(): def track(bunch, map_): for i in range(n_turns): for m in map_: m.track(bunch) def generate_bunch(n_macroparticles, alpha_x, alpha_y, beta_x, beta_y, alpha_0, Q_s, R): intensity = 1.05e11 sigma_z = 0.059958 gamma = 3730.26 eta = alpha_0 - 1. / gamma**2 gamma_t = 1. / np.sqrt(alpha_0) p0 = np.sqrt(gamma**2 - 1) * m_p * c beta_z = eta * R / Q_s epsn_x = 3.75e-6 # [m rad] epsn_y = 3.75e-6 # [m rad] epsn_z = 4 * np.pi * sigma_z**2 * p0 / (beta_z * e) bunch = generators.generate_Gaussian6DTwiss( macroparticlenumber=n_macroparticles, intensity=intensity, charge=e, gamma=gamma, mass=m_p, circumference=C, alpha_x=alpha_x, beta_x=beta_x, epsn_x=epsn_x, alpha_y=alpha_y, beta_y=beta_y, epsn_y=epsn_y, beta_z=beta_z, epsn_z=epsn_z) #print bunch.sigma_z() return bunch # In[4]: # Basic parameters. n_turns = 3 n_segments = 1 n_macroparticles = 10 Q_x = 64.28 Q_y = 59.31 Q_s = 0.0020443 C = 26658.883 R = C / (2. * np.pi) alpha_x_inj = 0. alpha_y_inj = 0. beta_x_inj = 66.0064 beta_y_inj = 71.5376 alpha_0 = 0.0003225 # In[5]: # Parameters for transverse map. s = np.arange(0, n_segments + 1) * C / n_segments alpha_x = alpha_x_inj * np.ones(n_segments) beta_x = beta_x_inj * np.ones(n_segments) D_x = np.zeros(n_segments) alpha_y = alpha_y_inj * np.ones(n_segments) beta_y = beta_y_inj * np.ones(n_segments) D_y = np.zeros(n_segments) # In[6]: # CASE I # With amplitude detuning (python implementation) # EXPECTED TUNE SPREADS AT THE GIVEN SETTINGS ARE 5e-4 FOR HORIZONTAL # AND VERTICAL. bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj, beta_x_inj, beta_y_inj, alpha_0, Q_s, R) ampl_det = AmplitudeDetuning.from_octupole_currents_LHC(i_focusing=400, i_defocusing=-400) trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, [ampl_det]) trans_one_turn = [m for m in trans_map] map_ = trans_one_turn track(bunch, map_) # In[7]: # CASE II # With first order Chromaticity (python implementation) bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj, beta_x_inj, beta_y_inj, alpha_0, Q_s, R) chroma = Chromaticity(Qp_x=[6], Qp_y=[3]) trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, [chroma]) trans_one_turn = [m for m in trans_map] map_ = trans_one_turn track(bunch, map_) # In[8]: # CASE III # With higher order Chromaticity (python implementation) bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj, beta_x_inj, beta_y_inj, alpha_0, Q_s, R) chroma = Chromaticity(Qp_x=[6., 4e4], Qp_y=[3., 0., 2e8]) trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, [chroma]) trans_one_turn = [m for m in trans_map] map_ = trans_one_turn track(bunch, map_)
def run(): def track_n_save(bunch, map_): mean_x = np.empty(n_turns) mean_y = np.empty(n_turns) sigma_z = np.empty(n_turns) for i in xrange(n_turns): mean_x[i] = bunch.mean_x() mean_y[i] = bunch.mean_y() sigma_z[i] = bunch.sigma_z() for m_ in map_: m_.track(bunch) return mean_x, mean_y, sigma_z def copy_bunch(bunch_source, bunch_target): bunch_target.x = bunch_source.x.copy() bunch_target.xp = bunch_source.xp.copy() bunch_target.y = bunch_source.y.copy() bunch_target.yp = bunch_source.yp.copy() bunch_target.z = bunch_source.z.copy() bunch_target.dp = bunch_source.dp.copy() def my_fft(data): t = np.arange(len(data)) fft = np.fft.rfft(data) fft_freq = np.fft.rfftfreq(t.shape[-1]) return fft_freq, np.abs(fft.real) def generate_bunch(n_macroparticles, alpha_x, alpha_y, beta_x, beta_y, linear_map): intensity = 1.05e11 sigma_z = 0.059958 gamma = 3730.26 gamma_t = 1. / np.sqrt(alpha_0) p0 = np.sqrt(gamma**2 - 1) * m_p * c beta_z = (linear_map.eta(dp=0, gamma=gamma) * linear_map.circumference / (2 * np.pi * linear_map.Q_s)) epsn_x = 3.75e-6 # [m rad] epsn_y = 3.75e-6 # [m rad] epsn_z = 4 * np.pi * sigma_z**2 * p0 / (beta_z * e) bunch = generators.generate_Gaussian6DTwiss( macroparticlenumber=n_macroparticles, intensity=intensity, charge=e, gamma=gamma, mass=m_p, circumference=C, alpha_x=alpha_x, beta_x=beta_x, epsn_x=epsn_x, alpha_y=alpha_y, beta_y=beta_y, epsn_y=epsn_y, beta_z=beta_z, epsn_z=epsn_z) return bunch # In[4]: # Basic parameters. n_turns = 3 n_segments = 5 n_macroparticles = 10 Q_x = 64.28 Q_y = 59.31 Q_s = 0.0020443 C = 26658.883 R = C / (2. * np.pi) alpha_0 = [0.0003225] # **Things tested** # - Instantiation of a TransverseMap (and therewith of several # TransverseSegmentMaps as we have more than 1 segment). # - With and without detuners, i.e. instantiation of Chromaticity and # AmplitudeDetuning DetunerCollections as well as the corresponding # SegmentDetuners. # - Are betatron tunes Q_{x,y} and detuning strengths correctly # scaled to segment lengths? # - TransverseSegmentMap.track(beam) method. # - Check spectrum of beam centroid motion. # - Betatron tune (implicitly checks the scaling to segment lengths) # - If chromaticity and linear synchro motion are on: synchrotron sidebands? # - If amplitude detuning is on and there is initial kick: decoherence? # - Does dispersion work correctly? # - Is exception risen when s[0] != 0 or s[-1] != C? # - Get optics at injection. # In[5]: # TEST CASE SETUP def gimme(D_x=None, D_y=None, *args, **kwargs): # Parameters for transverse map. s = np.arange(0, n_segments + 1) * C / n_segments alpha_x_inj = 0. alpha_y_inj = 0. beta_x_inj = 66.0064 beta_y_inj = 71.5376 alpha_x = alpha_x_inj * np.ones(n_segments) beta_x = beta_x_inj * np.ones(n_segments) if D_x is None: D_x = np.zeros(n_segments) alpha_y = alpha_y_inj * np.ones(n_segments) beta_y = beta_y_inj * np.ones(n_segments) if D_y is None: D_y = np.zeros(n_segments) if 'detuners' in kwargs: trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, kwargs['detuners']) else: trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, printer=SilentPrinter()) long_map = LinearMap(alpha_0, C, Q_s) bunch = generate_bunch(n_macroparticles, alpha_x_inj, alpha_y_inj, beta_x_inj, beta_y_inj, long_map) return bunch, trans_map, long_map # In[6]: # CASE I # Pure transverse tracking. Without detuners. bunch, trans_map, _ = gimme() map_ = trans_map mean_x, mean_y, sigma_z = track_n_save(bunch, map_) # In[7]: # CASE II # Without detuners. With linear synchrotron motion. bunch, trans_map, long_map = gimme() # This tests if TransverseMap is actually a sequence. trans_one_turn = [m for m in trans_map] map_ = trans_one_turn + [long_map] mean_x, mean_y, sigma_z = track_n_save(bunch, map_) # In[8]: # CASE III # With chromaticity in horizontal and vertical. With linear synchrotron motion. chroma = Chromaticity(Qp_x=6, Qp_y=10) bunch, trans_map, long_map = gimme(detuners=(chroma, )) trans_one_turn = [m for m in trans_map] map_ = trans_one_turn + [long_map] mean_x, mean_y, sigma_z = track_n_save(bunch, map_) # In[9]: # CASE IV # With amplitude detuning. With linear synchrotron motion. With initial kick. ampl_det = AmplitudeDetuning.from_octupole_currents_LHC(i_focusing=200, i_defocusing=-200) bunch, trans_map, long_map = gimme(detuners=(ampl_det, )) trans_one_turn = [m for m in trans_map] map_ = trans_one_turn + [long_map] bunch.x += 0.0003 bunch.y += 0.0005 mean_x, mean_y, sigma_z = track_n_save(bunch, map_) # In[10]: # CASE V # With amplitude detuning and chromaticity. With linear synchrotron motion. With initial kick. ampl_det = AmplitudeDetuning.from_octupole_currents_LHC(i_focusing=200, i_defocusing=-200) chroma = Chromaticity(Qp_x=10, Qp_y=6) bunch, trans_map, long_map = gimme(detuners=(ampl_det, chroma)) trans_one_turn = [m for m in trans_map] map_ = trans_one_turn + [long_map] bunch.x += 0.0003 bunch.y += 0.0005 mean_x, mean_y, sigma_z = track_n_save(bunch, map_) # In[11]: # CASE VI n_turns = 2 n_segments = 5 # Initial bunch distro to be used with and without dispersion # for direct comparison. bunch_orig, _, _ = gimme() # With dispersion. D_x = np.zeros(n_segments) D_y = np.zeros(n_segments) D_x[1] = 4.5 D_y[3] = 2.3 D_x[2] = 4.5 D_y[4] = 2.3 bunch_wD, trans_map_wD, _ = gimme(D_x, D_y) copy_bunch(bunch_orig, bunch_wD) # Add dispersion (manually for now). bunch_wD.x += D_x[0] * bunch_wD.dp bunch_wD.y += D_y[0] * bunch_wD.dp x_i_wD = np.zeros((n_segments * n_turns, n_macroparticles)) y_i_wD = np.zeros((n_segments * n_turns, n_macroparticles)) for j in xrange(n_segments * n_turns): x_i_wD[j, :] = bunch_wD.x y_i_wD[j, :] = bunch_wD.y trans_map_wD[j % n_segments].track(bunch_wD) # Without dispersion. bunch_woD, trans_map_woD, _ = gimme() copy_bunch(bunch_orig, bunch_woD) x_i_woD = np.zeros((n_segments * n_turns, n_macroparticles)) y_i_woD = np.zeros((n_segments * n_turns, n_macroparticles)) dp_i_woD = np.zeros((n_segments * n_turns, n_macroparticles)) for j in xrange(n_segments * n_turns): x_i_woD[j, :] = bunch_woD.x y_i_woD[j, :] = bunch_woD.y dp_i_woD[j, :] = bunch_woD.dp trans_map_woD[j % n_segments].track(bunch_woD) # In[12]: # Test how detuning parameters and betatron tunes are scaled # for the TransverseSegmentMaps. Qp_x = 8. Qp_y = 10. chroma = Chromaticity(Qp_x, Qp_y) bunch, trans_map, _ = gimme(detuners=(chroma, )) i = 1 # print 'Q_x' # print Q_x / float(n_segments) # print trans_map[i].dQ_x # print 'Q_y' # print Q_y / float(n_segments) # print trans_map[i].dQ_y # print 'Qp_x' # print Qp_x / float(n_segments) # print trans_map[i].segment_detuners[0].dQp_x # print 'Qp_y' # print Qp_y / float(n_segments) # print trans_map[i].segment_detuners[0].dQp_y app_x = 20. app_y = 12. app_xy = 30. ampl_det = AmplitudeDetuning(app_x, app_y, app_xy) bunch, trans_map, _ = gimme(detuners=(ampl_det, )) # print 'app_x' # print app_x / float(n_segments) # print trans_map[i].segment_detuners[0].dapp_x # print 'app_y' # print app_y / float(n_segments) # print trans_map[i].segment_detuners[0].dapp_y # print 'app_xy' # print app_xy / float(n_segments) # print trans_map[i].segment_detuners[0].dapp_xy # In[13]: # Test if optics at injection are correctly returned. s = np.arange(0, n_segments + 1) * C / n_segments alpha_x_inj = 0. alpha_y_inj = 0. beta_x_inj = 66.0064 beta_y_inj = 71.5376 alpha_x = alpha_x_inj * np.ones(n_segments) beta_x = beta_x_inj * np.ones(n_segments) alpha_y = alpha_y_inj * np.ones(n_segments) beta_y = beta_y_inj * np.ones(n_segments) trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, printer=SilentPrinter()) inj_opt_dict = trans_map.get_injection_optics() # In[14]: # Test if exception is risen when s[0] != 0 or s[-1] != C s = np.array([0., 4, 10, 12, C - 1]) try: trans_map = TransverseMap(s, alpha_x, beta_x, D_x, alpha_y, beta_y, D_y, Q_x, Q_y, printer=SilentPrinter()) # print ('test NOT passed. No error raised!') except ValueError as exc: # print ('test passed.\n') # print ('Error message:\n' + exc.message) pass
def generate_one_turn_map(original_map, list_of_optics_elements, list_of_PyHT_objects, optics_file, phase_x_col, beta_x_col, phase_y_col, beta_y_col, circumference, alpha_x=None, beta_x=None, D_x=None, alpha_y=None, beta_y=None, D_y=None, accQ_x=None, accQ_y=None, Qp_x=None, Qp_y=None, app_x=0, app_y=0, app_xy=0, other_detuners=[], use_cython=False): with open(optics_file) as f: content = f.readlines() content = [x.strip() for x in content] element_coordinates = np.zeros((len(list_of_optics_elements),5)) for j, element in enumerate(list_of_optics_elements): element_name = element[0] line_found = False for i, l in enumerate(content): s = l.split() if s[0] == ('"' + element_name + '"'): line_found = True element_coordinates[j,0] = j element_coordinates[j,1] = float(s[phase_x_col]) element_coordinates[j,2] = float(s[beta_x_col]) element_coordinates[j,3] = float(s[phase_y_col]) element_coordinates[j,4] = float(s[beta_y_col]) if line_found == False: raise ValueError('Element ' + element_name + ' not found from the optics file ' + optics_file) element_coordinates[:,3] = element_coordinates[:,3]*accQ_x/accQ_y map_locations = [] for j, element in enumerate(list_of_optics_elements): if element[1] == 'x': map_locations.append(element_coordinates[j,1]) elif element[1] == 'y': map_locations.append(element_coordinates[j,3]) else: raise ValueError('Unknown plane') map_locations = np.array(map_locations) order = np.argsort(map_locations) s = [] for i, j in enumerate(order): if i == 0: if map_locations[j]/accQ_x*circumference != 0.: s.append(0.) s.append(map_locations[j]/accQ_x*circumference) s.append(circumference) s = np.array(s) alpha_x = 0.*s beta_x = 0.*s+beta_x D_x = 0.*s+D_x alpha_y = 0.*s beta_y = 0.*s+beta_y D_y = 0.*s+D_y detuners = [] if any(np.atleast_1d(Qp_x) != 0) or \ any(np.atleast_1d(Qp_y) != 0): detuners.append(Chromaticity(Qp_x, Qp_y)) if app_x != 0 or app_y != 0 or app_xy != 0: detuners.append(AmplitudeDetuning(app_x, app_y, app_xy)) detuners += other_detuners transverse_map = TransverseMap( s=s, alpha_x=alpha_x, beta_x=beta_x, D_x=D_x, alpha_y=alpha_y, beta_y=beta_y, D_y=D_y, accQ_x=accQ_x, accQ_y=accQ_y, detuners=detuners) transverse_map.n_segments = len(s)-1 transverse_map.name = ['P_%d' % ip for ip in range(len(s)-1)] transverse_map.name.append('end_ring') for i_seg, m in enumerate(transverse_map): m.i0 = i_seg m.i1 = i_seg+1 m.s0 = transverse_map.s[i_seg] m.s1 = transverse_map.s[i_seg+1] m.name0 = transverse_map.name[i_seg] m.name1 = transverse_map.name[i_seg+1] m.beta_x0 = transverse_map.beta_x[i_seg] m.beta_x1 = transverse_map.beta_x[i_seg+1] m.beta_y0 = transverse_map.beta_y[i_seg] m.beta_y1 = transverse_map.beta_y[i_seg+1] # insert transverse map in the ring one_turn_map = [] for m in original_map: if type(m) != PyHEADTAIL.trackers.transverse_tracking.TransverseSegmentMap: one_turn_map.append(m) for i, m in enumerate(transverse_map): one_turn_map.append(m) if i < len(order): element_idx = order[i] one_turn_map.append(list_of_PyHT_objects[element_idx]) return one_turn_map