-
Notifications
You must be signed in to change notification settings - Fork 0
/
tune_suite.py
113 lines (94 loc) · 3.24 KB
/
tune_suite.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
#!/usr/bin/env python
import numpy as np
# suite of functions that can do different tune calculations
# calculate the simple x y z tunes just from fft
def basic_tunes(coords):
# coords has shape (6,n)
n = coords.shape[1]
df = 1.0/n
# get maximum amplitude location
maxn = int(n/2)
xt = np.abs(np.fft.fft(coords[0,:]))
# probably don't want 0 frequency
locmax = np.argmax(xt[0:maxn])
xtune = (locmax)*df
yt = np.abs(np.fft.fft(coords[2,:]))
locmax = np.argmax(yt[0:maxn])
ytune = (locmax)*df
zt = np.abs(np.fft.fft(coords[4,:]))
locmax = np.argmax(zt[0:maxn])
ztune = (locmax)*df
return (xtune,ytune,ztune)
# calculate tunes using the CFT algorithm
def cft_tunes(coords, search_range=((0.0,0.5),(0.0,0.5),(0.0,0.5))):
# coords has shape (6,n)
n = coords.shape[1]
#if n>100:
# print "warning! you have a lot of points. This will be slow"
# normal FFT precision is 1/n, CFT gives addition 1/n factor
df = 1.0/n**2
t = np.arange(n,dtype='d')
# loop over x,y,z
tunes = np.zeros(3)
for pln in range(3):
f = np.arange(search_range[pln][0],search_range[pln][1],df)
nf = len(f)
cft = np.zeros(nf,dtype='d')
idxcoord = 2*pln
for i in range(nf):
expfact = np.exp(-2.0*np.pi*(1.0j)*f[i]*t)
cft[i] = abs(np.dot(expfact, coords[idxcoord,:]))
cftmaxloc = np.argmax(cft)
tunes[pln] = f[cftmaxloc]
return tuple(tunes)
# get interpolated tunes
def interp_tunes(coords, search_range=((0.0,0.5),(0.0,0.5),(0.0,0.5))):
# coords has shape (6,n)
n = coords.shape[1]
# loop pver x. y, z
tunes = np.zeros(3)
maxn = int(n/2)
df = 1.0/n
f = np.arange(n,dtype='d')/n
for pln in range(0,6,2):
xt = abs(np.fft.fft(coords[pln, :]))
# cut off low end
xt[0:10] = 0.0
locmax = np.argmax(xt[0:maxn])
xtp = xt[locmax]
if xt[locmax-1] > xt[locmax+1]:
dir=-1.0
xtp2 = xt[locmax-1]
else:
dir=1.0
xtp2 = xt[locmax+1]
tune = f[locmax] + dir*(xtp2/(xtp+xtp2))/n
tunes[int(pln/2)] = tune
return tuple(tunes)
# get the (fractional) tunes of a set of coordinates from a single track
def refined_tunes(coords):
# coords has shape (6,n)
n = coords.shape[1]
df = 1.0/n
if n <= 100:
# really, if there are less than 100 points, just use
# the cft tunes
ctunes = cft_tunes(coords)
return ctunes
# first get the basic tunes
btunes = basic_tunes(coords)
if np.any(btunes == 0.0):
# if some of the basic tunes are 0.0, then there was
# some problem and it's probably not right. Get the CFT tunes
# for the first 100 points.
ctunes = cft_tunes(coords[:,:100])
xrange = (ctunes[0]-2.0*df, ctunes[0]+2.0*df)
yrange = (ctunes[1]-2.0*df, ctunes[1]+2.0*df)
zrange = (ctunes[2]-2.0*df, ctunes[2]+2.0*df)
else:
xrange = (btunes[0]-2.0*df, btunes[0]+2.0*df)
yrange = (btunes[1]-2.0*df, btunes[1]+2.0*df)
zrange = (btunes[2]-2.0*df, btunes[2]+2.0*df)
tunerange = (xrange,yrange,zrange)
ctunes = cft_tunes(coords, tunerange)
return ctunes