CbmAnaTreeContainer.cxx

Go to the documentation of this file.

/* Copyright (C) 2023-2023 Warsaw University of Technology, Warsaw
    SPDX-License-Identifier: GPL-3.0-only
    Authors: Daniel Wielanek [committer] */
#include "CbmAnaTreeContainer.h"
 
#include "AnalysisTree/Matching.hpp"
#include "HalCbmSource.h"
 
#include <TChain.h>
#include <TFile.h>
 
#include <AnalysisTree/Configuration.hpp>
#include <vector>
 
#include <Hal/Cout.h>
#include <Hal/InputDataInfo.h>
#include <Hal/RootIOManager.h>
 
void CbmAnaTreeRecoContainer::LoadFields(std::vector<TString> files)
{
  Bool_t found;
  AnalysisTree::BranchConfig conf = GetBranchConf(files, "VtxTracks", found);
  if (!found) return;
  GetFieldIds().vtx_px   = conf.GetFieldId("px");
  GetFieldIds().vtx_py   = conf.GetFieldId("py");
  GetFieldIds().vtx_pz   = conf.GetFieldId("pz");
  GetFieldIds().vtx_dcax = conf.GetFieldId("dcax");
  GetFieldIds().vtx_dcay = conf.GetFieldId("dcay");
  GetFieldIds().vtx_dcaz = conf.GetFieldId("dcaz");
 
  GetFieldIds().vtx_chi2    = conf.GetFieldId("chi2");
  GetFieldIds().vtx_vtxchi2 = conf.GetFieldId("vtx_chi2");
  GetFieldIds().vtx_q       = conf.GetFieldId("q");
  GetFieldIds().vtx_nhits   = conf.GetFieldId("nhits");
  GetFieldIds().vtx_mvdhits = conf.GetFieldId("nhits_mvd");
 
  GetFieldIds().vtx_x    = conf.GetFieldId("x");
  GetFieldIds().vtx_cx0  = conf.GetFieldId("cx0");
  GetFieldIds().vtx_cov1 = conf.GetFieldId("cov1");
  conf                   = GetBranchConf(files, "TofHits", found);
  if (!found) return;
  GetFieldIds().tof_mass2 = conf.GetFieldId("mass2");
}
 
void CbmAnaTreeMcContainer::LoadFields(std::vector<TString> files)
{
 
  Bool_t found;
  AnalysisTree::BranchConfig conf = GetBranchConf(files, "SimParticles", found);
  if (!found) return;
 
  GetFieldIds().px       = conf.GetFieldId("px");
  GetFieldIds().py       = conf.GetFieldId("py");
  GetFieldIds().pz       = conf.GetFieldId("pz");
  GetFieldIds().mass     = conf.GetFieldId("mass");
  GetFieldIds().pdg      = conf.GetFieldId("pid");
  GetFieldIds().motherId = conf.GetFieldId("mother_id");
 
  GetFieldIds().freezX = conf.GetFieldId("xfreez");
  GetFieldIds().freezY = conf.GetFieldId("yfreez");
  GetFieldIds().freezZ = conf.GetFieldId("zfreez");
  GetFieldIds().freezT = conf.GetFieldId("tfreez");
 
  conf = GetBranchConf(files, "SimEventHeader", found);
  if (!found) return;
 
  GetFieldIds().event_b   = conf.GetFieldId("b");
  GetFieldIds().event_psi = conf.GetFieldId("psi_RP");
}
 
AnalysisTree::BranchConfig CbmAnaTreeContainer::GetBranchConf(std::vector<TString> files, TString patternToFind,
                                                              Bool_t& found) const
{
  TFile* oldFile     = gFile;
  TDirectory* oldDir = gDirectory;
  for (auto file : files) {
    TFile* f                          = new TFile(file);
    AnalysisTree::Configuration* conf = (AnalysisTree::Configuration*) f->Get("Configuration");
    auto branches                     = conf->GetListOfBranches();
    for (auto branch : branches) {
      if (patternToFind.EqualTo(branch)) {
        found = kTRUE;
        f->Close();
        gFile      = oldFile;
        gDirectory = oldDir;
        return conf->GetBranchConfig(patternToFind.Data());
      }
    }
    f->Close();
  }
  gFile      = oldFile;
  gDirectory = oldDir;
  found      = kFALSE;
  return AnalysisTree::BranchConfig();
}
 
Bool_t CbmAnaTreeRecoContainer::ConnectToTree(TChain* tree)
{
 
  fEvent     = new AnalysisTree::EventHeader();
  fVtxTracks = new AnalysisTree::TrackDetector();
  fTofHits   = new AnalysisTree::HitDetector();
  fVtx2Tof   = new AnalysisTree::Matching();
  fVtx2Mc    = nullptr;
  if (tree->GetBranch("VtxTracks.") == nullptr) {
    Hal::Cout::PrintInfo("Lack of VtxTracks in AT", Hal::EInfo::kError);
    return kFALSE;
  }
  if (tree->GetBranch("RecEventHeader.") == nullptr) {
    Hal::Cout::PrintInfo("Lack of RecEventHeader in AT", Hal::EInfo::kError);
    return kFALSE;
  }
  if (tree->GetBranch("TofHits.") == nullptr) {
    Hal::Cout::PrintInfo("Lack of TofHits in AT", Hal::EInfo::kError);
    return kFALSE;
  }
  if (tree->GetBranch("VtxTracks2TofHits.") == nullptr) {
    Hal::Cout::PrintInfo("Lack of VtxTracks2TofHits tracks in AT", Hal::EInfo::kError);
    return kFALSE;
  }
  tree->SetBranchAddress("VtxTracks.", &fVtxTracks);
  tree->SetBranchAddress("RecEventHeader.", &fEvent);
  tree->SetBranchAddress("TofHits.", &fTofHits);
  tree->SetBranchAddress("VtxTracks2TofHits.", &fVtx2Tof);
  tree->SetBranchStatus("VtxTracks.", 1);
  tree->SetBranchStatus("RecEventHeader.", 1);
  tree->SetBranchStatus("TofHits.", 1);
  tree->SetBranchStatus("VtxTracks2TofHits.", 1);
  if (tree->GetBranch("VtxTracks2SimParticles.")) {
    tree->SetBranchAddress("VtxTracks2SimParticles.", &fVtx2Mc);
    tree->SetBranchStatus("VtxTracks2SimParticles.", 1);
  }
 
  return kTRUE;
}
 
Bool_t CbmAnaTreeMcContainer::ConnectToTree(TChain* tree)
{
  fEvent     = new AnalysisTree::EventHeader();
  fParticles = new AnalysisTree::Particles();
  if (tree->GetBranch("SimEventHeader.") == nullptr) return kFALSE;
  if (tree->GetBranch("SimParticles.") == nullptr) return kFALSE;
  tree->SetBranchAddress("SimEventHeader.", &fEvent);
  tree->SetBranchAddress("SimParticles.", &fParticles);
  tree->SetBranchStatus("SimEventHeader.", 1);
  tree->SetBranchStatus("SimParticles.", 1);
  return kTRUE;
}
 
void CbmAnaTreeV0Container::LoadFields(std::vector<TString> list)
{
  Bool_t found;
  AnalysisTree::BranchConfig conf = GetBranchConf(list, "Candidates", found);
  if (!found) return;
  GetFieldIds().chi2geo          = conf.GetFieldId("chi2_geo");
  GetFieldIds().chi2_prim_first  = conf.GetFieldId("chi2_prim_first");
  GetFieldIds().chi2_prim_second = conf.GetFieldId("chi2_prim_second");
  GetFieldIds().chi2_topo        = conf.GetFieldId("chi2_topo");
  GetFieldIds().cosine_first     = conf.GetFieldId("cosine_first");
  GetFieldIds().cosine_second    = conf.GetFieldId("cosine_second");
 
 
  GetFieldIds().cosine_topo = conf.GetFieldId("cosine_topo");
  GetFieldIds().distance    = conf.GetFieldId("distance");
 
  GetFieldIds().l         = conf.GetFieldId("l");
  GetFieldIds().l_over_dl = conf.GetFieldId("l_over_dl");
  GetFieldIds().px        = conf.GetFieldId("px");
  GetFieldIds().py        = conf.GetFieldId("py");
  GetFieldIds().pz        = conf.GetFieldId("pz");
  GetFieldIds().x         = conf.GetFieldId("x");
  GetFieldIds().y         = conf.GetFieldId("y");
  GetFieldIds().z         = conf.GetFieldId("z");
  GetFieldIds().dau1id    = conf.GetFieldId("daughter1_id");
  GetFieldIds().dau2id    = conf.GetFieldId("daughter2_id");
  GetFieldIds().mass      = conf.GetFieldId("mass");
 
#ifdef DEBUG_EXTR
  GetFieldIds().dau1px = conf.GetFieldId("dau1px");
  GetFieldIds().dau1py = conf.GetFieldId("dau1py");
  GetFieldIds().dau1pz = conf.GetFieldId("dau1pz");
  GetFieldIds().dau1x  = conf.GetFieldId("dau1x");
  GetFieldIds().dau1y  = conf.GetFieldId("dau1y");
  GetFieldIds().dau1z  = conf.GetFieldId("dau1z");
 
  GetFieldIds().dau2px = conf.GetFieldId("dau2px");
  GetFieldIds().dau2py = conf.GetFieldId("dau2py");
  GetFieldIds().dau2pz = conf.GetFieldId("dau2pz");
  GetFieldIds().dau2x  = conf.GetFieldId("dau2x");
  GetFieldIds().dau2y  = conf.GetFieldId("dau2y");
  GetFieldIds().dau2z  = conf.GetFieldId("dau2z");
 
#endif
 
 
  GetFieldIds().id  = conf.GetFieldId("id");
  GetFieldIds().pdg = conf.GetFieldId("pid");
  //  if (!found) return;
  // GetFieldIds().mass = conf.GetFieldId("dcaz");
}
 
Bool_t CbmAnaTreeV0Container::ConnectToTree(TChain* tree)
{
  fParticles   = new AnalysisTree::Particles();
  fParticlesMc = new AnalysisTree::Particles();
 
  auto vec   = Hal::RootIOManager::GetListOfBranches(tree, true);
  bool found = false;
  for (auto brName : vec) {
    brName.EqualTo("Candidates.");
    found = true;
  }
  if (!found) return kFALSE;
  tree->SetBranchAddress("Candidates.", &fParticles);
  tree->SetBranchAddress("Simulated.", &fParticlesMc);
  tree->SetBranchAddress("Candidates2Simulated.", &fV02Mc);
 
  tree->SetBranchStatus("Candidates.", 1);
  tree->SetBranchStatus("Simulated.", 1);
  tree->SetBranchStatus("Candidates2Simulated.", 1);
  return kTRUE;
}
 
void CbmAnaTreeContainers::LoadConf(Hal::InputDataInfo* info)
{
  auto listRow = info->GetListOfFiles();
  std::vector<TString> rec1;
  for (auto lev : listRow) {
    rec1.push_back(lev[0]);
  }
  for (int i = 0; i < GetSize(); i++) {
    fContainers[i]->LoadFields(rec1);
  }
}
 
CbmAnaTreeContainers::CbmAnaTreeContainers()
{
  fContainers[0] = new CbmAnaTreeMcContainer();
  fContainers[1] = new CbmAnaTreeRecoContainer();
  fContainers[2] = new CbmAnaTreeV0Container();
}
 
CbmAnaTreeContainer* CbmAnaTreeContainers::GetContainer(EContainerType type)
{
  switch (type) {
    case EContainerType::kMc: {
      return fContainers[0];
    } break;
    case EContainerType::kReco: {
      return fContainers[1];
    } break;
    case EContainerType::kV0: {
      return fContainers[2];
    } break;
  }
  return nullptr;
}
 
std::vector<TString> CbmAnaTreeContainers::GetNames()
{
  std::vector<TString> res;
  res.push_back("mc");
  res.push_back("reco");
  res.push_back("v0");
  return res;
}