TrackExtrapolateTool.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 ROOT
import shipunit as u
 
 
def cmp(a, b):
    return (a > b) - (a < b)
 
 
minNdf = 20
parallelToZ = ROOT.TVector3(0.0, 0.0, 1.0)
top = ROOT.gGeoManager.GetTopVolume()
if top.GetNode("SplitCalDetector_1"):
    z_ecal = top.GetNode("SplitCalDetector_1").GetMatrix().GetTranslation()[2]
else:
    print("TrackExtraploate tool: Error, no calo present")
    z_ecal = 100 * u.m
 
 
def extrapolateToPlane(fT, z):
    # etrapolate to a plane perpendicular to beam direction (z)
    rc, pos, mom = False, None, None
    fst = fT.getFitStatus()
    if fst.isFitConverged() and fst.getNdf() > minNdf:
        # test for fit status for each point
        if fT.getPoint(0).getFitterInfo() and fT.getPoint(1).getFitterInfo():
            fstate0, fstate1 = fT.getFittedState(0), fT.getFittedState(1)
            fPos0, fPos1 = fstate0.getPos(), fstate1.getPos()
            if abs(z - fPos0.z()) < abs(z - fPos1.z()):
                fstate = fstate0
            else:
                fstate = fstate1
            zs = min(z, z_ecal)
            NewPosition = ROOT.TVector3(0.0, 0.0, zs)
            rep = ROOT.genfit.RKTrackRep(13 * cmp(fstate.getPDG(), 0))
            state = ROOT.genfit.StateOnPlane(rep)
            pos, mom = fstate.getPos(), fstate.getMom()
            rep.setPosMom(state, pos, mom)
            try:
                rep.extrapolateToPlane(state, NewPosition, parallelToZ)
                pos, mom = state.getPos(), state.getMom()
                rc = True
            except Exception:
                # print 'error with extrapolation: z=',z/u.m,'m',pos.X(),pos.Y(),pos.Z(),mom.X(),mom.Y(),mom.Z()
                pass
            if not rc or z > z_ecal:
                # use linear extrapolation
                px, py, pz = mom.X(), mom.Y(), mom.Z()
                lam = (z - pos.Z()) / pz
                pos = ROOT.TVector3(pos.X() + lam * px, pos.Y() + lam * py, z)
    return rc, pos, mom