CbmKFV0FinderQa.cxx

Go to the documentation of this file.

/* Copyright (C) 2025 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Sergei Zharko [committer] */

 
#include "CbmKFV0FinderQa.h"
 
#include "CbmGlobalTrack.h"
#include "CbmKFV0FinderTask.h"
#include "CbmTofHit.h"
#include "algo/kfp/KfpV0FinderCutter.h"
 
#include <FairRootManager.h>
#include <Logger.h>
 
#include <TClonesArray.h>
#include <TFile.h>
 
using cbm::algo::kfp::ETrackPid;
using cbm::kfp::V0FinderQa;
using cbm::kfp::V0FinderTask;
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::ExecQa()
{
  // NOTE: reconstruction in timeslice mode
 
  FillTracks();
  FillParticles();
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
template<bool IsMixedEvent>
void V0FinderQa::FillV0Candidates(const KFParticle& v0cand)
{
  // Creates daughter 4-momentum vector
  constexpr auto CreatePVector = [](const auto& prt) -> PVector4_t {
    double px = prt.GetPx();
    double py = prt.GetPy();
    double pz = prt.GetPz();
    double m  = prt.NDaughters() == 1 ? prt.GetMassHypo() : prt.GetMass();
    // NOTE: KFParticle::GetE() has lower precision
    double e = std::sqrt(m * m + px * px + py * py + pz * pz);
    return PVector4_t(px, py, pz, e);
  };
 
  // Creates daughter 4-position vector
  constexpr auto CreateXVector = [](const auto& prtInfo) -> XVector4_t {
    const auto* h = prtInfo.pTofHit;
    return h ? XVector4_t(h->GetX(), h->GetY(), h->GetZ(), h->GetTime()) : XVector4_t();
  };
 
  const auto& particles = fpTopoReconstructor->GetParticles();
 
  //* Daughters
  size_t iDa = v0cand.DaughterIds()[0];  // daughter A (pi-)
  size_t iDb = v0cand.DaughterIds()[1];  // daughter B (proton)
  const auto& daughterA{particles[iDa]};
  const auto& daughterB{particles[iDb]};
  const auto& daughterAinfo{(*fpParticleInfo)[iDa]};
  const auto& daughterBinfo{(*fpParticleInfo)[iDb]};
 
  //* V0 properties
  Cutter_t::ParticleProperty property;
  KFParticleSIMD particleSIMD(const_cast<KFParticle&>(v0cand));  // yes, it is const_cast of reference
 
  // NOTE: in the mCBM lambda analysis the primary vertices are always set to the origin.
  //       Actually, the index of PV should be stored somewhere in the particleInfo and accessed here
  //       instead of "0"
  KFParticleSIMD primVertex(fpTopoReconstructor->GetPrimVertex(/*iPv=*/0));
  float_v decayLengthSIMD;
  float_v decayLengthErrSIMD;
  particleSIMD.GetDistanceToVertexLine(primVertex, decayLengthSIMD, decayLengthErrSIMD);
  property.daughterA.dcaRelToOrigin = daughterAinfo.dca;  // NOTE: from STS hits
  property.daughterB.dcaRelToOrigin = daughterBinfo.dca;  // NOTE: from STS hits
  property.decayLength              = decayLengthSIMD[0];
  property.dcaRelToPrimVertex       = particleSIMD.GetDistanceFromVertex(primVertex)[0];
  particleSIMD.SetProductionVertex(primVertex);
  property.chi2NdfTopo         = double(particleSIMD.Chi2()[0]) / double(particleSIMD.NDF()[0]);
  property.chi2NdfGeo          = v0cand.Chi2() / v0cand.NDF();
  property.dcaBetweenDaughters = daughterB.GetDistanceFromParticle(daughterA);
  property.daughterVtxZ        = v0cand.Z();
  // ---- Init 4-position and 4-momentum vectors of mother and daughters:
  // aX: position <- TOF hit - v0X
  // bX: position <- TOF hit - v0X
  auto aP      = CreatePVector(daughterA);
  auto bP      = CreatePVector(daughterB);
  auto v0P     = bP + aP;
  double v0Tof = property.decayLength * v0P.E() / (V0FinderTask::kSpeedOfLight * v0P.P());
  auto v0X     = XVector4_t(v0cand.X(), v0cand.Y(), v0cand.Z(), v0Tof);
  auto aX      = CreateXVector(daughterAinfo);
  auto bX      = CreateXVector(daughterBinfo);
  auto aR      = aX.Vect() - v0X.Vect();
  auto bR      = bX.Vect() - v0X.Vect();
  auto angle   = std::acos((aR.x() * bR.x() + aR.y() * bR.y() + aR.z() * bR.z()) / (aR.r() * bR.r()));
  auto angleKf = daughterA.GetAngle(daughterB);
  float v0Mass{0.f};
  float v0MassErr{0.f};
  v0cand.GetMass(v0Mass, v0MassErr);
  property.openingAngle   = angle;
  property.daughterA.beta = aP.P() / aP.E();  // NOTE: from KFParticle (TODO: compare with beta from TOF)
  property.daughterB.beta = bP.P() / bP.E();  // NOTE: from KFParticle (TODO: compare with beta from TOF)
 
  auto FillV0 = [&](RefHistogramsV0& h) {
    h.opening_angle->Fill(property.openingAngle);
    h.opening_angle_kf->Fill(angleKf);
    h.dca_rt_pv->Fill(property.dcaRelToPrimVertex);
    h.dca_daughters->Fill(property.dcaBetweenDaughters);
    h.decay_length->Fill(property.decayLength);
    h.decay_vtx_xy->Fill(v0cand.X(), v0cand.Y());
    h.decay_vtx_z->Fill(property.daughterVtxZ);
    h.chi2ndf_topo->Fill(property.chi2NdfTopo);
    h.chi2ndf_geo->Fill(property.chi2NdfGeo);
    h.mass->Fill(v0Mass);
  };
 
  auto FillSecondary = [&](RefHistogramsPrt& h, const auto& dtrProperty, const auto& dtrP) {
    h.beta->Fill(dtrProperty.beta);
    h.dca_rt_origin->Fill(dtrProperty.dcaRelToOrigin);
    h.p->Fill(dtrP.P());
    h.pt->Fill(dtrP.Pt());
  };
 
  // TODO: investigate
  // auto aPrefit  = aP;
  // auto bPrefit  = bP;
  // auto v0Prefit = v0P;
  // auto v0Xrefit = v0X;
  // Refit(aX, bX, v0Xrefit, aPrefit, bPrefit, v0Prefit);
  //LOG(info) << "m(before)=" << v0P.M() << ", m(after)=" << v0Prefit.M() << ", d=" << v0P.M() - v0Prefit.M();
 
  if constexpr (!IsMixedEvent) {
    FillSecondary(fhRefDaughtersBefore.all, property.daughterA, aP);
    FillSecondary(fhRefDaughtersBefore.all, property.daughterB, bP);
    FillSecondary(fhRefDaughtersBefore.pion, property.daughterA, aP);
    FillSecondary(fhRefDaughtersBefore.proton, property.daughterB, bP);
    FillV0(fhRefV0Before);
 
    if (fpCutter->SelectParticle(property)) {
      FillSecondary(fhRefDaughtersAfter.all, property.daughterA, aP);
      FillSecondary(fhRefDaughtersAfter.all, property.daughterB, bP);
      FillSecondary(fhRefDaughtersAfter.pion, property.daughterA, aP);
      FillSecondary(fhRefDaughtersAfter.proton, property.daughterB, bP);
      FillV0(fhRefV0After);
    }
  }
  else {  // Mixed event
    FillV0(fhRefMeBefore);
    if (fpCutter->SelectParticle(property)) {
      FillV0(fhRefMeAfter);
    }
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::FillSecondary(RefHistogramsPrt& h, const KFParticle& prt, const V0FinderTask::ParticleInfo& prtInfo)
{
  h.beta->Fill(prtInfo.beta);
  h.dca_rt_origin->Fill(prtInfo.dca);
  h.p->Fill(prt.GetP());
  h.pt->Fill(prt.GetPt());
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::FillParticles()
{
  const auto& particles = fpTopoReconstructor->GetParticles();
 
  // First loop: check, if the particle was attached to some mother
  // NOTE: STAR version of KFParticle should have the ParentID
  std::vector<char> vbUnused(particles.size(), true);  //< if particle not used to build a Lambda-candidate
  for (size_t iP = 0; iP < particles.size(); ++iP) {
    const auto& particle{particles[iP]};
    if (particle.NDaughters() > 1) {
      for (size_t iD : particle.DaughterIds()) {
        vbUnused[iD] = false;
      }
    }
  }
 
  //* Signal
  for (size_t iP = 0; iP < particles.size(); ++iP) {
    const auto& particle{particles[iP]};
    const auto& particleInfo{(*fpParticleInfo)[iP]};
 
    if (particle.NDaughters() == 1) {  // Secondary: pi- or proton
      bool bUnused = vbUnused[iP];
      FillSecondary(fhRefSecondaryStored.all, particle, particleInfo);
      if (bUnused) {
        FillSecondary(fhRefSecondaryUnused.all, particle, particleInfo);
      }
      if (particle.GetPDG() == -211) {
        FillSecondary(fhRefSecondaryStored.pion, particle, particleInfo);
        if (bUnused) {
          FillSecondary(fhRefSecondaryUnused.pion, particle, particleInfo);
        }
      }
      else if (particle.GetPDG() == 2212) {
        FillSecondary(fhRefSecondaryStored.proton, particle, particleInfo);
        if (bUnused) {
          FillSecondary(fhRefSecondaryUnused.proton, particle, particleInfo);
        }
      }
    }
    else if (particle.GetPDG() == Cutter_t::GetV0Pdg()) {
      FillV0Candidates</*IsMixedEvent=*/false>(particle);
    }
  }  // Signal
 
  //* Mixed-event background
  for (size_t iP = 0; iP < particles.size(); ++iP) {  // pions
    const auto& particleA{particles[iP]};
    if (particleA.GetPDG() != -211) {
      continue;
    }
    const auto& particleInfoA{(*fpParticleInfo)[iP]};
    for (size_t iQ = 0; iQ < particles.size(); ++iQ) {  // protons
      const auto& particleB{particles[iQ]};
      if (particleB.GetPDG() != 2212) {
        continue;
      }
      const auto& particleInfoB{(*fpParticleInfo)[iQ]};
      if (particleInfoA.eventID == particleInfoB.eventID) {  // Skip particles from the same event
        continue;
      }
 
      const KFParticle* pDaughters[2] = {&particleA, &particleB};
      KFParticle particle;
      particle.Construct(pDaughters, 2, nullptr);
      particle.AddDaughterId(iP);  // pi-
      particle.AddDaughterId(iQ);  // proton
      FillV0Candidates</*IsMixedEvent=*/true>(particle);
    }
  }
}
 
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::FillTracks()
{
  int nTracks = fpBrGlobalTracks->GetEntriesFast();
  for (int iTrk = 0; iTrk < nTracks; ++iTrk) {
    const auto* pTrk    = static_cast<CbmGlobalTrack*>(fpBrGlobalTracks->At(iTrk));
    const auto& trkInfo = (*fpTrackInfo)[iTrk];
    const auto* par     = pTrk->GetParamFirst();
    bool bMomDefined    = !std::isnan(par->GetQp());
    TVector3 mom;  // Momentum of the track
    par->Momentum(mom);
    double p  = bMomDefined ? mom.Mag() : -9999.;
    double pt = bMomDefined ? mom.Perp() : -9999.;
 
 
    int pdg = pTrk->GetPidHypo();
 
    auto Fill = [&p, &pt, &trkInfo](RefHistogramsPrt& h) -> void {
      h.beta->Fill(trkInfo.beta);
      h.dca_rt_origin->Fill(trkInfo.dca);
      h.p->Fill(p);
      h.pt->Fill(pt);
    };
 
    //* Fill histograms before cuts on the track
    Fill(fhRefTracksBefore.all);
    ETrackPid trkPid = ETrackPid::Undef;
    if (pdg == -211) {
      trkPid = ETrackPid::Pim;
      Fill(fhRefTracksBefore.pion);
    }
    else if (pdg == 2212) {
      trkPid = ETrackPid::P;
      Fill(fhRefTracksBefore.proton);
    }
 
    //* Select track with the same approach as in the V0FinderTask
    if (fpCutter->SelectTrack(Cutter_t::TrackProperty{                                 //
                                                      .dcaRelToOrigin = trkInfo.dca,   //
                                                      .beta           = trkInfo.beta,  //
                                                      .pid            = trkPid})) {
 
      Fill(fhRefTracksAfter.all);
      if (pdg == -211) {
        Fill(fhRefTracksAfter.pion);
      }
      else if (pdg == 2212) {
        Fill(fhRefTracksAfter.proton);
      }
    }
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
InitStatus V0FinderQa::InitQa()
{
  if (!fbDataBranchesInitialized) {
    LOG(error) << fName << ": data branches from V0FinderTask are not initialized. Make sure, that the SetDataBranches "
               << "was called from the macro";
    return kERROR;
  }
 
  auto* pFairManager = FairRootManager::Instance();
  if (!pFairManager) {
    LOG(error) << fName << ": FairRootManager not found";
    return kERROR;
  }
 
  auto InitBranch = [&](TClonesArray*& branch, const char* name) -> bool {
    branch = dynamic_cast<TClonesArray*>(pFairManager->GetObject(name));
    if (!branch) {
      LOG(error) << fName << ": branch \"" << name << "\" not found";
      return false;
    }
    return true;
  };
 
  bool bBranchesInitialized{true};
  bBranchesInitialized = InitBranch(fpBrGlobalTracks, "GlobalTrack") && bBranchesInitialized;
  if (!bBranchesInitialized) {
    return kERROR;
  }
 
  InitHistograms();
 
  return kSUCCESS;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::InitHistograms()
{
  auto InitRefPrt = [&](RefHistogramsPrt& h, const std::string& prefN, const std::string& prefT) {
    /* clang-format off */
    h.beta = MakeQaObject<TH1D>(
      Form("%s_beta", prefN.c_str()), 
      Form("#beta %s;#beta;Counts", prefT.c_str()), 
      kBbeta, kLbeta, kUbeta);
    h.dca_rt_origin = MakeQaObject<TH1D>(
      Form("%s_dca_rt_origin", prefN.c_str()),
      Form("DCA rel. to origin %s;DCA_{origin} [cm];Counts", prefT.c_str()),
      kBdca_rt_origin, kLdca_rt_origin, kUdca_rt_origin);
    h.p = MakeQaObject<TH1D>(
      Form("%s_p", prefN.c_str()), 
      Form("Momentum %s;p [GeV/c];Counts", prefT.c_str()), 
      kBp, kLp, kUp);
    h.pt = MakeQaObject<TH1D>(
      Form("%s_pt", prefN.c_str()),
      Form("Transverse momentum %s;p_{T} [GeV/c];Counts", prefT.c_str()),
      kBpt, kLpt, kUpt);
    /* clang-format on */
  };
 
  auto InitRef = [&](RefHistograms& h, const std::string& prefN, const std::string& prefT) {
    InitRefPrt(h.pion, Form("%s_pion", prefN.c_str()), Form("pions %s", prefT.c_str()));
    InitRefPrt(h.proton, Form("%s_proton", prefN.c_str()), Form("protons %s", prefT.c_str()));
    InitRefPrt(h.all, Form("%s_all", prefN.c_str()), Form("all %s", prefT.c_str()));
  };
 
  auto InitRefV0 = [&](RefHistogramsV0& h, const std::string& prefN, const std::string& prefT) {
    /* clang-format off */
    h.opening_angle    = MakeQaObject<TH1D>(
      Form("%s_opening_angle", prefN.c_str()),
      Form("Opening angle of daughters %s;#alpha [radians];Counts", prefT.c_str()),
      kBopening_angle, kLopening_angle, kUopening_angle);
    h.opening_angle_kf = MakeQaObject<TH1D>(
      Form("%s_opening_angle_kf", prefN.c_str()),
      Form("Opening angle of daughters (KFParticle) %s;#alpha [radians];Counts", prefT.c_str()),
      kBopening_angle, kLopening_angle, kUopening_angle);
    h.dca_rt_pv        = MakeQaObject<TH1D>(
      Form("%s_dca_rt_origin", prefN.c_str()),
      Form("DCA rel. to prim. vertex %s;DCA_{PV} [cm];Counts", prefT.c_str()),
      kBdca_rt_pv, kLdca_rt_pv, kUdca_rt_pv);
    h.dca_daughters    = MakeQaObject<TH1D>(
      Form("%s_dca_daughters", prefN.c_str()),
      Form("DCA between daughters %s;DCA_{daughters} [cm];Counts", prefT.c_str()),
      kBdca_daughters, kLdca_daughters, kUdca_daughters);
    h.decay_length     = MakeQaObject<TH1D>(
      Form("%s_decay_length", prefN.c_str()), 
      Form("Decay length %s;L [cm];Counts", prefT.c_str()),
      kBdecay_length, kLdecay_length, kUdecay_length);
    h.decay_vtx_z      = MakeQaObject<TH1D>(
      Form("%s_decay_vtx_z", prefN.c_str()),
      Form("Decay vertex z-coordinate %s;z_{v0-vertex} [cm];Counts", prefT.c_str()),
      kBdecay_vtx_z, kLdecay_vtx_z, kUdecay_vtx_z);
    h.decay_vtx_xy     = MakeQaObject<TH2D>(
      Form("%s_decay_vtx_xy", prefN.c_str()),
      Form("Decay vertex in xy plane %s;x_{v0-vertex} [cm];y_{v0-vertex} [cm];Counts", prefT.c_str()), 
      kBdecay_vtx_x, kLdecay_vtx_x, kUdecay_vtx_x, kBdecay_vtx_y, kLdecay_vtx_y, kUdecay_vtx_y);
    h.chi2ndf_topo     = MakeQaObject<TH1D>(
      Form("%s_chi2ndf_topo", prefN.c_str()),
      Form("#chi^{2}/NDF to PV %s; #chi^2/NDF;Counts", prefT.c_str()), 
      kBchi2ndf_topo, kLchi2ndf_topo, kUchi2ndf_topo);
    h.chi2ndf_geo      = MakeQaObject<TH1D>(
      Form("%s_chi2ndf_geo", prefN.c_str()),
      Form("#chi^{2}/NDF to origin %s; #chi^2/NDF;Counts", prefT.c_str()),
      kBchi2ndf_geo, kLchi2ndf_geo, kUchi2ndf_geo);
    h.mass             = MakeQaObject<TH1D>(
      Form("%s_mass", prefN.c_str()),
      Form("Invariant mass %s; m_{inv} [GeV/c^{2}];Counts", prefT.c_str()), 
      kBmass, kLmass, kUmass);
    /* clang-format on */
  };
 
  InitRef(fhRefTracksBefore, "trk_bef", "bef. track selection");
  InitRef(fhRefTracksAfter, "trk_aft", "aft. track selection");
  InitRef(fhRefSecondaryStored, "stored_prt", "sec. stored");
  InitRef(fhRefSecondaryUnused, "unused_prt", "sec. not attached to v0");
  InitRef(fhRefDaughtersBefore, "daughters_bef", "bef. v0 selection");
  InitRef(fhRefDaughtersAfter, "daughters_aft", "aft. v0 selection");
  InitRefV0(fhRefV0Before, "v0_bef", "v0 bef. selection");
  InitRefV0(fhRefV0After, "v0_aft", "v0 aft. selection");
  InitRefV0(fhRefMeBefore, "me_bef", "mixed-event bef. selection");
  InitRefV0(fhRefMeAfter, "me_aft", "mixed-event aft. selection");
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::Refit(const XVector4_t& aX, const XVector4_t& bX, XVector4_t& mX, PVector4_t& aP, PVector4_t& bP,
                       PVector4_t& mP)
{
  constexpr double MaxDifferenceV0Tof{0.001};
  constexpr size_t MaxNofIterations{25};
  double v0DecayLength{mX.r()};
  double v0Tof{mX.t()};
  double v0TofLast{100.};  // Some large initial time
  size_t iIter{0};         // Iterations
 
  auto aR = aX.Vect() - mX.Vect();  // direction of daughter A
  auto bR = bX.Vect() - mX.Vect();  // direction of daughter B
 
  auto ShiftMomentum = [&mX, &mP, &iIter](double tof, const Vector3_t& R, PVector4_t& P) {
    constexpr double MaxBeta{0.9999};
    double beta  = std::min(R.r() / ((tof - mX.t()) * V0FinderTask::kSpeedOfLight), MaxBeta);
    double mass  = P.M();
    double gamma = 1. / std::sqrt(1 - beta * beta);
    double mom   = beta * gamma * mass;
    Vector3_t p  = R * (mom / R.r());
    P            = PVector4_t(p.x(), p.y(), p.z(), std::sqrt(mom * mom + mass * mass));
    LOG(info) << "iter: " << iIter << ", mass=" << mass << ", beta=" << beta << ", mom=" << mom << ", v0R=" << mX.r();
  };
 
  while (std::fabs(v0Tof - v0TofLast) > MaxDifferenceV0Tof && iIter <= MaxNofIterations) {
    ShiftMomentum(aX.T(), aR, aP);
    ShiftMomentum(bX.T(), bR, bP);
    mP        = aP + bP;
    v0TofLast = v0Tof;
    v0Tof     = v0DecayLength * mP.E() / (V0FinderTask::kSpeedOfLight * mP.P());
    LOG(info) << "iter: " << iIter << ", v0Tof=" << v0TofLast << ", mom=" << mP.P() << ", mass=" << mP.M();
    ++iIter;
  }
 
  mX.SetE(v0Tof);
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderQa::SetFrom(const V0FinderTask* pFinder)
{
  if (!pFinder) {
    LOG(fatal) << fName << ": a nullptr V0FinderTask instance is passed to InitData() function";
  }
 
  fpTopoReconstructor = pFinder->GetTopoReconstructor();
  if (fpTopoReconstructor == nullptr) {
    LOG(fatal) << fName << ": the topology reconstructor is not defined in the finder task";
  }
 
  fpCutter = pFinder->GetCutter();
  if (fpCutter == nullptr) {
    LOG(fatal) << fName << ": the V0-cutter class is not defined in the finder task";
  }
 
  fpTrackInfo               = &pFinder->GetTrackInfo();
  fpParticleInfo            = &pFinder->GetParticleInfo();
  fbDataBranchesInitialized = true;
}