-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathUTILITY_linacPhaseAndAmplitude.py
125 lines (86 loc) · 5.49 KB
/
UTILITY_linacPhaseAndAmplitude.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
114
115
116
117
118
119
120
121
122
123
124
125
import numpy as np
# L1MatchStrings = []
# L2MatchStrings = []
# L3MatchStrings = []
# selectMarkers = []
def getLinacMatchStrings(tao):
"""Determine the strings required to match Bmad cavity elements"""
global L1MatchStrings, L2MatchStrings, L3MatchStrings, selectMarkers
#These more general match strings might cause problems by including both lords and slaves!
L1MatchStrings = ["K11_1*", "K11_2*"]
L2MatchStrings = ["K11_4*", "K11_5*", "K11_6*", "K11_7*", "K11_8*", "K12_*", "K13_*", "K14_*"]
L3MatchStrings = ["K15_*", "K16_*", "K17_*", "K18_*", "K19_*"]
#Therefore, use lat_list and these general search strings to get specific elements
#Recall that, by default, lat_list has the flag "-track_only" so lords are not included
#But some of the tracking elements are super_slaves which cannot be `set`
#"-array_out" flag is on by default but if flags are edited, need to re-add manually to get numeric values back
L1MatchStrings = np.concatenate( [ tao.lat_list(i,"ele.name", flags="-no_slaves") for i in L1MatchStrings ] )
L2MatchStrings = np.concatenate( [ tao.lat_list(i,"ele.name", flags="-no_slaves") for i in L2MatchStrings ] )
L3MatchStrings = np.concatenate( [ tao.lat_list(i,"ele.name", flags="-no_slaves") for i in L3MatchStrings ] )
selectMarkers = np.array(["ENDDL10", "BEGBC11_1", "BC11CEND", "BEGBC14_1", "ENDBC14_2", "BEGBC20", "ENDBC20", "PENT"])
return [L1MatchStrings, L2MatchStrings, L3MatchStrings, selectMarkers]
def getEnergyChangeFromElements(tao, activeMatchStrings):
"""Calculate the energy change imparted by selected elements
activeMatchStrings may be a string of a common name ("L1", "L2", or "L3") or an actual list of strings to match
"""
activeMatchStrings = matchStringWrapper(tao, activeMatchStrings)
#VOLTAGE is just gradient times length; need to manually include phase info
voltagesArr = [tao.lat_list(i, "ele.VOLTAGE", flags="-no_slaves -array_out") for i in activeMatchStrings]
voltagesArr = np.concatenate(voltagesArr)
angleMultArr = [tao.lat_list(i, "ele.PHI0", flags="-no_slaves -array_out") for i in activeMatchStrings]
angleMultArr = np.concatenate(angleMultArr)
angleMultArr = [np.cos(i * (2*3.1415) ) for i in angleMultArr] #Recall Bmad uses units of 1/2pi
return( np.dot(voltagesArr, angleMultArr) )
def setLinacGradientAuto(tao, activeMatchStrings, targetVoltage):
"""Set all elements to a constant gradient required to achieve target voltage change
activeMatchStrings may be a string of a common name ("L1", "L2", or "L3") or an actual list of strings to match
"""
activeMatchStrings = matchStringWrapper(tao, activeMatchStrings)
#Set to a fixed gradient so everything is the same. Not exactly physical but probably close
baseGradient = 1.0e7
for i in activeMatchStrings: tao.cmd(f'set ele {i} GRADIENT = {baseGradient}')
#See the resulting voltage gain
voltageSum = getEnergyChangeFromElements(tao, activeMatchStrings)
#print(voltageSum/1e6)
#Uniformly scale gradient to hit target voltage
for i in activeMatchStrings: tao.cmd(f'set ele {i} GRADIENT = {baseGradient*targetVoltage/voltageSum}')
voltageSum = getEnergyChangeFromElements(tao, activeMatchStrings)
#print(voltageSum/1e6)
def setLinacPhase(tao, activeMatchStrings, phi0Deg):
"""Set all elements to a phase
activeMatchStrings may be a string of a common name ("L1", "L2", or "L3") or an actual list of strings to match
"""
activeMatchStrings = matchStringWrapper(tao, activeMatchStrings)
for i in activeMatchStrings: tao.cmd(f'set ele {i} PHI0 = {phi0Deg / 360.}')
def matchStringWrapper(tao, activeMatchStrings):
"""Translate common names ("L1", "L2", or "L3") into relevant match strings
This is just to make life a little simpler for people.
Usually activeMatchStrings should be an array
However, if it's a string that's a common name ("L1", "L2", or "L3"), return the relevant match strings
Otherwise just spit out whatever was provided
"""
if isinstance(activeMatchStrings, str):
[L1MatchStrings, L2MatchStrings, L3MatchStrings, selectMarkers] = getLinacMatchStrings(tao)
if activeMatchStrings == "L1":
return L1MatchStrings
if activeMatchStrings == "L2":
return L2MatchStrings
if activeMatchStrings == "L3":
return L3MatchStrings
return activeMatchStrings
#Broken! If you comma separate too many values it just starts ignoring them...
#def printETOTValues(activeElements):
# printThing = tao.cmd(f"show lat {','.join(activeElements)} -att E_TOT");
# [print(row) for row in printThing];
#Broken! If you comma separate too many values it just starts ignoring them...
#def printArbValues(activeElements, attString):
# printThing = tao.cmd(f"show lat {','.join(activeElements)} -att {attString}");
# [print(row) for row in printThing];
def printArbValues(tao, activeElements, attString):
#The .tolist() is because of issues when activeElements is a numpy array
namesList = [ tao.lat_list(i, "ele.name", flags="-no_slaves -array_out") for i in activeElements.tolist() ]
namesList = np.concatenate(namesList)
valuesList = [ tao.lat_list(i, f"ele.{attString}",flags="-no_slaves -array_out") for i in activeElements.tolist() ]
valuesList = np.concatenate(valuesList) / 1e9
printThing = np.transpose([namesList, valuesList])
display(pd.DataFrame(printThing))