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
