add AT->KFPF interface

analysis/CMakeLists.txt

@@ -13,6 +13,7 @@ add_subdirectory (PWGCHA/jpsi)
 add_subdirectory (PWGHAD/hadron)
 
 add_subdirectory (common/analysis_tree_converter)
+add_subdirectory (common/at_kfpf_interface)
 
 add_subdirectory (detectors)
 

analysis/common/at_kfpf_interface/ATKFParticleFinder.cxx

+#include "ATKFParticleFinder.h"
+
+void ATKFParticleFinder::InitInput(const std::string& file_name, const std::string& tree_name)
+{
+  std::cout << "ATKFParticleFinder::InitInput()\n";
+  
+//   topo_reconstructor_ = new KFParticleTopoReconstructor;
+//   topo_reconstructor_->Clear();
+  
+  in_file_ = TFile::Open(file_name.c_str(), "read");
+  config_ = (AnalysisTree::Configuration*) in_file_->Get("Configuration");
+
+  in_chain_ =  new TChain(tree_name.c_str());
+  in_chain_ -> Add(file_name.c_str());
+  in_chain_ -> SetBranchAddress("VtxTracks", &kf_tracks_);
+  in_chain_ -> SetBranchAddress("RecEventHeader", &rec_event_header_);
+
+  auto branch_conf_kftr = config_->GetBranchConfig( "VtxTracks" );
+  q_field_id_ = branch_conf_kftr.GetFieldId("q");
+
+  par_field_id_ = branch_conf_kftr.GetFieldId("x");   // par0
+  mf_field_id_ = branch_conf_kftr.GetFieldId("cx0");  // magnetic field par0
+  cov_field_id_ = branch_conf_kftr.GetFieldId("cov1"); // cov matrix 0
+  
+  passcuts_field_id_ = branch_conf_kftr.GetFieldId("pass_cuts");
+  pdg_field_id_ = branch_conf_kftr.GetFieldId("mc_pdg");
+  nhits_field_id_ = branch_conf_kftr.GetFieldId("nhits");
+  nhits_mvd_field_id_ = branch_conf_kftr.GetFieldId("nhits_mvd");
+  vtx_chi2_field_id_ = branch_conf_kftr.GetFieldId("vtx_chi2");
+}
+
+void ATKFParticleFinder::InitOutput(const std::string& file_name)
+{
+  std::cout << "ATKFParticleFinder::InitOutput()\n";
+  
+  out_file_ = TFile::Open(file_name.c_str(), "recreate");
+  
+  AnalysisTree::BranchConfig ParticlesRecoBranch("ParticlesReconstructed", AnalysisTree::DetType::kParticle);
+  
+  ParticlesRecoBranch.AddField<int>("daughter1id");
+  ParticlesRecoBranch.AddField<int>("daughter2id");
+  
+  out_config_.AddBranchConfig( ParticlesRecoBranch );
+  particles_reco_ = new AnalysisTree::Particles(out_config_.GetLastId());
+  
+  out_tree_ = new TTree("aTree", "AnalysisTree ParticlesReco");
+  out_tree_ -> Branch("ParticlesReconstructed", "AnalysisTree::Particles", &particles_reco_);
+  out_config_.Write("Configuration");
+  
+  daughter1_id_field_id_ = out_config_.GetBranchConfig( particles_reco_->GetId() ).GetFieldId("daughter1id");
+  daughter2_id_field_id_ = out_config_.GetBranchConfig( particles_reco_->GetId() ).GetFieldId("daughter2id");  
+}
+
+void ATKFParticleFinder::Finish()
+{
+  std::cout << "ATKFParticleFinder::Finish()\n";
+  
+  out_tree_->Write();  
+  out_file_->Close();
+}
+
+void ATKFParticleFinder::Run(int n_events)
+{
+  if(n_events<0 || n_events>in_chain_->GetEntries())
+    n_events = in_chain_->GetEntries();
+  
+  std::cout << "ATKFParticleFinder::Run() " << n_events << " events\n";
+  
+  for(int i_event=0; i_event<n_events; i_event++)
+  {
+    std::cout << "eveNo = " << i_event << "\n";
+    in_chain_->GetEntry(i_event);
+    KFParticleTopoReconstructor* eventTopoReconstructor = CreateTopoReconstructor();
+    
+//     const KFPTrackVector* tv = eventTopoReconstructor->GetTracks();
+//     KFPTrackVector tvv = *tv;
+//     tvv.Print();
+    
+    eventTopoReconstructor->SortTracks();
+    eventTopoReconstructor->ReconstructParticles();
+       
+    WriteCandidates(eventTopoReconstructor);
+  }
+  Finish();
+}
+
+KFParticleTopoReconstructor* ATKFParticleFinder::CreateTopoReconstructor()
+{
+// 
+// Creates the pointer on the KFParticleTopoReconstructor object
+// with all necessary input information in order to perform particle selection using
+// non-simplified "standard" KFParticle algorithm. 
+// 
+  auto* TR = new KFParticleTopoReconstructor;
+  
+  // cuts setting
+  TR -> GetKFParticleFinder() -> SetChiPrimaryCut2D(cuts_.GetCutChi2Prim());
+  TR -> GetKFParticleFinder() -> SetMaxDistanceBetweenParticlesCut(cuts_.GetCutDistance());
+  TR -> GetKFParticleFinder() -> SetChi2Cut2D(cuts_.GetCutChi2Geo());
+  TR -> GetKFParticleFinder() -> SetLCut(cuts_.GetCutLDown());
+  TR -> GetKFParticleFinder() -> SetLdLCut2D(cuts_.GetCutLdL());
+  
+  int n_good_tracks=0;
+  
+  for(int i_track=0; i_track < kf_tracks_->GetNumberOfChannels(); i_track++)
+  {
+    const AnalysisTree::Track& rec_track = kf_tracks_->GetChannel(i_track);
+    if(rec_track.GetField<bool>(passcuts_field_id_)==0) continue;
+    if(rec_track.GetField<int>(pdg_field_id_)==-2 && pid_mode_==1) continue;
+    n_good_tracks++;
+  }
+    
+  KFPTrackVector track_vector1, track_vector2;
+  track_vector1.Resize(n_good_tracks);
+  
+  int j_track = 0;
+    
+  for(int i_track=0; i_track < kf_tracks_->GetNumberOfChannels(); i_track++)
+  {
+    const AnalysisTree::Track& rec_track = kf_tracks_->GetChannel(i_track);
+    
+    if(rec_track.GetField<bool>(passcuts_field_id_)==0) continue;
+    if(rec_track.GetField<int>(pdg_field_id_)==-2 && pid_mode_==1) continue;
+    
+    for(Int_t iP=0; iP<3; iP++)
+      track_vector1.SetParameter(rec_track.GetField<float>(par_field_id_+iP), iP, j_track);
+    track_vector1.SetParameter(rec_track.GetPx(), 3, j_track);
+    track_vector1.SetParameter(rec_track.GetPy(), 4, j_track);
+    track_vector1.SetParameter(rec_track.GetPz(), 5, j_track);
+    auto cov_matrix = GetCovMatrixCbm(rec_track);
+    for(Int_t iC=0; iC<21; iC++)
+      track_vector1.SetCovariance(cov_matrix.at(iC), iC, j_track);
+    for(Int_t iF=0; iF<10; iF++)
+      track_vector1.SetFieldCoefficient(rec_track.GetField<float>(mf_field_id_+iF), iF, j_track);
+    if(pid_mode_==0)
+      track_vector1.SetPDG(-1, j_track);
+    else if (pid_mode_==1)
+      track_vector1.SetPDG(rec_track.GetField<int>(pdg_field_id_), j_track);
+    track_vector1.SetQ(rec_track.GetField<int>(q_field_id_), j_track);
+    if(rec_track.GetField<float>(vtx_chi2_field_id_)<3.)
+      track_vector1.SetPVIndex(0, j_track);
+    else
+      track_vector1.SetPVIndex(-1, j_track);
+    track_vector1.SetNPixelHits(rec_track.GetField<int>(nhits_mvd_field_id_), j_track);
+    track_vector1.SetId(rec_track.GetId(), j_track);  
+    j_track++;
+  }
+  TR->Init(track_vector1, track_vector2);
+  
+  KFPVertex primVtx_tmp;
+  primVtx_tmp.SetXYZ(rec_event_header_->GetVertexX(), rec_event_header_->GetVertexY(), rec_event_header_->GetVertexZ());
+  primVtx_tmp.SetCovarianceMatrix( 0,0,0,0,0,0 );
+  primVtx_tmp.SetNContributors( 0 );
+  primVtx_tmp.SetChi2( -100 );
+  std::vector<int> pvTrackIds;
+  KFVertex pv(primVtx_tmp);
+  TR->AddPV(pv, pvTrackIds);
+  
+  std::cout << track_vector1.Size() << "\n";
+
+  return TR; 
+}
+
+void ATKFParticleFinder::WriteCandidates(const KFParticleTopoReconstructor* eventTR)
+{
+  particles_reco_->ClearChannels();
+  
+  for(unsigned int iP=0;iP<eventTR->GetParticles().size();iP++)
+  {
+    const KFParticle& particle = eventTR->GetParticles()[iP];
+    
+    auto* particlerec = particles_reco_->AddChannel();
+    particlerec->Init(out_config_.GetBranchConfig(particles_reco_->GetId()));
+    
+    float mass, masserr;
+    particle.GetMass(mass, masserr);
+    particlerec->SetMass(mass);
+    particlerec->SetField(particle.DaughterIds()[0], daughter1_id_field_id_);
+    particlerec->SetField(particle.DaughterIds()[1], daughter2_id_field_id_);    
+    particlerec->SetMomentum(particle.GetPx(), particle.GetPy(), particle.GetPz());
+    particlerec->SetPid ( particle.GetPDG() );
+        
+//     topo_reconstructor_->AddParticle(particle);
+  }
+  out_tree_->Fill();
+}
+
+std::vector<float> ATKFParticleFinder::GetCovMatrixCbm(const AnalysisTree::Track& track) const
+{
+  const double tx = track.GetField<float>(par_field_id_+3); 
+  const double ty = track.GetField<float>(par_field_id_+4); 
+  const double qp = track.GetField<float>(par_field_id_+5);
+  const Int_t q = track.GetField<int>(q_field_id_);
+  
+  //calculate covariance matrix
+  const double t = sqrt(1.f + tx*tx + ty*ty);
+  const double t3 = t*t*t;
+  const double dpxdtx = q/qp*(1.f+ty*ty)/t3;
+  const double dpxdty = -q/qp*tx*ty/t3;
+  const double dpxdqp = -q/(qp*qp)*tx/t;
+  const double dpydtx = -q/qp*tx*ty/t3;
+  const double dpydty = q/qp*(1.f+tx*tx)/t3;
+  const double dpydqp = -q/(qp*qp)*ty/t;
+  const double dpzdtx = -q/qp*tx/t3;
+  const double dpzdty = -q/qp*ty/t3;
+  const double dpzdqp = -q/(qp*qp*t3);
+  
+  const double F[6][5] = { {1.f, 0.f, 0.f,    0.f,    0.f},
+  {0.f, 1.f, 0.f,    0.f,    0.f},
+  {0.f, 0.f, 0.f,    0.f,    0.f},
+  {0.f, 0.f, dpxdtx, dpxdty, dpxdqp},
+  {0.f, 0.f, dpydtx, dpydty, dpydqp},
+  {0.f, 0.f, dpzdtx, dpzdty, dpzdqp} };
+    
+  double VFT[5][6];
+  for(int i=0; i<5; i++)
+    for(int j=0; j<6; j++)
+    {
+      VFT[i][j] = 0;
+      for(int k=0; k<5; k++)
+      {
+        if (k<=i)
+          VFT[i][j] += track.GetField<float>(cov_field_id_ + k + i*(i+1)/2) * F[j][k];   //parameters->GetCovariance(i,k) * F[j][k];
+        else
+          VFT[i][j] += track.GetField<float>(cov_field_id_ + i + k*(k+1)/2) * F[j][k];   //parameters->GetCovariance(i,k) * F[j][k];
+      }
+    }
+
+    
+  std::vector<float> cov(21, 0);
+  for(int i=0, l=0; i<6; i++)
+    for(int j=0; j<=i; j++, l++)
+    {
+      cov[l] = 0;
+      for(int k=0; k<5; k++)
+      {
+        cov[l] += F[i][k] * VFT[k][j];
+      }
+    }
+  return cov;
+}
\ No newline at end of file

analysis/common/at_kfpf_interface/ATKFParticleFinder.h

+#ifndef ATKFParticleFinder_HH
+#define ATKFParticleFinder_HH
+
+#include <string>
+#include <utility>
+
+#include "TFile.h"
+#include "TChain.h"
+#include "TTree.h"
+
+#include "AnalysisTree/Detector.hpp"
+#include "AnalysisTree/Configuration.hpp"
+#include "AnalysisTree/EventHeader.hpp"
+
+#include "KFParticleTopoReconstructor.h"
+#include "CutsContainer.h"
+
+class ATKFParticleFinder
+{
+public:
+  
+  void InitInput(const std::string& file_name, const std::string& tree_name);
+  void InitOutput(const std::string& file_name="KFPF_output.root");
+  void Finish();
+  void SetCuts(CutsContainer cuts) {cuts_ = cuts;};
+//   const KFParticleTopoReconstructor* GetTopoReconstructor() const {return topo_reconstructor_;};
+  void SetPIDMode(int pid){pid_mode_ = pid;};
+  void WriteCandidates(const KFParticleTopoReconstructor* TR);
+  void Run(int n_events=-1);
+  
+private:
+  
+  KFParticleTopoReconstructor* CreateTopoReconstructor();
+  
+  std::string in_file_name_;
+  std::string in_tree_name_;
+  TFile* in_file_{nullptr};
+  TChain* in_chain_{nullptr};
+  AnalysisTree::Configuration* config_{nullptr};
+  AnalysisTree::EventHeader* rec_event_header_{nullptr};
+  AnalysisTree::TrackDetector* kf_tracks_{nullptr};
+  
+  TFile* out_file_{nullptr};
+  TTree* out_tree_{nullptr};
+  AnalysisTree::Configuration out_config_;
+  AnalysisTree::Particles* particles_reco_{nullptr};
+    
+  int q_field_id_{-999};
+  int par_field_id_ {-999};
+  int mf_field_id_ {-999};
+  int cov_field_id_{-999};
+  int passcuts_field_id_{-999};
+  int pdg_field_id_{-999};
+  int vtx_chi2_field_id_{-999};
+  int nhits_field_id_{-999};
+  int nhits_mvd_field_id_{-999};
+  
+  int daughter1_id_field_id_{-999};
+  int daughter2_id_field_id_{-999};
+  
+  CutsContainer cuts_;
+//   KFParticleTopoReconstructor* topo_reconstructor_{nullptr};
+  
+  std::vector<float> GetCovMatrixCbm(const AnalysisTree::Track& track) const;
+  
+  int pid_mode_{1};
+};
+#endif //ATKFParticleFinder_HH
\ No newline at end of file

analysis/common/at_kfpf_interface/AnalysisTreeKfpfInterfaceLinkDef.h

+#ifdef __CINT__
+
+#pragma link off all globals;
+#pragma link off all classes;
+#pragma link off all functions;
+
+#pragma link C++ class ATKFParticleFinder+;
+#pragma link C++ class CutsContainer+;
+
+#endif

analysis/common/at_kfpf_interface/CMakeLists.txt

+Set(LIBRARY_NAME AnalysisTreeKfpfInterface)
+
+set(SRCS
+        ATKFParticleFinder.cxx
+        CutsContainer.cxx
+    )
+
+Set(INCLUDE_DIRECTORIES
+  ${CMAKE_CURRENT_SOURCE_DIR}
+
+  ${KFParticle_INCLUDE_DIR}
+  ${AnalysisTree_INCLUDE_DIR}
+)
+
+Include_Directories (${INCLUDE_DIRECTORIES})
+
+Set(SYSTEM_INCLUDE_DIRECTORIES
+  ${VC_INCLUDE_DIRS}
+  ${BASE_INCLUDE_DIRECTORIES}
+  ${Boost_INCLUDE_DIR} 
+)
+
+Include_Directories (SYSTEM ${SYSTEM_INCLUDE_DIRECTORIES})
+
+set (LINK_DIRECTORIES
+  ${ROOT_LIBRARY_DIR}
+  ${FAIRROOT_LIBRARY_DIR}
+  ${Boost_LIBRARY_DIRS}
+  ${Vc_LIB_DIR}
+  ${KFParticle_LIB_DIR}
+  ${AnalysisTree_LIBRARY_DIR}
+)
+ 
+link_directories(${LINK_DIRECTORIES})
+
+IF (SSE_FOUND)
+  Message(STATUS "${LIBRARY_NAME} will be compiled with SSE support")
+  ADD_DEFINITIONS(-DHAVE_SSE)
+  SET_SOURCE_FILES_PROPERTIES(${SRCS} PROPERTIES COMPILE_FLAGS "-msse -O3")
+ELSE (SSE_FOUND)
+  MESSAGE(STATUS "${LIBRARY_NAME} will be compiled without SSE support")
+  SET_SOURCE_FILES_PROPERTIES(${SRCS} PROPERTIES COMPILE_FLAGS "-msse -O3")
+ENDIF (SSE_FOUND)
+
+Set(LINKDEF AnalysisTreeKfpfInterfaceLinkDef.h)
+
+If(UNIX AND NOT APPLE)
+  Set(_AnalysisTree_LIB  AnalysisTreeBase.so AnalysisTreeInfra.so)
+  Else()
+  Set(_AnalysisTree_LIB  AnalysisTreeBase.dylib AnalysisTreeInfra.dylib)
+EndIf()
+
+Set(DEPENDENCIES 
+	${_AnalysisTree_LIB} 
+	KFParticle
+	Vc.a)
+Set(DEFINITIONS -DDO_TPCCATRACKER_EFF_PERFORMANCE -DNonhomogeneousField -DCBM -DUSE_TIMERS)
+
+ADD_DEFINITIONS(${DEFINITIONS})
+
+GENERATE_LIBRARY()
+
+# Install(FILES ../../../macro/analysis/common/analysis_tree_converter/run_treemaker.C
+#         DESTINATION share/cbmroot/macro/analysis_tree
+#        )
+
+Add_Dependencies(AnalysisTreeKfpfInterface ANALYSISTREE KFPARTICLE)
\ No newline at end of file

analysis/common/at_kfpf_interface/CutsContainer.cxx

+#include "CutsContainer.h"

analysis/common/at_kfpf_interface/CutsContainer.h

+//
+// @class CutsContainer
+// @brief Container with values of cuts.
+// @authors Oleksii Lubynets, Viktor Klochkov, Ilya Selyuzhenkov
+//
+// The meaning of quantities to be cut is described in the OutputContainer.h
+//
+
+
+#ifndef CutsContainer_H
+#define CutsContainer_H
+
+class CutsContainer
+{
+  
+ public:
+   
+  CutsContainer() = default;
+  virtual ~CutsContainer() = default;  
+  
+  //  lambda candidate parameters setters
+  void SetCutChi2Prim(float value){cut_chi2_prim_ = value;};
+  void SetCutDistance(float value){cut_distance_ = value;};
+  void SetCutChi2Geo(float value){cut_chi2_geo_ = value;};
+  void SetCutLDown(float value){cut_l_down_ = value;};
+  void SetCutLdL(float value){cut_ldl_ = value;};
+    
+  //  lambda candidate parameters getters
+  float GetCutChi2Prim() const {return cut_chi2_prim_;};
+  float GetCutDistance() const {return cut_distance_;};
+  float GetCutChi2Geo() const {return cut_chi2_geo_;};
+  float GetCutLDown() const {return cut_l_down_;};
+  float GetCutLdL() const {return cut_ldl_;};
+    
+ protected:
+   
+  // Cuts with their default values
+  float cut_chi2_prim_{18.4207};
+  float cut_distance_{1.};
+  float cut_chi2_geo_{3.};
+  float cut_l_down_{-5.};;
+  float cut_ldl_{5.};   
+    
+};
+#endif//CutsContainer_H
\ No newline at end of file