test/visualization/graph.py

0001 import ROOT
0002 import sys
0003 import numpy as np
0004
0005 def generate_graph(filename, dataset, title, param_no, db, plane):
0006     file = ROOT.TFile.Open("visualization/"+filename+".root", "RECREATE")
0007
0008     arr = np.loadtxt("visualization/"+dataset+".dat")
0009     rows, columns = arr.shape
0010     x_line = []
0011     y_line = []
0012
0013     for r in range(rows):
0014         if int(arr[r,0]) == int(param_no) and int(arr[r,1]) == int(db) and int(arr[r,2]) == int(plane):
0015             x_line.append(arr[r,3])
0016             y_line.append(arr[r,5])
0017
0018     size = len(x_line)
0019
0020     x = ROOT.std.vector('double')()
0021     for i in x_line: x.push_back(float(i))
0022     y = ROOT.std.vector('double')()
0023     for i in y_line: y.push_back(float(i))
0024
0025     c1 = ROOT.TCanvas()
0026     gr = ROOT.TGraph(size, x.data(), y.data())
0027     gr.SetTitle(title)
0028     gr.SetMarkerColor(4)
0029     gr.SetMarkerStyle(21)
0030     #gr.Draw("AL")
0031     file.WriteObject(gr, title)
0032
0033 if __name__ == '__main__':
0034     generate_graph(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], sys.argv[6])