Exemplo n.º 1
0
    def __init__(self, fout, fgeo):
        self.fn = ROOT.TFile.Open(fout, 'update')
        self.sTree = self.fn.cbmsim
        if self.sTree.GetBranch("FitTracks"):
            print "remove RECO branches and rerun reconstruction"
            self.fn.Close()
            # make a new file without reco branches
            f = ROOT.TFile(fout)
            sTree = f.cbmsim
            if sTree.GetBranch("FitTracks"):
                sTree.SetBranchStatus("FitTracks", 0)
            if sTree.GetBranch("goodTracks"):
                sTree.SetBranchStatus("goodTracks", 0)
            if sTree.GetBranch("VetoHitOnTrack"):
                sTree.SetBranchStatus("VetoHitOnTrack", 0)
            if sTree.GetBranch("Particles"):
                sTree.SetBranchStatus("Particles", 0)
            if sTree.GetBranch("fitTrack2MC"):
                sTree.SetBranchStatus("fitTrack2MC", 0)
            if sTree.GetBranch("EcalClusters"):
                sTree.SetBranchStatus("EcalClusters", 0)
            if sTree.GetBranch("EcalReconstructed"):
                sTree.SetBranchStatus("EcalReconstructed", 0)
            if sTree.GetBranch("Pid"): sTree.SetBranchStatus("Pid", 0)
            if sTree.GetBranch("Digi_StrawtubesHits"):
                sTree.SetBranchStatus("Digi_StrawtubesHits", 0)
            if sTree.GetBranch("Digi_SBTHits"):
                sTree.SetBranchStatus("Digi_SBTHits", 0)
            if sTree.GetBranch("digiSBT2MC"):
                sTree.SetBranchStatus("digiSBT2MC", 0)
            if sTree.GetBranch("Digi_StrawtubesHits"):
                sTree.SetBranchStatus("Digi_StrawtubesHits", 0)
            rawFile = fout.replace("_rec.root", "_raw.root")
            recf = ROOT.TFile(rawFile, "recreate")
            newTree = sTree.CloneTree(0)
            for n in range(sTree.GetEntries()):
                sTree.GetEntry(n)
                rc = newTree.Fill()
            sTree.Clear()
            newTree.AutoSave()
            f.Close()
            recf.Close()
            os.system('cp ' + rawFile + ' ' + fout)
            self.fn = ROOT.TFile(fout, 'update')
            self.sTree = self.fn.cbmsim
#  check that all containers are present, otherwise create dummy version
        self.dummyContainers = {}
        branch_class = {"vetoPoint":"vetoPoint","ShipRpcPoint":"ShipRpcPoint","TargetPoint":"TargetPoint",\
                        "strawtubesPoint":"strawtubesPoint","EcalPointLite":"ecalPoint","HcalPointLite":"hcalPoint"}
        for x in branch_class:
            if not self.sTree.GetBranch(x):
                self.dummyContainers[x + "_array"] = ROOT.TClonesArray(
                    branch_class[x])
                self.dummyContainers[x] = self.sTree.Branch(
                    x, self.dummyContainers[x + "_array"], 32000, -1)
                setattr(self.sTree, x, self.dummyContainers[x + "_array"])
                self.dummyContainers[x].Fill()
#
        if self.sTree.GetBranch("GeoTracks"):
            self.sTree.SetBranchStatus("GeoTracks", 0)
        # prepare for output
        # event header
        self.header = ROOT.FairEventHeader()
        self.eventHeader = self.sTree.Branch("ShipEventHeader", self.header,
                                             32000, -1)
        # fitted tracks
        self.fGenFitArray = ROOT.TClonesArray("genfit::Track")
        self.fGenFitArray.BypassStreamer(ROOT.kFALSE)
        self.fitTrack2MC = ROOT.std.vector('int')()
        self.goodTracksVect = ROOT.std.vector('int')()
        self.mcLink = self.sTree.Branch("fitTrack2MC", self.fitTrack2MC, 32000,
                                        -1)
        self.fitTracks = self.sTree.Branch("FitTracks", self.fGenFitArray,
                                           32000, -1)
        self.goodTracksBranch = self.sTree.Branch("goodTracks",
                                                  self.goodTracksVect, 32000,
                                                  -1)
        #
        self.digiStraw = ROOT.TClonesArray("strawtubesHit")
        self.digiStrawBranch = self.sTree.Branch("Digi_StrawtubesHits",
                                                 self.digiStraw, 32000, -1)
        self.digiSBT = ROOT.TClonesArray("vetoHit")
        self.digiSBTBranch = self.sTree.Branch("Digi_SBTHits", self.digiSBT,
                                               32000, -1)
        self.vetoHitOnTrackArray = ROOT.TClonesArray("vetoHitOnTrack")
        self.vetoHitOnTrackBranch = self.sTree.Branch("VetoHitOnTrack",
                                                      self.vetoHitOnTrackArray,
                                                      32000, -1)
        self.digiSBT2MC = ROOT.std.vector('std::vector< int >')()
        self.mcLinkSBT = self.sTree.Branch("digiSBT2MC", self.digiSBT2MC,
                                           32000, -1)
        # for the digitizing step
        self.v_drift = modules["Strawtubes"].StrawVdrift()
        self.sigma_spatial = modules["Strawtubes"].StrawSigmaSpatial()

        # setup ecal reconstruction
        self.caloTasks = []
        if self.sTree.GetBranch("EcalPoint"):
            # Creates. exports and fills calorimeter structure
            dflag = 0
            if debug: dflag = 10
            ecalGeo = ecalGeoFile + 'z' + str(ShipGeo.ecal.z) + ".geo"
            if not ecalGeo in os.listdir(os.environ["FAIRSHIP"] + "/geometry"):
                shipDet_conf.makeEcalGeoFile(ShipGeo.ecal.z, ShipGeo.ecal.File)
            ecalFiller = ROOT.ecalStructureFiller("ecalFiller", dflag, ecalGeo)
            ecalFiller.SetUseMCPoints(ROOT.kTRUE)
            ecalFiller.StoreTrackInformation()
            self.caloTasks.append(ecalFiller)
            #GeV -> ADC conversion
            ecalDigi = ROOT.ecalDigi("ecalDigi", 0)
            self.caloTasks.append(ecalDigi)
            #ADC -> GeV conversion
            ecalPrepare = ROOT.ecalPrepare("ecalPrepare", 0)
            self.caloTasks.append(ecalPrepare)
            # Maximums locator
            ecalMaximumFind = ROOT.ecalMaximumLocator("maximumFinder", dflag)
            self.caloTasks.append(ecalMaximumFind)
            # Cluster calibration
            ecalClusterCalib = ROOT.ecalClusterCalibration(
                "ecalClusterCalibration", 0)
            #4x4 cm cells
            ecalCl3PhS = ROOT.TFormula("ecalCl3PhS",
                                       "[0]+x*([1]+x*([2]+x*[3]))")
            ecalCl3PhS.SetParameters(6.77797e-04, 5.75385e+00, 3.42690e-03,
                                     -1.16383e-04)
            ecalClusterCalib.SetStraightCalibration(3, ecalCl3PhS)
            ecalCl3Ph = ROOT.TFormula(
                "ecalCl3Ph", "[0]+x*([1]+x*([2]+x*[3]))+[4]*x*y+[5]*x*y*y")
            ecalCl3Ph.SetParameters(0.000750975, 5.7552, 0.00282783,
                                    -8.0025e-05, -0.000823651, 0.000111561)
            ecalClusterCalib.SetCalibration(3, ecalCl3Ph)
            #6x6 cm cells
            ecalCl2PhS = ROOT.TFormula("ecalCl2PhS",
                                       "[0]+x*([1]+x*([2]+x*[3]))")
            ecalCl2PhS.SetParameters(8.14724e-04, 5.67428e+00, 3.39030e-03,
                                     -1.28388e-04)
            ecalClusterCalib.SetStraightCalibration(2, ecalCl2PhS)
            ecalCl2Ph = ROOT.TFormula(
                "ecalCl2Ph", "[0]+x*([1]+x*([2]+x*[3]))+[4]*x*y+[5]*x*y*y")
            ecalCl2Ph.SetParameters(0.000948095, 5.67471, 0.00339177,
                                    -0.000122629, -0.000169109, 8.33448e-06)
            ecalClusterCalib.SetCalibration(2, ecalCl2Ph)
            self.caloTasks.append(ecalClusterCalib)
            # Cluster finder
            ecalClusterFind = ROOT.ecalClusterFinder("clusterFinder", dflag)
            self.caloTasks.append(ecalClusterFind)
            # Calorimeter reconstruction
            ecalReco = ROOT.ecalReco('ecalReco', 0)
            self.caloTasks.append(ecalReco)
            # Match reco to MC
            ecalMatch = ROOT.ecalMatch('ecalMatch', 0)
            self.caloTasks.append(ecalMatch)
            if EcalDebugDraw:
                # ecal drawer: Draws calorimeter structure, incoming particles, clusters, maximums
                ecalDrawer = ROOT.ecalDrawer("clusterFinder", 10)
                self.caloTasks.append(ecalDrawer)

    # add pid reco
            import shipPid
            self.caloTasks.append(shipPid.Task(self))
# prepare vertexing
        self.Vertexing = shipVertex.Task(h, self.sTree)
        # setup random number generator
        self.random = ROOT.TRandom()
        ROOT.gRandom.SetSeed(13)
        self.PDG = ROOT.TDatabasePDG.Instance()
        # access ShipTree
        self.sTree.GetEvent(0)
        if len(self.caloTasks) > 0:
            print "** initialize Calo reconstruction **"
            self.ecalStructure = ecalFiller.InitPython(
                self.sTree.EcalPointLite)
            ecalDigi.InitPython(self.ecalStructure)
            ecalPrepare.InitPython(self.ecalStructure)
            self.ecalMaximums = ecalMaximumFind.InitPython(self.ecalStructure)
            self.ecalCalib = ecalClusterCalib.InitPython()
            self.ecalClusters = ecalClusterFind.InitPython(
                self.ecalStructure, self.ecalMaximums, self.ecalCalib)
            self.EcalClusters = self.sTree.Branch("EcalClusters",
                                                  self.ecalClusters, 32000, -1)
            self.ecalReconstructed = ecalReco.InitPython(
                self.sTree.EcalClusters, self.ecalStructure, self.ecalCalib)
            self.EcalReconstructed = self.sTree.Branch("EcalReconstructed",
                                                       self.ecalReconstructed,
                                                       32000, -1)
            ecalMatch.InitPython(self.ecalStructure, self.ecalReconstructed,
                                 self.sTree.MCTrack)
            if EcalDebugDraw:
                ecalDrawer.InitPython(self.sTree.MCTrack, self.sTree.EcalPoint,
                                      self.ecalStructure, self.ecalClusters)
        else:
            ecalClusters = ROOT.TClonesArray("ecalCluster")
            ecalReconstructed = ROOT.TClonesArray("ecalReconstructed")
            self.EcalClusters = self.sTree.Branch("EcalClusters", ecalClusters,
                                                  32000, -1)
            self.EcalReconstructed = self.sTree.Branch("EcalReconstructed",
                                                       ecalReconstructed,
                                                       32000, -1)
#
        self.geoMat = ROOT.genfit.TGeoMaterialInterface()
        # init geometry and mag. field
        gMan = ROOT.gGeoManager
        #
        self.bfield = ROOT.genfit.BellField(ShipGeo.Bfield.max,
                                            ShipGeo.Bfield.z, 2,
                                            ShipGeo.Yheight / 2. * u.m)
        self.fM = ROOT.genfit.FieldManager.getInstance()
        self.fM.init(self.bfield)
        ROOT.genfit.MaterialEffects.getInstance().init(self.geoMat)

        # init fitter, to be done before importing shipPatRec
        #fitter          = ROOT.genfit.KalmanFitter()
        #fitter          = ROOT.genfit.KalmanFitterRefTrack()
        self.fitter = ROOT.genfit.DAF()
        if debug: self.fitter.setDebugLvl(1)  # produces lot of printout
        #set to True if "real" pattern recognition is required also
        if debug == True: shipPatRec.debug = 1

        # for 'real' PatRec
        shipPatRec.initialize(fgeo)
Exemplo n.º 2
0
    def __init__(self, fout):
        self.fn = ROOT.TFile(fout, 'update')
        self.sTree = self.fn.cbmsim
        if self.sTree.GetBranch("SmearedHits"):
            print "remove RECO branches and rerun reconstruction"
            self.fn.Close()
            # make a new file without reco branches
            f = ROOT.TFile(fout)
            sTree = f.cbmsim
            sTree.SetBranchStatus("SmearedHits", 0)
            if sTree.GetBranch("FitTracks"):
                sTree.SetBranchStatus("FitTracks", 0)
            if sTree.GetBranch("Particles"):
                sTree.SetBranchStatus("Particles", 0)
            if sTree.GetBranch("fitTrack2MC"):
                sTree.SetBranchStatus("fitTrack2MC", 0)
            if sTree.GetBranch("FitTracks_PR"):
                sTree.SetBranchStatus("FitTracks_PR", 0)
            if sTree.GetBranch("Particles_PR"):
                sTree.SetBranchStatus("Particles_PR", 0)
            if sTree.GetBranch("fitTrack2MC_PR"):
                sTree.SetBranchStatus("fitTrack2MC_PR", 0)
            if sTree.GetBranch("EcalClusters"):
                sTree.SetBranchStatus("EcalClusters", 0)
            rawFile = fout.replace("_rec.root", "_raw.root")
            recf = ROOT.TFile(rawFile, "recreate")
            newTree = sTree.CloneTree(0)
            for n in range(sTree.GetEntries()):
                sTree.GetEntry(n)
                rc = newTree.Fill()
            sTree.Clear()
            newTree.AutoSave()
            f.Close()
            recf.Close()
            os.system('cp ' + rawFile + ' ' + fout)
            self.fn = ROOT.TFile(fout, 'update')
            self.sTree = self.fn.cbmsim
#  check that all containers are present, otherwise create dummy version
        self.dummyContainers = {}
        branch_class = {"vetoPoint":"vetoPoint","ShipRpcPoint":"ShipRpcPoint","TargetPoint":"TargetPoint",\
                        "strawtubesPoint":"strawtubesPoint","EcalPointLite":"ecalPoint","HcalPointLite":"hcalPoint"}
        for x in branch_class:
            if not self.sTree.GetBranch(x):
                self.dummyContainers[x + "_array"] = ROOT.TClonesArray(
                    branch_class[x])
                self.dummyContainers[x] = self.sTree.Branch(
                    x, self.dummyContainers[x + "_array"], 32000, -1)
                setattr(self.sTree, x, self.dummyContainers[x + "_array"])
                self.dummyContainers[x].Fill()
#
        if self.sTree.GetBranch("GeoTracks"):
            self.sTree.SetBranchStatus("GeoTracks", 0)
        self.nEvents = min(self.sTree.GetEntries(), nEvents)
        # prepare for output
        self.fGenFitArray = ROOT.TClonesArray("genfit::Track")
        self.fGenFitArray.BypassStreamer(ROOT.kFALSE)
        self.fitTrack2MC = ROOT.std.vector('int')()
        self.mcLink = self.sTree.Branch("fitTrack2MC" + realPR,
                                        self.fitTrack2MC, 32000, -1)
        self.fitTracks = self.sTree.Branch("FitTracks" + realPR,
                                           self.fGenFitArray, 32000, -1)
        #
        self.SmearedHits = ROOT.TClonesArray("TVectorD")
        self.SHbranch = self.sTree.Branch("SmearedHits", self.SmearedHits,
                                          32000, -1)
        #
        self.random = ROOT.TRandom()
        ROOT.gRandom.SetSeed(13)
        #
        self.Vertexing = shipVertex.Task(h, self)
        self.Pid = shipPid.Task(h, self)