#divY = 0.05*np.random.choice([-1,0,1])
#divX = 0.05
E = 150
E = np.random.uniform(E-20,E+20)
#E = 160 + 15*np.random.choice([-1,0,1])
p = EtoP(E)
dponp = (p-ref_p)/ref_p
beam[it_z,:,i] = np.array([z,.00001,divX,.00001,divY,0,dponp])
L = 0.3
[beam[:,:,i],it_z] = drift(L,beam[:,:,i],refE,it_z,1)
#it_z = it_z + 1
L=0.8
B=7
Brho = 1/300*math.sqrt(E**2+2*938*E)
if Brho/B < L*10 :
# need more segments
N_segments = max(int(L/(Brho/B))*10,1)
else:
N_segments = 1
[beam[:,:,i],it_z] = solenoid(L,B,0.1,beam[:,:,i],it_z,N_segments,refE)
for i in range(0, nb_part):
it_z = 0
z = 0
divX = np.random.normal(0, 0.05)
divY = np.random.normal(0, 0.05)
beam[it_z, :, i] = np.array([z, .00001, divX, .00001, divY, 0, 160])
it_z = it_z + 1
length = 0.1
z = z + length
beam[it_z, :, i] = z
beam[it_z, 1:7, i] = drift(length, beam[it_z - 1, 1:7, i])
use_wedge = False
L = 0.01
if use_wedge:
it_z = it_z + 1
L_min = 0.0011
L_max = L
z = z + L_max
beam[it_z, 0, i] = z
beam[it_z, 1:7, i] = wedge(L_min, .004, L_max, 0.035, 'tantalum', 10,
beam[it_z - 1, 1:7, i])
else:
it_z = it_z + 1
length = L
z = z + length
p = EtoP(E)
dponp = (p - ref_p) / ref_p
divX = 0.0
divY = 0.0
#E = 160
beam[0, :, i] = np.array([z, 0, divX, .0, divY, 0, dponp])
it_z = 0
#ref_x[it_z] = 0
# filter directly partcles above 50mrad
L = 0.1
#[beam[:,:,i],it_z] = collimator(L,0.005*sqrt(2),'tantalum',beam[:,:,i],it_z,10)
[beam[:, :, i], it_z] = drift(L, beam[:, :, i], refE, it_z, 1)
L = 0.2
[beam[:, :, i], it_z] = drift(L, beam[:, :, i], refE, it_z, 1)
L = 0.2
B = 1.5
pole_in = -4.5
pole_out = -4.5
gap = 0.1
N_segments = 10
#[beam[:,:,i],it_z] = sec_bend(L,B,0,beam[:,:,i],160,True,it_z,N_segments)
[beam[:, :, i], it_z] = bending(L, B, 0, pole_in, pole_out, gap, 0.5, 0,
beam[:, :, i], refE, it_z, N_segments)
[beam[:, :, i], it_z] = drift(0.1, beam[:, :, i], refE, it_z, 1)
for i in range(0, iterations):
z = 0
divX = np.random.normal(0, 0.05)
divY = np.random.normal(0, 0.05)
#beam = np.array([.00001,divX,.00001,divY,0,160])
beam[i, :] = np.array([z, .00001, divX, .00001, divY, 0, 160])
#beam = np.array([.00001,.05,.00001,.05,0,160])
#beam = np.array([.01,.0,.01,.0,0,160])
#x_vs_z = np.array([z,beam[0]][i])
length = 0.1
z = z + length
beam[i, 0] = z
beam[i, 1:7] = drift(length, beam[i, 1:7])
#x_vs_z = np.vstack((x_vs_z,np.array([z,beam[0][i]])))
#length=0.15
#z=z+length
#beam[0:6] = quad(length,-0.964,0.015,beam[0:6])
#x_vs_z = np.vstack((x_vs_z,np.array([z,beam[0]])))
L_min = 0.001
L_max = 0.01
z = z + L_max
beam[i, 0] = z
beam[i, 1:7] = wedge(L_min, .005, L_max, .02, 'tantalum', 10, beam[i, 1:7])
#x_vs_z = np.vstack((x_vs_z,np.array([z,beam[0][i]])))
# =============================================================================
# px = refp*divX
# py = refp*divY
# pz = sqrt(refp**2 - px**2 - py**2)
#
x = 0
y = 0
z = 0
# vx = PtoV(px)
# vy = PtoV(py)
# vz = PtoV(pz)
beam[0,it_p,:] = [z,x,divX,y,divY,0,dponp]
L = 0.15
[beam[:,it_p,:],it_z] = drift(L,beam[:,it_p,:],refE,it_z,1)
[beam[:,it_p,:],it_z] = particle_in_B_map(field_map,beam[:,it_p,:],Delta_z,it_z,nb_it_z,it_p,refp,0.002,0.002,spatial_res = 10)
L = 0.55
[beam[:,it_p,:],it_z] = drift(L,beam[:,it_p,:],refE,it_z,1)
print("--- execution time: %s seconds ---" % (time.perf_counter() - start_time))
plt.figure(0)
plt.plot(beam[0:it_z+1,:,0],beam[0:it_z+1,:,1])
plt.figure(1)
# filter directly partcles above 50mrad
L = 0.11
[beam[:, :, i], it_z] = collimator(L, 0.05 * sqrt(2), 'tantalum',
beam[:, :, i], it_z, 10, refE)
#[beam[:,:,i],it_z] = drift(L,beam[:,:,i],it_z,1)
L = 0.15
B = -0.972 * Brho_factor
a = 0.015
N_segments = 10
[beam[:, :, i], it_z] = quad(L, B, a, beam[:, :, i], it_z, refE,
N_segments)
L = 0.15
[beam[:, :, i], it_z] = drift(L, beam[:, :, i], refE, it_z, 1)
L = 0.33
B = 0.996 * Brho_factor
a = 0.04
N_segments = 10
[beam[:, :, i], it_z] = quad(L, B, a, beam[:, :, i], it_z, refE,
N_segments)
L = 0.2
[beam[:, :, i], it_z] = drift(L, beam[:, :, i], refE, it_z, 1)
L = 0.33
B = -0.586 * Brho_factor
a = 0.04
N_segments = 10
