def test_ntuple(): with TemporaryFile(): ntuple = Ntuple(('a', 'b', 'c'), name='test') for i in range(100): ntuple.Fill(gauss(.3, 2.), gauss(0, 1.), gauss(-1., 5)) ntuple.Write()
def add_seed(intuple): outtuple = Ntuple( ('id', 'px', 'py', 'pz', 'x', 'y', 'z', 'opx', 'opy', 'opz', 'ox', 'oy', 'oz', 'pythiaid', 'parentid', 'w', 'ecut', 'seed'), name='pythia8-Geant4') for i, muon in enumerate(intuple): a = array('f', [y for x in muon.values() for y in x]) a.append(float(i)) outtuple.fill(a) return outtuple
#!/usr/bin/env python """ ======================= A simple Ntuple example ======================= This example demonstrates how to create a simple Ntuple. """ print __doc__ from rootpy.tree import Ntuple from rootpy.io import root_open from random import gauss f = root_open("test.root", "recreate") # create an ntuple with three float fields: a, b, c ntuple = Ntuple(('a', 'b', 'c'), name="test") # fill the ntuple with fake data for i in xrange(20): ntuple.Fill(gauss(.5, 1.), gauss(.3, 2.), gauss(13., 42.)) ntuple.write() # write as CSV ntuple.csv() f.close()
def analyse(tree, outputfile): """Analyse tree to find hit positions and create histograms. Parameters ---------- tree Tree or Chain of trees with the usual `cbmsim` format outputfile : str Filename for the file in which the histograms are saved, will be overwritten Returns ------- std::vector<double> Vector of hit x-positions [cm] """ # TODO handle addition of new output files! r.gROOT.SetBatch(True) maxpt = 6.5 maxp = 360. h = {} ut.bookHist(h, 'mu_pos', '#mu- hits;x[cm];y[cm]', 100, -1000, +1000, 100, -800, 1000) ut.bookHist(h, 'anti-mu_pos', '#mu+ hits;x[cm];y[cm]', 100, -1000, +1000, 100, -800, 1000) ut.bookHist(h, 'mu_w_pos', '#mu- hits;x[cm];y[cm]', 100, -1000, +1000, 100, -800, 1000) ut.bookHist(h, 'anti-mu_w_pos', '#mu+ hits;x[cm];y[cm]', 100, -1000, +1000, 100, -800, 1000) ut.bookHist(h, 'mu_p', '#mu+-;p[GeV];', 100, 0, maxp) ut.bookHist(h, 'mu_p_original', '#mu+-;p[GeV];', 100, 0, maxp) ut.bookHist(h, 'mu_pt_original', '#mu+-;p_t[GeV];', 100, 0, maxpt) ut.bookHist(h, 'mu_ppt_original', '#mu+-;p[GeV];p_t[GeV];', 100, 0, maxp, 100, 0, maxpt) ut.bookHist(h, 'smear', '#mu+- initial vertex;x[cm];y[cm]', 100, -10, +10, 100, -10, 10) xs = r.std.vector('double')() i, n = 0, tree.GetEntries() print '0/{}\r'.format(n), fout = r.TFile.Open('gallery.root', 'recreate') fout2 = r.TFile.Open('misis.root', 'recreate') fout2.cd() outtuple = Ntuple( ('px', 'py', 'pz', 'x', 'y', 'z', 'opx', 'opy', 'opz', 'w'), name='MISIS') histz = Hist(100, -8000, -3000, title='#mu z position at origin') mom = r.TVector3() for event in tree: i += 1 if i % 1000 == 0: print '{}/{}\r'.format(i, n), original_muon = event.MCTrack[1] h['smear'].Fill(original_muon.GetStartX(), original_muon.GetStartY()) for hit in event.vetoPoint: draw = False if hit: if not hit.GetEnergyLoss() > 0: continue pid = hit.PdgCode() if hit.GetZ() > 2597 and hit.GetZ() < 2599 and abs(pid) == 13: hit.Momentum(mom) P = mom.Mag() / u.GeV y = hit.GetY() x = hit.GetX() if pid == 13: h['mu_pos'].Fill(x, y) else: h['anti-mu_pos'].Fill(x, y) x *= pid / 13. if (P > 1 and abs(y) < 5 * u.m and (x < 2.6 * u.m and x > -3 * u.m)): xs.push_back(x) w = np.sqrt((560. - (x + 300.)) / 560.) h['mu_p'].Fill(P) original_muon = event.MCTrack[1] h['mu_p_original'].Fill(original_muon.GetP()) h['mu_pt_original'].Fill(original_muon.GetPt()) h['mu_ppt_original'].Fill(original_muon.GetP(), original_muon.GetPt()) if pid == 13: h['mu_w_pos'].Fill(x, y, w) else: h['anti-mu_w_pos'].Fill(-x, y, w) if (P > 1) and (x < 2.6 * u.m) and (abs(y) < 5 * u.m): draw = 4 if pid > 0 else 6 if draw: graph_x, graph_y = graph_tracks(event) graph_x.SetLineColor(draw) graph_y.SetLineColor(draw) c = Canvas(1600, 900, name='c{}'.format(i)) multigraph = r.TMultiGraph( 'tracks_{}'.format(i), 'Tracks in acceptance and/or on wrong side;z [cm];x/y [cm]') graph_x.SetLineStyle(1) graph_x.SetMarkerStyle(20) graph_x.SetTitle('x-projection') graph_x.SetFillStyle(0) graph_y.SetTitle('y-projection') graph_y.SetLineStyle(2) graph_y.SetMarkerStyle(20) graph_y.SetFillStyle(0) multigraph.Add(graph_x, 'lp') multigraph.Add(graph_y, 'lp') multigraph.Draw('Alp') c.BuildLegend() fout.cd() c.Write() fout2.cd() hitlist = {} for hit in event.vetoPoint: if abs(hit.PdgCode()) != 13: continue lp = hit.LastPoint() lm = hit.LastMom() if lp.z() < 2597: hitlist[lp.z()] = [lp.x(), lp.y(), lm.Px(), lm.Py(), lm.Pz()] zs = hitlist.keys() if len(zs) > 1: zs.sort() if not hitlist: continue last_hit = hitlist[zs[-1]] if last_hit[4] > 0: a = array('f', (last_hit[2], last_hit[3], last_hit[4], last_hit[0], last_hit[1], zs[-1], original_muon.GetPx(), original_muon.GetPy(), original_muon.GetPz(), original_muon.GetWeight())) outtuple.fill(a) histz.fill(original_muon.GetStartZ(), original_muon.GetWeight()) print 'Loop done' ut.writeHists(h, outputfile) fout2.cd() outtuple.Write() histz.Write() return xs