1
|
import os
|
2
|
import subprocess
|
3
|
import json
|
4
|
import operator
|
5
|
import glob
|
6
|
import sys
|
7
|
|
8
|
import numpy
|
9
|
import ROOT
|
10
|
|
11
|
import Configuration
|
12
|
import maus_cpp.globals
|
13
|
|
14
|
import xboa.Common as common
|
15
|
from xboa.Bunch import Bunch
|
16
|
from xboa.Hit import Hit
|
17
|
from xboa.tracking import MAUSTracking
|
18
|
|
19
|
class Tracking:
|
20
|
|
21
|
def amp(self, ellipse_inv, hit):
|
22
|
arr = numpy.matrix([hit['x'], hit['px'], hit['y'], hit['py']])
|
23
|
amplitude = arr*ellipse_inv*arr.T
|
24
|
return amplitude[0, 0]
|
25
|
|
26
|
def amp_list(self, bunch, beta, alpha, bz):
|
27
|
ellipse = Bunch.build_penn_ellipse(1., self.m_mu, beta, alpha, self.p, 0., bz, 1.)
|
28
|
ellipse_inv = numpy.linalg.inv(ellipse)
|
29
|
amplitude_list = [self.amp(ellipse_inv, hit) for hit in bunch]
|
30
|
return amplitude_list
|
31
|
|
32
|
def amp_list_6d(self, bunch):
|
33
|
bunch.set_covariance_matrix(True)
|
34
|
amp_list = [bunch.get_amplitude(bunch, hit, ['x', 'y', 'ct']) for hit in bunch]
|
35
|
bunch.set_covariance_matrix(False)
|
36
|
return amp_list
|
37
|
|
38
|
def scatter_plot_mean_rms(self, x_values, y_values, xmin, xmax, xstep):
|
39
|
nbins = 9
|
40
|
bin_low = 0.
|
41
|
bin_step = 10.
|
42
|
hist_data = []
|
43
|
for x, y in zip(x_values, y_values):
|
44
|
bin = int((x-bin_low)/bin_step)
|
45
|
while len(hist_data) - 1 < bin:
|
46
|
hist_data.append([])
|
47
|
hist_data[bin].append(y)
|
48
|
hist = ROOT.TH1D("Mean", "", nbins, bin_low, bin_low+(nbins+1)*bin_step)
|
49
|
for i, data in enumerate(hist_data):
|
50
|
mean, sigmaSqu = 0., 0.
|
51
|
if len(data) > 0:
|
52
|
mean = sum(data)/len(data)
|
53
|
if len(data) > 1:
|
54
|
sigmaSqu = sum([x**2 for x in data])/len(data)-mean**2
|
55
|
hist.SetBinContent(i+1, mean)
|
56
|
hist.SetBinError(i+1, sigmaSqu**0.5)
|
57
|
common._hist_persistent.append(hist)
|
58
|
return hist
|
59
|
|
60
|
def amplitude_scatter_plot(self, bunch_in, bunch_out, predicate):
|
61
|
"""predicate takes arguments (spill_index, spills_in_bunch_out)"""
|
62
|
spills_out = set([(hit['spill'], hit['event_number']) for hit in bunch_out])
|
63
|
hits_in = [hit for hit in bunch_in if (hit['spill'], hit['event_number']) in spills_out]
|
64
|
hits_out = [hit for hit in bunch_out if (hit['spill'], hit['event_number']) in spills_out]
|
65
|
amp_in = self.amp_list(hits_in, self.beta, 0., self.bz_us)
|
66
|
amp_out = self.amp_list(hits_out, self.beta, 0., self.bz_ds)
|
67
|
delta_amp = [(amp_out[i]-amp_in[i])/amp_in[i] for i, a_in in enumerate(amp_in)]
|
68
|
if len(amp_in) == 0 or len(amp_in) != len(amp_out):
|
69
|
print "Failed to make graph:", len(amp_in), len(amp_out)
|
70
|
hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'x-y amplitude out [mm]')
|
71
|
else:
|
72
|
hist, graph = common.make_root_graph('amplitude x y', amp_in, 'x-y amplitude in [mm]', delta_amp, '(A_{out} - A_{in})/A_{in} [mm]', xmin=0., xmax=200.)
|
73
|
hist_means = self.scatter_plot_mean_rms(amp_in, delta_amp, 0., 200., 10.)
|
74
|
graph.SetMarkerStyle(7)
|
75
|
return hist, hist_means, graph
|
76
|
|
77
|
def amplitude_p_scatter_plot(self, bunch_in):
|
78
|
"""predicate takes arguments (spill_index, spills_in_bunch_out)"""
|
79
|
hits_in = [hit for hit in bunch_in]
|
80
|
amp_in = self.amp_list(hits_in, self.beta, 0., self.bz_us)
|
81
|
p_in = [hit['p'] for hit in hits_in]
|
82
|
if len(amp_in) == 0:
|
83
|
print "Failed to make graph:", len(amp_in)
|
84
|
hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'P [MeV/c]', xmin=0., xmax=200.)
|
85
|
else:
|
86
|
hist, graph = common.make_root_graph('amplitude x y', amp_in, 'x-y amplitude in [mm]', p_in, 'P [MeV/c]', xmin=0., xmax=200.)
|
87
|
hist_means = self.scatter_plot_mean_rms(amp_in, p_in, 0., 200., 10.)
|
88
|
graph.SetMarkerStyle(7)
|
89
|
return hist, hist_means, graph
|
90
|
|
91
|
def amplitude_t_scatter_plot(self, bunch_in, bunch_out):
|
92
|
"""predicate takes arguments (spill_index, spills_in_bunch_out)"""
|
93
|
spills_out = set([(hit['spill'], hit['event_number']) for hit in bunch_out])
|
94
|
hits_in = [hit for hit in bunch_in if (hit['spill'], hit['event_number']) in spills_out]
|
95
|
hits_out = [hit for hit in bunch_out if (hit['spill'], hit['event_number']) in spills_out]
|
96
|
amp_in = self.amp_list_6d(hits_in, self.beta, 0., self.bz_us)
|
97
|
amp_out = self.amp_list(hits_out, self.beta, 0., self.bz_ds)
|
98
|
delta_t = [h_o['t']-h_i['t'] for h_i, h_o in zip(hits_in, hits_out)]
|
99
|
if len(amp_in) == 0:
|
100
|
print "Failed to make graph:", len(amp_in)
|
101
|
hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'dt [ns]', xmin=0., xmax=200.)
|
102
|
else:
|
103
|
hist, graph = common.make_root_graph('amplitude x y', amp_in, 'x-y amplitude in [mm]', delta_t, 'dt [ns]', xmin=0., xmax=200.)
|
104
|
hist_means = self.scatter_plot_mean_rms(amp_in, delta_t, 0., 200., 10.)
|
105
|
graph.SetMarkerStyle(7)
|
106
|
return hist, hist_means, graph
|
107
|
|
108
|
def amplitude_6d_scatter_plot(self, bunch_in, bunch_out):
|
109
|
"""predicate takes arguments (spill_index, spills_in_bunch_out)"""
|
110
|
amp_in = self.amp_list_6d(bunch_in)
|
111
|
amp_out = self.amp_list_6d(bunch_out)
|
112
|
hist_in = common.make_root_histogram('A_{6D} in', amp_in, 'A_{6D} [mm]', 20, xmin=0., xmax=200.)
|
113
|
hist_out = common.make_root_histogram('A_{6D} out', amp_out, 'A_{6D} [mm]', 20, xmin=0., xmax=200.)
|
114
|
hist_in.SetName('A_{6D} in')
|
115
|
hist_out.SetName('A_{6D} out')
|
116
|
return hist_in, hist_out
|
117
|
|
118
|
def bin_data(data_list, bin_edge_list):
|
119
|
bin = 0
|
120
|
#for amp_in in
|
121
|
|
122
|
def amplitude_6d_capture_plot(self, bunch_in, bunch_out):
|
123
|
amp_in = sorted(self.amp_list_6d(bunch_in))
|
124
|
amp_out = sorted(self.amp_list_6d(bunch_out))
|
125
|
bins = [i*10. for i in range(11)]+[i*50. for i in range(2, 5)]
|
126
|
contents_in = []
|
127
|
contents_out = []
|
128
|
|
129
|
|
130
|
def amplitude_scatter(self, bunch_in, bunch_out):
|
131
|
canvas = common.make_root_canvas('amplitude 6d hist')
|
132
|
hist_in, hist_out = self.amplitude_6d_scatter_plot(bunch_in, bunch_out)
|
133
|
hist_out.SetLineColor(4)
|
134
|
hist_out.Draw()
|
135
|
hist_in.SetLineColor(2)
|
136
|
hist_in.Draw("SAME")
|
137
|
common.make_root_legend(canvas, [hist_in, hist_out])
|
138
|
canvas.Update()
|
139
|
self.print_canvas(canvas, 'acceptance_6d_hist')
|
140
|
"""
|
141
|
spills_out = set([hit['spill'] for hit in bunch_out])
|
142
|
canvas = common.make_root_canvas('amplitude x y scatter')
|
143
|
hist, hist_means, graph = self.amplitude_scatter_plot(bunch_in, bunch_out, lambda x: True)
|
144
|
hist.Draw()
|
145
|
hist_means.Draw("same")
|
146
|
graph.Draw('p')
|
147
|
canvas.Update()
|
148
|
self.print_canvas(canvas, 'acceptance')
|
149
|
canvas = common.make_root_canvas('amplitude x y p scatter')
|
150
|
hist, hist_means, graph = self.amplitude_p_scatter_plot(bunch_in)
|
151
|
hist.Draw()
|
152
|
hist_means.Draw("same")
|
153
|
graph.Draw('p')
|
154
|
canvas.Update()
|
155
|
self.print_canvas(canvas, 'acceptance_p')
|
156
|
canvas = common.make_root_canvas('amplitude x y t scatter')
|
157
|
hist, hist_means, graph = self.amplitude_t_scatter_plot(bunch_in, bunch_out)
|
158
|
hist.Draw()
|
159
|
hist_means.Draw("same")
|
160
|
graph.Draw('p')
|
161
|
canvas.Update()
|
162
|
self.print_canvas(canvas, 'acceptance_t')
|
163
|
"""
|
164
|
pass
|
165
|
|
166
|
def two_d_plots(self, bunch_list, colour, canvas_list = None):
|
167
|
n_graphs = 4
|
168
|
if canvas_list == None:
|
169
|
canvas_list = [None]*n_graphs
|
170
|
graph_list = [None]*n_graphs
|
171
|
|
172
|
canvas_list[0], hist, graph_list[0] = bunch_list[0].root_scatter_graph("t", "energy", "ns", "MeV", include_weightless=False, canvas=canvas_list[0])
|
173
|
canvas_list[1], hist, graph_list[1] = bunch_list[-1].root_scatter_graph("t", "energy", "ns", "MeV", include_weightless=False, canvas=canvas_list[1])
|
174
|
canvas_list[2], hist, graph_list[2] = bunch_list[0].root_scatter_graph("r", "pt", "mm", "MeV/c", include_weightless=False, canvas=canvas_list[2])
|
175
|
canvas_list[3], hist, graph_list[3] = bunch_list[-1].root_scatter_graph("r", "pt", "mm", "MeV/c", include_weightless=False, canvas=canvas_list[3])
|
176
|
for i in range(n_graphs):
|
177
|
canvas_list[i].cd()
|
178
|
graph_list[i].SetMarkerStyle(6)
|
179
|
graph_list[i].SetMarkerColor(colour)
|
180
|
graph_list[i].Draw("p")
|
181
|
canvas_list[i].Update()
|
182
|
self.print_canvas(canvas_list[0], "time_energy_at_start")
|
183
|
self.print_canvas(canvas_list[1], "time_energy_at_end")
|
184
|
self.print_canvas(canvas_list[2], "r_pt_at_start")
|
185
|
self.print_canvas(canvas_list[3], "r_pt_at_end")
|
186
|
return canvas_list
|
187
|
|
188
|
def cut(self, bunch_list):
|
189
|
print " weights in "
|
190
|
print " no cut:", bunch_list[0].bunch_weight(), bunch_list[-1].bunch_weight()
|
191
|
for bunch in bunch_list:
|
192
|
bunch.transmission_cut(bunch_list[-1], global_cut=True)
|
193
|
bunch.transmission_cut(bunch_list[-1], global_cut=True, test_variable = ['spill', 'event_number', 'particle_number', 'pid'])
|
194
|
bunch.cut({'pid':-13}, operator.ne, global_cut=True)
|
195
|
print ' transmission cut:', bunch_list[0].bunch_weight(), bunch_list[-1].bunch_weight(), "**"
|
196
|
bunches = [bunch_list[0], bunch_list[-1]]
|
197
|
cut_stream = ['upstream', 'downstream']
|
198
|
for i in [0, 1]:
|
199
|
print ' ', cut_stream[i]
|
200
|
print ' ', 'cut ', ' u/s', ' d/s'
|
201
|
bunches[i].cut({'r':self.r_cut[i]}, operator.gt, global_cut=True)
|
202
|
print ' cut r: ', self.r_cut[i], str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)
|
203
|
bunches[i].cut({'amplitude x y':self.amplitude_trans_cut[i]}, operator.gt, global_cut=True)
|
204
|
print ' cut amp_trans: ', str(self.amplitude_trans_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)
|
205
|
bunches[i].cut({'amplitude x y':self.amplitude_trans_cut[i]}, operator.gt, global_cut=True)
|
206
|
print ' cut amp_trans: ', str(self.amplitude_trans_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)
|
207
|
if self.do_long_cut:
|
208
|
bunches[i].cut({'amplitude ct':self.amplitude_ct_cut[i]}, operator.gt, global_cut=True)
|
209
|
print ' cut amp_long: ', str(self.amplitude_ct_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)
|
210
|
bunches[i].cut({'amplitude ct':self.amplitude_ct_cut[i]}, operator.gt, global_cut=True)
|
211
|
print ' cut amp_long: ', str(self.amplitude_ct_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)
|
212
|
mean_p = bunches[i].mean(['p'])['p']
|
213
|
bunches[i].cut({'p':mean_p+self.p_cut[i][1]}, operator.gt, global_cut=True)
|
214
|
bunches[i].cut({'p':mean_p-self.p_cut[i][0]}, operator.lt, global_cut=True)
|
215
|
print ' cut', round(mean_p-self.p_cut[i][0], 1), '< p <', round(mean_p+self.p_cut[i][1], 1), ' ', str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)
|
216
|
|
217
|
def z_plot(self, bunch_list):
|
218
|
print "Start MOMENTS"
|
219
|
print bunch_list[0].moment(['t', 't'])
|
220
|
print bunch_list[0].moment(['pz', 't'])
|
221
|
print bunch_list[0].moment(['pz', 'pz'])
|
222
|
"""
|
223
|
for item in ['t', 'pz', 'energy']:
|
224
|
canvas = common.make_root_canvas('sigma_'+item)
|
225
|
z_list = [bunch.mean(['z'])['z'] for bunch in bunch_list]
|
226
|
item_list = [bunch.moment([item, item])**0.5 for bunch in bunch_list]
|
227
|
hist, graph = common.make_root_graph('sigma_'+item, z_list, 'mean z [mm]', item_list, '#sigma('+item+')')
|
228
|
hist.Draw()
|
229
|
graph.Draw('l')
|
230
|
self.print_canvas(canvas, 's_'+item+'_vs_z')
|
231
|
"""
|
232
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['p'], 'mm', 'MeV/c')
|
233
|
self.print_canvas(canvas, 'momentum_vs_z')
|
234
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['x', 'y'], 'mm', 'mm')
|
235
|
self.print_canvas(canvas, 'beta_trans_vs_z')
|
236
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['x', 'y'], 'mm', 'mm')
|
237
|
self.print_canvas(canvas, 'emittance_trans_vs_z')
|
238
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['t'], 'mm', 'mm')
|
239
|
self.print_canvas(canvas, 'beta_long_vs_z')
|
240
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['ct'], 'mm', 'mm')
|
241
|
self.print_canvas(canvas, 'emittance_long_vs_z')
|
242
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['ct', 'x', 'y'], 'mm', 'mm')
|
243
|
self.print_canvas(canvas, 'emittance_6d_vs_z')
|
244
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['bz'])
|
245
|
self.print_canvas(canvas, 'bz_vs_z')
|
246
|
|
247
|
def grok_filename(self, filename):
|
248
|
base = os.path.basename(filename)[12:-5]
|
249
|
emittance = base.split('_')[0]
|
250
|
emittance = float(emittance.split('=')[1])
|
251
|
a_dir = os.path.dirname(filename).split('/')
|
252
|
index = a_dir[-1].split('_')[-1]
|
253
|
base += '_'+index
|
254
|
a_dir = "/".join(a_dir[:-1])
|
255
|
print " directory", a_dir, "name", base
|
256
|
self.current = {
|
257
|
"dir":a_dir,
|
258
|
"name":base,
|
259
|
"full":filename,
|
260
|
"emittance":emittance,
|
261
|
"index":index
|
262
|
}
|
263
|
return self.current
|
264
|
|
265
|
def print_canvas(self, canvas, name):
|
266
|
for format in ["png", "root"]:
|
267
|
canvas.Print(self.current["dir"]+"/"+name+"_"+self.current["name"]+"."+format)
|
268
|
|
269
|
def single_analysis(self, filename):
|
270
|
print "Loading data", filename
|
271
|
metadata = self.grok_filename(filename)
|
272
|
bunch_list = Bunch.new_list_from_read_builtin('maus_root_virtual_hit', filename)
|
273
|
bunch_list = [bunch for bunch in bunch_list if bunch.bunch_weight() > 100]
|
274
|
if self.z_max != None:
|
275
|
bunch_list = [bunch for bunch in bunch_list if bunch[0]['z'] < self.z_max]
|
276
|
print " Printing transmission (before cut)"
|
277
|
canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'bunch_weight', [], 'mm', '')
|
278
|
self.print_canvas(canvas, 'weight_vs_z')
|
279
|
print " Making cuts"
|
280
|
canvas_list = self.two_d_plots(bunch_list, 2)
|
281
|
self.cut(bunch_list)
|
282
|
self.amplitude_scatter(bunch_list[0], bunch_list[-1])
|
283
|
print " Making 2d plots"
|
284
|
self.two_d_plots(bunch_list, 1, canvas_list)
|
285
|
print " Making z plots"
|
286
|
self.z_plot(bunch_list)
|
287
|
print " Making acceptance plots"
|
288
|
self.data.append((metadata, bunch_list))
|
289
|
Bunch.clear_global_weights()
|
290
|
|
291
|
def multi_analysis(self):
|
292
|
multi_data = open(self.data[0][0]["dir"]+"/multi_analysis_"+self.current["index"]+".json", "w")
|
293
|
delta_transmission = []
|
294
|
delta_emittance_trans = []
|
295
|
nominal_emittance_trans = []
|
296
|
actual_emittance_trans = []
|
297
|
actual_emittance_long = []
|
298
|
delta_emittance_long = []
|
299
|
actual_emittance_6d = []
|
300
|
delta_emittance_6d = []
|
301
|
sigma_energy = []
|
302
|
for item in self.data:
|
303
|
self.cut(item[1])
|
304
|
try:
|
305
|
delta_trans = float(len(item[1][-1]))/float(len(item[1][0]))
|
306
|
emit_in = item[1][0].get_emittance(['x', 'y'])
|
307
|
emit_out = item[1][-1].get_emittance(['x', 'y'])
|
308
|
delta_emit_t = emit_out/emit_in-1.
|
309
|
s_energy = (item[1][0].moment(['energy', 'energy']))**0.5
|
310
|
emit_long_in = item[1][0].get_emittance(['ct'])
|
311
|
emit_long_out = item[1][-1].get_emittance(['ct'])
|
312
|
delta_emit_l = emit_long_out/emit_long_in-1.
|
313
|
emit_6d_in = item[1][0].get_emittance(['x', 'y', 'ct'])
|
314
|
emit_6d_out = item[1][-1].get_emittance(['x', 'y', 'ct'])
|
315
|
delta_emit_6d = emit_6d_out/emit_6d_in-1.
|
316
|
sigma_energy.append(s_energy)
|
317
|
actual_emittance_long.append(emit_long_in)
|
318
|
delta_emittance_long.append(delta_emit_l)
|
319
|
actual_emittance_trans.append(emit_in)
|
320
|
delta_emittance_trans.append(delta_emit_t)
|
321
|
actual_emittance_6d.append(emit_6d_in)
|
322
|
delta_emittance_6d.append(delta_emit_6d)
|
323
|
delta_transmission.append(delta_trans)
|
324
|
nominal_emittance_trans.append(item[0]['emittance'])
|
325
|
except Exception:
|
326
|
sys.excepthook(*sys.exc_info())
|
327
|
Bunch.clear_global_weights()
|
328
|
json_doc = {
|
329
|
"delta_transmission":delta_transmission,
|
330
|
"nominal_emittance_trans":nominal_emittance_trans,
|
331
|
"delta_emittance_trans":delta_emittance_trans,
|
332
|
"actual_emittance_trans":actual_emittance_trans,
|
333
|
"delta_emittance_long":delta_emittance_long,
|
334
|
"actual_emittance_long":actual_emittance_long,
|
335
|
"delta_emittance_6d":delta_emittance_6d,
|
336
|
"actual_emittance_6d":actual_emittance_6d,
|
337
|
"sigma_energy":sigma_energy,
|
338
|
}
|
339
|
print >> multi_data, json.dumps(json_doc)
|
340
|
|
341
|
|
342
|
def __init__(self, momentum, amplitude):
|
343
|
self.data = []
|
344
|
self.current = {}
|
345
|
self.p = momentum
|
346
|
self.m_mu = common.pdg_pid_to_mass[13]
|
347
|
self.bz_us = 4.e-3
|
348
|
self.bz_ds = -4.e-3
|
349
|
self.beta = self.p/self.bz_us*2./common.constants['c_light']
|
350
|
self.z_max = 1090.1
|
351
|
self.scraped = {'__any__':[]}
|
352
|
self.r_cut = [150., 150.]
|
353
|
self.amplitude_trans_cut = [72., 72.]
|
354
|
self.do_long_cut = True
|
355
|
self.amplitude_ct_cut = [50., 150.]
|
356
|
self.p_cut = [[10., 15.], [15., 30.]]
|
357
|
file_list = glob.glob("output/output_bross/output_++--_FC=40.0_p=200_21/*.root")
|
358
|
print file_list
|
359
|
for filename in sorted(file_list):
|
360
|
try:
|
361
|
self.single_analysis(filename)
|
362
|
except Exception:
|
363
|
sys.excepthook(*sys.exc_info())
|
364
|
#self.multi_analysis()
|
365
|
|
366
|
def main():
|
367
|
batch = True
|
368
|
ROOT.gROOT.SetBatch(batch)
|
369
|
Tracking(200., 1.)
|
370
|
print "Done\n\n"
|
371
|
if not batch:
|
372
|
raw_input()
|
373
|
|
374
|
if __name__ == "__main__":
|
375
|
main()
|