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_Adetuning(self): ''' Track the GPU bunch through a TransverseMap with amplitude detuning Check if results on CPU and GPU are the same Special test since it requires the PyCERNmachines module to create the LHC beam. Skip if not available. ''' # Use a real bunch (not all set to 1), otherwise the Action/detuning # gets too big and the numbers blow up. Dx = np.append(np.linspace(0., 20., self.nsegments), [0]) # add some dispersion/alpha lhc = m.LHC(n_segments=self.nsegments, machine_configuration='450GeV', app_x=1e-9, app_y=2e-9, app_xy=-1.5e-11, chromaticity_on=False, amplitude_detuning_on=True, alpha_x=1.2 * np.ones(self.nsegments), D_x=Dx, printer=SilentPrinter()) # create pure Python map, PyCERNmachine uses Cython. adetuner = AmplitudeDetuning(lhc.app_x, lhc.app_y, lhc.app_xy) transverse_map = tt.TransverseMap(lhc.s, lhc.alpha_x, lhc.beta_x, lhc.D_x, lhc.alpha_y, lhc.beta_y, lhc.D_y, lhc.Q_x, lhc.Q_y, adetuner, printer=SilentPrinter()) bunch_cpu = self.create_lhc_bunch(lhc) bunch_gpu = self.create_lhc_bunch(lhc) self.assertTrue( self._track_cpu_gpu(transverse_map, bunch_cpu, bunch_gpu), 'Transverse tracking with Adetuning CPU/GPU differs')
def test_tracking_with_detuner(self): '''Tests whether the cython and pure python version of the TransverseMap yield the same results (x,y,xp,yp) when tracking a beam with detuners ''' adetuner = AmplitudeDetuning(1e-2, 5e-2, 1e-3) pure_python_map = pure_py.TransverseMap(self.s, self.alpha_x, self.beta_x, self.Dx, self.alpha_y, self.beta_y, self.Dy, self.Qx, self.Qy, adetuner, printer=SilentPrinter()) beam_p = self.create_bunch() for s in pure_python_map: s.track(beam_p)
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) file2.close()
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) Jy = 1 / 2 * ((1 + pp.alpha_y**2) / pp.beta_y * bunch.y**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)
convert_hanarmonicities = True if convert_hanarmonicities: scale_factor = 2 * bunch.p0 app_x = pp.app_x * scale_factor app_y = pp.app_y * scale_factor app_xy = pp.app_xy * scale_factor else: app_x = pp.app_x app_y = pp.app_y app_xy = pp.app_xy 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(app_x, app_y, app_xy)]) longitudinal_map = LinearMap([pp.alpha], pp.circumference, pp.Q_s) one_turn_map = [transverse_map[0]] + [longitudinal_map] # 3. CREATE LISTS TO SAVE THE TBT DATA X = [] Y = [] DP = [] # 4. SET UP THE ACCELERATOR AND START TRACKING for i in range(n_turns): # bunch.yp += ampKicks[i] * np.sin(2 * np.pi * 400e6 / (bunch.beta * c) * bunch.z)
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