-
Notifications
You must be signed in to change notification settings - Fork 0
/
evaluate.py
193 lines (145 loc) · 5.85 KB
/
evaluate.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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
#!/usr/bin/env python
#
import synergia
import rr_tune_survey
import rr_setup
from rr_options import opts
import rrnova_qt60x
import tune_suite
import h5py
import numpy as np
DEBUG=False
RR_template_file = "RR2020V0922_TEMPLATE_fixed"
#RR_template_file = "RR2020V0922FLAT_fixed"
RR_ring_name = "ring605_fodo"
"""
Given input dict with parameters { 'kxl_even': xxx , ...
'kxl_odd': xxx, ...
}
given x in [1, 2, 3]
propagate particles at a set of momenta in the Recycler ring and
return the spectra or tunes.
"""
#----------------------------------------------------------------------
def generate_lattice(kxl_values, adjust_tunes=True):
lattice = rr_tune_survey.get_rr_lattice_for_opt(RR_template_file, RR_ring_name, kxl_values)
# replacing old rr_setup.setup() do-it-all,
if opts.start_element:
lattice_tmp1 = rr_setup.reorder_lattice(lattice, opts.start_element)
else:
lattice_tmp1 = lattice
lattice_tmp1a = rr_setup.convert_rbends_to_sbends(lattice_tmp1)
if opts.lattice_simplify:
lattice_tmp2 = rr_setup.keep_qt(lattice_tmp1a)
else:
lattice_tmp2 = lattice_tmp1a
(xtune, ytune, cdt) = synergia.simulation.Lattice_simulator.calculate_tune_and_cdt(lattice_tmp2)
print('generate_lattice, initial xtune: ', xtune, ', ytune: ', ytune)
if adjust_tunes and (opts.xtune_adjust or opts.ytune_adjust):
print("Adjusting tunes to:")
print("xtune: ", opts.xtune_adjust)
print("ytune: ", opts.ytune_adjust)
if opts.xtune_adjust:
delta_xtune = opts.xtune_adjust - xtune
print('delta_xtune: ', delta_xtune)
else:
delta_xtune = 0.0
if opts.ytune_adjust:
delta_ytune = opts.ytune_adjust - ytune
print('delta_ytune: ', delta_ytune)
else:
delta_ytune = 0.0
rrnova_qt60x.adjust_rr60_trim_quads(lattice_tmp2, delta_xtune, delta_ytune)
else:
print('skipping tune adjustment')
lattice_tmp2.set_all_string_attribute("extractor_type", "libff")
harmno = 588 # harmonic number
# cavities in lattice_tmp2 are modified in-place
rr_setup.setup_rf_cavities(lattice_tmp2, opts.rf_voltage, harmno)
synergia.simulation.Lattice_simulator.tune_circular_lattice(lattice_tmp2)
return lattice_tmp2
#----------------------------------------------------------------------
def save_lattice_txt(lattice, filename):
if synergia.utils.Commxx().get_rank() == 0:
with open(filename, 'w') as f:
print(lattice, file=f)
#----------------------------------------------------------------------
def run_particles(lattice):
# We're only going to propagate a small number of particles
# each at a different momentum to determine their tunes so I
# don't really need the grid stuff.
rr_tune_survey.run_rr(lattice, opts.turns)
#----------------------------------------------------------------------
# calculate the x and y tunes for each particle
def analyze_propagation():
h5 = h5py.File('tracks.h5', 'r')
trks = h5.get('track_coords')[()]
npart = trks.shape[1]
xtunes = np.zeros(npart)
ytunes = np.zeros(npart)
for n in range(npart):
# calculate tunes this particle
t = tune_suite.interp_tunes(trks[:, n, 0:6].transpose())
xtunes[n] = t[0]
ytunes[n] = t[1]
return xtunes, ytunes
#----------------------------------------------------------------------
def evaluate(kxl_values, chatty=False, adjust_tunes=True):
try:
lattice = generate_lattice(kxl_values, adjust_tunes=adjust_tunes)
except:
if DEBUG:
print("exception generated from generate_lattice")
print("kxl_values: ", kxl_values, flush=True)
return None
if chatty:
print("read lattice, ", len(lattice.get_elements()), ", length: ", lattice.get_length())
save_lattice_txt(lattice, 'rr_lattice.out')
try:
(nux, nuy, cdT) = synergia.simulation.Lattice_simulator.calculate_tune_and_cdt(lattice)
except:
if DEBUG:
print("exception generated from calculate_tunes_and_cdt")
print("kxl_values: ", kxl_values, flush=True)
return None
if chatty:
print('tune x: ', nux)
print('tune y: ', nuy)
print('cdT: ', cdT)
print('T: ', cdT/synergia.foundation.pconstants.c)
try:
chrom_t = synergia.simulation.Lattice_simulator.get_chromaticities(lattice)
except:
if DEBUG:
print("exception generated from calculate_get_chromaticities")
print("kxl_values: ", kxl_values, flush=True)
return None
if chatty:
print('horizontal chromaticity: ', chrom_t.horizontal_chromaticity)
print('vertical chromaticity: ', chrom_t.vertical_chromaticity)
print('compaction factor: ', chrom_t.momentum_compaction)
print('slip factor: ', chrom_t.slip_factor)
run_particles(lattice)
return analyze_propagation()
#----------------------------------------------------------------------
# main() run evaluate() for representative set of KxL values
def main():
# sample settings
#kxl_values = {} # this one uses the unmodified multipole settings
# from the lattice file
kxl_values = { # this one has slight changes
'K1L_EVEN': 0.01,
'K1L_ODD': -0.005,
'K2L_EVEN': -0.005,
'K2L_ODD': -0.00033,
'K3L_EVEN': 0.0000129,
'K3L_ODD': -0.0000333,
'K4L_EVEN': 0.0,
'K4L_ODD': 0.0,
'K5L_EVEN': 0.0,
'K5L_ODD': 0.0
}
return evaluate(kxl_values, True)
#----------------------------------------------------------------------
if __name__ == "__main__":
main()