Vertex-track association tool implementation

src/appl/emrec/emvertex.cpp

@@ -11,6 +11,8 @@
 #include "EdbCombGen.h"
 #include "EdbVertexComb.h"
 
+#include <TROOT.h>
+
 using namespace std;
 using namespace TMath;
 
@@ -20,14 +22,24 @@ EdbID        idset;
 EdbPVRec     gAli;
 EdbScanProc  gSproc;
 EdbVertexRec gEVR;
+EdbVertexRec rfEVR;
+
+bool do_vtxrefit = false;
+float tr_pfit = 1000;
+float tr_mfit = 0.1390;
+int last_trkID = -1;
+
 
 void VertexRec(EdbID id, TEnv &cenv);
 void ReadVertex(EdbID id,TEnv &env);
 void MakeScanCondBT(EdbScanCond &cond, TEnv &env);
 void SetTracksErrors(TObjArray &tracks, EdbScanCond &cond, float p, float m);
 void do_vertex(TEnv &env);
-void AddCompatibleTracks(EdbPVRec &v_trk, EdbPVRec &v_vtx);
-bool IsCompatible(EdbVertex &v, EdbTrackP &t);
+void AddCompatibleTracks(TEnv &env, EdbPVRec &v_trk, EdbPVRec &v_vtx, float r2max, float dzmax, TObjArray &v_out, TObjArray &v_out2, TNtuple* outTree);
+bool IsCompatible(EdbVertex &v, EdbTrackP &t, float r2max, float dzmax, float *r2, float *dz);
+void SplitTrack(EdbTrackP *t, EdbTrackP *&t_in, EdbTrackP *&t_out, Int_t zsplit);
+void ExecuteVTA(EdbVertex *vtx, EdbTrackP *track);
+int SetSegmentsP(EdbTrackP t, float p) {for(int i=0; i<t.N(); i++) t.GetSegment(i)->SetP(p); return t.N();}
 void Display( const char *dsname,  EdbVertexRec *evr, TEnv &env );
 
 //----------------------------------------------------------------------------------------
@@ -68,6 +80,9 @@ void set_default(TEnv &env)
   env.SetValue("emvertex.trfit.doit"     ,  1 );
   env.SetValue("emvertex.trfit.P"        , 10 );
   env.SetValue("emvertex.trfit.M"        ,  0.139);
+  env.SetValue("emvertex.trfit.r2max", 5. );
+  env.SetValue("emvertex.trfit.dzmax", 4000. );
+
   env.SetValue("emvertex.bt.Sigma0", "0.2 0.2 0.002 0.002" );
   env.SetValue("emvertex.bt.Degrad", 5. );
 }
@@ -164,7 +179,7 @@ int main(int argc, char* argv[])
   bool        do_set     = false;
   bool        do_display = false;
   bool        do_read    = false;
-  
+
   for(int i=1; i<argc; i++ ) {
     char *key  = argv[i];
     if(!strncmp(key,"-set=",5))
@@ -183,6 +198,10 @@ int main(int argc, char* argv[])
     {
       do_display=true;
     }
+    else if(!strncmp(key,"-fit", 6))
+    {
+      do_vtxrefit=true;
+    }
   } 
   cenv.WriteFile("vertex.save.rootrc");
 
@@ -229,6 +248,16 @@ void ReadVertex(EdbID id, TEnv &env)
   gEVR.eQualityMode= env.GetValue("emvertex.vtx.QualityMode"   , 0);  // (0:=Prob/(sigVX^2+sigVY^2); 1:= inverse average track-vertex distance)
   TCut cutvtx      = env.GetValue("emvertex.vtx.cutvtx"        , "(flag==0||flag==3)&&n>4");
 
+  float r2max      = env.GetValue("emvertex.trfit.r2max"        , 5. );
+  float dzmax      = env.GetValue("emvertex.trfit.dzmax"        , 4000. );
+  TEnv trenv("trenv");
+  trenv.ReadFile("track.rootrc", kEnvLocal);
+  tr_pfit   = trenv.GetValue("fedra.track.momentum"     , 1000);
+  tr_mfit   = trenv.GetValue("fedra.track.mass"     , 0.1390);
+
+  TObjArray v_out;
+  TObjArray v_out2;
+
   EdbDataProc *dproc = new EdbDataProc();
   TString name;
   gSproc.MakeFileName(name,id,"vtx.root",false);
@@ -241,7 +270,16 @@ void ReadVertex(EdbID id, TEnv &env)
       EdbPVRec *vtr = new EdbPVRec();
       vtr->SetScanCond( new EdbScanCond(gCond) );
       gSproc.ReadTracksTree( idset,*vtr, cuttr);
-      AddCompatibleTracks( *vtr, gAli );  // assign to the vertices of gAli additional tracks from vtr if any
+      TNtuple *outTree = new TNtuple("tracks","Tree of matched tracks","chosen:n:vid:tid:nseg:npl:tx:ty:firstp:lastp:r2:dz");
+      AddCompatibleTracks(env, *vtr, gAli, r2max, dzmax, v_out, v_out2, outTree);  // assign to the vertices of gAli additional tracks from vtr if any
+      EdbDataProc::MakeVertexTree(v_out,"flag0.vtx.root");
+      EdbDataProc::MakeVertexTree(v_out2,"flag1.vtx.root");
+      TFile *outFile = new TFile("found_tracks.root","RECREATE");
+      outTree->Write();
+      outFile->Write();
+      outFile->Close();
+      delete outTree;
+      delete outFile;
     }
   }
 }
@@ -315,37 +353,194 @@ void MakeScanCondBT(EdbScanCond &cond, TEnv &env)
   cond.SetName("SND_basetrack");
 }
 
-void AddCompatibleTracks(EdbPVRec &v_trk, EdbPVRec &v_vtx)
+void AddCompatibleTracks(TEnv &env, EdbPVRec &v_trk, EdbPVRec &v_vtx, float r2max, float dzmax, TObjArray &v_out, TObjArray &v_out2, TNtuple* outTree)
 {
   int ntr  = v_trk.Ntracks();
   int nvtx = v_vtx.Nvtx();
+  rfEVR.eEdbTracks = gAli.eTracks;
+  rfEVR.SetPVRec(&gAli);
+  rfEVR.eDZmax      = env.GetValue("emvertex.vtx.DZmax"         , 3000.);
+  rfEVR.eProbMin    = env.GetValue("emvertex.vtx.ProbMinV"      , 0.001);
+  rfEVR.eImpMax     = env.GetValue("emvertex.vtx.ImpMax"        , 10.);
+  rfEVR.eUseMom     = env.GetValue("emvertex.vtx.UseMom"        , false);
+  rfEVR.eUseSegPar  = env.GetValue("emvertex.vtx.UseSegPar"     , false);
+  rfEVR.eQualityMode= env.GetValue("emvertex.vtx.QualityMode"   , 0);  // (0:=Prob/(sigVX^2+sigVY^2); 1:= inverse average track-vertex distance)
   Log(1,"AddCompatibleTracks", "%d tracks, %d vertex", ntr,nvtx );
+  if (do_vtxrefit){
+    TFile *trk_f = TFile::Open(Form("b%06d.0.0.0.trk.root", idset.eBrick));
+    TTree *trackstree = (TTree*) trk_f->Get("tracks");
+    last_trkID = trackstree->GetEntries() - 1;
+    trk_f->Close();
+  }
+  Log(3, "AddCompatibleTracks", "Last track ID from trk file: %d", last_trkID);
   for(int iv=0; iv<nvtx; iv++)
   {
     EdbVertex *v = v_vtx.GetVertex(iv);
+    if (v->Flag()==1) {rfEVR.AddVertex(v); continue;}
+    bool flag1 = false;
+    std::vector<int> trackids;
+    Log(1,"\nAddCompatibleTracks","Looking for parent tracks of vtx: %i",v->ID());
+    for(int i=0; i<v->N(); i++){
+      EdbTrackP *t = (EdbTrackP*)v->GetTrack(i);
+      int trid = t->ID();
+      trackids.push_back(trid);
+    }
+    EdbTrackP *t_chosen = 0;
+    float r2, dz;
+    float r2_ini = r2max;
+    float dz_ini = dzmax;
+    int founds=0;
     for(int it=0; it<ntr; it++) 
     {
       EdbTrackP *t = v_trk.GetTrack(it);
-      if( IsCompatible(*v,*t) ) {
-	t->SetFlag(999999);
-	v_vtx.AddTrack(t);
+      //Log(2, "AddCompatibleTracks", "Here getting track n. %d from v_trk", it);
+      int trid = t->ID();
+      if (std::find(trackids.begin(), trackids.end(), trid)!=trackids.end()) continue; //Maybe here can be changed to EdbVertex::TrackInVertex(EdbTrackP *t)
+      if( IsCompatible(*v, *t, r2max, dzmax, &r2, &dz) ) {
+        flag1 = true;
+        t->SetFlag(999999);
+        if( r2 < r2_ini ) { r2_ini=r2; dz_ini=dz; t_chosen=t; }
+        // v_vtx.AddTrack(t);
+        founds++;
+        outTree->Fill(0, 1, v->ID(), t->ID() ,t->N(), t->Npl(), t->TX(), t->TY(), t->GetSegmentFirst()->Plate(), t->GetSegmentLast()->Plate(), r2, dz);
+      }
+    }
+    if (flag1){
+      //v_vtx.AddTrack(t_chosen);
+      outTree->Fill(1, founds, v->ID(), t_chosen->ID(), t_chosen->N(), t_chosen->Npl(), t_chosen->TX(), t_chosen->TY(), t_chosen->GetSegmentFirst()->Plate(), t_chosen->GetSegmentLast()->Plate(), r2max, dzmax);
+      Log(2,"AddCompatibleTracks","Closest track found at r2=%.4f dz=%.2f",r2_ini,dz_ini);
+      v_out2.Add(v);
+      if (do_vtxrefit)
+      {
+        Log(2,"AddCompatibleTracks","Executing VTA on vertex vID=%d",v->ID());
+	      ExecuteVTA(v, t_chosen);
       }
     }
+    else 
+    {
+      if (do_vtxrefit) {rfEVR.AddVertex(v); Log(2, "AddCompatibleTracks", "Adding vertex %d with flag %d", v->ID(), v->Flag());}
+      v_out.Add(v);
+    }
+    //r2max = r2_ini; dzmax = dz_ini;
   }
+  if (do_vtxrefit)
+  {
+  TString name;
+  gSproc.MakeFileName(name,idset,"vtx.refit.root",false);
+  EdbDataProc::MakeVertexTree(*(rfEVR.eVTX),name.Data());
+}
 }
 
-bool IsCompatible(EdbVertex &v, EdbTrackP &t)
+bool IsCompatible(EdbVertex &v, EdbTrackP &t, float r2max, float dzmax, float *r2, float *dz)
 {
   EdbSegP ss;
+  EdbSegP *tst = t.GetSegmentFirst();
+  float tz = tst->Z();
+  if (tz > v.VZ()) return false;   //only tracks starting upstream of the vtx
   t.EstimatePositionAt(v.VZ(),ss);
   float dx=ss.X()-v.VX();
   float dy=ss.Y()-v.VY();
-  float r2 = Sqrt(dx*dx+dy*dy);
-  float dz = Abs(ss.DZ());
-  if(r2<5&&dz<4000) { printf("r2=%.4f dz=%.2f\n",r2,ss.DZ()); return true;}
+  *r2 = Sqrt(dx*dx+dy*dy);
+  *dz = Abs(ss.DZ());
+
+  if(*r2<r2max&&*dz<dzmax) {
+    Log(3,"IsCompatible","r2=%.4f dz=%.2f\n",*r2,ss.DZ());
+    return true;
+  }
   return false;
 }
 
+void SplitTrack(EdbTrackP *t, EdbTrackP *&t_in, EdbTrackP *&t_out, Float_t zsplit)
+{
+  EdbSegP *sbest = (EdbSegP *) t->GetSegmentWithClosestZ(zsplit, 5000);
+  if  (!sbest) {Log(1, "EdbTrackP::GetSegmentWithClosestZ", "closest segment not found!"); return;}
+  Int_t cutplate = -1;
+  if (sbest->Z() < zsplit) cutplate = sbest->Plate();
+  else {cutplate = sbest->Plate() - 1;}
+  for (Int_t iseg = 0; iseg < t->N(); iseg++)
+  {
+    EdbSegP *seg = (EdbSegP *) t->GetSegment(iseg);
+    if (seg->Plate() <= cutplate) t_in->AddSegment(seg);
+    else {t_out->AddSegment(seg);}
+  }
+  Log(3, "SplitTrack", "Track found is %d, print follows", t->ID());
+  if (gEDBDEBUGLEVEL == 3) t->PrintNice();
+  //SetSegmentsP(t_in, tr_pfit);
+  t_in->SetP(tr_pfit);
+  t_in->SetM(tr_mfit);
+  t_in->SetCounters();
+  t_in->SetMC(t->MCEvt(), t->MCTrack());
+  t_in->SetID(last_trkID+1);
+  t_in->SetTrack(last_trkID+1);
+  t_in->SetSegmentsTrack(last_trkID+1);
+  t_in->SetFlag(999999);
+  t_in->FitTrackKFS(false, 3504); // using segments with min Z and W radlen
+  if (t_out->N() != 0)
+  {
+  //SetSegmentsP(t_out, tr_pfit);
+  t_out->SetP(tr_pfit);
+  t_out->SetM(tr_mfit);
+  t_out->SetCounters();
+  t_out->SetMC(t->MCEvt(), t->MCTrack());
+  t_out->SetID(last_trkID+2);
+  t_out->SetTrack(last_trkID+2);
+  t_out->SetSegmentsTrack(last_trkID+2);
+  t_out->SetFlag(999999);
+  t_out->FitTrackKFS(false, 3504);
+  last_trkID+=2;
+  }
+  else {last_trkID+=1;Log(1, "SplitTrack", "Out-track has 0 segments!");}
+  Log(3, "SplitTrack", "Track %d is splitted in %d and %d", t->ID(), t_in->ID(), t_out->ID());
+  if (gEDBDEBUGLEVEL >= 3) {t_in->PrintNice();t_out->PrintNice();}
+}
+void ExecuteVTA(EdbVertex *vtx, EdbTrackP *track)
+{
+  // Make a new EdbVertex object in order to not change the original EdbVertex obj
+  EdbVertex *vtx_new = new EdbVertex();
+  Log(2, "ExecuteVTA", "Multiplicity of original vertex is %d", vtx->N());
+  //vtx_new->SetV(vtx->V());  // this was causing a bug when accessing vtx->VZ()
+  for (int t=0;t<vtx->N();t++)
+  { 
+	  EdbVTA *vta = nullptr;
+	  vta = rfEVR.AddTrack(*vtx_new, (EdbTrackP*)vtx->GetTrack(t), true);
+  }
+  EdbTrackP *intrack = new EdbTrackP();
+  EdbTrackP *outtrack = new EdbTrackP();
+  if (vtx->N() != vtx_new->N())
+  {
+	  Log(2, "ExecuteVTA", "Multiplicities between vtx and vtx_new mismatch before VTA, printing them");
+	  if( gEDBDEBUGLEVEL >= 2) {vtx->Print();vtx_new->Print();}
+  }
+  Log(2, "ExecuteVTA", "Attaching track %d to vertex %d at z=%f", track->ID(), vtx->ID(), vtx->VZ());
+  SplitTrack(track, intrack, outtrack, vtx->VZ());
+  EdbVTA *vta_in = new EdbVTA(intrack, vtx_new);
+  vta_in->SetFlag(2);
+  vtx_new->AddVTA(vta_in);
+  vta_in->SetZpos(0);
+  intrack->AddVTA(vta_in);
+  EdbVTA *vta_out;
+  if (outtrack->N()){ 
+  vta_out = new EdbVTA(outtrack, vtx_new);
+  vta_out->SetFlag(2);
+  vtx_new->AddVTA(vta_out);
+  vta_out->SetZpos(1);
+  outtrack->AddVTA(vta_out);
+  }
+  else {Log(1, "ExecuteVTA", "WARNING: Out-track invalid, not attaching it to the vertex");}
+  vtx_new->EstimateVertexFlag();
+  if (vtx_new->Flag() != 1) {Log(2, "ExecuteVTA", "Warning, vtx new %d has flag %d", vtx->ID(), vtx_new->Flag());if (gEDBDEBUGLEVEL == 2) {vtx_new->Print();}}
+  vtx_new->SetID(vtx->ID());
+  if (rfEVR.MakeV(*vtx_new) ) {rfEVR.AddVertex(vtx_new);}
+  else {Log(1, "ExecuteVTA", "VERTEX %d NOT ADDED TO THE VERTEXREC", vtx_new->ID());}
+  Log(2, "ExecuteVTA", "New vertex created vID=%d, prob is %.3f, the original one was %.3f", vtx_new->ID(), vtx_new->V()->prob(), vtx->V()->prob());
+  if (outtrack->N()) {Log(2, "SplitTrack", "Track, attached to vertex vID=%d, trid=%d splitted in trid=%d (in) and in trid=%d (out)", vtx->ID(), track->ID(), intrack->ID(), outtrack->ID());}
+  else {Log(2, "ExecuteVTA", "Track, attached to vertex vID=%d, trid=%d splitted in trid=%d", vtx->ID(), track->ID(), intrack->ID());}
+  rfEVR.ePVR->AddTrack(intrack);
+  if (outtrack->N()) rfEVR.ePVR->AddTrack(outtrack);
+  //SafeDelete(vta_in);
+  //SafeDelete(vta_out);
+}
+
 //-----------------------------------------------------------------------------
 void SetTracksErrors(TObjArray &tracks, EdbScanCond &cond, float p, float m)
 {

src/libEdr/EdbPattern.cxx

@@ -1096,7 +1096,7 @@ int EdbTrackP::EstimatePositionAt( Float_t z, EdbSegP &ss )
   // TODO: dz=0: mean?
   if( N()<2 ) return 0;
   EdbSegP *s1=0,*s2=0;
-  float   dz1,dz2;   dz1=dz2=kMaxInt;
+  float   dz1,dz2;   dz1=dz2=1e9f; //kMaxInt into float;
   for(int i=0; i<N(); i++)
   {
     EdbSegP *s = GetSegment(i);
@@ -1114,12 +1114,14 @@ int EdbTrackP::EstimatePositionAt( Float_t z, EdbSegP &ss )
   float dx0 = s2->X()-s1->X();
   float dy0 = s2->Y()-s1->Y();
   float dz  = z-s1->Z();
-  ss.SetX( s1->X() + dz*dx0/dz0 );
-  ss.SetY( s1->Y() + dz*dy0/dz0 );
+  float tx = dx0/dz0;
+  float ty = dy0/dz0;
+  ss.SetX( s1->X() + dz*tx );
+  ss.SetY( s1->Y() + dz*ty );
   ss.SetZ( z );
   ss.SetDZ( dz );           // keep dz distance
-  ss.SetTX( dx0/dz0 );
-  ss.SetTY( dy0/dz0 );
+  ss.SetTX( tx );
+  ss.SetTY( ty );
   return 1;
 }
 

src/libEdr/EdbVertex.cxx

@@ -160,7 +160,7 @@ EdbSegP   *EdbVertex::GetTrackV(int i, bool usesegpar)
 {
 	EdbTrackP *t = GetTrack(i); if(!t) return 0;
 	EdbSegP *s = t->TrackExtremity(Zpos(i), usesegpar);
-  if( s->P()>=0 && (s->P() != t->P()) ) Log(1,"GetTrackV","Warning! segment momentum=%f is not equal to the track momentum=%f",s->P(), t->P() );
+  //if( s->P()>=0 && (s->P() != t->P()) ) Log(1,"GetTrackV","Warning! segment momentum=%f is not equal to the track momentum=%f",s->P(), t->P() );
 	return s;
 }