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 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 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 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 # dE here is in eV
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 # Convert dE from eV to GeV
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_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 # dE is already in GeV - convert to eV
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 # dE already converted to GeV
# ~ 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
def generate_initial_distribution(
parameters,
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'])
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])
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.
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_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'])))
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
def saveBunchAsMatfile(bunch, filename=None):
b = bunch
#take the MPI Communicator from bunch: it could be different from MPI_COMM_WORLD
comm = b.getMPIComm()
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
main_rank = 0
# n_parts_arr - array of size of the number of CPUs,
# and have the number of macroparticles on each CPU
n_parts_arr = [0] * size
n_parts_arr[rank] = b.getSize()
n_parts_arr = orbit_mpi.MPI_Allreduce(n_parts_arr, mpi_datatype.MPI_INT,
mpi_op.MPI_SUM, comm)
mp_array = range(n_parts_arr[rank])
particles = {}
particles['x'] = map(b.x, mp_array)
particles['xp'] = map(b.xp, mp_array)
particles['y'] = map(b.y, mp_array)
particles['yp'] = map(b.yp, mp_array)
particles['z'] = map(b.z, mp_array)
particles['dE'] = map(b.dE, mp_array)
phase_space_keys = particles.keys()
for attribute in b.getPartAttrNames():
particles[attribute] = [[]
for i in range(b.getPartAttrSize(attribute))]
for j in xrange(b.getPartAttrSize(attribute)):
particles[attribute][j] += map(
lambda i: b.partAttrValue(attribute, i, j), mp_array)
#This is just for case. Actually, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
for i_cpu in range(1, size):
for key in phase_space_keys:
if (rank == main_rank):
#get the particle coordinates and attributes
bunch_size_remote = orbit_mpi.MPI_Recv(mpi_datatype.MPI_INT,
i_cpu, 222, comm)
if bunch_size_remote:
particles[key] += list(
np.atleast_1d(
orbit_mpi.MPI_Recv(mpi_datatype.MPI_DOUBLE, i_cpu,
222, comm)))
elif (rank == i_cpu):
#send the coordinate array if there are any particles ...
bunch_size_local = bunch.getSize()
orbit_mpi.MPI_Send(bunch_size_local, mpi_datatype.MPI_INT,
main_rank, 222, comm)
if bunch_size_local:
orbit_mpi.MPI_Send(particles[key], mpi_datatype.MPI_DOUBLE,
main_rank, 222, comm)
for i_cpu in range(1, size):
for attribute in b.getPartAttrNames():
if (rank == main_rank):
bunch_size_remote = orbit_mpi.MPI_Recv(mpi_datatype.MPI_INT,
i_cpu, 222, comm)
if bunch_size_remote:
#get the particle coordinates and attributes
for j in xrange(b.getPartAttrSize(attribute)):
particles[attribute][j] += list(
np.atleast_1d(
orbit_mpi.MPI_Recv(mpi_datatype.MPI_DOUBLE,
i_cpu, 222, comm)))
elif (rank == i_cpu):
bunch_size_local = bunch.getSize()
orbit_mpi.MPI_Send(bunch_size_local, mpi_datatype.MPI_INT,
main_rank, 222, comm)
if bunch_size_local:
#send the coordinate array if there are any particles ...
for j in xrange(b.getPartAttrSize(attribute)):
orbit_mpi.MPI_Send(particles[attribute][j],
mpi_datatype.MPI_DOUBLE, main_rank,
222, comm)
bunchparameters = {'classical_radius': bunch.classicalRadius(), \
'charge': bunch.charge(),
'mass': bunch.mass(), \
'momentum': bunch.getSyncParticle().momentum(), \
'beta': bunch.getSyncParticle().beta(), \
'gamma': bunch.getSyncParticle().gamma(), \
'time': bunch.getSyncParticle().time()}
if filename:
if rank == main_rank:
sio.savemat(filename + '.mat', {
'particles': particles,
'bunchparameters': bunchparameters
},
do_compression=True)
orbit_mpi.MPI_Barrier(comm)
def save_to_matfile(self, filename):
rank = orbit_mpi.MPI_Comm_rank(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
if not rank:
sio.savemat(filename, self.output_dict)
orbit_mpi.MPI_Barrier(orbit_mpi.mpi_comm.MPI_COMM_WORLD)
def BunchGather(bunch, turn, p, plot_footprint=False):
b = bunch
verbose = False
# take the MPI Communicator from bunch: it could be different from MPI_COMM_WORLD
comm = b.getMPIComm()
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
main_rank = 0
# n_parts_arr - array of size of the number of CPUs,
# and have the number of macroparticles on each CPU
n_parts_arr = [0] * size
n_parts_arr[rank] = b.getSize()
n_parts_arr = orbit_mpi.MPI_Allreduce(n_parts_arr, mpi_datatype.MPI_INT,
mpi_op.MPI_SUM, comm)
if verbose:
print 'BunchMoments:: bunch_size on MPI Rank: ', rank, ' = ', n_parts_arr[
rank]
print 'BunchMoments:: n_parts_arr on MPI Rank: ', rank, ' = ', n_parts_arr
mp_array = range(n_parts_arr[rank])
particles = {}
particles['x'] = map(b.x, mp_array)
particles['xp'] = map(b.xp, mp_array)
particles['y'] = map(b.y, mp_array)
particles['yp'] = map(b.yp, mp_array)
particles['z'] = map(b.z, mp_array)
particles['dE'] = map(b.dE, mp_array)
phase_space_keys = particles.keys()
for attribute in b.getPartAttrNames():
particles[attribute] = [[]
for i in range(b.getPartAttrSize(attribute))]
for j in xrange(b.getPartAttrSize(attribute)):
particles[attribute][j] += map(
lambda i: b.partAttrValue(attribute, i, j), mp_array)
# This is just in case. Actually, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
for i_cpu in range(1, size):
for key in phase_space_keys:
if (rank == main_rank):
# get the particle coordinates and attributes
bunch_size_remote = orbit_mpi.MPI_Recv(mpi_datatype.MPI_INT,
i_cpu, 222, comm)
if bunch_size_remote:
particles[key] += list(
np.atleast_1d(
orbit_mpi.MPI_Recv(mpi_datatype.MPI_DOUBLE, i_cpu,
222, comm)))
elif (rank == i_cpu):
# send the coordinate array if there are any particles ...
bunch_size_local = bunch.getSize()
orbit_mpi.MPI_Send(bunch_size_local, mpi_datatype.MPI_INT,
main_rank, 222, comm)
if bunch_size_local:
orbit_mpi.MPI_Send(particles[key], mpi_datatype.MPI_DOUBLE,
main_rank, 222, comm)
for i_cpu in range(1, size):
for attribute in b.getPartAttrNames():
if (rank == main_rank):
bunch_size_remote = orbit_mpi.MPI_Recv(mpi_datatype.MPI_INT,
i_cpu, 222, comm)
if bunch_size_remote:
# get the particle coordinates and attributes
for j in xrange(b.getPartAttrSize(attribute)):
particles[attribute][j] += list(
np.atleast_1d(
orbit_mpi.MPI_Recv(mpi_datatype.MPI_DOUBLE,
i_cpu, 222, comm)))
elif (rank == i_cpu):
bunch_size_local = bunch.getSize()
orbit_mpi.MPI_Send(bunch_size_local, mpi_datatype.MPI_INT,
main_rank, 222, comm)
if bunch_size_local:
# send the coordinate array if there are any particles ...
for j in xrange(b.getPartAttrSize(attribute)):
orbit_mpi.MPI_Send(particles[attribute][j],
mpi_datatype.MPI_DOUBLE, main_rank,
222, comm)
bunchparameters = {'classical_radius': bunch.classicalRadius(), \
'charge': bunch.charge(),
'mass': bunch.mass(), \
'momentum': bunch.getSyncParticle().momentum(), \
'beta': bunch.getSyncParticle().beta(), \
'gamma': bunch.getSyncParticle().gamma(), \
'time': bunch.getSyncParticle().time()}
########################################################################
# Plot tune footprint with histograms #
########################################################################
if rank is main_rank:
# ~ print 'Rank: ', rank
if turn >= 0:
# ~ print 'Turn: ', turn
if plot_footprint:
if verbose:
print 'BunchGather:: Plot tune footprint on rank', rank
tunex = str(p['tunex'][0] + '.' + p['tunex'][1:])
tuney = str(p['tuney'][0] + '.' + p['tuney'][1:])
tunex_sav = str(p['tunex'][0] + 'p' + p['tunex'][1:])
tuney_sav = str(p['tuney'][0] + 'p' + p['tuney'][1:])
fontsize = 15
qx = np.array(particles['ParticlePhaseAttributes'][2])
qy = np.array(particles['ParticlePhaseAttributes'][3])
qx[np.where(qx > 0.5)] -= 1
qy[np.where((qy > 0.6) & (qx < 0.25))] -= 1
print 'resonances'
resonances = resonance_lines((5.75, 6.25), (5.75, 6.25),
(1, 2, 3, 4), 10)
fontsize = 17
f, ax = plt.subplots(1, figsize=(6, 6))
gridspec.GridSpec(3, 3)
#f.subplots_adjust(hspace = 0) # Horizontal spacing between subplots
f.subplots_adjust(
wspace=0) # Vertical spacing between subplots
my_cmap = plt.cm.jet
my_cmap.set_under('w', 1)
r = resonances
print 'title'
title = str(tunex_sav + ' ' + tuney_sav + ' turn ' + str(turn))
# First subplot
print 'plot1'
plt.subplot2grid((3, 3), (0, 0), colspan=2, rowspan=1)
plt.hist(6 + qx, bins=1000,
range=(r.Qx_min,
r.Qx_max)) #, norm=mcolors.PowerNorm(gamma))
plt.ylabel('Frequency')
plt.grid(which='both')
plt.title(title, fontsize=fontsize)
# Main plot
print 'plot2'
plt.subplot2grid((3, 3), (1, 0), colspan=2, rowspan=2)
plt.hist2d(6 + qx,
6 + qy,
bins=1000,
cmap=my_cmap,
vmin=1,
range=[[r.Qx_min, r.Qx_max], [r.Qy_min, r.Qy_max]
]) #, norm=mcolors.PowerNorm(gamma))
plt.xlabel(r'Q$_x$')
plt.ylabel(r'Q$_y$')
print 'plot_resonance'
resonances.plot_resonance(f)
# Second subplot
print 'plot3'
plt.subplot2grid((3, 3), (1, 2), colspan=1, rowspan=2)
plt.hist(6 + qy,
bins=1000,
range=(r.Qy_min, r.Qy_max),
orientation=u'horizontal'
) #, norm=mcolors.PowerNorm(gamma))
plt.xlabel('Frequency')
plt.grid(which='both')
current_axis = plt.gca()
#current_axis.axes.get_yaxis().set_visible(False)
ax.xaxis.label.set_size(fontsize)
ax.yaxis.label.set_size(fontsize)
ax.tick_params(labelsize=fontsize)
plt.tight_layout()
savename = str('Tune_Footprints/' + tunex_sav + '_' +
tuney_sav + '_turn_' + str(turn) + '_hist.png')
print 'savefig'
f.savefig(savename, dpi=100)
plt.close(f)
outputs = dict()
if rank is main_rank:
x = np.array(particles['x'])
xp = np.array(particles['xp'])
y = np.array(particles['y'])
yp = np.array(particles['yp'])
z = np.array(particles['z'])
dE = np.array(particles['dE'])
mu_x = moment(x, 2)
mu_xp = moment(xp, 2)
mu_y = moment(y, 2)
mu_yp = moment(yp, 2)
mu_z = moment(z, 2)
mu_dE = moment(dE, 2)
sig_x = np.sqrt(mu_x)
sig_xp = np.sqrt(mu_xp)
sig_y = np.sqrt(mu_y)
sig_yp = np.sqrt(mu_yp)
sig_z = np.sqrt(mu_z)
sig_dE = np.sqrt(mu_dE)
x_6_sig = x[np.where((x >= -6 * sig_x) & (x <= 6 * sig_x))]
xp_6_sig = xp[np.where((xp >= -6 * sig_xp) & (xp <= 6 * sig_xp))]
y_6_sig = y[np.where((y >= -6 * sig_y) & (y <= 6 * sig_y))]
yp_6_sig = yp[np.where((yp >= -6 * sig_yp) & (yp <= 6 * sig_yp))]
z_6_sig = z[np.where((z >= -6 * sig_z) & (z <= 6 * sig_z))]
dE_6_sig = dE[np.where((dE >= -6 * sig_dE) & (dE <= 6 * sig_dE))]
# Later add something to cut large amplitude particles to reduce noise for kurtosis calculation
outputs = {
'Mu_x': mu_x,
'Mu_xp': mu_xp,
'Mu_y': mu_y,
'Mu_yp': mu_yp,
'Mu_z': mu_z,
'Mu_dE': mu_dE,
'Sig_x': sig_x,
'Sig_xp': sig_xp,
'Sig_y': sig_y,
'Sig_yp': sig_yp,
'Sig_z': sig_z,
'Sig_dE': sig_dE,
'Max_x': np.max(x),
'Max_xp': np.max(xp),
'Max_y': np.max(y),
'Max_yp': np.max(yp),
'Max_z': np.max(z),
'Max_dE': np.max(dE),
'Min_x': np.min(x),
'Min_xp': np.min(xp),
'Min_y': np.min(y),
'Min_yp': np.min(yp),
'Min_z': np.min(z),
'Min_dE': np.min(dE),
'Kurtosis_x': kurtosis(x, fisher=True, nan_policy='omit'),
'Kurtosis_xp': kurtosis(xp, fisher=True, nan_policy='omit'),
'Kurtosis_y': kurtosis(y, fisher=True, nan_policy='omit'),
'Kurtosis_yp': kurtosis(yp, fisher=True, nan_policy='omit'),
'Kurtosis_z': kurtosis(z, fisher=True, nan_policy='omit'),
'Kurtosis_dE': kurtosis(dE, fisher=True, nan_policy='omit'),
'Kurtosis_x_6sig': kurtosis(x_6_sig,
fisher=True,
nan_policy='omit'),
'Kurtosis_xp_6sig': kurtosis(xp_6_sig,
fisher=True,
nan_policy='omit'),
'Kurtosis_y_6sig': kurtosis(y_6_sig,
fisher=True,
nan_policy='omit'),
'Kurtosis_yp_6sig': kurtosis(yp_6_sig,
fisher=True,
nan_policy='omit'),
'Kurtosis_z_6sig': kurtosis(z_6_sig,
fisher=True,
nan_policy='omit'),
'Kurtosis_dE_6sig': kurtosis(dE_6_sig,
fisher=True,
nan_policy='omit')
}
return outputs
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 bunch_pyorbit_to_orbit_nHarm(ringLength, nHarm, pyOrbitBunch, \
name_of_orbit_mpi_bunch_file):
"""
Translates pyORBIT bunch to ORBIT_MPI bunch, incorporating RF
harmonic number, and dumps it into a file.
The ring length should be defined in the input (in meters).
Lines in bunch files:
ORBIT_MPI: x[mm] xp[mrad] y[mm] yp[mrad] phi[rad] dE[GeV].
pyORBIT: x[m] xp[rad] y[m] yp[rad] z[m] dE[GeV]
"""
pi2 = 2.0 * math.pi
zfac = pi2 * nHarm / ringLength
b = pyOrbitBunch
# Take the MPI Communicator from bunch: it could be different
# from MPI_COMM_WORLD
comm = pyOrbitBunch.getMPIComm()
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
main_rank = 0
# n_parts_arr - array of size number of CPUs,
# Contains the number of macroparticles on each CPU
n_parts_arr = [0] * size
n_parts_arr[rank] = b.getSize()
n_parts_arr = orbit_mpi.MPI_Allreduce(n_parts_arr, \
mpi_datatype.MPI_INT, mpi_op.MPI_SUM, comm)
file_out = None
if (rank == main_rank):
file_out = open(name_of_orbit_mpi_bunch_file, "w")
if (rank == main_rank):
for i in range(n_parts_arr[rank]):
x = b.x(i) * 1000.
px = b.px(i) * 1000.
y = b.y(i) * 1000.
py = b.py(i) * 1000.
z = -(math.fmod(b.z(i) * zfac, pi2))
if (z > math.pi):
z = z - 2 * math.pi
if (z < -math.pi):
z = z + 2 * math.pi
dE = b.dE(i)
file_out.write(str(x) + " " + str(px) + " " + \
str(y) + " " + str(py) + " "+ \
str(z) + " " + str(dE) + "\n")
# MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
val_arr = (0., 0., 0., 0., 0., 0.)
for i_cpu in range(1, size):
#Again, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
for i in range(n_parts_arr[i_cpu]):
if (rank == main_rank):
# Get the coordinate array
(x, px, y, py, z, dE) = orbit_mpi.MPI_Recv(\
mpi_datatype.MPI_DOUBLE, \
i_cpu, 222, comm)
file_out.write(str(x) + " " + str(px) + \
" " + str(y) + " " + str(py) + \
" " + str(z) + " " + str(dE) + "\n")
elif (rank == i_cpu):
#send the coordinate array
x = b.x(i) * 1000.
px = b.px(i) * 1000.
y = b.y(i) * 1000.
py = b.py(i) * 1000.
z = -(math.fmod(b.z(i) * zfac, pi2))
if (z > math.pi):
z = z - 2 * math.pi
if (z < -math.pi):
z = z + 2 * math.pi
dE = b.dE(i)
val_arr = (x, px, y, py, z, dE)
orbit_mpi.MPI_Send(val_arr, \
mpi_datatype.MPI_DOUBLE, \
main_rank, 222, comm)
if (rank == main_rank):
file_out.close()
def generate_initial_5mm_distribution(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] = (5E-3 - 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] = (5E-3 - 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
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
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
def profiles(Bunch, coord, histogram, steps=100, Min=1.0, Max=-1.0):
"""
Returns a profile for one of the following Bunch coordinates:
x[m] xp[rad] y[m] yp[rad] z[m] dE[GeV]
"""
b = Bunch
# Take the MPI Communicator from bunch: It could be
# different from MPI_COMM_WORLD
comm = Bunch.getMPIComm()
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
main_rank = 0
# n_parts_arr - array of size of the number of CPUs,
# contains the number of macroparticles on each CPU
n_parts_arr = [0] * size
n_parts_arr[rank] = b.getSize()
n_parts_arr = orbit_mpi.MPI_Allreduce(n_parts_arr, \
mpi_datatype.MPI_INT,mpi_op.MPI_SUM,comm)
partdat = []
if (rank == main_rank):
for i in range(n_parts_arr[rank]):
if coord == "x":
partdat.append(b.x(i))
if coord == "px":
partdat.append(b.px(i))
if coord == "y":
partdat.append(b.y(i))
if coord == "py":
partdat.append(b.py(i))
if coord == "z":
partdat.append(b.z(i))
if coord == "dE":
partdat.append(b.dE(i))
# That is just for case.
# Actually, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
val_arr = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
for i_cpu in range(1, size):
# Again, that is just for case.
# Actually, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
for i in range(n_parts_arr[i_cpu]):
if (rank == main_rank):
#get the coordinate array
(x, px, y, py, z, dE) = \
orbit_mpi.MPI_Recv(mpi_datatype.MPI_DOUBLE, \
i_cpu, 222, comm)
if coord == "x":
partdat.append(x)
if coord == "px":
partdat.append(px)
if coord == "y":
partdat.append(y)
if coord == "py":
partdat.append(py)
if coord == "z":
partdat.append(z)
if coord == "dE":
partdat.append(dE)
elif (rank == i_cpu):
#send the coordinate array
x = b.x(i)
px = b.px(i)
y = b.y(i)
py = b.py(i)
z = b.z(i)
dE = b.dE(i)
val_arr = (x, px, y, py, z, dE)
orbit_mpi.MPI_Send(val_arr, \
mpi_datatype.MPI_DOUBLE, main_rank, 222, comm)
l = len(partdat)
m = min(partdat)
M = max(partdat)
c = (M + m) / 2.0
d = (M - m) * 1.1 / 2.0
M = c + d
m = c - d
if Max > M:
M = Max
if Min < m:
m = Min
dx = (M - m) / steps
grid = [m]
prof = [0]
for i in range(1, steps + 1):
x = m + i * dx
grid.append(x)
prof.append(0)
grid.append(M)
prof.append(0)
for n in range(l):
i = (partdat[n] - m) / dx
i = int(i)
if i < 0:
pass
elif i > range(steps):
pass
else:
frac = (partdat[n] - m) / dx % 1
prof[i] = prof[i] + (1.0 - frac)
prof[i + 1] = prof[i + 1] + frac
sum = 0.0
for i in range(steps + 1):
sum = sum + prof[i]
file_out = histogram
if (rank == main_rank):
file_out = open(histogram, "w")
file_out.write("Min = " + str(m) + " Max = " + \
str(M) + " steps = " + str(steps) + "\n")
file_out.write("nParts = " + str(l) + " HistSum = " + \
str(sum) + "\n\n")
for i in range(steps + 1):
file_out.write(str(grid[i]) + " " + \
str(prof[i]) + "\n")
if (rank == main_rank):
file_out.close()
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)
#----------------------------------------------
# Initialize a Teapot-Style PTC lattice
#----------------------------------------------
PTC_File = "SIS_18_BENCHMARK.flt"
Lattice = PTC_Lattice("MACHINE")
Lattice.readPTC(PTC_File)
readScriptPTC('Input/time.ptc')
paramsDict = {}
paramsDict["length"]=Lattice.getLength()/Lattice.nHarm
#----------------------------------------------
# Add apertures
#----------------------------------------------
def bunch_pyorbit_to_orbit(ringLength, pyOrbitBunch,
name_of_orbit_mpi_bunch_file):
"""
Translates pyORBIT bunch to ORBIT_MPI bunch and dumps it into the file.
The ring length should be defined in the input (in meters).
ORBIT_MPI file has lines: x[mm] xp[mrad] y[mm] yp[mrad] phi[rad] dE[GeV].
pyORBIT: x[m] xp[rad] y[m] yp[rad] z[m] dE[GeV]
"""
pi2 = 2.0 * math.pi
L = ringLength
b = pyOrbitBunch
#take the MPI Communicator from bunch: it could be different from MPI_COMM_WORLD
comm = pyOrbitBunch.getMPIComm()
rank = orbit_mpi.MPI_Comm_rank(comm)
size = orbit_mpi.MPI_Comm_size(comm)
main_rank = 0
# n_parts_arr - array of size of the number of CPUs,
# and have the number of macroparticles on each CPU
n_parts_arr = [0] * size
n_parts_arr[rank] = b.getSize()
n_parts_arr = orbit_mpi.MPI_Allreduce(n_parts_arr, mpi_datatype.MPI_INT,
mpi_op.MPI_SUM, comm)
file_out = None
if (rank == main_rank):
file_out = open(name_of_orbit_mpi_bunch_file, "w")
if (rank == main_rank):
for i in range(n_parts_arr[rank]):
x = b.x(i) * 1000.
px = b.px(i) * 1000.
y = b.y(i) * 1000.
py = b.py(i) * 1000.
z = -(math.fmod(b.z(i) * pi2 / L, pi2))
if (z > math.pi):
z = z - 2 * math.pi
if (z < -math.pi):
z = z + 2 * math.pi
dE = b.dE(i)
file_out.write(
str(x) + " " + str(px) + " " + str(y) + " " + str(py) + " " +
str(z) + " " + str(dE) + "\n")
#That is just for case. Actually, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
val_arr = (0., 0., 0., 0., 0., 0.)
for i_cpu in range(1, size):
#Again, that is just for case. Actually, MPI_Barrier command is not necessary.
orbit_mpi.MPI_Barrier(comm)
for i in range(n_parts_arr[i_cpu]):
if (rank == main_rank):
#get the coordinate array
(x, px, y, py, z,
dE) = orbit_mpi.MPI_Recv(mpi_datatype.MPI_DOUBLE, i_cpu, 222,
comm)
file_out.write(
str(x) + " " + str(px) + " " + str(y) + " " + str(py) +
" " + str(z) + " " + str(dE) + "\n")
elif (rank == i_cpu):
#send the coordinate array
x = b.x(i) * 1000.
px = b.px(i) * 1000.
y = b.y(i) * 1000.
py = b.py(i) * 1000.
z = -(math.fmod(b.z(i) * pi2 / L, pi2))
if (z > math.pi):
z = z - 2 * math.pi
if (z < -math.pi):
z = z + 2 * math.pi
dE = b.dE(i)
val_arr = (x, px, y, py, z, dE)
orbit_mpi.MPI_Send(val_arr, mpi_datatype.MPI_DOUBLE, main_rank,
222, comm)
if (rank == main_rank):
file_out.close()