def run_from_transport(input_file = "C:/TRANS/for001.dat", nb_part=1000, \
N_segments = 10, kill_lost_particles = True, \
refE = 160, old_refE = 160, DeltaE=0, E_dist='uniform', \
DeltaX = 10**-5, DeltaY = 10**-5, size_dist='cst', \
DeltaDivX = 0.05, DeltaDivY = 0.05, div_dist='uniform', \
gap = 0.03, k1 = 0.5, k2 = 0, \
paraxial_correction = False, dpzTolerance = 10**-4, \
plot_results = True, output_results=True):
"""
Compute trajectories of protons in beamline defined in the Transport file 'input_file'
Definition of inputs:
input_file = Transport input file.
nb_part = number of particles to generate
N_segments = number of points that represent the particle trajectory in each magnet
kill_lost_particles = boolean to determine if paticle trajectories are stopped when they hit an element in the beamline
refE = reference energy of the particle
old_refE = reference energy in the Transport file (in case magnetic fields need to be scaled)
E_dist = Statistical distribution of particle energies (uniform', 'normal' or 'cst')
DeltaX = beam extention in X plane
DeltaY = beam extention in Y plane
size_dist = Statistical distribution of the beam size (uniform', 'normal' or 'cst')
DeltaDivX = divergence in X plane
DeltaDivY = divergence in Y plane
div_dist = Statistical distribution of the beam divergence (uniform', 'normal' or 'cst')
gap = vertical gap of dipoles (useful if kill_lost_particles = True)
k1, k2 = fringe field parameters (see Transport manual)
paraxial_correction = Boolean to correct the effective field for diverging particles (increases computation time)
dpzTolerance = tolerance on normalized magnetic field in paraxial correction
plot_results = boolean
output_results = boolean
Definition of outputs:
sigmaX, sigmaY = spot size at isocenter
eff_ESS_dEonE_1pc, eff_GTR_dEonE_1pc = ESS and GTR effieciency for a dE/E=1% (if DeltaE high enough in input)
"""
split_transport_file(input_file)
ref_p = EtoP(refE)
#Brho = 1/300*sqrt(refE**2+2*938*refE)
Brho_factor = Brho_scaling(old_refE, refE)
gapX = gap # case of CCT magnets
########################################
nb_pts_z = transport_count_lines(input_file, N_segments)
###############################################################################
# initial conditions
beam = np.empty(shape=[nb_pts_z, 7, nb_part])
for i in range(0, nb_part):
# initialize particle properties
z = 0
it_z = 0
if size_dist == 'uniform':
sizeX = DeltaX * np.random.uniform(-1, 1)
sizeY = DeltaY * np.random.uniform(-1, 1)
elif size_dist == 'normal':
sizeX = DeltaX * np.random.normal(0, 1)
sizeY = DeltaY * np.random.normal(0, 1)
elif size_dist == 'cst':
sizeX = DeltaX * (i - nb_part / 2) / nb_part * 2
sizeY = DeltaY * (i - nb_part / 2) / nb_part * 2
else:
raise Exception('Distribution chosen for "size_dist" is not valid')
if div_dist == 'uniform':
divX = DeltaDivX * np.random.uniform(-1, 1)
divY = DeltaDivY * np.random.uniform(-1, 1)
elif div_dist == 'normal':
divX = DeltaDivX * np.random.normal(0, 1)
divY = DeltaDivY * np.random.normal(0, 1)
elif div_dist == 'cst':
divX = DeltaDivX * (i - nb_part / 2) / nb_part * 2
divY = DeltaDivY * (i - nb_part / 2) / nb_part * 2
else:
raise Exception('Distribution chosen for "div_dist" is not valid')
if E_dist == 'uniform':
E = refE + DeltaE * np.random.uniform(-1, 1)
elif E_dist == 'normal':
E = refE + DeltaE * np.random.normal(0, 1)
elif E_dist == 'cst':
E = refE + DeltaE * (i - nb_part / 2) / nb_part * 2
else:
raise Exception('Distribution chosen for "E_dist" is not valid')
p = EtoP(E)
dponp = (p - ref_p) / ref_p
beam[0, :, i] = np.array([z, sizeX, divX, sizeY, divY, 0, dponp])
###############################################################################
# extraction section
last_z = beam[it_z, 0, :]
last_itz = it_z
for i in range(0, nb_part):
it_z = last_itz
[beam, it_z] = transport_input('transport_file_ESS.txt',
beam,
refE,
i,
N_segments,
gap,
k1,
k2,
last_z,
last_itz,
Brho_factor,
kill_lost_particles,
gap_X=gapX,
paraxial_correction=paraxial_correction,
dpz_tolerance=dpzTolerance)
E_list_in = np.vectorize(PtoE)(ref_p * (1 + beam[0, 6, :]))
dEonE_1pc = refE / 100
nb_part_in_ESS_dEonE1pc = ((E_list_in < refE + dEonE_1pc) &
(E_list_in > refE - dEonE_1pc)).sum()
nb_part_in_ESS_dEonE2pc = ((E_list_in < refE + 2 * dEonE_1pc) &
(E_list_in > refE - 2 * dEonE_1pc)).sum()
E_list_out_ESS = np.vectorize(PtoE)(ref_p * (1 + beam[it_z, 6, :]))
E_list_out_ESS[np.isnan(E_list_out_ESS)] = 0
nb_part_out_ESS_dEonE1pc = ((E_list_out_ESS < refE + dEonE_1pc) &
(E_list_out_ESS > refE - dEonE_1pc)).sum()
###############################################################################
# gantry section
last_z = beam[it_z, 0, :]
last_itz = it_z
it_z_GTR = it_z
for i in range(0, nb_part):
it_z = last_itz
[beam, it_z] = transport_input('transport_file_GTR.txt',
beam,
refE,
i,
N_segments,
gap,
k1,
k2,
last_z,
last_itz,
Brho_factor,
kill_lost_particles,
gap_X=gapX,
dpz_tolerance=dpzTolerance)
# make plots
if plot_results:
[sigmaX, sigmaY] = plot_beam(input_file, beam, it_z, it_z_GTR, ref_p)
else:
[sigmaX, sigmaY] = get_spot_size(it_z, beam)
# output results
if output_results:
eff_ESS_dEonE_1pc = nb_part_out_ESS_dEonE1pc / max(
nb_part_in_ESS_dEonE1pc, 1) * 100
print('ESS efficiency within E range [ %0.2f , %0.2f ] = %0.2f %%' %
(refE - dEonE_1pc, refE + dEonE_1pc, eff_ESS_dEonE_1pc))
E_list_out_GTR = np.vectorize(PtoE)(ref_p * (1 + beam[it_z, 6, :]))
E_list_out_GTR[np.isnan(E_list_out_GTR)] = 0
nb_part_out_GTR_dEonE1pc = ((E_list_out_GTR < refE + dEonE_1pc) &
(E_list_out_GTR > refE - dEonE_1pc)).sum()
eff_GTR_dEonE_1pc = nb_part_out_GTR_dEonE1pc / max(
nb_part_out_ESS_dEonE1pc, 1) * 100
print('GTR efficiency within E range [ %0.2f , %0.2f ] = %0.2f %%' %
(refE - dEonE_1pc, refE + dEonE_1pc, eff_GTR_dEonE_1pc))
print('Total efficiency within E range [ %0.2f , %0.2f ] = %0.2f %%' %
(refE - dEonE_1pc, refE + dEonE_1pc, nb_part_out_GTR_dEonE1pc /
max(nb_part_in_ESS_dEonE1pc, 1) * 100))
nb_part_out_GTR_dEonE2pc = (
(E_list_out_GTR < refE + 2 * dEonE_1pc) &
(E_list_out_GTR > refE - 2 * dEonE_1pc)).sum()
eff_tot_dEonE_2pc = nb_part_out_GTR_dEonE2pc / max(
nb_part_in_ESS_dEonE2pc, 1) * 100
print('Total efficiency within E range [ %0.2f , %0.2f ] = %0.2f %%' %
(refE - 2 * dEonE_1pc, refE + 2 * dEonE_1pc, eff_tot_dEonE_2pc))
return [sigmaX, sigmaY, eff_ESS_dEonE_1pc, eff_GTR_dEonE_1pc]
k1 = 0.5
k2 = 0
old_refE = 160 # default energy considered for fields bleow
refE = 160
ref_p = EtoP(refE)
Brho = 1 / 300 * sqrt(refE**2 + 2 * 938 * refE)
Brho_factor = Brho_scaling(old_refE, refE)
kill_lost_particles = True
gapX = gap # case of CCT magnets
########################################
N_segments = 10
nb_pts_z = transport_count_lines(input_file, N_segments)
###############################################################################
# initial conditions
beam = np.empty(shape=[nb_pts_z, 7, nb_part])
for i in range(0, nb_part):
z = 0
it_z = 0
sizeX = np.random.uniform(-0.00001, 0.00001)
sizeY = np.random.uniform(-0.00001, 0.00001)
#sizeX = 0.00001
def plot_beam(input_file, beam, it_z_ISO, it_z_GTR, ref_p):
"""
Plot beam properties
Retun spot size parameters at isocenter
"""
refE = PtoE(ref_p)
#######################################
# compute GTR drawing
nb_pts_z = transport_count_lines(input_file, 1)
layout = GTR_layout_from_transport(input_file, nb_pts_z, refE)
plt.figure('Gantry layout', figsize=(6, 4))
plt.scatter(layout[0:nb_pts_z - 1, 0], layout[0:nb_pts_z - 1, 1])
plt.title('Gantry layout')
plt.xlabel('length [m]')
plt.ylabel('heigth [m]')
#######################################
# Particle traces
fig = plt.figure('Particle traces', figsize=(9, 6))
fig.suptitle('Particle traces', fontweight="bold")
ax0 = fig.add_subplot(
211) # add subplot 1 (211 = 2 rows, 1 column, first plot)
ax1 = fig.add_subplot(212) # add subplot 2
ax0.plot(beam[0:it_z_ISO, 0, :], beam[0:it_z_ISO,
1, :]) # plot X on left plot
plt.ylim((-0.05, 0.05))
ax0.set_xlabel('z [m]')
ax0.set_ylabel('x [m]')
#plt.legend(np.vectorize(PtoE)(ref_p*(1+beam[it_z_ISO,6,:])))
ax0.grid(which='major')
ax1.plot(beam[0:it_z_ISO, 0, :], beam[0:it_z_ISO,
3, :]) # plot Y on right plot
plt.ylim((-0.05, 0.05))
ax1.set_xlabel('z [m]')
ax1.set_ylabel('y [m]')
ax1.grid(which='major')
#plt.legend(np.vectorize(PtoE)(ref_p*(1+beam[it_z_ISO,6,:])))
################################################
plt.figure('Spot size')
plt.title('Spot size at isocenter')
# filter NaN values
X_iso_filtered = beam[it_z_ISO, 1, :][~np.isnan(beam[it_z_ISO, 1, :])]
Y_iso_filtered = beam[it_z_ISO, 3, :][~np.isnan(beam[it_z_ISO, 3, :])]
## add filter on central values for the fit
#maskX = np.logical_and(X_iso_filtered < 0.01,X_iso_filtered>-0.01)
#maskY = np.logical_and(Y_iso_filtered < 0.01,Y_iso_filtered>-0.01)
#
#(muX, sigmaX) = norm.fit(X_iso_filtered[maskX])
#print('mu X = ',muX,' , sigma X = ',sigmaX)
#(muY, sigmaY) = norm.fit(Y_iso_filtered[maskY])
#print('mu Y = ',muY,' , sigma Y = ',sigmaY)
nb_bins = 100
n, bins, patches = plt.hist(X_iso_filtered,
nb_bins,
range=(-0.03, 0.03),
alpha=0.5,
label="X")
X_spot = [bins[:-1], n]
#hist_area = np.sum(n)*(max(bins)-min(bins))/len(bins)
## add a 'best fit' line
#y = norm.pdf(bins, loc=muX, scale=sigmaX)*hist_area
#plt.plot(bins, y, 'b--', linewidth=2,label='X fit: \u03BC = {:.2f} mm, \u03C3 = {:.2f} mm'.format(muX*1000,sigmaX*1000))
# make a Gaussian fit
#tplInitial contains the "first guess" of the parameters
param_initial = [100, 0, 0.003]
param_final, success = leastsq(errorFuncGaussian,
param_initial[:],
args=(X_spot[:][0], X_spot[:][1]))
(muX, sigmaX) = param_final[1:3]
print('mu X = %0.2e mm, sigma X = %0.2e mm' % (muX * 1000, sigmaX * 1000))
plt.plot(X_spot[:][0],
funcGaussian(param_final, X_spot[:][0]),
'b--',
label='X fit: \u03BC = {:.2f} mm, \u03C3 = {:.2f} mm'.format(
muX * 1000, sigmaX * 1000))
# double Gaussian fit
param_initial = [100, 0, 0.003, 10, 0, 0.01]
param_final, success = leastsq(errorFuncDoubleGaussian,
param_initial[:],
args=(X_spot[:][0], X_spot[:][1]))
print('double Gaussian X: ', param_final)
#plt.plot(X_spot[:][0],funcDoubleGaussian(param_final,X_spot[:][0]),'b:',label='double Gaussian')
n, bins, patches = plt.hist(Y_iso_filtered,
nb_bins,
range=(-0.03, 0.03),
alpha=0.5,
label="Y")
Y_spot = [bins[:-1], n]
##hist_area = np.sum(n)*(max(bins)-min(bins))/len(bins)
## add a 'best fit' line
#y = norm.pdf(bins, loc=muY, scale=sigmaY)*hist_area
#plt.plot(bins, y, 'r--', linewidth=2,label='Y fit: \u03BC = {:.2f} mm, \u03C3 = {:.2f} mm'.format(muY*1000,sigmaY*1000))
param_initial = [100, 0, 0.003]
param_final, success = leastsq(errorFuncGaussian,
param_initial[:],
args=(Y_spot[:][0], Y_spot[:][1]))
(muY, sigmaY) = param_final[1:3]
print('mu Y = %0.2e mm, sigma Y = %0.2e mm' % (muY * 1000, sigmaY * 1000))
plt.plot(Y_spot[:][0],
funcGaussian(param_final, Y_spot[:][0]),
'r--',
label='Y fit: \u03BC = {:.2f} mm, \u03C3 = {:.2f} mm'.format(
muY * 1000, sigmaY * 1000))
# double Gaussian fit
param_initial = [100, 0, 0.003, 10, 0, 0.01]
param_final, success = leastsq(errorFuncDoubleGaussian,
param_initial[:],
args=(Y_spot[:][0], Y_spot[:][1]))
print('double Gaussian Y: ', param_final)
#plt.plot(Y_spot[:][0],funcDoubleGaussian(param_final,Y_spot[:][0]),'r:',label='double Gaussian')
plt.xlabel('X/Y [m]')
plt.ylabel('counts')
plt.legend(loc='upper right')
################################################
plt.figure('Spot size vs enrgy')
plt.scatter(beam[it_z_ISO, 1, :],
beam[it_z_ISO, 3, :],
c=np.vectorize(PtoE)(ref_p * (1 + beam[0, 6, :])))
plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.colorbar()
plt.xlim((-0.05, 0.05))
plt.ylim((-0.05, 0.05))
plt.figure('Energy transmission')
plt.hist(np.vectorize(PtoE)(ref_p * (1 + beam[0, 6, :])),
100,
range=(refE - 10, refE + 10),
alpha=0.3,
label="E source")
plt.hist(np.vectorize(PtoE)(ref_p * (1 + beam[it_z_GTR, 6, :])),
100,
range=(refE - 10, refE + 10),
alpha=0.3,
label="E_in GTR")
plt.hist(np.vectorize(PtoE)(ref_p * (1 + beam[it_z_ISO, 6, :])),
100,
range=(refE - 10, refE + 10),
alpha=0.3,
label="E_out GTR")
plt.title('Energy transmission')
plt.legend(loc='upper right')
plt.xlabel('E [MeV]')
plt.ylabel('nb of protons (arb. units)')
################################################
fig = plt.figure('Emittance', figsize=(9, 6))
ax0 = fig.add_subplot(
121) # add subplot 1 (121 = 1 row, 2 columns, first plot)
ax1 = fig.add_subplot(122) # add subplot 2
ax0.scatter(beam[it_z_ISO, 1, :],
beam[it_z_ISO, 2, :],
c=np.vectorize(PtoE)(ref_p * (1 + beam[0, 6, :])))
ax0.set_xlabel('x [m]')
ax0.set_ylabel('divx [mrad]')
im1 = ax1.scatter(beam[it_z_ISO, 3, :],
beam[it_z_ISO, 4, :],
c=np.vectorize(PtoE)(ref_p * (1 + beam[0, 6, :])))
ax1.set_xlabel('y [m]')
ax1.set_ylabel('divy [mrad]')
divider = make_axes_locatable(ax1)
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')
plt.subplots_adjust(wspace=0.4) # put more space between plts
return [sigmaX, sigmaY]