muDIS_mergeFiles.py

Go to the documentation of this file.

# SPDX-License-Identifier: LGPL-3.0-or-later
# SPDX-FileCopyrightText: Copyright CERN for the benefit of the SHiP Collaboration
 
import os
 
import ROOT
import rootUtils as ut
 
h = {}
 
 
def merge() -> None:
    sTree = ROOT.TChain("DIS")
    for i in range(0, 12):
        fn = "muonDis_" + str(i) + ".root"
        sTree.AddFile(fn)
    fm = ROOT.TFile("test.root", "recreate")
    nTree = ROOT.TTree("DIS", "muon DIS")
    iMuon = ROOT.TClonesArray("TVectorD")
    iMuonBranch = nTree.Branch("InMuon", iMuon, 32000, -1)
    dPart = ROOT.TClonesArray("TVectorD")
    dPartBranch = nTree.Branch("Particles", dPart, 32000, -1)
    for n in range(sTree.GetEntries()):
        sTree.GetEvent(n)
        dPart.Clear()
        iMuon.Clear()
        tca_vec = iMuon.ConstructedAt(0)
        tca_vec.ResizeTo(sTree.InMuon[0])
        ROOT.std.swap(tca_vec, sTree.InMuon[0])
        for part in sTree.Particles:
            nPart = len(dPart)
            if dPart.GetSize() == nPart:
                dPart.Expand(nPart + 10)
            tca_vec = dPart.ConstructedAt(nPart)
            tca_vec.ResizeTo(part)
            ROOT.std.swap(tca_vec, part)
        dPartBranch.Fill()
        iMuonBranch.Fill()
        nTree.Fill()
    fm.cd()
    nTree.Write()
    fm.Close()
 
 
def hadd() -> None:
    cmd = "hadd -f muonDIS.root "
    for i in range(0, 12):
        cmd += " muonDis_" + str(i) + ".root "
    os.system(cmd)
 
 
def makePlots(sTree) -> None:
    ut.bookHist(h, "muP", "muon mom", 100, 0.0, 400.0)
    ut.bookHist(h, "nOut", "outgoing part", 100, -0.5, 99.5)
    ut.bookHist(h, "pos", "position", 100, -25.0, 100.0, 100, -12.0, 12.0, 100, -13.0, 13.0)
    for n in range(sTree.GetEntries()):
        sTree.GetEvent(n)
        inMu = sTree.InMuon[0]
        w = inMu[8]
        P = ROOT.TMath.Sqrt(inMu[1] ** 2 + inMu[2] ** 2 + inMu[3] ** 2)
        h["muP"].Fill(P, w)
        h["nOut"].Fill(len(sTree.Particles))
        h["pos"].Fill(inMu[7], inMu[5], inMu[6])
 
 
def test(fn: str = "test.root") -> None:
    fm = ROOT.TFile(fn)
    sTree = fm.DIS
    makePlots(sTree)