Skip to content

Commit

Permalink
Merge pull request #119 from EmanuelPerez/update_vertex
Browse files Browse the repository at this point in the history 
Update vertex
  • Loading branch information
@clementhelsens
clementhelsens authored Feb 21, 2022
2 parents bd4b0c4 + 56a1813 commit a04f9a4
Show file tree
Hide file tree
Showing 3 changed files with 167 additions and 12 deletions.
105 changes: 105 additions & 0 deletions analyzers/dataframe/VertexFitterSimple.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -722,3 +722,108 @@ VertexingUtils::FCCAnalysesVertex VertexFitterSimple::VertexFitter_Tk( int Prim


////////////////////////////////////////////////////



ROOT::VecOps::RVec<edm4hep::TrackState> VertexFitterSimple::get_PrimaryTracks( VertexingUtils::FCCAnalysesVertex initialVertex,
ROOT::VecOps::RVec<edm4hep::TrackState> tracks,
bool BeamSpotConstraint,
double bsc_sigmax, double bsc_sigmay, double bsc_sigmaz,
double bsc_x, double bsc_y, double bsc_z,
int ipass ) {


// iterative procedure to determine the primary vertex - and the primary tracks
// Start from a vertex reconstructed from all tracks, remove the one with the highest chi2, fit again etc

// tracks = the collection of tracks that was used in the first step

//bool debug = true ;
bool debug = false;
float CHI2MAX = 25 ;

if (debug) {
if (ipass == 0) std::cout << " \n --------------------------------------------------------\n" << std::endl;
std::cout << " ... enter in VertexFitterSimple::get_PrimaryTracks ipass = " << ipass << std::endl;
if (ipass == 0) std::cout << " initial number of tracks = " << tracks.size() << std::endl;
}

ROOT::VecOps::RVec<edm4hep::TrackState> seltracks = tracks;
ROOT::VecOps::RVec<float> reco_chi2 = initialVertex.reco_chi2;

if ( seltracks.size() <= 1 ) return seltracks;

int isPrimaryVertex = initialVertex.vertex.primary ;

int maxElementIndex = std::max_element(reco_chi2.begin(),reco_chi2.end()) - reco_chi2.begin();
auto minmax = std::minmax_element(reco_chi2.begin(), reco_chi2.end());
float chi2max = *minmax.second ;

if ( chi2max < CHI2MAX ) {
if (debug) {
std::cout << " --- DONE, all tracks have chi2 < CHI2MAX " << std::endl;
std::cout << " number of primary tracks selected = " << seltracks.size() << std::endl;

}
return seltracks ;
}

if (debug) {
std::cout << " remove a track that has chi2 = " << chi2max << std::endl;
}

seltracks.erase( seltracks.begin() + maxElementIndex );
ipass ++;

VertexingUtils::FCCAnalysesVertex vtx = VertexFitterSimple::VertexFitter_Tk( isPrimaryVertex,
seltracks,
BeamSpotConstraint,
bsc_sigmax, bsc_sigmay, bsc_sigmaz,
bsc_x, bsc_y, bsc_z ) ;

return VertexFitterSimple::get_PrimaryTracks( vtx, seltracks, BeamSpotConstraint, bsc_sigmax, bsc_sigmay, bsc_sigmaz,
bsc_x, bsc_y, bsc_z, ipass ) ;



}


ROOT::VecOps::RVec<edm4hep::TrackState> VertexFitterSimple::get_NonPrimaryTracks( ROOT::VecOps::RVec<edm4hep::TrackState> allTracks,
ROOT::VecOps::RVec<edm4hep::TrackState> primaryTracks ) {

ROOT::VecOps::RVec<edm4hep::TrackState> result;
for (auto & track: allTracks) {
bool isInPrimary = false;
for ( auto & primary: primaryTracks) {
if ( track.D0 == primary.D0 && track.Z0 == primary.Z0 && track.phi == primary.phi && track.omega == primary.omega && track.tanLambda == primary.tanLambda ) {
isInPrimary = true;
break;
}
}
if ( !isInPrimary) result.push_back( track );
}

return result;
}


ROOT::VecOps::RVec<bool> VertexFitterSimple::IsPrimary_forTracks( ROOT::VecOps::RVec<edm4hep::TrackState> allTracks,
ROOT::VecOps::RVec<edm4hep::TrackState> primaryTracks ) {

ROOT::VecOps::RVec<bool> result;
for (auto & track: allTracks) {
bool isInPrimary = false;
for ( auto & primary: primaryTracks) {
if ( track.D0 == primary.D0 && track.Z0 == primary.Z0 && track.phi == primary.phi && track.omega == primary.omega && track.tanLambda == primary.tanLambda ) {
isInPrimary = true;
break;
}
}
result.push_back( isInPrimary );
}
return result;
}



18 changes: 18 additions & 0 deletions analyzers/dataframe/VertexFitterSimple.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,24 @@ namespace VertexFitterSimple{
double sigmax=0., double sigmay=0., double sigmaz=0.,
double bsc_x=0., double bsc_y=0., double bsc_z=0. ) ;

/// Return the tracks that are flagged as coming from the primary vertex
ROOT::VecOps::RVec<edm4hep::TrackState> get_PrimaryTracks( VertexingUtils::FCCAnalysesVertex initialVertex,
ROOT::VecOps::RVec<edm4hep::TrackState> tracks,
bool BeamSpotConstraint,
double bsc_sigmax, double bsc_sigmay, double bsc_sigmaz,
double bsc_x, double bsc_y, double bsc_z,
int ipass = 0 ) ;


/// Return the tracks that are NOT flagged as coming from the primary vertex
ROOT::VecOps::RVec<edm4hep::TrackState> get_NonPrimaryTracks( ROOT::VecOps::RVec<edm4hep::TrackState> allTracks,
ROOT::VecOps::RVec<edm4hep::TrackState> primaryTracks ) ;

/// for an input vector of tracks, return a vector of bools that tell if the track was identified as a primary track
ROOT::VecOps::RVec<bool> IsPrimary_forTracks( ROOT::VecOps::RVec<edm4hep::TrackState> allTracks,
ROOT::VecOps::RVec<edm4hep::TrackState> primaryTracks ) ;



Double_t FastRv(TVectorD p1, TVectorD p2) ;
TMatrixDSym RegInv3(TMatrixDSym &Smat0) ;
Expand Down
56 changes: 44 additions & 12 deletions examples/FCCee/vertex/analysis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,14 @@ def __init__(self, inputlist, outname, ncpu):
if ".root" not in outname:
self.outname+=".root"

ROOT.ROOT.EnableImplicitMT(ncpu)
#ROOT.ROOT.EnableImplicitMT(ncpu)

self.df = ROOT.RDataFrame("events", inputlist)
print (" done")
#__________________________________________________________
def run(self):
#df2 = (self.df.Range(0,5000)
df2 = (self.df
df2 = (self.df.Range(0,1000)
#df2 = (self.df

# MC event primary vertex
.Define("MC_PrimaryVertex", "MCParticle::get_EventPrimaryVertex(21)( Particle )" )
Expand All @@ -53,28 +53,55 @@ def run(self):
# gen-level primary vertex is not (0,0,0)
.Alias("MCRecoAssociations0", "MCRecoAssociations#0.index")
.Alias("MCRecoAssociations1", "MCRecoAssociations#1.index")
.Define("PrimaryTracks", "VertexingUtils::SelPrimaryTracks(MCRecoAssociations0,MCRecoAssociations1,ReconstructedParticles,Particle, MC_PrimaryVertex)" )
.Define("nPrimaryTracks", "ReconstructedParticle::get_n(PrimaryTracks)")
# the recoParticles corresponding to the tracks that are primaries, according to MC-matching :
.Define("MC_PrimaryTracks_RP", "VertexingUtils::SelPrimaryTracks(MCRecoAssociations0,MCRecoAssociations1,ReconstructedParticles,Particle, MC_PrimaryVertex)" )
# and the corresponding tracks :
.Define("MC_PrimaryTracks", "ReconstructedParticle2Track::getRP2TRK( MC_PrimaryTracks_RP, EFlowTrack_1)" )

.Define("nPrimaryTracks", "ReconstructedParticle::get_n(MC_PrimaryTracks_RP)")

# Reconstruct the vertex from these primary tracks :
.Define("VertexObject_primaryTracks", "VertexFitterSimple::VertexFitter ( 1, PrimaryTracks, EFlowTrack_1) ")
.Define("VertexObject_primaryTracks", "VertexFitterSimple::VertexFitter ( 1, MC_PrimaryTracks_RP, EFlowTrack_1) ")
.Define("Vertex_primaryTracks", "VertexingUtils::get_VertexData( VertexObject_primaryTracks )") # primary vertex, in mm

# Idem, but adding the beam-spot constraint to the vertex fitter
# At the Z peak, the beam-spot size is ( 4.5 mum, 20 nm, 0.3 mm)
# The beam-spot dimensions are passed in mum :
.Define("VertexObject_primaryTracks_BSC", "VertexFitterSimple::VertexFitter ( 1, PrimaryTracks, EFlowTrack_1, true, 4.5, 20e-3, 300) ")
.Define("VertexObject_primaryTracks_BSC", "VertexFitterSimple::VertexFitter ( 1, MC_PrimaryTracks_RP, EFlowTrack_1, true, 4.5, 20e-3, 300) ")
.Define("Vertex_primaryTracks_BSC", "VertexingUtils::get_VertexData( VertexObject_primaryTracks_BSC )") # primary vertex, in mm


#Run the Acts AMVF vertex finder
.Define("VertexObject_actsFinder","VertexFinderActs::VertexFinderAMVF( EFlowTrack_1)")
.Define("Vertex_actsFinder", "VertexingUtils::get_VertexData( VertexObject_actsFinder )") # primary vertex, in mm
.Define("nPrimaryTracks_actsFinder", "VertexingUtils::get_VertexNtrk(VertexObject_actsFinder)")
#.Define("VertexObject_actsFinder","VertexFinderActs::VertexFinderAMVF( EFlowTrack_1)")
#.Define("Vertex_actsFinder", "VertexingUtils::get_VertexData( VertexObject_actsFinder )") # primary vertex, in mm
#.Define("nPrimaryTracks_actsFinder", "VertexingUtils::get_VertexNtrk(VertexObject_actsFinder)")

#Run the Acts full Billoir vertex fitter
.Define("VertexObject_primaryTracks_actsFitter","VertexFitterActs::VertexFitterFullBilloir(PrimaryTracks, EFlowTrack_1)")
.Define("Vertex_primaryTracks_actsFitter", "VertexingUtils::get_VertexData( VertexObject_primaryTracks_actsFitter )") # primary vertex, in mm
#.Define("VertexObject_primaryTracks_actsFitter","VertexFitterActs::VertexFitterFullBilloir(PrimaryTracks, EFlowTrack_1)")
#.Define("Vertex_primaryTracks_actsFitter", "VertexingUtils::get_VertexData( VertexObject_primaryTracks_actsFitter )") # primary vertex, in mm


# --- now, determime the primary (and secondary) tracks without using the MC-matching:

# First, reconstruct a vertex from all tracks
.Define("VertexObject_allTracks", "VertexFitterSimple::VertexFitter_Tk ( 1, EFlowTrack_1, true, 4.5, 20e-3, 300)")
# Select the tracks that are reconstructed as primaries
.Define("RecoedPrimaryTracks", "VertexFitterSimple::get_PrimaryTracks( VertexObject_allTracks, EFlowTrack_1, true, 4.5, 20e-3, 300, 0., 0., 0., 0)")
.Define("n_RecoedPrimaryTracks", "ReconstructedParticle2Track::getTK_n( RecoedPrimaryTracks )")
# the final primary vertex :
.Define("FinalVertexObject", "VertexFitterSimple::VertexFitter_Tk ( 1, RecoedPrimaryTracks, true, 4.5, 20e-3, 300) ")
.Define("FinalVertex", "VertexingUtils::get_VertexData( FinalVertexObject )")

# the secondary tracks
.Define("SecondaryTracks", "VertexFitterSimple::get_NonPrimaryTracks( EFlowTrack_1, RecoedPrimaryTracks )")
.Define("n_SecondaryTracks", "ReconstructedParticle2Track::getTK_n( SecondaryTracks )" )

# which of the tracks are primary according to the reco algprithm
.Define("IsPrimary_based_on_reco", "VertexFitterSimple::IsPrimary_forTracks( EFlowTrack_1, RecoedPrimaryTracks )")
# which of the tracks are primary according to the MC-matching
.Define("IsPrimary_based_on_MC", "VertexFitterSimple::IsPrimary_forTracks( EFlowTrack_1, MC_PrimaryTracks )")





Expand All @@ -87,8 +114,10 @@ def run(self):
"MC_PrimaryVertex",
"ntracks",
"nPrimaryTracks",
# primary vertex, seeded by the tracks that are MC-matched to MC-primaries
"Vertex_primaryTracks",
"Vertex_primaryTracks_BSC",
"IsPrimary_based_on_MC",

#"nPrimaryTracks_actsFinder",
#
Expand All @@ -97,6 +126,9 @@ def run(self):
#"Vertex_actsFinder", # on Zuds: both track selections lead to very similar results for the vertex
#"Vertex_primaryTracks_actsFitter", # on Zuds: both track selections lead to very similar results for the vertex

# primary vertex and primary tracks w/o any MC-matching :
"IsPrimary_based_on_reco",
"FinalVertex",

]:
branchList.push_back(branchName)
Expand Down

0 comments on commit a04f9a4

Please sign in to comment.