Project

General

Profile

Feature #1543 » simulation_analysis.py

Rogers, Chris, 29 October 2014 14:34

 
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()
(31-31/57)