Skip to content

Commit 6ab5d3e

Browse files
rdemariardemaria
authored
Merge pull request #56 from kparasch/feature/closed-orbit_and_linear_normal_form
Feature/closed orbit and linear normal form
2 parents 8e1d432 + 52e38a3 commit 6ab5d3e

File tree

3 files changed

+247
-0
lines changed

3 files changed

+247
-0
lines changed

‎pysixtrack/closed_orbit.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
import numpy as np
2+
from .particles import Particles
3+
4+
5+
def get_init_particles_for_linear_map(closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum):
6+
part = Particles(p0c=p0c)
7+
# part = Particles()
8+
# part.p0c = p0c
9+
part.x = closed_orbit[0] + np.array([0.0, 1.0 * d, 0.0, 0.0, 0.0, 0.0, 0.0])
10+
part.px = closed_orbit[1] + np.array([0.0, 0.0, 1.0 * d, 0.0, 0.0, 0.0, 0.0])
11+
part.y = closed_orbit[2] + np.array([0.0, 0.0, 0.0, 1.0 * d, 0.0, 0.0, 0.0])
12+
part.py = closed_orbit[3] + np.array([0.0, 0.0, 0.0, 0.0, 1.0 * d, 0.0, 0.0])
13+
14+
longi = closed_orbit[4] + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 1.0 * d, 0.0])
15+
plongi = closed_orbit[5] + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 * d])
16+
17+
setattr(part, longitudinal_coordinate, longi)
18+
setattr(part, longitudinal_momentum, plongi)
19+
20+
return part
21+
22+
23+
def part_to_array(part, longitudinal_coordinate, longitudinal_momentum):
24+
X_array = np.empty([6, 7])
25+
X_array[0, :] = part.x
26+
X_array[1, :] = part.px
27+
X_array[2, :] = part.y
28+
X_array[3, :] = part.py
29+
X_array[4, :] = getattr(part, longitudinal_coordinate)
30+
X_array[5, :] = getattr(part, longitudinal_momentum)
31+
32+
return X_array
33+
34+
35+
def linearize_around_closed_orbit(line, closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum):
36+
37+
part = get_init_particles_for_linear_map(closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum)
38+
X_init = part_to_array(part, longitudinal_coordinate, longitudinal_momentum)
39+
40+
line.track(part)
41+
X_fin = part_to_array(part, longitudinal_coordinate, longitudinal_momentum)
42+
43+
m = X_fin[:, 0] - X_init[:, 0]
44+
M = np.empty([6, 6])
45+
for j in range(6):
46+
M[:, j] = (X_fin[:, j + 1] - X_fin[:, 0]) / d
47+
48+
X_CO = X_init[:, 0] + np.matmul(np.linalg.inv(np.identity(6) - M), m.T)
49+
50+
return X_CO, M
51+
52+
53+
def healy_symplectify(M):
54+
# https://accelconf.web.cern.ch/e06/PAPERS/WEPCH152.PDF
55+
print("Symplectifying linear One-Turn-Map...")
56+
57+
print("Before symplectifying: det(M) = {}".format(np.linalg.det(M)))
58+
I = np.identity(6)
59+
60+
S = np.array(
61+
[
62+
[0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
63+
[-1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
64+
[0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
65+
[0.0, 0.0, -1.0, 0.0, 0.0, 0.0],
66+
[0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
67+
[0.0, 0.0, 0.0, 0.0, -1.0, 0.0],
68+
]
69+
)
70+
71+
V = np.matmul(S, np.matmul(I - M, np.linalg.inv(I + M)))
72+
W = (V + V.T) / 2
73+
if np.linalg.det(I - np.matmul(S, W)) != 0:
74+
M_new = np.matmul(I + np.matmul(S, W), np.linalg.inv(I - np.matmul(S, W)))
75+
else:
76+
print("WARNING: det(I - SW) = 0!")
77+
V_else = np.matmul(S, np.matmul(I + M, np.linalg.inv(I - M)))
78+
W_else = (V_else + V_else.T) / 2
79+
M_new = -np.matmul(
80+
I + np.matmul(S, W_else), np.linalg(I - np.matmul(S, W_else))
81+
)
82+
83+
print("After symplectifying: det(M) = {}".format(np.linalg.det(M_new)))
84+
return M_new

‎pysixtrack/line.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@
66

77
from .loader_sixtrack import _expand_struct
88
from .loader_mad import iter_from_madx_sequence
9+
from .closed_orbit import linearize_around_closed_orbit
10+
from .closed_orbit import healy_symplectify
11+
from .linear_normal_form import _linear_normal_form
12+
913

1014
# missing access to particles._m:
1115
deg2rad = np.pi / 180.
@@ -192,6 +196,46 @@ def get_elements_of_type(self, types):
192196

193197
return elements, names
194198

199+
def linear_normal_form(self, M):
200+
return _linear_normal_form(M)
201+
202+
def find_closed_orbit_and_linear_OTM(
203+
self, p0c, guess=None, d=1.e-7, tol=1.e-10, max_iterations=20, longitudinal_coordinate='zeta'
204+
):
205+
if guess is None:
206+
guess = [0., 0., 0., 0., 0., 0.]
207+
208+
assert len(guess) == 6
209+
210+
closed_orbit = np.array(guess).copy()
211+
212+
canonical_conjugate_momentum = {'tau' : 'ptau', 'zeta' : 'delta', 'sigma' : 'psigma'}
213+
214+
if longitudinal_coordinate not in ['tau', 'zeta', 'sigma']:
215+
raise Exception('Longitudinal variable not recognized in search of closed orbit')
216+
217+
longitudinal_momentum = canonical_conjugate_momentum[longitudinal_coordinate]
218+
219+
for i in range(max_iterations):
220+
new_closed_orbit, M = linearize_around_closed_orbit(
221+
self, closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum
222+
)
223+
224+
error = np.linalg.norm( new_closed_orbit - closed_orbit )
225+
226+
closed_orbit = new_closed_orbit
227+
if error < tol:
228+
print('Converged with approximate distance: {}'.format(error))
229+
_, M = linearize_around_closed_orbit(
230+
self, closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum
231+
)
232+
return closed_orbit, healy_symplectify(M)
233+
234+
print ('Closed orbit search iteration: {}'.format(i))
235+
236+
print('WARNING!: Search did not converge, approximate distance: {}'.format(error))
237+
return closed_orbit, healy_symplectify(M)
238+
195239
def find_closed_orbit(
196240
self, p0c, guess=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], method="Nelder-Mead"
197241
):

‎pysixtrack/linear_normal_form.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import numpy as np
2+
3+
S = np.array([[0., 1., 0., 0., 0., 0.],
4+
[-1., 0., 0., 0., 0., 0.],
5+
[ 0., 0., 0., 1., 0., 0.],
6+
[ 0., 0.,-1., 0., 0., 0.],
7+
[ 0., 0., 0., 0., 0., 1.],
8+
[ 0., 0., 0., 0.,-1., 0.]])
9+
10+
######################################################
11+
### Implement Normalization of fully coupled motion ##
12+
######################################################
13+
14+
def Rot2D(mu):
15+
return np.array([[ np.cos(mu), np.sin(mu)],
16+
[-np.sin(mu), np.cos(mu)]])
17+
18+
def _linear_normal_form(M):
19+
w0, v0 = np.linalg.eig(M)
20+
21+
a0 = np.real(v0)
22+
b0 = np.imag(v0)
23+
24+
index_list = [0,5,1,2,3,4]
25+
26+
##### Sort modes in pairs of conjugate modes #####
27+
28+
conj_modes = np.zeros([3,2], dtype=np.int)
29+
for j in [0,1]:
30+
conj_modes[j,0] = index_list[0]
31+
del index_list[0]
32+
33+
min_index = 0
34+
min_diff = abs(np.imag(w0[conj_modes[j,0]] + w0[index_list[min_index]]))
35+
for i in range(1,len(index_list)):
36+
diff = abs(np.imag(w0[conj_modes[j,0]] + w0[index_list[i]]))
37+
if min_diff > diff:
38+
min_diff = diff
39+
min_index = i
40+
41+
conj_modes[j,1] = index_list[min_index]
42+
del index_list[min_index]
43+
44+
conj_modes[2,0] = index_list[0]
45+
conj_modes[2,1] = index_list[1]
46+
47+
##################################################
48+
#### Select mode from pairs with positive (real @ S @ imag) #####
49+
50+
modes = np.empty(3, dtype=np.int)
51+
for ii,ind in enumerate(conj_modes):
52+
if np.matmul(np.matmul(a0[:,ind[0]], S), b0[:,ind[0]]) > 0:
53+
modes[ii] = ind[0]
54+
else:
55+
modes[ii] = ind[1]
56+
57+
##################################################
58+
#### Sort modes such that (1,2,3) is close to (x,y,zeta) ####
59+
for i in [1,2]:
60+
if abs(v0[:,modes[0]])[0] < abs(v0[:,modes[i]])[0]:
61+
modes[0], modes[i] = modes[i], modes[0]
62+
63+
if abs(v0[:,modes[1]])[2] < abs(v0[:,modes[2]])[2]:
64+
modes[2], modes[1] = modes[1], modes[2]
65+
66+
##################################################
67+
#### Rotate eigenvectors to the Courant-Snyder parameterization ####
68+
phase0 = np.log(v0[0,modes[0]]).imag
69+
phase1 = np.log(v0[2,modes[1]]).imag
70+
phase2 = np.log(v0[4,modes[2]]).imag
71+
72+
v0[:,modes[0]] *= np.exp(-1.j*phase0)
73+
v0[:,modes[1]] *= np.exp(-1.j*phase1)
74+
v0[:,modes[2]] *= np.exp(-1.j*phase2)
75+
76+
##################################################
77+
#### Construct W #################################
78+
79+
a1 = v0[:,modes[0]].real
80+
a2 = v0[:,modes[1]].real
81+
a3 = v0[:,modes[2]].real
82+
b1 = v0[:,modes[0]].imag
83+
b2 = v0[:,modes[1]].imag
84+
b3 = v0[:,modes[2]].imag
85+
86+
n1 = 1./np.sqrt(np.matmul(np.matmul(a1, S), b1))
87+
n2 = 1./np.sqrt(np.matmul(np.matmul(a2, S), b2))
88+
n3 = 1./np.sqrt(np.matmul(np.matmul(a3, S), b3))
89+
90+
a1 *= n1
91+
a2 *= n2
92+
a3 *= n3
93+
94+
b1 *= n1
95+
b2 *= n2
96+
b3 *= n3
97+
98+
W = np.array([a1,b1,a2,b2,a3,b3]).T
99+
W[abs(W) < 1.e-14] = 0. # Set very small numbers to zero.
100+
invW = np.matmul(np.matmul(S.T, W.T), S)
101+
102+
##################################################
103+
#### Get tunes and rotation matrix in the normalized coordinates ####
104+
105+
mu1 = np.log(w0[modes[0]]).imag
106+
mu2 = np.log(w0[modes[1]]).imag
107+
mu3 = np.log(w0[modes[2]]).imag
108+
109+
#q1 = mu1/(2.*np.pi)
110+
#q2 = mu2/(2.*np.pi)
111+
#q3 = mu3/(2.*np.pi)
112+
113+
R = np.zeros_like(W)
114+
R[0:2,0:2] = Rot2D(mu1)
115+
R[2:4,2:4] = Rot2D(mu2)
116+
R[4:6,4:6] = Rot2D(mu3)
117+
##################################################
118+
119+
return W, invW, R

0 commit comments

Comments
 (0)