-
Notifications
You must be signed in to change notification settings - Fork 0
/
TestModules.py
91 lines (73 loc) · 2.84 KB
/
TestModules.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#theta in degrees
def calc_quadrupole(r, theta, p, rotation_angle=0):
import numpy as np
#ep_0 = 8.857187*10**(-12)
return 3*p*np.cos(theta-rotation_angle)*np.sin(theta-rotation_angle)/(r**3.)
def calc_dipole_onaxis(r, theta, p):
import numpy as np
#ep_0 = 8.857187*10**(-12)
return p*(3*np.cos(theta)**2. - 1)/(r**3.)
def calc_dipole(r, theta, p, rotation_angle=0):
import numpy as np
return 2*p*np.sin(theta-rotation_angle) / r**3.
def dipole_field(beam_map, p, rotation_angle=0):
import numpy as np
#square map
N = len(beam_map)
#coordinate system
ones = np.ones(N)
inds = (np.arange(N) - N/2.)
X = np.outer(ones,inds)
Y = -np.transpose(X)
theta_map = np.arctan2(Y,X)
for pix_col in range(len(beam_map)):
for pix_row in range(len(beam_map[pix_col,:])):
theta = theta_map[pix_col,pix_row]
r = np.sqrt( (pix_col - int(N/2))**2. + (pix_row - int(N/2))**2. )
# try:
# theta = np.arctan((pix_col- int(N/2)) / (pix_row-int(N/2)))
# except ZeroDivisionError:
# if pix_col < 512:
# theta = np.pi/2
# else:
# theta = 3*np.pi/2
if r == 0:
dipole_component = 1
else:
#dipole_component = calc_dipole_onaxis(r,theta,p)
dipole_component = calc_dipole(r,theta,p,rotation_angle=rotation_angle)
if dipole_component > 1:
beam_map[pix_col,pix_row] = 1
else:
beam_map[pix_col,pix_row] = dipole_component
return beam_map
def quadrupole_field(beam_map, Q, rotation_angle=0):
import numpy as np
#square map
N = len(beam_map)
#coordinate system
ones = np.ones(N)
inds = (np.arange(N) - N/2.)
X = np.outer(ones,inds)
Y = -np.transpose(X)
theta_map = np.arctan2(Y,X)
for pix_col in range(len(beam_map)):
for pix_row in range(len(beam_map[pix_col,:])):
theta = theta_map[pix_col,pix_row]
r = np.sqrt( (pix_col - int(N/2))**2. + (pix_row - int(N/2))**2. )
# try:
# theta = np.arctan( (pix_col-int(N/2)) / (pix_row-int(N/2)) )
# except ZeroDivisionError:
# if pix_col < 512:
# theta = np.pi/2
# else:
# theta = 3*np.pi/2
if r == 0:
quadrupole_component = 0
else:
quadrupole_component = calc_quadrupole(r,theta,Q,rotation_angle=rotation_angle)
if quadrupole_component > 1:
beam_map[pix_col,pix_row] = 1
else:
beam_map[pix_col,pix_row] = quadrupole_component
return beam_map