## Feature #1543 » simulation_analysis.py

 1 ```import os ``` ```import subprocess ``` ```import json ``` ```import operator ``` ```import glob ``` ```import sys ``` ```import numpy ``` ```import ROOT ``` ```import Configuration ``` ```import maus_cpp.globals ``` ```import xboa.Common as common ``` ```from xboa.Bunch import Bunch ``` ```from xboa.Hit import Hit ``` ```from xboa.tracking import MAUSTracking ``` ```class Tracking: ``` ``` def amp(self, ellipse_inv, hit): ``` ``` arr = numpy.matrix([hit['x'], hit['px'], hit['y'], hit['py']]) ``` ``` amplitude = arr*ellipse_inv*arr.T ``` ``` return amplitude[0, 0] ``` ``` def amp_list(self, bunch, beta, alpha, bz): ``` ``` ellipse = Bunch.build_penn_ellipse(1., self.m_mu, beta, alpha, self.p, 0., bz, 1.) ``` ``` ellipse_inv = numpy.linalg.inv(ellipse) ``` ``` amplitude_list = [self.amp(ellipse_inv, hit) for hit in bunch] ``` ``` return amplitude_list ``` ``` def amp_list_6d(self, bunch): ``` ``` bunch.set_covariance_matrix(True) ``` ``` amp_list = [bunch.get_amplitude(bunch, hit, ['x', 'y', 'ct']) for hit in bunch] ``` ``` bunch.set_covariance_matrix(False) ``` ``` return amp_list ``` ``` def scatter_plot_mean_rms(self, x_values, y_values, xmin, xmax, xstep): ``` ``` nbins = 9 ``` ``` bin_low = 0. ``` ``` bin_step = 10. ``` ``` hist_data = [] ``` ``` for x, y in zip(x_values, y_values): ``` ``` bin = int((x-bin_low)/bin_step) ``` ``` while len(hist_data) - 1 < bin: ``` ``` hist_data.append([]) ``` ``` hist_data[bin].append(y) ``` ``` hist = ROOT.TH1D("Mean", "", nbins, bin_low, bin_low+(nbins+1)*bin_step) ``` ``` for i, data in enumerate(hist_data): ``` ``` mean, sigmaSqu = 0., 0. ``` ``` if len(data) > 0: ``` ``` mean = sum(data)/len(data) ``` ``` if len(data) > 1: ``` ``` sigmaSqu = sum([x**2 for x in data])/len(data)-mean**2 ``` ``` hist.SetBinContent(i+1, mean) ``` ``` hist.SetBinError(i+1, sigmaSqu**0.5) ``` ``` common._hist_persistent.append(hist) ``` ``` return hist ``` ``` def amplitude_scatter_plot(self, bunch_in, bunch_out, predicate): ``` ``` """predicate takes arguments (spill_index, spills_in_bunch_out)""" ``` ``` spills_out = set([(hit['spill'], hit['event_number']) for hit in bunch_out]) ``` ``` hits_in = [hit for hit in bunch_in if (hit['spill'], hit['event_number']) in spills_out] ``` ``` hits_out = [hit for hit in bunch_out if (hit['spill'], hit['event_number']) in spills_out] ``` ``` amp_in = self.amp_list(hits_in, self.beta, 0., self.bz_us) ``` ``` amp_out = self.amp_list(hits_out, self.beta, 0., self.bz_ds) ``` ``` delta_amp = [(amp_out[i]-amp_in[i])/amp_in[i] for i, a_in in enumerate(amp_in)] ``` ``` if len(amp_in) == 0 or len(amp_in) != len(amp_out): ``` ``` print "Failed to make graph:", len(amp_in), len(amp_out) ``` ``` hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'x-y amplitude out [mm]') ``` ``` else: ``` ``` 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.) ``` ``` hist_means = self.scatter_plot_mean_rms(amp_in, delta_amp, 0., 200., 10.) ``` ``` graph.SetMarkerStyle(7) ``` ``` return hist, hist_means, graph ``` ``` def amplitude_p_scatter_plot(self, bunch_in): ``` ``` """predicate takes arguments (spill_index, spills_in_bunch_out)""" ``` ``` hits_in = [hit for hit in bunch_in] ``` ``` amp_in = self.amp_list(hits_in, self.beta, 0., self.bz_us) ``` ``` p_in = [hit['p'] for hit in hits_in] ``` ``` if len(amp_in) == 0: ``` ``` print "Failed to make graph:", len(amp_in) ``` ``` hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'P [MeV/c]', xmin=0., xmax=200.) ``` ``` else: ``` ``` 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.) ``` ``` hist_means = self.scatter_plot_mean_rms(amp_in, p_in, 0., 200., 10.) ``` ``` graph.SetMarkerStyle(7) ``` ``` return hist, hist_means, graph ``` ``` def amplitude_t_scatter_plot(self, bunch_in, bunch_out): ``` ``` """predicate takes arguments (spill_index, spills_in_bunch_out)""" ``` ``` spills_out = set([(hit['spill'], hit['event_number']) for hit in bunch_out]) ``` ``` hits_in = [hit for hit in bunch_in if (hit['spill'], hit['event_number']) in spills_out] ``` ``` hits_out = [hit for hit in bunch_out if (hit['spill'], hit['event_number']) in spills_out] ``` ``` amp_in = self.amp_list_6d(hits_in, self.beta, 0., self.bz_us) ``` ``` amp_out = self.amp_list(hits_out, self.beta, 0., self.bz_ds) ``` ``` delta_t = [h_o['t']-h_i['t'] for h_i, h_o in zip(hits_in, hits_out)] ``` ``` if len(amp_in) == 0: ``` ``` print "Failed to make graph:", len(amp_in) ``` ``` hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'dt [ns]', xmin=0., xmax=200.) ``` ``` else: ``` ``` hist, graph = common.make_root_graph('amplitude x y', amp_in, 'x-y amplitude in [mm]', delta_t, 'dt [ns]', xmin=0., xmax=200.) ``` ``` hist_means = self.scatter_plot_mean_rms(amp_in, delta_t, 0., 200., 10.) ``` ``` graph.SetMarkerStyle(7) ``` ``` return hist, hist_means, graph ``` ``` def amplitude_6d_scatter_plot(self, bunch_in, bunch_out): ``` ``` """predicate takes arguments (spill_index, spills_in_bunch_out)""" ``` ``` amp_in = self.amp_list_6d(bunch_in) ``` ``` amp_out = self.amp_list_6d(bunch_out) ``` ``` hist_in = common.make_root_histogram('A_{6D} in', amp_in, 'A_{6D} [mm]', 20, xmin=0., xmax=200.) ``` ``` hist_out = common.make_root_histogram('A_{6D} out', amp_out, 'A_{6D} [mm]', 20, xmin=0., xmax=200.) ``` ``` hist_in.SetName('A_{6D} in') ``` ``` hist_out.SetName('A_{6D} out') ``` ``` return hist_in, hist_out ``` ``` def bin_data(data_list, bin_edge_list): ``` ``` bin = 0 ``` ``` #for amp_in in ``` ``` def amplitude_6d_capture_plot(self, bunch_in, bunch_out): ``` ``` amp_in = sorted(self.amp_list_6d(bunch_in)) ``` ``` amp_out = sorted(self.amp_list_6d(bunch_out)) ``` ``` bins = [i*10. for i in range(11)]+[i*50. for i in range(2, 5)] ``` ``` contents_in = [] ``` ``` contents_out = [] ``` ``` ``` ``` def amplitude_scatter(self, bunch_in, bunch_out): ``` ``` canvas = common.make_root_canvas('amplitude 6d hist') ``` ``` hist_in, hist_out = self.amplitude_6d_scatter_plot(bunch_in, bunch_out) ``` ``` hist_out.SetLineColor(4) ``` ``` hist_out.Draw() ``` ``` hist_in.SetLineColor(2) ``` ``` hist_in.Draw("SAME") ``` ``` common.make_root_legend(canvas, [hist_in, hist_out]) ``` ``` canvas.Update() ``` ``` self.print_canvas(canvas, 'acceptance_6d_hist') ``` ``` """ ``` ``` spills_out = set([hit['spill'] for hit in bunch_out]) ``` ``` canvas = common.make_root_canvas('amplitude x y scatter') ``` ``` hist, hist_means, graph = self.amplitude_scatter_plot(bunch_in, bunch_out, lambda x: True) ``` ``` hist.Draw() ``` ``` hist_means.Draw("same") ``` ``` graph.Draw('p') ``` ``` canvas.Update() ``` ``` self.print_canvas(canvas, 'acceptance') ``` ``` canvas = common.make_root_canvas('amplitude x y p scatter') ``` ``` hist, hist_means, graph = self.amplitude_p_scatter_plot(bunch_in) ``` ``` hist.Draw() ``` ``` hist_means.Draw("same") ``` ``` graph.Draw('p') ``` ``` canvas.Update() ``` ``` self.print_canvas(canvas, 'acceptance_p') ``` ``` canvas = common.make_root_canvas('amplitude x y t scatter') ``` ``` hist, hist_means, graph = self.amplitude_t_scatter_plot(bunch_in, bunch_out) ``` ``` hist.Draw() ``` ``` hist_means.Draw("same") ``` ``` graph.Draw('p') ``` ``` canvas.Update() ``` ``` self.print_canvas(canvas, 'acceptance_t') ``` ``` """ ``` ``` pass ``` ``` def two_d_plots(self, bunch_list, colour, canvas_list = None): ``` ``` n_graphs = 4 ``` ``` if canvas_list == None: ``` ``` canvas_list = [None]*n_graphs ``` ``` graph_list = [None]*n_graphs ``` ``` canvas_list[0], hist, graph_list[0] = bunch_list[0].root_scatter_graph("t", "energy", "ns", "MeV", include_weightless=False, canvas=canvas_list[0]) ``` ``` canvas_list[1], hist, graph_list[1] = bunch_list[-1].root_scatter_graph("t", "energy", "ns", "MeV", include_weightless=False, canvas=canvas_list[1]) ``` ``` 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]) ``` ``` 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]) ``` ``` for i in range(n_graphs): ``` ``` canvas_list[i].cd() ``` ``` graph_list[i].SetMarkerStyle(6) ``` ``` graph_list[i].SetMarkerColor(colour) ``` ``` graph_list[i].Draw("p") ``` ``` canvas_list[i].Update() ``` ``` self.print_canvas(canvas_list[0], "time_energy_at_start") ``` ``` self.print_canvas(canvas_list[1], "time_energy_at_end") ``` ``` self.print_canvas(canvas_list[2], "r_pt_at_start") ``` ``` self.print_canvas(canvas_list[3], "r_pt_at_end") ``` ``` return canvas_list ``` ``` def cut(self, bunch_list): ``` ``` print " weights in " ``` ``` print " no cut:", bunch_list[0].bunch_weight(), bunch_list[-1].bunch_weight() ``` ``` for bunch in bunch_list: ``` ``` bunch.transmission_cut(bunch_list[-1], global_cut=True) ``` ``` bunch.transmission_cut(bunch_list[-1], global_cut=True, test_variable = ['spill', 'event_number', 'particle_number', 'pid']) ``` ``` bunch.cut({'pid':-13}, operator.ne, global_cut=True) ``` ``` print ' transmission cut:', bunch_list[0].bunch_weight(), bunch_list[-1].bunch_weight(), "**" ``` ``` bunches = [bunch_list[0], bunch_list[-1]] ``` ``` cut_stream = ['upstream', 'downstream'] ``` ``` for i in [0, 1]: ``` ``` print ' ', cut_stream[i] ``` ``` print ' ', 'cut ', ' u/s', ' d/s' ``` ``` bunches[i].cut({'r':self.r_cut[i]}, operator.gt, global_cut=True) ``` ``` print ' cut r: ', self.r_cut[i], str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10) ``` ``` bunches[i].cut({'amplitude x y':self.amplitude_trans_cut[i]}, operator.gt, global_cut=True) ``` ``` 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) ``` ``` bunches[i].cut({'amplitude x y':self.amplitude_trans_cut[i]}, operator.gt, global_cut=True) ``` ``` 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) ``` ``` if self.do_long_cut: ``` ``` bunches[i].cut({'amplitude ct':self.amplitude_ct_cut[i]}, operator.gt, global_cut=True) ``` ``` 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) ``` ``` bunches[i].cut({'amplitude ct':self.amplitude_ct_cut[i]}, operator.gt, global_cut=True) ``` ``` 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) ``` ``` mean_p = bunches[i].mean(['p'])['p'] ``` ``` bunches[i].cut({'p':mean_p+self.p_cut[i][1]}, operator.gt, global_cut=True) ``` ``` bunches[i].cut({'p':mean_p-self.p_cut[i][0]}, operator.lt, global_cut=True) ``` ``` 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) ``` ``` def z_plot(self, bunch_list): ``` ``` print "Start MOMENTS" ``` ``` print bunch_list[0].moment(['t', 't']) ``` ``` print bunch_list[0].moment(['pz', 't']) ``` ``` print bunch_list[0].moment(['pz', 'pz']) ``` ``` """ ``` ``` for item in ['t', 'pz', 'energy']: ``` ``` canvas = common.make_root_canvas('sigma_'+item) ``` ``` z_list = [bunch.mean(['z'])['z'] for bunch in bunch_list] ``` ``` item_list = [bunch.moment([item, item])**0.5 for bunch in bunch_list] ``` ``` hist, graph = common.make_root_graph('sigma_'+item, z_list, 'mean z [mm]', item_list, '#sigma('+item+')') ``` ``` hist.Draw() ``` ``` graph.Draw('l') ``` ``` self.print_canvas(canvas, 's_'+item+'_vs_z') ``` ``` """ ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['p'], 'mm', 'MeV/c') ``` ``` self.print_canvas(canvas, 'momentum_vs_z') ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['x', 'y'], 'mm', 'mm') ``` ``` self.print_canvas(canvas, 'beta_trans_vs_z') ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['x', 'y'], 'mm', 'mm') ``` ``` self.print_canvas(canvas, 'emittance_trans_vs_z') ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['t'], 'mm', 'mm') ``` ``` self.print_canvas(canvas, 'beta_long_vs_z') ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['ct'], 'mm', 'mm') ``` ``` self.print_canvas(canvas, 'emittance_long_vs_z') ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['ct', 'x', 'y'], 'mm', 'mm') ``` ``` self.print_canvas(canvas, 'emittance_6d_vs_z') ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['bz']) ``` ``` self.print_canvas(canvas, 'bz_vs_z') ``` ``` def grok_filename(self, filename): ``` ``` base = os.path.basename(filename)[12:-5] ``` ``` emittance = base.split('_')[0] ``` ``` emittance = float(emittance.split('=')[1]) ``` ``` a_dir = os.path.dirname(filename).split('/') ``` ``` index = a_dir[-1].split('_')[-1] ``` ``` base += '_'+index ``` ``` a_dir = "/".join(a_dir[:-1]) ``` ``` print " directory", a_dir, "name", base ``` ``` self.current = { ``` ``` "dir":a_dir, ``` ``` "name":base, ``` ``` "full":filename, ``` ``` "emittance":emittance, ``` ``` "index":index ``` ``` } ``` ``` return self.current ``` ``` def print_canvas(self, canvas, name): ``` ``` for format in ["png", "root"]: ``` ``` canvas.Print(self.current["dir"]+"/"+name+"_"+self.current["name"]+"."+format) ``` ``` def single_analysis(self, filename): ``` ``` print "Loading data", filename ``` ``` metadata = self.grok_filename(filename) ``` ``` bunch_list = Bunch.new_list_from_read_builtin('maus_root_virtual_hit', filename) ``` ``` bunch_list = [bunch for bunch in bunch_list if bunch.bunch_weight() > 100] ``` ``` if self.z_max != None: ``` ``` bunch_list = [bunch for bunch in bunch_list if bunch[0]['z'] < self.z_max] ``` ``` print " Printing transmission (before cut)" ``` ``` canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'bunch_weight', [], 'mm', '') ``` ``` self.print_canvas(canvas, 'weight_vs_z') ``` ``` print " Making cuts" ``` ``` canvas_list = self.two_d_plots(bunch_list, 2) ``` ``` self.cut(bunch_list) ``` ``` self.amplitude_scatter(bunch_list[0], bunch_list[-1]) ``` ``` print " Making 2d plots" ``` ``` self.two_d_plots(bunch_list, 1, canvas_list) ``` ``` print " Making z plots" ``` ``` self.z_plot(bunch_list) ``` ``` print " Making acceptance plots" ``` ``` self.data.append((metadata, bunch_list)) ``` ``` Bunch.clear_global_weights() ``` ``` def multi_analysis(self): ``` ``` multi_data = open(self.data[0][0]["dir"]+"/multi_analysis_"+self.current["index"]+".json", "w") ``` ``` delta_transmission = [] ``` ``` delta_emittance_trans = [] ``` ``` nominal_emittance_trans = [] ``` ``` actual_emittance_trans = [] ``` ``` actual_emittance_long = [] ``` ``` delta_emittance_long = [] ``` ``` actual_emittance_6d = [] ``` ``` delta_emittance_6d = [] ``` ``` sigma_energy = [] ``` ``` for item in self.data: ``` ``` self.cut(item[1]) ``` ``` try: ``` ``` delta_trans = float(len(item[1][-1]))/float(len(item[1][0])) ``` ``` emit_in = item[1][0].get_emittance(['x', 'y']) ``` ``` emit_out = item[1][-1].get_emittance(['x', 'y']) ``` ``` delta_emit_t = emit_out/emit_in-1. ``` ``` s_energy = (item[1][0].moment(['energy', 'energy']))**0.5 ``` ``` emit_long_in = item[1][0].get_emittance(['ct']) ``` ``` emit_long_out = item[1][-1].get_emittance(['ct']) ``` ``` delta_emit_l = emit_long_out/emit_long_in-1. ``` ``` emit_6d_in = item[1][0].get_emittance(['x', 'y', 'ct']) ``` ``` emit_6d_out = item[1][-1].get_emittance(['x', 'y', 'ct']) ``` ``` delta_emit_6d = emit_6d_out/emit_6d_in-1. ``` ``` sigma_energy.append(s_energy) ``` ``` actual_emittance_long.append(emit_long_in) ``` ``` delta_emittance_long.append(delta_emit_l) ``` ``` actual_emittance_trans.append(emit_in) ``` ``` delta_emittance_trans.append(delta_emit_t) ``` ``` actual_emittance_6d.append(emit_6d_in) ``` ``` delta_emittance_6d.append(delta_emit_6d) ``` ``` delta_transmission.append(delta_trans) ``` ``` nominal_emittance_trans.append(item[0]['emittance']) ``` ``` except Exception: ``` ``` sys.excepthook(*sys.exc_info()) ``` ``` Bunch.clear_global_weights() ``` ``` json_doc = { ``` ``` "delta_transmission":delta_transmission, ``` ``` "nominal_emittance_trans":nominal_emittance_trans, ``` ``` "delta_emittance_trans":delta_emittance_trans, ``` ``` "actual_emittance_trans":actual_emittance_trans, ``` ``` "delta_emittance_long":delta_emittance_long, ``` ``` "actual_emittance_long":actual_emittance_long, ``` ``` "delta_emittance_6d":delta_emittance_6d, ``` ``` "actual_emittance_6d":actual_emittance_6d, ``` ``` "sigma_energy":sigma_energy, ``` ``` } ``` ``` print >> multi_data, json.dumps(json_doc) ``` ``` ``` ``` def __init__(self, momentum, amplitude): ``` ``` self.data = [] ``` ``` self.current = {} ``` ``` self.p = momentum ``` ``` self.m_mu = common.pdg_pid_to_mass[13] ``` ``` self.bz_us = 4.e-3 ``` ``` self.bz_ds = -4.e-3 ``` ``` self.beta = self.p/self.bz_us*2./common.constants['c_light'] ``` ``` self.z_max = 1090.1 ``` ``` self.scraped = {'__any__':[]} ``` ``` self.r_cut = [150., 150.] ``` ``` self.amplitude_trans_cut = [72., 72.] ``` ``` self.do_long_cut = True ``` ``` self.amplitude_ct_cut = [50., 150.] ``` ``` self.p_cut = [[10., 15.], [15., 30.]] ``` ``` file_list = glob.glob("output/output_bross/output_++--_FC=40.0_p=200_21/*.root") ``` ``` print file_list ``` ``` for filename in sorted(file_list): ``` ``` try: ``` ``` self.single_analysis(filename) ``` ``` except Exception: ``` ``` sys.excepthook(*sys.exc_info()) ``` ``` #self.multi_analysis() ``` ```def main(): ``` ``` batch = True ``` ``` ROOT.gROOT.SetBatch(batch) ``` ``` Tracking(200., 1.) ``` ``` print "Done\n\n" ``` ``` if not batch: ``` ``` raw_input() ``` ```if __name__ == "__main__": ``` ``` main() ```
(31-31/57)