def PrintAllParticles(self, filename=None):
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
if filename is None:
filename = 'Particles_all.dat'
# if file exists then append
if os.path.exists(filename):
f = open(filename, "a+")
# if file doesn't exist create and add header
else:
f = open(filename, "w+")
f.write("#ParticleID\tTurn\tx[m]\txp\ty[m]\typ\tz[m]\tdE[GeV]")
for n in self.particle_list:
for t in self.turn_list:
f.write("\n%i\t%i\t%f\t%f\t%f\t%f\t%f\t%f" % ( \
n, t, \
self.particles[str(n)][str(t)]['x'], \
self.particles[str(n)][str(t)]['xp'], \
self.particles[str(n)][str(t)]['y'], \
self.particles[str(n)][str(t)]['yp'], \
self.particles[str(n)][str(t)]['z'], \
self.particles[str(n)][str(t)]['dE'] ))
f.close()
def LinearRestoringForce(bunch, force):
rank = 0
numprocs = 1
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
numprocs = orbit_mpi.MPI_Comm_size(comm)
nparts_arr_local = []
for i in range(numprocs):
nparts_arr_local.append(0)
nparts_arr_local[rank] = bunch.getSize()
data_type = mpi_datatype.MPI_INT
op = mpi_op.MPI_SUM
nparts_arr = orbit_mpi.MPI_Allreduce(nparts_arr_local, data_type, op, comm)
for i in range(bunch.getSize()):
en = bunch.dE(i)
en = en + bunch.z(i) * force
bunch.dE(i, en)
return
def getBunch(self,
nParticles=0,
distributorClass=WaterBagDist3D,
cut_off=-1.):
"""
Returns the pyORBIT bunch with particular number of particles.
"""
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
data_type = mpi_datatype.MPI_DOUBLE
main_rank = 0
bunch = Bunch()
self.bunch.copyEmptyBunchTo(bunch)
macrosize = (self.beam_current * 1.0e-3 / self.bunch_frequency)
macrosize /= (math.fabs(bunch.charge()) * self.si_e_charge)
distributor = distributorClass(self.twiss[0], self.twiss[1],
self.twiss[2], cut_off)
bunch.getSyncParticle().time(0.)
for i in range(nParticles):
(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 (i % size == rank):
bunch.addParticle(x, xp, y, yp, z, dE)
nParticlesGlobal = bunch.getSizeGlobal()
bunch.macroSize(macrosize / nParticlesGlobal)
return bunch
def writeStatLats(self, s, bunch, lattlength=0):
self.bunchtwissanalysis.analyzeBunch(bunch)
emitx = self.bunchtwissanalysis.getEmittance(0)
betax = self.bunchtwissanalysis.getBeta(0)
alphax = self.bunchtwissanalysis.getAlpha(0)
betay = self.bunchtwissanalysis.getBeta(1)
alphay = self.bunchtwissanalysis.getAlpha(1)
emity = self.bunchtwissanalysis.getEmittance(1)
dispersionx = self.bunchtwissanalysis.getDispersion(0)
ddispersionx = self.bunchtwissanalysis.getDispersionDerivative(0)
#dispersiony = self.bunchtwissanalysis.getDispersion(1, bunch)
#ddispersiony = self.bunchtwissanalysis.getDispersionDerivative(1, bunch)
sp = bunch.getSyncParticle()
time = sp.time()
if lattlength > 0:
time = sp.time() / (lattlength / (sp.beta() * speed_of_light))
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.file_out.write(
str(s) + "\t" + str(time) + "\t" + str(emitx) + "\t" +
str(emity) + "\t" + str(betax) + "\t" + str(betay) + "\t" +
str(alphax) + "\t" + str(alphay) + "\t" + str(dispersionx) +
"\t" + str(ddispersionx) + "\n")
def PrintParticleForTurn(self, turn, n, filename=None):
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
if filename is None:
filename = 'Particle_' + str(n) + '_turn_' + str(turn) + '.dat'
# Check that the particle exists
if n not in self.particle_list:
print "Particle_output_dictionary::print_particle_for_turn: Particle not stored, use AddNewParticle(n) before tracking."
else:
# if file exists then append
if os.path.exists(filename):
f = open(filename, "a+")
# if file doesn't exist create and add header
else:
f = open(filename, "w+")
f.write(
"#ParticleID\tTurn\tx[m]\txp\ty[m]\typ\tz[m]\tdE[GeV]")
f.write("\n%i\t%i\t%f\t%f\t%f\t%f\t%f\t%f" % ( \
n, turn, \
self.particles[str(n)][str(turn)]['x'], \
self.particles[str(n)][str(turn)]['xp'], \
self.particles[str(n)][str(turn)]['y'], \
self.particles[str(n)][str(turn)]['yp'], \
self.particles[str(n)][str(turn)]['z'], \
self.particles[str(n)][str(turn)]['dE'] ))
f.close()
def Update(self, bunch, turn, verbose=False):
self.update_flag = 1
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
for n in self.particle_list:
# Create the turn dictionary
self.particles[str(n)][str(turn)] = {
} # Second level : N-2 : Turn
# self.particles[index][turn]['x'] = bunch.x(index)
self.particles[str(n)][str(turn)]['x'] = bunch.x(n)
self.particles[str(n)][str(turn)]['xp'] = bunch.xp(n)
self.particles[str(n)][str(turn)]['y'] = bunch.y(n)
self.particles[str(n)][str(turn)]['yp'] = bunch.yp(n)
self.particles[str(n)][str(turn)]['z'] = bunch.z(n)
self.particles[str(n)][str(turn)]['dE'] = bunch.dE(n)
self.turn_list.append(turn)
if verbose:
print "Particle_output_dictionary::update: Added turn %i" % (turn)
print "Dictionary now:"
print self.particles
def writeMoments(self, s, bunch, lattlength=0):
sp = bunch.getSyncParticle()
time = sp.time()
if lattlength > 0:
time = sp.time() / (lattlength / (sp.beta() * speed_of_light))
self.bunchtwissanalysis.computeBunchMoments(bunch, self.order,
self.dispterm,
self.emitnormterm)
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.file_out.write(str(s) + "\t" + str(time) + "\t")
for i in range(0, self.order + 1):
for j in range(0, i + 1):
self.file_out.write(
str(self.bunchtwissanalysis.getBunchMoment(i - j, j)) +
"\t")
self.file_out.write("\n")
def generate_initial_distribution_FMA(parameters, output_file = 'Input/ParticleDistribution.in', summary_file = 'Input/ParticleDistribution_summary.txt', outputFormat='Orbit', triangular_grid = True):
# 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']); phi.fill(parameters['phi_s'])
dE = np.zeros(parameters['n_macroparticles'])
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_sqrt = np.floor(np.sqrt(parameters['n_macroparticles']))
Jx = np.linspace(emittance_x/gamma_x/n_macroparticles_sqrt/10, emittance_x/gamma_x, n_macroparticles_sqrt)
Jy = np.linspace(emittance_y/gamma_y, emittance_y/gamma_y/n_macroparticles_sqrt/10, n_macroparticles_sqrt)
x, y = np.meshgrid(np.sqrt(Jx)*parameters['TransverseCut'], np.sqrt(Jy)*parameters['TransverseCut'])
if triangular_grid:
indcs = np.tril_indices(len(x))
x = x[indcs]
y = y[indcs]
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 bunch_from_matfile(matfile):
d = sio.loadmat(matfile, squeeze_me=True)
p = dict((key, value) for (key, value) in map(
lambda k: (k, d['particles'][k][()]), d['particles'].dtype.names))
attributes = list(set(p) - set(['x', 'xp', 'y', 'yp', 'z', 'dE']))
attributes.sort(key=str.lower)
bunch = Bunch()
bunch.classicalRadius(d['bunchparameters']['classical_radius'])
bunch.charge(d['bunchparameters']['charge'])
bunch.mass(d['bunchparameters']['mass'])
bunch.getSyncParticle().momentum(d['bunchparameters']['momentum'])
bunch.getSyncParticle().time(d['bunchparameters']['time'])
x = np.atleast_1d(d['particles']['x'][()])
xp = np.atleast_1d(d['particles']['xp'][()])
y = np.atleast_1d(d['particles']['y'][()])
yp = np.atleast_1d(d['particles']['yp'][()])
z = np.atleast_1d(d['particles']['z'][()])
dE = np.atleast_1d(d['particles']['dE'][()])
n_part = len(x)
import orbit_mpi
comm = bunch.getMPIComm()
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
count = n_part / size
remainder = n_part % size
if (rank < remainder):
i_start = rank * (count + 1)
i_stop = i_start + count + 1
else:
i_start = rank * count + remainder
i_stop = i_start + count
# print rank, i_start, i_stop
map(lambda i: bunch.addParticle(x[i], xp[i], y[i], yp[i], z[i], dE[i]),
xrange(i_start, i_stop))
orbit_mpi.MPI_Barrier(comm)
for a in attributes:
bunch.addPartAttr(a)
a_size = bunch.getPartAttrSize(a)
if a_size > 1:
for j in xrange(a_size):
map(
lambda
(ip, i): bunch.partAttrValue(a, ip, j,
np.atleast_1d(p[a][j])[i]),
enumerate(xrange(i_start, i_stop)))
else:
map(
lambda (ip, i): bunch.partAttrValue(a, ip, 0,
np.atleast_1d(p[a])[i]),
enumerate(xrange(i_start, i_stop)))
orbit_mpi.MPI_Barrier(comm)
return bunch
def call(*args, **kwargs):
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
if not rank:
result = func(*args, **kwargs)
else:
result = None
orbit_mpi.MPI_Barrier(comm)
return result
def write_SextupoleRamp_files(target_file, pattern, ptc_source_table):
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
if not rank:
with open(target_file, 'w') as fid:
fid.write('SET ORBIT RAMPING \n')
fid.write(' ramp\n %s\t"%s"\t%1.9f \n' %
(pattern, ptc_source_table, 1.0))
fid.write('return')
orbit_mpi.MPI_Barrier(comm)
def update(self):
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
map(
lambda key: self.output_dict[key].append(
self.parameter_functions[key]()), self.entries)
if self.print_header_flag:
self.print_header()
self.print_header_flag = 0
self.print_last()
def __init__(self,f_name,title):
self.f_name = f_name
rank = orbit_mpi.MPI_Comm_rank(mpi_comm.MPI_COMM_WORLD)
if (rank == 0):
f_out = open(self.f_name,"w")
for i in range(len(title)):
f_out.write("%19s"%title[i])
f_out.write("\n")
f_out.close()
def UpdatePTCTwiss(self, Lattice, turn, verbose=False):
self.update_flag = 1
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
# Create the turn dictionary
self.twiss_dict[int(turn)] = {} # Second level : N-2 : Turn
# Third level: twiss
self.twiss_dict[int(turn)]['beta_x'] = ([
n.getParamsDict()['betax'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['beta_y'] = ([
n.getParamsDict()['betay'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['alpha_x'] = ([
n.getParamsDict()['alphax'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['alpha_y'] = ([
n.getParamsDict()['alphay'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['eta_x'] = ([
n.getParamsDict()['etax'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['eta_y'] = ([
n.getParamsDict()['etay'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['eta_px'] = ([
n.getParamsDict()['etapx'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['eta_py'] = ([
n.getParamsDict()['etapy'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['orbit_x'] = ([
n.getParamsDict()['orbitx'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['orbit_px'] = ([
n.getParamsDict()['orbitpx'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['orbit_y'] = ([
n.getParamsDict()['orbity'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['orbit_py'] = ([
n.getParamsDict()['orbitpy'] for n in Lattice.getNodes()
])
self.twiss_dict[int(turn)]['s'] = np.cumsum(
[n.getLength() for n in Lattice.getNodes()])
self.turn_list.append(turn)
if verbose:
print "PTCLatticeFunctionsDictionary::update: Added turn %i" % (
turn)
def orbitFinalize(message = None): """ Method. Finalizes the execution of the ORBIT script. """ import orbit_mpi import sys import traceback if(orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD) == 0): print "ORBIT message: ", message traceback.print_stack() print "STOP" sys.exit(1)
def call(*args, **kwargs):
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
#print 'rank %i before executing the function'%rank
if not rank:
#print 'rank %i executing the function'%rank
result = func(*args, **kwargs)
else:
result = None
#print 'rank %i before the barrier'%rank
orbit_mpi.MPI_Barrier(comm)
#print 'rank %i after the barrier'%rank
return result
def write_RFtable(filename, harmonic_factors, time, E_kin, RF_voltage, RF_phase):
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
if not rank:
n_lines = len(time)
n_harmonics = len(harmonic_factors)
arr = np.vstack((time,E_kin, np.dstack((RF_voltage,RF_phase)).flatten().reshape(n_lines, 2*n_harmonics).T)).T
with open(filename, 'w') as fid:
fid.write('%d 1 1 0 %d\n'%(n_lines, n_harmonics))
fid.write(' '.join(map(lambda i: '%d'%i, harmonic_factors))+'\n')
for j in xrange(n_lines):
fid.write('\t'.join(map(lambda i: '%1.8f'%i, arr[j, :]))+'\n')
orbit_mpi.MPI_Barrier(comm)
def fdata(self,data):
rank = orbit_mpi.MPI_Comm_rank(mpi_comm.MPI_COMM_WORLD)
if (rank == 0):
f_out = open(self.f_name,"a")
for i in range(len(data)):
if (type(data[i]) == types.StringType):
f_out.write("%19s"%data[i])
if (type(data[i]) == types.IntType):
f_out.write("%19i"%data[i])
if (type(data[i]) == types.FloatType):
f_out.write("%19.4f"%data[i])
f_out.write("\n")
f_out.close()
def write_QuadRamp_files(target_file, twissfile, pattern, ptc_source_table):
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
if not rank:
t = metaclass.twiss(twissfile)
q_i = [i for i, n in enumerate(t.NAME) if pattern in n]
with open(target_file, 'w') as fid:
fid.write('SET ORBIT RAMPING \n')
for i in q_i:
fid.write(' ramp\n %s\t"%s"\t%1.9f \n' %
(t.NAME[i], ptc_source_table, (t.K1L[i] / t.L[i]) /
(t.K1L[q_i[0]] / t.L[q_i[0]])))
fid.write('return')
orbit_mpi.MPI_Barrier(comm)
def write_PTCtable(filename, multipole_orders, time, normal_components, skew_components):
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
factor = 1./np.math.factorial(multipole_orders-1) # the factorial factor is needed to be consistent with MADX
if not rank:
n_lines = len(time)
n_multipoles = 1 # number of multipole orders to be changed (for the moment only 1 is implemented)
arr = np.vstack((time,normal_components*factor,skew_components*factor)).T
with open(filename, 'w') as fid:
fid.write('%d 1 %d\n'%(n_lines, n_multipoles))
fid.write(' %d\n'%multipole_orders)
for j in xrange(n_lines):
fid.write('\t'.join(map(lambda i: '%1.11f'%i, arr[j, :]))+'\n')
orbit_mpi.MPI_Barrier(comm)
def PrintOrbitExtrema(self, lattice_folder='.'):
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
filename = lattice_folder + '/Orbit_Extrema.dat'
f = open(filename, "w")
f.write(
'# Turn\tOrbit_Min_x\tOrbit_Max_x\tOrbit_Min_y\tOrbit_Max_y')
for turn in self.turn_list:
f.write("\n%i\t%f\t%f\t%f\t%f" % (turn, \
np.min(self.twiss_dict[int(turn)]['orbit_x']), \
np.max(self.twiss_dict[int(turn)]['orbit_x']), \
np.min(self.twiss_dict[int(turn)]['orbit_y']), \
np.max(self.twiss_dict[int(turn)]['orbit_y'])))
f.close()
def getAutoionizationBunch(self, mult, b, time_par):
rank = orbit_mpi.MPI_Comm_rank(mpi_comm.MPI_COMM_WORLD)
random.seed((rank + 1) * 12571 + time_par * int(time.time()))
N_evol = int(b.getPartAttrDicts()['Evolution']['size']) - 5
bunch = Bunch()
bunch.charge(1)
bunch.mass(0.938256)
bunch_unstr = Bunch()
bunch_unstr.charge(0)
bunch_unstr.mass(b.mass())
for i in range(b.getSize()):
dt = b.partAttrValue("Evolution", i, N_evol + 1)
x0 = b.partAttrValue("Evolution", i, N_evol + 2)
y0 = b.partAttrValue("Evolution", i, N_evol + 3)
z0 = b.partAttrValue("Evolution", i, N_evol + 4)
(x1, y1, z1, px, py, pz) = (b.x(i), b.y(i), b.z(i), b.px(i),
b.py(i), b.pz(i))
p1 = b.partAttrValue("Evolution", i, N_evol)
for j in range(mult):
ran = random.random()
if (p1 <= ran):
bunch_unstr.addParticle(x1, px, y1, py, z1, pz)
else:
for k in range(N_evol):
f1 = b.partAttrValue("Evolution", i, k)
f2 = b.partAttrValue("Evolution", i, k + 1)
if (f1 <= ran and ran < f2):
coeff = (k + (ran - f1) / (f2 - f1)) / N_evol
(x, y, z) = (x0 + (x1 - x0) * coeff,
y0 + (y1 - y0) * coeff,
z0 + (z1 - z0) * coeff)
bunch.addParticle(x, px, y, py, z, pz)
break
return bunch, bunch_unstr
def analyzeSignal(self, bunch):
self.bunchtwissanalysis.analyzeBunch(bunch)
# if mpi operations are enabled, this section of code will
# determine the rank of the present node
rank = 0 # default is primary node
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
# only the primary node needs to output the calculated information
if (rank == 0):
self.xAvg = self.bunchtwissanalysis.getAverage(0)
self.xpAvg = self.bunchtwissanalysis.getAverage(1)
self.yAvg = self.bunchtwissanalysis.getAverage(2)
self.ypAvg = self.bunchtwissanalysis.getAverage(3)
def addParticleIdNumbers(b, fixedidnumber=-1, part_ind=0):
rank = 0
numprocs = 1
mpi_init = orbit_mpi.MPI_Initialized()
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
if (mpi_init):
rank = orbit_mpi.MPI_Comm_rank(comm)
numprocs = orbit_mpi.MPI_Comm_size(comm)
nparts_arr_local = []
for i in range(numprocs):
nparts_arr_local.append(0)
nparts_arr_local[rank] = b.getSize()
data_type = mpi_datatype.MPI_INT
op = mpi_op.MPI_SUM
nparts_arr = orbit_mpi.MPI_Allreduce(nparts_arr_local, data_type, op,
comm)
if (b.hasPartAttr("ParticleIdNumber") == 0):
b.addPartAttr("ParticleIdNumber")
if (fixedidnumber >= 0):
for i in range(part_ind, b.getSize()):
b.partAttrValue("ParticleIdNumber", i, 0, fixedidnumber)
else:
istart = 0
if (rank == 0):
istart = 0
else:
for i in range(rank):
istart = istart + nparts_arr[i]
for i in range(b.getSize()):
idnumber = istart + i
b.partAttrValue("ParticleIdNumber", i, 0, idnumber)
def update(self, bunch, turn):
self.update_flag = 1
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
for n in self.particle_list:
# Create the turn dictionary
self.particles[str(n)][str(turn)] = {
} # Second level : N-2 : Turn
# self.particles[index][turn]['x'] = bunch.x(index)
self.particles[str(n)][str(turn)]['x'] = bunch.x(n)
self.particles[str(n)][str(turn)]['xp'] = bunch.xp(n)
self.particles[str(n)][str(turn)]['y'] = bunch.y(n)
self.particles[str(n)][str(turn)]['yp'] = bunch.yp(n)
self.particles[str(n)][str(turn)]['z'] = bunch.z(n)
self.particles[str(n)][str(turn)]['dE'] = bunch.dE(n)
self.turn_list.append(turn)
def getBunch(self, time_par):
rank = orbit_mpi.MPI_Comm_rank(mpi_comm.MPI_COMM_WORLD)
random.seed((rank + 1) * 12571 + time_par * int(time.time()))
E = self.mass + self.TK
P = math.sqrt(E * E - self.mass * self.mass)
#vz = 299792458*P/E
#beta = P/E
gamma = E / self.mass
#bg = beta*gamma
bunch = Bunch()
bunch.charge(self.charge)
bunch.mass(self.mass)
for i in range(self.N_part):
x, xp = self.getCoords(self.alphaX, self.betaX, self.emtX,
math.sqrt(self.emtX * self.betaX) * 5.0)
y, yp = self.getCoords(self.alphaY, self.betaY, self.emtY,
math.sqrt(self.emtY * self.betaY) * 5.0)
z, zp = self.getCoords(self.alphaZ, self.betaZ, self.emtZ,
math.sqrt(self.emtZ * self.betaZ) * 5.0)
#zp = dbeta/beta XAL definition
zp = zp * gamma * gamma
#zp = dp/p
x += self.dispD * zp
xp += self.dispDP * zp
px = xp * P
py = yp * P
pz = (1 + zp) * P
bunch.addParticle(x, px, y, py, z, pz)
return bunch
def PrintPTCTwissForTurn(self, turn, lattice_folder='.', filename=None):
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
if filename is None:
filename = lattice_folder + '/PTC_Twiss_turn_' + str(
turn) + '.dat'
# Check that the particle exists
if turn not in self.turn_list:
print "PTCLatticeFunctionsDictionary::PrintPTCTwissForTurn: Turn not stored, use UpdatePTCTwiss function on this turn to store."
else:
# if file exists then overwrite
if os.path.exists(filename):
f = open(filename, "w")
# if file doesn't exist create and add header
else:
f = open(filename, "w")
f.write(
'# s\tbeta_x\tbeta_y\talpha_x\talpha_y\tD_x\tD_y\tD_px\tD_py\torbit_x\torbit_px\torbit_y\torbit_py'
)
for i in range(0, len(self.twiss_dict[int(turn)]['s']), 1):
f.write("\n%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f" % ( \
self.twiss_dict[int(turn)]['s'][i], \
self.twiss_dict[int(turn)]['beta_x'][i], \
self.twiss_dict[int(turn)]['beta_y'][i], \
self.twiss_dict[int(turn)]['alpha_x'][i], \
self.twiss_dict[int(turn)]['alpha_y'][i], \
self.twiss_dict[int(turn)]['eta_x'][i], \
self.twiss_dict[int(turn)]['eta_y'][i], \
self.twiss_dict[int(turn)]['eta_px'][i], \
self.twiss_dict[int(turn)]['eta_py'][i], \
self.twiss_dict[int(turn)]['orbit_x'][i], \
self.twiss_dict[int(turn)]['orbit_px'][i], \
self.twiss_dict[int(turn)]['orbit_y'][i], \
self.twiss_dict[int(turn)]['orbit_py'][i]))
f.close()
def getBunch(self, nParticles, twissX, twissY, twissZ, cut_off = -1.): """ Returns the pyORBIT bunch with particular number of particles. """ (x0,xp0,y0,yp0,z0,dE0) = self.init_coords comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD rank = orbit_mpi.MPI_Comm_rank(comm) size = orbit_mpi.MPI_Comm_size(comm) data_type = mpi_datatype.MPI_DOUBLE main_rank = 0 bunch = Bunch() self.bunch.copyEmptyBunchTo(bunch) macrosize = (self.beam_current*1.0e-3/self.bunch_frequency) macrosize /= (math.fabs(bunch.charge())*self.si_e_charge) distributor = GaussDist3D(twissX,twissY,twissZ, cut_off) bunch.getSyncParticle().time(0.) for i in range(nParticles): (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(i%size == rank): bunch.addParticle(x+x0,xp+xp0,y+y0,yp+yp0,z+z0,dE+dE0) nParticlesGlobal = bunch.getSizeGlobal() bunch.macroSize(macrosize/nParticlesGlobal) return bunch
def addAxisField(cls,fl_name,dir_location = ""): """ This method add to the store the axis RF field for the RF gap node. The dir_location string variable will be added to the fl_name to get the file name. Returns the axis RF field function. """ if(cls.static_axis_field_dict.has_key(fl_name)): return cls.static_axis_field_dict[fl_name] comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD data_type = mpi_datatype.MPI_DOUBLE rank = orbit_mpi.MPI_Comm_rank(comm) main_rank = 0 x_arr = [] y_arr = [] if(rank == 0): fl_in = open(dir_location + fl_name,"r") lns = fl_in.readlines() fl_in.close() for ln in lns: res_arr = ln.split() if(len(res_arr) == 2): x = float(res_arr[0]) y = float(res_arr[1]) x_arr.append(x) y_arr.append(y) x_arr = orbit_mpi.MPI_Bcast(x_arr,data_type,main_rank,comm) y_arr = orbit_mpi.MPI_Bcast(y_arr,data_type,main_rank,comm) function = Function() for ind in range(len(x_arr)): function.add(x_arr[ind],y_arr[ind]) #---- setting the const step (if function will allow it) #---- will speed up function calculation later function.setConstStep(1) cls.static_axis_field_dict[fl_name] = function return function
if slicebyslice:
#PIC
from orbit.space_charge.sc2p5d import scAccNodes, scLatticeModifications
from spacecharge import SpaceChargeCalcAnalyticGaussian
from spacecharge import InterpolatedLineDensityProfile
from lib.output_dictionary import *
from lib.particle_output_dictionary import *
from lib.pyOrbit_GenerateInitialDistribution2 import *
from lib.save_bunch_as_matfile import *
from lib.pyOrbit_LinearRestoringForce import *
print "Start ..."
comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
rank = orbit_mpi.MPI_Comm_rank(comm)
#----------------------------------------------
# Create folder structure
#----------------------------------------------
from lib.mpi_helpers import mpi_mkdir_p
mpi_mkdir_p('input')
mpi_mkdir_p('output')
mpi_mkdir_p('lost')
#----------------------------------------------
# Generate Lattice (MADX + PTC)
#----------------------------------------------
if not rank:
os.system("/afs/cern.ch/eng/sl/MAD-X/pro/releases/5.02.00/madx-linux64 < Input/SIS18.madx")
orbit_mpi.MPI_Barrier(comm)