def generate_initial_distribution( parameters, Lattice, output_file='Input/ParticleDistribution.in', summary_file='Input/ParticleDistribution_summary.txt', outputFormat='Orbit'): parameters['alphax0'] = Lattice.alphax0 parameters['betax0'] = Lattice.betax0 parameters['alphay0'] = Lattice.alphay0 parameters['betay0'] = Lattice.betay0 parameters['etax0'] = Lattice.etax0 parameters['etapx0'] = Lattice.etapx0 parameters['etay0'] = Lattice.etay0 parameters['etapy0'] = Lattice.etapy0 parameters['x0'] = Lattice.orbitx0 parameters['xp0'] = Lattice.orbitpx0 parameters['y0'] = Lattice.orbity0 parameters['yp0'] = Lattice.orbitpy0 parameters['gamma_transition'] = Lattice.gammaT parameters['circumference'] = Lattice.getLength() parameters['length'] = Lattice.getLength() / Lattice.nHarm # twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['etax0'], 'etapx0': parameters['etapx0'] } dispersiony = { 'etay0': parameters['etay0'], 'etapy0': parameters['etapy0'] } # ~ dispersionx = {'etax0': parameters['etax0'], 'etapx0': parameters['etapx0']} # ~ dispersiony = {'etay0': parameters['etay0'], 'etapy0': parameters['etapy0']} closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # initialize particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) # building the distributions Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=parameters['TransverseCut']) Longitudinal_distribution = LongitudinalJohoDistributionSingleHarmonic( parameters, parameters['LongitudinalJohoParameter']) # only the main CPU is actually writing its distribution to a file ... comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file, "w") as fid: csv_writer = csv.writer(fid, delimiter=' ') for i in range(parameters['n_macroparticles']): (phi[i], dE[i]) = Longitudinal_distribution.getCoordinates() (x[i], xp[i], y[i], yp[i]) = Transverse_distribution.getCoordinates() x[i] += closedOrbitx['x0'] xp[i] += closedOrbitx['xp0'] y[i] += closedOrbity['y0'] yp[i] += closedOrbity['yp0'] dpp = dE[i] / (parameters['energy']) / parameters['beta']**2 x[i] += dpp * dispersionx['etax0'] xp[i] += dpp * dispersionx['etapx0'] y[i] += dpp * dispersiony['etay0'] yp[i] += dpp * dispersiony['etapy0'] if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. dE[i] /= 1.e9 csv_writer.writerow( [x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) #csv_writer.writerow([x[i], xp[i], y[i], yp[i], z[i], dE[i]]) if summary_file: with open(summary_file, 'w') as fid: map( lambda key: fid.write(key + ' = ' + str(parameters[key]) + '\n'), parameters) print '\nCreated particle distribution with ' + str( parameters['n_macroparticles'] ) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
def generate_initial_distribution_tomo_old( parameters, matfile=0, Lattice=None, output_file='ParticleDistribution.in', outputFormat='pyOrbit', summary_file='ParticleDistribution_summary.txt', summary_mat_file=None): assert outputFormat in ['Orbit', 'pyOrbit'] p = parameters beta = p['beta'] gamma = p['gamma'] if Lattice: p['alphax0'] = Lattice.alphax0 p['betax0'] = Lattice.betax0 p['alphay0'] = Lattice.alphay0 p['betay0'] = Lattice.betay0 p['etax0'] = Lattice.etax0 p['etapx0'] = Lattice.etapx0 p['etay0'] = Lattice.etay0 p['etapy0'] = Lattice.etapy0 p['x0'] = Lattice.orbitx0 p['xp0'] = Lattice.orbitpx0 p['y0'] = Lattice.orbity0 p['yp0'] = Lattice.orbitpy0 p['gamma_transition'] = Lattice.gammaT p['circumference'] = Lattice.getLength() # building the distributions # eta = 1/p['gamma_transition']**2 - 1/p['gamma']**2 # R = p['circumference']/2/np.pi # beta = p['beta'] # energy = p['energy'] # phi_rf = p['phi_s'] # h = p['harmonic_number'] # h_main = np.atleast_1d(p['harmonic_number'])[0] # rf_voltage = p['rf_voltage'] # RF = DoubleRF(R, eta, beta, energy, phi_rf, h, rf_voltage) # Longitudinal_distribution = LongitudinalBinomialDistribution(RF, p['LongitudinalDistribution_z_max'], p['LongitudinalJohoParameter']) # z, dpp = Longitudinal_distribution.getCoordinates(p['n_macroparticles']) # building the distributions beta = p['beta'] try: noise_level = p['noise_level'] except KeyError: noise_level = 0 # ~ Longitudinal_distribution = LongitudinalDistributionFromTomoscope(p['tomo_file']) Longitudinal_distribution = LongitudinalDistributionFromTomoscope( p['tomo_file'], matfile) # ~ Longitudinal_distribution.plot_Tomoscope_data() # ~ Longitudinal_distribution.plot_generated_distribution() t_rand, dE_rand = Longitudinal_distribution.getCoordinates( p['n_macroparticles'], noise_level) z = t_rand * speed_of_light * beta * 1e-9 # convert ns to s and then m # ~ z = (t_rand * 1e-9) * speed_of_light * beta * 0.075 # convert ns to s and then m dE = dE_rand * 1e-3 # convert from MeV to GeV dpp = dE / p['energy'] / 1.e-9 / beta**2 # ~ dpp = dE / p['energy'] / beta**2 # Not sure which dpp definition is correct # h_main = np.atleast_1d(p['harmonic_number'])[0] # R = p['circumference']/2/np.pi # phi = - z * h_main / R # z_arr, z_profile, z_rms, dp, dp_profile, dpp_rms = Longitudinal_distribution.getBunchProfile() # p['dpp_sigma'] = _GaussianFit(dp, dp_profile)[0][2] # p['dpp_sigma_from_FWHM'] = _Gaussian_sigma_from_FWHM(dp, dp_profile) # p['dpp_profile'] = np.array([dp, dp_profile]) # p['dpp_rms'] = dpp_rms # p['linedensity_profile'] = np.array([z_arr, z_profile]) # phi = - z * h_main / R # dE = dpp * p['energy'] * beta**2 * 1.e-9 # transverse coordinates x, xp, y, yp = [], [], [], [] for epsn_x, epsn_y, intensity in zip(np.atleast_1d(p['epsn_x']), np.atleast_1d(p['epsn_y']), np.atleast_1d(p['intensity'])): # twiss containers twissX = TwissContainer(alpha=p['alphax0'], beta=p['betax0'], emittance=epsn_x / gamma / beta) twissY = TwissContainer(alpha=p['alphay0'], beta=p['betay0'], emittance=epsn_y / gamma / beta) Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=p['TransverseCut']) n_macroparticles_tmp = int(p['n_macroparticles'] * (intensity / np.sum(p['intensity']))) Transverse_coords = np.array( map(lambda i: Transverse_distribution.getCoordinates(), xrange(n_macroparticles_tmp))) x.extend(Transverse_coords[:, 0].tolist()) xp.extend(Transverse_coords[:, 1].tolist()) y.extend(Transverse_coords[:, 2].tolist()) yp.extend(Transverse_coords[:, 3].tolist()) # in case x has not yet a length of n_macroparticles # ~ while len(x)<p['n_macroparticles']: # ~ Transverse_coords = Transverse_distribution.getCoordinates() # ~ x.append(Transverse_coords[0]) # ~ xp.append(Transverse_coords[1]) # ~ y.append(Transverse_coords[2]) # ~ yp.append(Transverse_coords[3]) # Dispersion and closed orbit x = np.array(x) + p['x0'] + dpp * p['etax0'] xp = np.array(xp) + p['xp0'] + dpp * p['etapx0'] y = np.array(y) + p['y0'] + dpp * p['etay0'] yp = np.array(yp) + p['yp0'] + dpp * p['etapy0'] # only the main CPU is actually writing its distribution to a file ... comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file, "w") as fid: csv_writer = csv.writer(fid, delimiter=' ') if outputFormat == 'Orbit': x *= 1000. xp *= 1000. y *= 1000. yp *= 1000. # ~ dE[i] /= 1.e9 # Already in the correct units map( lambda i: csv_writer.writerow( [x[i], xp[i], y[i], yp[i], phi[i], dE[i]]), range(p['n_macroparticles'])) elif outputFormat == 'pyOrbit': map( lambda i: csv_writer.writerow( [x[i], xp[i], y[i], yp[i], z[i], dE[i]]), range(p['n_macroparticles'])) if summary_file: with open(summary_file, 'w') as fid: map(lambda key: fid.write(key + ' = ' + str(p[key]) + '\n'), p) if summary_mat_file: with open(summary_mat_file, 'w') as fid: sio.savemat(fid, parameters) print '\nCreated particle distribution with ' + str( p['n_macroparticles']) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
def generate_initial_distribution_from_tomo( parameters, matfile=0, Lattice=None, output_file='ParticleDistribution.in', outputFormat='pyOrbit', summary_file='ParticleDistribution_summary.txt', summary_mat_file=None): # Get parameters from the lattice parameters['alphax0'] = Lattice.alphax0 parameters['betax0'] = Lattice.betax0 parameters['alphay0'] = Lattice.alphay0 parameters['betay0'] = Lattice.betay0 parameters['etax0'] = Lattice.etax0 parameters['etapx0'] = Lattice.etapx0 parameters['etay0'] = Lattice.etay0 parameters['etapy0'] = Lattice.etapy0 parameters['x0'] = Lattice.orbitx0 parameters['xp0'] = Lattice.orbitpx0 parameters['y0'] = Lattice.orbity0 parameters['yp0'] = Lattice.orbitpy0 parameters['gamma_transition'] = Lattice.gammaT parameters['circumference'] = Lattice.getLength() parameters['length'] = Lattice.getLength() / Lattice.nHarm # Create Twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['etax0'], 'etapx0': parameters['etapx0'] } dispersiony = { 'etay0': parameters['etay0'], 'etapy0': parameters['etapy0'] } closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # Initialize empty particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) z = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) # Instatiate the classes for longitudinal and transverse distns Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=parameters['TransverseCut']) Longitudinal_distribution = LongitudinalDistributionFromTomoscope( parameters['tomo_file'], matfile) try: noise_level = parameters['noise_level'] except KeyError: noise_level = 0 t_rand, dE_rand = Longitudinal_distribution.getCoordinates( parameters['n_macroparticles'], noise_level) z = (t_rand * 1e-9 ) * speed_of_light * parameters['beta'] # convert ns to s and then m dE = dE_rand * 1e-3 # convert from MeV to GeV # We need to convert z into phi h_main = np.atleast_1d(parameters['harmonic_number'])[0] R = parameters['circumference'] / 2 / np.pi phi = -z * h_main / R # Write the distn to a file only on one CPU comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file, "w") as fid: csv_writer = csv.writer(fid, delimiter=' ') for i in range(parameters['n_macroparticles']): # ~ (z[i], dE[i]) = Longitudinal_distribution.getCoordinates() # ~ z[i] = z[i] * speed_of_light * parameters['beta'] * 1e-9 # convert ns to s and then m # ~ dE[i] = dE[i] * 1e-3 # convert from MeV to GeV (x[i], xp[i], y[i], yp[i]) = Transverse_distribution.getCoordinates() x[i] += closedOrbitx['x0'] xp[i] += closedOrbitx['xp0'] y[i] += closedOrbity['y0'] yp[i] += closedOrbity['yp0'] dpp = dE[i] / ( parameters['energy']) / parameters['beta']**2 * 1E9 print '\n dpp = ', dpp x[i] += dpp * dispersionx['etax0'] xp[i] += dpp * dispersionx['etapx0'] y[i] += dpp * dispersiony['etay0'] yp[i] += dpp * dispersiony['etapy0'] # ~ if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. # ~ dE[i] /= 1.e9 # ~ if outputFormat == 'Orbit': map( lambda i: csv_writer.writerow( [x[i], xp[i], y[i], yp[i], phi[i], dE[i]]), range(parameters['n_macroparticles'])) # ~ elif outputFormat == 'pyOrbit': # ~ map(lambda i: csv_writer.writerow([x[i], xp[i], y[i], yp[i], z[i], dE[i]]), range(parameters['n_macroparticles'])) if summary_file: with open(summary_file, 'w') as fid: map( lambda key: fid.write(key + ' = ' + str(parameters[key]) + '\n'), parameters) print '\nCreated particle distribution with ' + str( parameters['n_macroparticles'] ) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
print "relat. gamma=",gamma print "relat. beta=",beta frequency = 402.5e+6 v_light = 2.99792458e+8 # in [m/sec] #------ emittances are in 1e-6 (alphaX,betaX,emittX) = ( -3.1105 , 10.1926 , 0.5469 ) (alphaY,betaY,emittY) = ( 1.3880 , 6.2129 , 0.5353 ) (alphaZ,betaZ,emittZ) = ( 0.1691 , 12.5589 , 0.2149 ) print " ========= PyORBIT Twiss ===========" print " aplha beta emitt[mm*mrad] X= %6.4f %6.4f %6.4f "%(alphaX,betaX,emittX) print " aplha beta emitt[mm*mrad] Y= %6.4f %6.4f %6.4f "%(alphaY,betaY,emittY) print " aplha beta emitt[mm*MeV] Z= %6.4f %6.4f %6.4f "%(alphaZ,betaZ,emittZ) twissX = TwissContainer(alphaX,betaX,emittX*1.0e-6) twissY = TwissContainer(alphaY,betaY,emittY*1.0e-6) twissZ = TwissContainer(alphaZ,betaZ,emittZ*1.0e-6) print "Start Bunch Generation." bunch_gen = SNS_Linac_BunchGenerator(twissX,twissY,twissZ) #set the initial kinetic energy in GeV bunch_gen.setKinEnergy(e_kin_ini) #set the beam peak current in mA bunch_gen.setBeamCurrent(50.0) #bunch_in = bunch_gen.getBunch(nParticles = 20000, distributorClass = WaterBagDist3D) bunch_in = bunch_gen.getBunch(nParticles = 20000, distributorClass = GaussDist3D) #bunch_in = bunch_gen.getBunch(nParticles = 20000, distributorClass = KVDist3D)
#-------------------------------------------------------- # The script will test the gerenrators from bunch_generators #-------------------------------------------------------- import math import random import sys from orbit.bunch_generators import TwissContainer, TwissAnalysis from orbit.bunch_generators import WaterBagDist2D, GaussDist2D, KVDist2D n = 10000 #--------------------------------------------- # KV 2D #--------------------------------------------- twissX = TwissContainer(alpha=1., beta=2., emittance=3.) twissY = TwissContainer(alpha=2., beta=3., emittance=4.) dist = KVDist2D(twissX, twissY) twiss_analysis = TwissAnalysis(2) for i in range(n): (x, xp, y, yp) = dist.getCoordinates() twiss_analysis.account((x, xp, y, yp)) print "=================================================" print "KV 2D - done!" print " alpha beta [m/rad] gamma emitt[m*rad] " print "Twiss X %12.5g %12.5g %12.5g %12.5g " % twissX.getAlphaBetaGammaEmitt( ) print "Generated X %12.5g %12.5g %12.5g %12.5g " % twiss_analysis.getTwiss( 0) print "......................................................................"
b = Bunch() b.mass(0.93827231) b.macroSize(macrosize) b.getSyncParticle().kinEnergy(energy) #----------------------------------------------------------------------------- # Twiss Parameters for Distribution #----------------------------------------------------------------------------- alphaX = 0.0 alphaY = 0.0 betaX = Lattice_Length / (2 * pi) betaY = Lattice_Length / (2 * pi) emittanceX = 2.e-5 emittanceY = 2.e-5 twissX = TwissContainer(alphaX, betaX, emittanceX) twissY = TwissContainer(alphaY, betaY, emittanceY) #----------------------------------------------------------------------------- # Making Distribution #----------------------------------------------------------------------------- # dist = SelfConDist2D(twissX, twissY) #Code wrote by Austin H. dist = SelfConDist2D_2(twissX, twissY) #Code wrote by J. Holmes adapted to python if SingleParticle==True: PlotBeta = False x, xp = 0.0795, 0. y, yp = 0., -0.01 z, zp = 0., 0. b.addParticle(x,xp,y,yp,z,zp)
#-------------------------------------------------------- # The script will test the gerenrators from bunch_generators #-------------------------------------------------------- import math import random import sys from orbit.bunch_generators import TwissContainer, TwissAnalysis from orbit.bunch_generators import WaterBagDist3D, GaussDist3D, KVDist3D n = 10000 #--------------------------------------------- # KV 3D #--------------------------------------------- twissX = TwissContainer(alpha = 1., beta = 2., emittance = 3.) twissY = TwissContainer(alpha = 2., beta = 3., emittance = 4.) twissZ = TwissContainer(alpha = 3., beta = 4., emittance = 5.) dist = KVDist3D(twissX,twissY,twissZ) twiss_analysis = TwissAnalysis(3) for i in range(n): (x,xp,y,yp,z,zp) = dist.getCoordinates() twiss_analysis.account((x,xp,y,yp,z,zp)) print "=================================================" print "KV 3D - done!" print " alpha beta [m/rad] gamma emitt[m*rad] " print "Twiss X %12.5g %12.5g %12.5g %12.5g "%twissX.getAlphaBetaGammaEmitt() print "Generated X %12.5g %12.5g %12.5g %12.5g "%twiss_analysis.getTwiss(0) print "......................................................................" print "Twiss Y %12.5g %12.5g %12.5g %12.5g "%twissY.getAlphaBetaGammaEmitt()
scLatticeModifications.setSC2DSliceBySliceAccNodes(teapot_latt, sc_path_length_min, sc_calc) #SPACE CHARGE ''' #------------------------------------------------------------------------------------------------------------- #---------------------------------------------ORBIT ERRORS------------------------------------------- # get lattice function from transfer matrices matrix_lattice = TEAPOT_MATRIX_Lattice(teapot_latt, b) (muX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (muY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() #----------------------------------Bunch-Distribusion------------------------------------------------ # machted beam emittance_x, emittance_y = 35 * 10**(-6), 15 * 10**(-6) twissX = TwissContainer(arrPosAlphaX[0][1], arrPosBetaX[0][1], emittance_x) twissY = TwissContainer(arrPosAlphaY[0][1], arrPosBetaY[0][1], emittance_y) #twissZ = TwissContainer(alpha = 0., beta = 1., emittance = 10**(-6)) # not matched #---------------------------------------------CALCULATE ORBIT------------------------------ OrbitX, OrbitY = orbit(teapot_latt, b).get_orbit() x0, px0 = OrbitX[0][1], OrbitX[0][2] y0, py0 = OrbitY[0][1], OrbitY[0][2] print('Closed x-orbit at s_0 ({},{})'.format(x0, px0)) #--------------------------------------------------------------------------------------------------- #x0 = 0.001 n = 10**(2) # num of particles #dist = KVDist3D(twissX,twissY,twissZ)
def generate_initial_distribution_3DGaussian( parameters, Lattice, output_file='Input/ParticleDistribution.in', summary_file='Input/ParticleDistribution_summary.txt', outputFormat='Orbit'): parameters['alphax0'] = Lattice.alphax0 parameters['betax0'] = Lattice.betax0 parameters['alphay0'] = Lattice.alphay0 parameters['betay0'] = Lattice.betay0 parameters['etax0'] = Lattice.etax0 parameters['etapx0'] = Lattice.etapx0 parameters['etay0'] = Lattice.etay0 parameters['etapy0'] = Lattice.etapy0 parameters['x0'] = Lattice.orbitx0 parameters['xp0'] = Lattice.orbitpx0 parameters['y0'] = Lattice.orbity0 parameters['yp0'] = Lattice.orbitpy0 parameters['gamma_transition'] = Lattice.gammaT parameters['circumference'] = Lattice.getLength() parameters['length'] = Lattice.getLength() / Lattice.nHarm # twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['etax0'], 'etapx0': parameters['etapx0'] } dispersiony = { 'etay0': parameters['etay0'], 'etapy0': parameters['etapy0'] } closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # initialize particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) # ~ z = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) # building the distributions Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=parameters['TransverseCut']) # We need to convert z into phi h_main = np.atleast_1d(parameters['harmonic_number'])[0] R = parameters['circumference'] / 2 / np.pi sig_E = (parameters['dpp_rms'] * parameters['energy'] * parameters['beta']**2) if orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD) == 0: fid = open(output_file, "w") csv_writer = csv.writer(fid, delimiter=' ') for i in range(parameters['n_macroparticles']): # Longitudinal distn - use 5 sigma as cut-off (manual) outside_limits_E = True while outside_limits_E: dE[i] = random.gauss(0., sig_E) # Energy in eV if abs(dE[i]) < (5 * sig_E): print '\n\tdE = ', dE[i] outside_limits_E = False outside_limits_z = True while outside_limits_z: z_temp = random.gauss(0., parameters['blength_rms']) if abs(z_temp) < (5 * parameters['blength_rms']): print '\n\tz_temp = ', z_temp phi[i] = -z_temp * h_main / R outside_limits_z = False (x[i], xp[i], y[i], yp[i]) = Transverse_distribution.getCoordinates() x[i] += closedOrbitx['x0'] xp[i] += closedOrbitx['xp0'] y[i] += closedOrbity['y0'] yp[i] += closedOrbity['yp0'] dpp = dE[i] / (parameters['energy']) / parameters['beta']**2 x[i] += dpp * dispersionx['etax0'] xp[i] += dpp * dispersionx['etapx0'] y[i] += dpp * dispersiony['etay0'] yp[i] += dpp * dispersiony['etapy0'] if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000 dE[i] /= 1.e9 csv_writer.writerow([x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) # else: # still need to convert from phi to z!! #csv_writer.writerow([x[i], xp[i], y[i], yp[i], z[i], dE[i]]) fid.close() fid = open(summary_file, 'w') parameter_list = [ 'circumference', 'rf_voltage', 'phi_s', 'harmonic_number', 'gamma_transition', 'n_macroparticles', 'energy', 'gamma', 'bunch_length', 'LongitudinalCut', 'LongitudinalJohoParameter', 'x0', 'xp0', 'betax0', 'alphax0', 'etax0', 'etapx0', 'y0', 'yp0', 'betay0', 'alphay0', 'etay0', 'etapy0', 'epsn_x', 'epsn_y', 'TransverseCut' ] for key in parameter_list: fid.write(key + ' = ' + str(parameters[key]) + '\n') fid.close() print '\nCreated particle distribution with ' + str( parameters['n_macroparticles'] ) + ' macroparticles into file: ', output_file return output_file
print "Start." #=====Make a Teapot style lattice====== lattice = teapot.TEAPOT_Lattice() print "Read MAD." lattice.readMAD("SNSring_pyOrbitBenchmark.LAT", "RING") print "Lattice=", lattice.getName(), " length [m] =", lattice.getLength( ), " nodes=", len(lattice.getNodes()) #------------------------------ # Main Bunch init #------------------------------ n_particles = 100000 twissX = TwissContainer(alpha=0.0046902, beta=10.207, emittance=3.e-5) twissY = TwissContainer(alpha=0.056823, beta=10.639, emittance=3.e-5) twissZ = TwissContainer(alpha=0., beta=100000., emittance=0.008) dist = GaussDist3D(twissX, twissY, twissZ) b = Bunch() for i in range(n_particles): (x, xp, y, yp, z, zp) = dist.getCoordinates() b.addParticle(x, xp, y, yp, z, zp) total_macroSize = 1.e+14 b.mass(0.93827231) energy = 1.0 #Gev #b.readBunch(distribution_file, n_particles) print "Bunch Generated." b.getSyncParticle().kinEnergy(energy)
def generate_initial_5mm_distribution( start, half_range, parameters, Lattice, horizontal=1, output_file='Input/ParticleDistribution.in', summary_file='Input/ParticleDistribution_summary.txt', outputFormat='Orbit'): parameters['alphax0'] = Lattice.alphax0 parameters['betax0'] = Lattice.betax0 parameters['alphay0'] = Lattice.alphay0 parameters['betay0'] = Lattice.betay0 parameters['etax0'] = Lattice.etax0 parameters['etapx0'] = Lattice.etapx0 parameters['etay0'] = Lattice.etay0 parameters['etapy0'] = Lattice.etapy0 parameters['x0'] = Lattice.orbitx0 parameters['xp0'] = Lattice.orbitpx0 parameters['y0'] = Lattice.orbity0 parameters['yp0'] = Lattice.orbitpy0 parameters['gamma_transition'] = Lattice.gammaT parameters['circumference'] = Lattice.getLength() parameters['length'] = Lattice.getLength() / Lattice.nHarm # twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['beta'] * parameters['etax0'], 'etapx0': parameters['beta'] * parameters['etapx0'] } dispersiony = { 'etay0': parameters['beta'] * parameters['etay0'], 'etapy0': parameters['beta'] * parameters['etapy0'] } closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # initialize particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) Longitudinal_distribution = LongitudinalJohoDistributionSingleHarmonic( parameters, parameters['LongitudinalJohoParameter']) # only the main CPU is actually writing its distribution to a file ... comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file, "w") as fid: csv_writer = csv.writer(fid, delimiter=' ') for i in range(parameters['n_macroparticles']): # ~ (phi[i], dE[i]) = Longitudinal_distribution.getCoordinates() # EQUAL STEPS if horizontal: x[i] = (start - half_range) + ( (i * (2 * half_range)) / float(parameters['n_macroparticles'])) # z = (-phi*L)/(2*pi) # phi = (-2*pi*z)/L phi[i] = (-2 * np.pi * 2.5 * parameters['blength_rms']) / parameters["length"] print x[i], phi[i] elif not horizontal: y[i] = (start - half_range) + ( (i * (2 * half_range)) / float(parameters['n_macroparticles'])) print y[i] if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. dE[i] /= 1.e9 csv_writer.writerow( [x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) if summary_file: with open(summary_file, 'w') as fid: map( lambda key: fid.write(key + ' = ' + str(parameters[key]) + '\n'), parameters) print '\nCreated particle distribution with ' + str( parameters['n_macroparticles'] ) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
# The script will test the gerenrators from bunch_generators #-------------------------------------------------------- import math import random import sys from bunch import Bunch from orbit.bunch_generators import TwissContainer, TwissAnalysis from orbit.bunch_generators import WaterBagDist3D, GaussDist3D, KVDist3D n = 300000 #--------------------------------------------- # Water Bag 3D #--------------------------------------------- twissX = TwissContainer(alpha = 0.045, beta = 12.355, emittance = 30e-6) twissY = TwissContainer(alpha = 0.046, beta = 12.044, emittance = 30e-6) twissZ = TwissContainer(alpha = 0, beta = 100000., emittance = .01) #dist = WaterBagDist3D(twissX,twissY,twissZ) dist = KVDist3D(twissX,twissY,twissZ) #dist = GaussDist3D(twissX,twissY,twissZ) twiss_analysis = TwissAnalysis(3) twiss_analysis.init() file_out = open("KV.dat","w") for i in range(n): (x,xp,y,yp,z,zp) = dist.getCoordinates() file_out.write(str(x) + " " + str(xp) + " " + str(y) + " " + str(yp) + " "+ str(z) + " " + str(zp) + "\n") twiss_analysis.account((x,xp,y,yp,z,zp)) file_out.close()
def generate_initial_distribution_from_BLonD_manual_Twiss(parameters, TwissDict, Lattice=None, output_file='ParticleDistribution.in', outputFormat='pyOrbit', summary_file='ParticleDistribution_summary.txt', summary_mat_file=None): # Get parameters from the TwissDict dictionary parameters['alphax0'] = TwissDict['alpha_x'] parameters['betax0'] = TwissDict['beta_x'] parameters['alphay0'] = TwissDict['alpha_y'] parameters['betay0'] = TwissDict['beta_y'] parameters['etax0'] = TwissDict['D_x'] parameters['etapx0'] = TwissDict['D_xp'] parameters['etay0'] = TwissDict['D_y'] parameters['etapy0'] = TwissDict['D_yp'] parameters['x0'] = TwissDict['x0'] parameters['xp0'] = TwissDict['xp0'] parameters['y0'] = TwissDict['y0'] parameters['yp0'] = TwissDict['yp0'] parameters['gamma_transition'] = TwissDict['gamma_transition'] parameters['circumference'] = TwissDict['circumference'] parameters['length'] = TwissDict['length'] # Create Twiss containers twissX = TwissContainer(alpha = parameters['alphax0'], beta = parameters['betax0'], emittance = parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha = parameters['alphay0'], beta = parameters['betay0'], emittance = parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = {'etax0': parameters['etax0'], 'etapx0': parameters['etapx0']} dispersiony = {'etay0': parameters['etay0'], 'etapy0': parameters['etapy0']} closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # Initialize empty particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) z = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) # Instatiate the classes for longitudinal and transverse distns Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=parameters['TransverseCut']) # Open BLonD file BLonD_data = np.load(parameters['BLonD_file']) # Old iterative method # ~ for i in range(parameters['n_macroparticles']): # ~ try: # ~ # Set co-ordinates # ~ z[i] = BLonD_data['dz'][i] # ~ phi[i] = -1 * z[i] * h_main / R # ~ dE[i] = (BLonD_data['dE'][i] / 1E9) # in eV # ~ print i, ': ', z[i] # ~ except IndexError: # ~ print 'ERROR: pyOrbit_GenerateInitialDistribution::generate_initial_distribution_from_BLonD' # ~ print parameters['BLonD_file'], ' does not contain enough particles to fill the bunch co-ordinates' # ~ exit(0) if len(BLonD_data['dz']) <= (parameters['n_macroparticles']-1): print 'generate_initial_distribution_from_BLonD::Error: input array length', len(BLonD_data['dz']), ' does not meet number of requested particles', parameters['n_macroparticles'] exit(0) if len(BLonD_data['dE']) <= (parameters['n_macroparticles']-1): print 'generate_initial_distribution_from_BLonD::Error: input file length', len(BLonD_data['dE']), ' does not meet number of requested particles', parameters['n_macroparticles'] exit(0) z = BLonD_data['dz'] dE = (BLonD_data['dE']/ 1E9) # We need to convert z into phi h_main = np.atleast_1d(parameters['harmonic_number'])[0] R = parameters['circumference'] / 2 / np.pi phi = - z * h_main / R # Write the distn to a file only on one CPU comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file,"w") as fid: csv_writer = csv.writer(fid, delimiter=' ') for i in range(parameters['n_macroparticles']): phi[i] = -1 * z[i] * h_main / R # ~ (z[i], dE[i]) = Longitudinal_distribution.getCoordinates() # ~ z[i] = z[i] * speed_of_light * parameters['beta'] * 1e-9 # convert ns to s and then m # ~ dE[i] = dE[i] * 1e-3 # convert from MeV to GeV (x[i], xp[i], y[i], yp[i]) = Transverse_distribution.getCoordinates() x[i] += closedOrbitx['x0'] xp[i] += closedOrbitx['xp0'] y[i] += closedOrbity['y0'] yp[i] += closedOrbity['yp0'] dpp = dE[i] / (parameters['energy']) / parameters['beta']**2 * 1E9 #print '\n dpp = ', dpp x[i] += dpp * dispersionx['etax0'] xp[i] += dpp * dispersionx['etapx0'] y[i] += dpp * dispersiony['etay0'] yp[i] += dpp * dispersiony['etapy0'] # ~ if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. # ~ dE[i] /= 1.e9 # ~ if outputFormat == 'Orbit': map(lambda i: csv_writer.writerow([x[i], xp[i], y[i], yp[i], phi[i], dE[i]]), range(parameters['n_macroparticles'])) # ~ elif outputFormat == 'pyOrbit': # ~ map(lambda i: csv_writer.writerow([x[i], xp[i], y[i], yp[i], z[i], dE[i]]), range(parameters['n_macroparticles'])) if summary_file: with open(summary_file, 'w') as fid: map(lambda key: fid.write(key + ' = ' + str(parameters[key]) + '\n'), parameters) print '\nCreated particle distribution with ' + str(parameters['n_macroparticles']) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
positioni, positionf, setDict, paramsDict, seed_value=40) #ESet = AddErrorSet(lattice, positioni, positionf, setDict, paramsDict) # Random print "FIELD ERRORS IN THE QUADRUPOLES/MULTIPOLES ARE APPLIED" matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, b) TwissDataX1, TwissDataY1 = matrix_lattice.getRingTwissDataX( ), matrix_lattice.getRingTwissDataY() #-------------------------------------------------------------------------------- n_particles = 5000 twissZ = TwissContainer(alpha=0., beta=100000000., emittance=sigma_p) b.macroSize(total_macroSize / n_particles) #---------------------------------------------SPLIT LONG ELEMENTS-------------------------------------- print("split long elements") for node in lattice.getNodes(): if node.getLength() > 1.0: node.setnParts(int(node.getLength() // 1 + 1)) #----------------------------Add Space Charge nodes---------------------------------------------------- sc_path_length_min = 0.000001 sizeX = 32 #number of grid points in horizontal direction sizeY = 32 #number of grid points in vertical direction sizeZ = 16 #number of longitudinal slices
def generate_initial_distribution_y0( parameters, output_file='Input/ParticleDistribution.in', summary_file='Input/ParticleDistribution_summary.txt', outputFormat='Orbit', orientation='x'): # twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['beta'] * parameters['etax0'], 'etapx0': parameters['beta'] * parameters['etapx0'] } dispersiony = { 'etay0': parameters['beta'] * parameters['etay0'], 'etapy0': parameters['beta'] * parameters['etapy0'] } closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} emittance_x = parameters['epsn_x'] / parameters['gamma'] / parameters[ 'beta'] emittance_y = parameters['epsn_y'] / parameters['gamma'] / parameters[ 'beta'] gamma_x = (1. + parameters['alphax0']**2) / parameters['betax0'] gamma_y = (1. + parameters['alphay0']**2) / parameters['betay0'] n_macroparticles = (parameters['n_macroparticles']) # initialize particle arrays x = np.zeros(parameters['n_macroparticles']) x.fill(100 * emittance_x / gamma_x / n_macroparticles) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) y.fill(100 * emittance_y / gamma_y / n_macroparticles) yp = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) phi.fill(parameters['phi_s']) dE = np.zeros(parameters['n_macroparticles']) if orientation == 'y': y = np.linspace(emittance_y / gamma_y / n_macroparticles, 300 * emittance_y / gamma_y, n_macroparticles) * parameters['TransverseCut'] else: x = np.linspace(emittance_x / gamma_x / n_macroparticles, 300 * emittance_x / gamma_x, n_macroparticles) * parameters['TransverseCut'] x = x.flatten() y = y.flatten() if orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD) == 0: fid = open(output_file, "w") csv_writer = csv.writer(fid, delimiter=' ') for i in range(len(x)): if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. dE[i] /= 1.e9 csv_writer.writerow([x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) fid.close() fid = open(summary_file, 'w') parameter_list = [ 'circumference', 'rf_voltage', 'phi_s', 'harmonic_number', 'gamma_transition', 'n_macroparticles', 'energy', 'gamma', 'bunch_length', 'LongitudinalCut', 'LongitudinalJohoParameter', 'x0', 'xp0', 'betax0', 'alphax0', 'etax0', 'etapx0', 'y0', 'yp0', 'betay0', 'alphay0', 'etay0', 'etapy0', 'epsn_x', 'epsn_y', 'TransverseCut' ] for key in parameter_list: fid.write(key + ' = ' + str(parameters[key]) + '\n') fid.close() print '\nCreated particle distribution with ' + str( len(x)) + ' macroparticles into file: ', output_file return output_file
def generate_initial_poincare_distribution( n_sigma, parameters, Lattice, horizontal=1, output_file='Input/ParticleDistribution.in', summary_file='Input/ParticleDistribution_summary.txt', outputFormat='Orbit'): parameters['alphax0'] = Lattice.alphax0 parameters['betax0'] = Lattice.betax0 parameters['alphay0'] = Lattice.alphay0 parameters['betay0'] = Lattice.betay0 parameters['etax0'] = Lattice.etax0 parameters['etapx0'] = Lattice.etapx0 parameters['etay0'] = Lattice.etay0 parameters['etapy0'] = Lattice.etapy0 parameters['x0'] = Lattice.orbitx0 parameters['xp0'] = Lattice.orbitpx0 parameters['y0'] = Lattice.orbity0 parameters['yp0'] = Lattice.orbitpy0 parameters['gamma_transition'] = Lattice.gammaT parameters['circumference'] = Lattice.getLength() parameters['length'] = Lattice.getLength() / Lattice.nHarm # twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['beta'] * parameters['etax0'], 'etapx0': parameters['beta'] * parameters['etapx0'] } dispersiony = { 'etay0': parameters['beta'] * parameters['etay0'], 'etapy0': parameters['beta'] * parameters['etapy0'] } closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # initialize particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) # building the distributions # ~ Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=parameters['TransverseCut']) Transverse_distribution = KVDist1D(twissX) Longitudinal_distribution = LongitudinalJohoDistributionSingleHarmonic( parameters, parameters['LongitudinalJohoParameter']) # only the main CPU is actually writing its distribution to a file ... comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file, "w") as fid: csv_writer = csv.writer(fid, delimiter=' ') for i in range(parameters['n_macroparticles']): # RANDOM UNIFORM # ~ x[i] = random.uniform(0., n_sigma) * np.sqrt(parameters['betax0'] * parameters['epsn_x']) # EQUAL STEPS if horizontal: # ~ print 'beta = ',parameters['beta'], ' gamma = ', parameters['gamma'] x[i] = i * float(n_sigma / float( parameters['n_macroparticles'])) * np.sqrt( float(parameters['betax0']) * (parameters['epsn_x'] / (parameters['beta'] * parameters['gamma']))) # ~ x[i] = i * float(n_sigma/float(parameters['n_macroparticles'])) # ~ print 'Generate Poincare Distn: maximum sigma = ', n_sigma, ', current particle = ', i * float(n_sigma/float(parameters['n_macroparticles'])), ' sigma' # ~ print 'eps_geo_x = ', ( parameters['epsn_x'] / (parameters['beta'] * parameters['gamma'])) # ~ print 'beta_x0 = ', float(parameters['betax0']) # ~ print 'sigma_x = ', np.sqrt(float(parameters['betax0']) * ( parameters['epsn_x'] / (parameters['beta'] * parameters['gamma']))) elif not horizontal: # ~ print '\nVERTICAL BUNCH: n_sigma = ',n_sigma, ', sigma = ', (np.sqrt(parameters['betay0'] * parameters['epsn_y'])) # ~ print '\ty =', i * (n_sigma/parameters['n_macroparticles']) * np.sqrt(parameters['betay0'] * parameters['epsn_y']) # ~ print '\ti = ', i, ', betay0 = ', parameters['betay0'], ', epsn_y = ', parameters['epsn_y'], ', macroparticles = ', parameters['n_macroparticles'] # ~ print '\tsqrt(bet*eps) = ', np.sqrt(parameters['betay0'] * parameters['epsn_y']) # ~ print '\tn_sigma/macro = ', float(n_sigma/float(parameters['n_macroparticles'])) y[i] = i * float(n_sigma / float( parameters['n_macroparticles'])) * np.sqrt( float(parameters['betay0']) * (parameters['epsn_y'] / (parameters['beta'] * parameters['gamma']))) if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. dE[i] /= 1.e9 csv_writer.writerow( [x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) #csv_writer.writerow([x[i], xp[i], y[i], yp[i], z[i], dE[i]]) if summary_file: with open(summary_file, 'w') as fid: map( lambda key: fid.write(key + ' = ' + str(parameters[key]) + '\n'), parameters) print '\nCreated particle distribution with ' + str( parameters['n_macroparticles'] ) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
beam_current = 0.038 N_particles = 10000 macrosize = (beam_current / frequency) macrosize /= (math.fabs(bunch.charge()) * si_e_charge) macrosize /= N_particles (alphaX, betaX, emittX) = (-1.39, 0.126, 3.67 * 1.0e-6) (alphaY, betaY, emittY) = (2.92, 0.281, 3.74 * 1.0e-6) (alphaZ, betaZ, emittZ) = (0.0, 117.0, 0.0166 * 1.0e-6) #---- we artificially increase the emittances to see apertures effects emittX *= 5. emittY *= 10. twissX = TwissContainer(alphaX, betaX, emittX) twissY = TwissContainer(alphaY, betaY, emittY) twissZ = TwissContainer(alphaZ, betaZ, emittZ) distributor = WaterBagDist3D(twissX, twissY, twissZ) for ind in range(N_particles): (x, xp, y, yp, z, dE) = distributor.getCoordinates() (x, xp, y, yp, z, dE) = orbit_mpi.MPI_Bcast((x, xp, y, yp, z, dE), data_type, main_rank, comm) if (ind % size == rank): bunch.addParticle(x, xp, y, yp, z, dE) nParticlesGlobal = bunch.getSizeGlobal() if (rank == 0): print "total number of particles =", nParticlesGlobal
def generate_initial_long_poincare_distribution( z_offset, parameters, Lattice, zero_particle=True, output_file='Input/ParticleDistribution.in', summary_file='Input/ParticleDistribution_summary.txt', outputFormat='Orbit'): parameters['alphax0'] = Lattice.alphax0 parameters['betax0'] = Lattice.betax0 parameters['alphay0'] = Lattice.alphay0 parameters['betay0'] = Lattice.betay0 parameters['etax0'] = Lattice.etax0 parameters['etapx0'] = Lattice.etapx0 parameters['etay0'] = Lattice.etay0 parameters['etapy0'] = Lattice.etapy0 parameters['x0'] = Lattice.orbitx0 parameters['xp0'] = Lattice.orbitpx0 parameters['y0'] = Lattice.orbity0 parameters['yp0'] = Lattice.orbitpy0 parameters['gamma_transition'] = Lattice.gammaT parameters['circumference'] = Lattice.getLength() parameters['length'] = Lattice.getLength() / Lattice.nHarm # twiss containers twissX = TwissContainer(alpha=parameters['alphax0'], beta=parameters['betax0'], emittance=parameters['epsn_x'] / parameters['gamma'] / parameters['beta']) twissY = TwissContainer(alpha=parameters['alphay0'], beta=parameters['betay0'], emittance=parameters['epsn_y'] / parameters['gamma'] / parameters['beta']) dispersionx = { 'etax0': parameters['beta'] * parameters['etax0'], 'etapx0': parameters['beta'] * parameters['etapx0'] } dispersiony = { 'etay0': parameters['beta'] * parameters['etay0'], 'etapy0': parameters['beta'] * parameters['etapy0'] } closedOrbitx = {'x0': parameters['x0'], 'xp0': parameters['xp0']} closedOrbity = {'y0': parameters['y0'], 'yp0': parameters['yp0']} # initialize particle arrays x = np.zeros(parameters['n_macroparticles']) xp = np.zeros(parameters['n_macroparticles']) y = np.zeros(parameters['n_macroparticles']) yp = np.zeros(parameters['n_macroparticles']) phi = np.zeros(parameters['n_macroparticles']) dE = np.zeros(parameters['n_macroparticles']) # building the distributions # ~ Transverse_distribution = GaussDist2D(twissX, twissY, cut_off=parameters['TransverseCut']) # ~ Transverse_distribution = KVDist1D(twissX) # ~ Longitudinal_distribution = LongitudinalJohoDistributionSingleHarmonic(parameters, parameters['LongitudinalJohoParameter']) # only the main CPU is actually writing its distribution to a file ... comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD if orbit_mpi.MPI_Comm_rank(comm) == 0: with open(output_file, "w") as fid: csv_writer = csv.writer(fid, delimiter=' ') h_main = np.atleast_1d(parameters['harmonic_number'])[0] R = parameters['circumference'] / 2 / np.pi #phi = - z * h_main / R for i in range(parameters['n_macroparticles']): if zero_particle: if i == 0: phi[i] = -z_offset * h_main / R else: phi[i] = i * -z_offset * h_main / R if outputFormat == 'Orbit': x[i] *= 1000. xp[i] *= 1000. y[i] *= 1000. yp[i] *= 1000. dE[i] /= 1.e9 # ~ csv_writer.writerow([x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) csv_writer.writerow([x[i], xp[i], y[i], yp[i], phi[i], dE[i]]) # ~ if summary_file: # ~ with open(summary_file, 'w') as fid: # ~ map(lambda key: fid.write(key + ' = ' + str(parameters[key]) + '\n'), parameters) print '\nCreated particle distribution with ' + str( parameters['n_macroparticles'] ) + ' macroparticles into file: ', output_file orbit_mpi.MPI_Barrier(comm) return output_file
matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) (muX, arrPosAlphaX, arrPosBetaX) = matrix_lattice.getRingTwissDataX() (muY, arrPosAlphaY, arrPosBetaY) = matrix_lattice.getRingTwissDataY() (arrDispX, arrDispPrimeX) = matrix_lattice.getRingDispersionDataX() (arrDispY, arrDispPrimeY) = matrix_lattice.getRingDispersionDataY() alpax = arrPosAlphaX[0][1] betax = arrPosBetaX[0][1] alpay = arrPosAlphaY[0][1] betay = arrPosBetaY[0][1] #============add particles to beam=============== # ini. 3D match distribution twissX = TwissContainer(alpha=alpax, beta=betax, emittance=emittance_x / 2) twissY = TwissContainer(alpha=alpay, beta=betay, emittance=emittance_y / 2) dist = KVDist2D(twissX, twissY) for i in range(NPIC): (x, xp, y, yp) = dist.getCoordinates() bunch.addParticle(x + offset, xp, y, yp, 0.0, 0.0) #============add particles to beam=============== xAvg = [] yAvg = [] s = [] #============Get BPMs signal=============== nodes = lattice.getNodes() for i in range(1):
#===================== m a t c h i n g =====================# solve = EnvelopeSolver(beamline) Ksc = Optics().getPerveance(b, beamline[-1].data['s'], emitx, emity) twiss = solve.match_twiss_matrix(emitx, emity, 0.0, 0.0) twiss_sc = solve.match_twiss_matrix(emitx, emity, sigma_p, Ksc) s_tw = [x.data['s'] for x in beamline] #===================== m a t c h i n g =====================# n_particles = 1000 twissX = TwissContainer(alpha=arrPosAlphaX[0][1], beta=arrPosBetaX[0][1], emittance=emitx) #twissX = TwissContainer(alpha = twiss[0,1], beta = twiss[0,0], emittance = emitx) twissY = TwissContainer(alpha=arrPosAlphaY[0][1], beta=arrPosBetaY[0][1], emittance=emity) twissZ = TwissContainer(alpha=0., beta=100000000., emittance=0.001) dist = GaussDist3D(twissX, twissY, twissZ) dist = KVDist3D(twissX, twissY, twissZ) for i in range(n_particles): (x, xp, y, yp, z, zp) = dist.getCoordinates() b.addParticle(x, xp, y, yp, z, zp) x_space0 = [(b.x(i), b.px(i)) for i in range(n_particles)]
#-------------------------------------------------------- import math import random import sys from orbit.bunch_generators import TwissContainer, KVDist1D, WaterBagDist1D, GaussDist1D from orbit.bunch_generators import TwissAnalysis # number of particles generated n = 10000 #--------------------------------------------- # KV 1D #--------------------------------------------- twiss = TwissContainer(1.,2.,3.) dist = KVDist1D(twiss) twiss_analysis = TwissAnalysis(1) for i in range(n): (x,y) = dist.getCoordinates() twiss_analysis.account((x,y)) print "=================================================" print "KV 1D - done!" print " alpha beta [m/rad] gamma emitt[m*rad] " print "Twiss X %12.5g %12.5g %12.5g %12.5g "%twiss.getAlphaBetaGammaEmitt() print "Generated X %12.5g %12.5g %12.5g %12.5g "%twiss_analysis.getTwiss(0) print "=================================================" #--------------------------------------------- # Waterbag 1D
"""WaterBagDist3D Parmila input ;input -2 100000 -1.9619 18.3140 0.0021824 ! Toutatis 38 mA emittance to 7mm into MEBT ; 1.7681 16.1030 0.0021856 ; 0.0196 58.2172 0.003088 ;input -2 100000 -1.9899 19.6360 0.0017573 ! exact match to rms properties of ReadDist distribution ; 1.92893 17.778 0.0017572 ; 0.015682 67.0939 0.002420 """ #--------------------------------------------- # Set up Twiss for X,Y,Z #--------------------------------------------- twissX = TwissContainer(alpha = alpha_x, beta = beta_x, emittance = emitt_x) twissY = TwissContainer(alpha = alpha_y, beta = beta_y, emittance = emitt_y) twissZ = TwissContainer(alpha = alpha_z, beta = beta_z, emittance = emitt_z) print "-------------------input parameters----------------------------------------" print "X alpha= %12.5g beta [cm/rad] =%12.5g emitt[cm*rad] = %12.5g "%(alpha_x,beta_x,emitt_x) print "Y alpha= %12.5g beta [cm/rad] =%12.5g emitt[cm*rad] = %12.5g "%(alpha_y,beta_y,emitt_y) print "Z alpha= %12.5g beta [deg/MeV] =%12.5g emitt[deg*MeV] = %12.5g "%(alpha_z,beta_z,emitt_z) #distributor = GaussDist3D(twissX,twissY,twissZ,cut_off = 4.0) distributor = WaterBagDist3D(twissX,twissY,twissZ) #distributor = KVDist3D(twissX,twissY,twissZ) twiss_analysis = TwissAnalysis(3) #---------------------------------- #open file for Parmila