def make_shot_aligned(p_array_init, cell, meth=1):
method = oclt.MethodTM()
if meth > 1: method.global_method = oclt.SecondTM
p_array = oclt.deepcopy(p_array_init)
lat = oclt.MagneticLattice(cell, method=method)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
return p_array, np.array(
[p_array.x().mean() * 1e3,
p_array.y().mean() * 1e3])
def make_shot(p_array_init, QAP_T, cell, meth=1):
D1, QAP1, D2, QAP2, D3, QAP3, D4 = cell
method = oclt.MethodTM()
if meth > 1: method.global_method = oclt.SecondTM
p_array = oclt.deepcopy(p_array_init)
QAP1_Tx, QAP1_Tz, QAP2_Tx, QAP2_Tz, QAP3_Tx, QAP3_Tz = QAP_T
#print QAP1_Tx,QAP1_Tz,QAP2_Tx,QAP2_Tz,QAP3_Tx,QAP3_Tz
lat = oclt.MagneticLattice(cell, method=method, start=D1, stop=D1)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
p_array.particles[::6] -= QAP1_Tx * 1e-3
p_array.particles[2::6] -= QAP1_Tz * 1e-3
lat = oclt.MagneticLattice(cell, method=method, start=QAP1, stop=QAP1)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
p_array.particles[::6] += QAP1_Tx * 1e-3
p_array.particles[2::6] += QAP1_Tz * 1e-3
lat = oclt.MagneticLattice(cell, method=method, start=D2, stop=D2)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
p_array.particles[::6] -= QAP2_Tx * 1e-3
p_array.particles[2::6] -= QAP2_Tz * 1e-3
lat = oclt.MagneticLattice(cell, method=method, start=QAP2, stop=QAP2)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
p_array.particles[::6] += QAP2_Tx * 1e-3
p_array.particles[2::6] += QAP2_Tz * 1e-3
lat = oclt.MagneticLattice(cell, method=method, start=D3, stop=D3)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
p_array.particles[::6] -= QAP3_Tx * 1e-3
p_array.particles[2::6] -= QAP3_Tz * 1e-3
lat = oclt.MagneticLattice(cell, method=method, start=QAP3, stop=QAP3)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
p_array.particles[::6] += QAP3_Tx * 1e-3
p_array.particles[2::6] += QAP3_Tz * 1e-3
lat = oclt.MagneticLattice(cell, method=method, start=D4)
navi = oclt.Navigator(lat)
tw, p_array = oclt.track(lat, p_array, navi, print_progress=False)
return p_array, np.array(
[p_array.x().mean() * 1e3,
p_array.y().mean() * 1e3])
def print_R(latt_objs, BeamEnergy):
"""
Prints some elements of the transport matrices
Parameters
----------
lat: LattObjs
Magnetic lattice objects retuned by make_line
BeamEnergy : float
Electron energy in GEVs
"""
lat = oclt.MagneticLattice(latt_objs,method=method)
optics.lattice_transfer_map_R(lat,BeamEnergy)
print(R_par_str.format(*lat.R.flatten()))
def print_latt_chars(latt_objs, BeamEnergy):
"""
Prints some elements of the transport matrices
Parameters
----------
lat: LattObjs
Magnetic lattice objects retuned by make_line
BeamEnergy : float
Electron energy in GEVs
"""
lat = oclt.MagneticLattice(latt_objs,method=method)
optics.lattice_transfer_map_RT(lat,BeamEnergy)
params = [lat.R[0,0],lat.R[2,2],lat.R[4,5]*1000,\
lat.T[1,1,5], lat.T[3,3,5],lat.T[0,1,5],lat.T[2,3,5]]
print(latt_par_string.format(*params))
def supermatch_matrix(latt_objs, nrg, varz, MAG=True, printout=True):
lat = oclt.MagneticLattice(latt_objs, method=method)
def error_func(x):
for i in range(len(varz)):
if varz[i].__class__ == oclt.Quadrupole:
varz[i].k1 = x[i]
varz[i].transfer_map = lat.method.create_tm(varz[i])
optics.lattice_transfer_map_RT(lat, nrg)
err = np.abs( lat.R[0,1])**2 + np.abs(lat.R[2,3] )**2 \
+ np.abs(lat.T[1,1,5])**2 + np.abs(lat.T[3,3,5])**2 \
+ (np.abs(lat.T[0,1,5])-np.abs(lat.T[2,3,5]))**2
if MAG is True:
err += (np.abs(lat.R[0, 0]) - np.abs(lat.R[2, 2]))**2
elif MAG is not True and MAG is not None:
err += (np.abs(lat.R[0, 0]) - MAG[0])**2 + (np.abs(lat.R[2, 2]) -
MAG[1])**2
return err
x = [0.0] * len(varz)
for i in range(len(varz)):
if varz[i].__class__ == oclt.Quadrupole:
x[i] = varz[i].k1
if printout: print("initial value: x = ", x)
fmin(error_func, x, xtol=1e-9, maxiter=20000, maxfun=20000, disp=printout)
for i in range(len(varz)):
if varz[i].__class__ == oclt.Quadrupole:
x[i] = varz[i].k1
if printout: print("final value: x = ", x)
return x
def get_envs(QAP1=1., QAP2=1., QAP3=1., \
SRC=Drifts['LPWA-QAP1']*1e3,TARG=Drifts['QAP3-IMG1']*1e2, ):
emin,emax = 0.04,0.25
energies = np.r_[emin:emax:300j]
mags = []
for BeamEnergy in energies:
indx_stop = np.nonzero(np.array(cell_keys)=='IMG1')[0][0]
QuadGradients1 = oclt.deepcopy(QuadGradients)
QuadGradients1['QAP1'] *= QAP1
QuadGradients1['QAP2'] *= QAP2
QuadGradients1['QAP3'] *= QAP3
Drifts1 = oclt.deepcopy(Drifts)
Drifts1['LPWA-QAP1'] = SRC*1e-3
Drifts1['QAP3-IMG1'] = TARG*1e-2
latt_elmts = cox.make_line( \
BeamEnergy=BeamEnergy,QuadGradients=QuadGradients1,Drifts=Drifts1 )
lat = oclt.MagneticLattice( \
[latt_elmts[key] for key in cell_keys[:indx_stop+1]], \
method=method,stop=latt_elmts['IMG1'])
oclt.cpbd.optics.lattice_transfer_map_R(lat,BeamEnergy)
mags.append([lat.R[0,1],lat.R[2,3],lat.R[0,0],lat.R[2,2]])
mags = np.array(mags)
sx = np.abs(mags[:,0]*3e-3 + mags[:,2]*1e-6)*1e3
sy = np.abs(mags[:,1]*3e-3 + mags[:,3]*1e-6)*1e3
plt.plot(energies*1e3,sx,lw=1.5)
plt.plot(energies*1e3,sy,lw=1.5)
plt.plot(energies*1e3,np.sqrt(sx*sy),c='k',lw=1.5)
plt.legend(('$\sigma_x$','$\sigma_z$','$\sqrt{\sigma_x\sigma_z}$'),loc=1)
plt.ylabel('Beam sizes (mm)')
plt.xlabel('Electron energy (MeV)')
plt.ylim(0,4.);
def aligh_slices(p_arrays, lattice_elements, cell_keys,\
BeamEnergy_ref, start_key=None, stop_key=None,\
method=method,verbose=False):
"""
Aligns the slices of the spectrally sliced beam at the end of
the transport simulation.
Parameters
----------
p_arrays: list of oclt.ParticleArray objects
lattice_elements: line parameters as defined in make_line and make_shot
stop_key : str
key-name of lattice element at which to stop the transport (exclusive)
method : oclt.methodTM
transport method (works with first or second orders)
Returns
-------
p_arrays : list or single oclt.ParticleArray object
See Also
--------
make_line, make_shot
"""
if verbose:
sys.stdout.write('\nAligning '+str(len(p_arrays))+' slices')
sys.stdout.flush()
r16 = []
r36 = []
r56 = []
method_ord1 = oclt.MethodTM()
for i in range(len(p_arrays)):
latt_objs = make_line(lattice_elements, cell_keys, \
BeamEnergy=p_arrays[i].E, BeamEnergy_ref=BeamEnergy_ref)
if stop_key!=None and start_key!=None:
stop_indx = cell_keys.index(stop_key)-1
start_indx = cell_keys.index(start_key)
lat = oclt.MagneticLattice(latt_objs, method=method_ord1,\
start=latt_objs[start_indx],stop=latt_objs[stop_indx])
elif stop_key!=None and start_key==None:
stop_indx = cell_keys.index(stop_key)-1
lat = oclt.MagneticLattice(latt_objs, method=method_ord1,\
stop=latt_objs[stop_indx])
elif stop_key==None and start_key!=None:
start_indx = cell_keys.index(start_key)
lat = oclt.MagneticLattice(latt_objs, method=method_ord1,\
start=latt_objs[start_indx])
else:
lat = oclt.MagneticLattice(latt_objs, method=method_ord1)
optics.lattice_transfer_map_R(lat,p_arrays[i].E)
r16.append(lat.R[0,5])
r36.append(lat.R[2,5])
r56.append(lat.R[4,5])
dx = [0,]
dy = [0,]
ds = [0,]
x_loc = 0.0
y_loc = 0.0
s_loc = 0.0
for i in range(1,len(p_arrays)):
de = 1-p_arrays[i].E/p_arrays[i-1].E
x_loc += de*r16[i]
y_loc += de*r36[i]
s_loc += de*r56[i]
dx.append(x_loc)
dy.append(y_loc)
ds.append(s_loc)
ee = np.array([p_array.E for p_array in p_arrays])
# ee_centr = 0.5*(ee.max()+ee.min())
ee_centr = BeamEnergy_ref
indx_centr = (ee_centr-ee>0).sum()
dx = np.array(dx)
dy = np.array(dy)
ds = np.array(ds)
dx -= dx[indx_centr]
dy -= dy[indx_centr]
ds -= ds[indx_centr]
for i in range(len(p_arrays)):
p_arrays[i].x_c = -dx[i]
p_arrays[i].y_c = -dy[i]
p_arrays[i].z_c = ds[i]
return p_arrays
def make_shot(p_arrays, lattice_elements, cell_keys, \
BeamEnergy_ref,start_key=None, stop_key=None, \
method=method, Nit = 1,output = None,verbose=False):
"""
Performs the transport simulation through a COXINEL-type transport line.
Parameters
----------
p_arrays: list or single oclt.ParticleArray object
lattice_elements: line parameters as defined in make_line
start_key : str
key-name of lattice element at which to stop the transport (inclusive)
stop_key : str
key-name of lattice element at which to stop the transport (exclusive)
method : oclt.methodTM
transport method (works with first or second orders)
Nit : int
number of steps to perform (for output and damping resolutions)
output : dictionary or None
outputs to perform as defined in beam_diags
Returns
-------
p_arrays : list of oclt.ParticleArray objects
outputs : dictionary with arrays
See Also
--------
make_line, beam_diags, damp_particles
"""
latt_objs = make_line(lattice_elements, cell_keys, \
BeamEnergy=BeamEnergy_ref, BeamEnergy_ref=BeamEnergy_ref,)
if stop_key!=None and start_key!=None:
stop_indx = cell_keys.index(stop_key)-1
start_indx = cell_keys.index(start_key)
lat = oclt.MagneticLattice(latt_objs, method=method, \
start=latt_objs[start_indx], stop=latt_objs[stop_indx])
elif stop_key!=None and start_key==None:
stop_indx = cell_keys.index(stop_key)-1
lat = oclt.MagneticLattice(latt_objs, method=method, \
stop=latt_objs[stop_indx])
elif stop_key==None and start_key!=None:
start_indx = cell_keys.index(start_key)
lat = oclt.MagneticLattice(latt_objs, method=method,\
start=latt_objs[start_indx])
else:
lat = oclt.MagneticLattice(latt_objs, method=method)
dz = lat.totalLen/Nit
if type(p_arrays)!=list: p_arrays = [p_arrays,]
outputs = {}
if output!=None:
for key in output: outputs[key] = np.zeros(Nit+1)
for i in range(len(p_arrays)):
latt_objs = make_line(lattice_elements, cell_keys, \
BeamEnergy=p_arrays[i].E, BeamEnergy_ref=BeamEnergy_ref)
if stop_key!=None and start_key!=None:
stop_indx = cell_keys.index(stop_key)-1
start_indx = cell_keys.index(start_key)
lat = oclt.MagneticLattice(latt_objs, method=method,\
start=latt_objs[start_indx],stop=latt_objs[stop_indx])
elif stop_key!=None and start_key==None:
stop_indx = cell_keys.index(stop_key)-1
lat = oclt.MagneticLattice(latt_objs, method=method,\
stop=latt_objs[stop_indx])
elif stop_key==None and start_key!=None:
start_indx = cell_keys.index(start_key)
lat = oclt.MagneticLattice(latt_objs,method=method,\
start=latt_objs[start_indx])
else:
lat = oclt.MagneticLattice(latt_objs, method=method)
navi = oclt.Navigator(lat)
sss = '\r'+'Transporing slice '+str(i+1)+' of '+str(len(p_arrays))+':'
for j in range(Nit+1):
if verbose: sys.stdout.write(sss+'step '+str(j+1)+' of '+str(Nit))
oclt.tracking_step(lat, p_arrays[i], dz,navi)
if output!=None:
for key in output:
outputs[key][j] += beam_diags(p_arrays[i],key)
return p_arrays, outputs