CbmKFV0FinderTask.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 "CbmKFV0FinderTask.h"
 
#include "CbmEvent.h"
#include "CbmEventTriggers.h"
#include "CbmGlobalTrack.h"
#include "CbmKfTarget.h"
#include "CbmStsHit.h"
#include "CbmStsTrack.h"
#include "CbmTofAddress.h"
#include "CbmTofHit.h"
#include "CbmTofTrack.h"
#include "CbmTrdHit.h"
#include "CbmTrdTrack.h"
#include "CbmYaml.h"
#include "FairRunAna.h"
#include "KFPTrackVector.h"
#include "KFVertex.h"
#include "Logger.h"
#include "TClonesArray.h"
#include "TMatrixTSym.h"
#include "algo/kfp/KfpV0FinderConfig.h"
#include "algo/kfp/KfpV0FinderCutter.h"
 
#include <boost/filesystem.hpp>
 
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <limits>
#include <sstream>
 
// Log macros:
#define ERR_ LOG(error) << fName << ": "
#define WRN_ LOG(warning) << fName << ": "
#define LOG_(SEVERITY, VERBOSITY) LOG_IF(SEVERITY, fVerbose >= VERBOSITY) << fName << ": "
 
using cbm::algo::kfp::ETrackPid;
using cbm::algo::kfp::EV0Type;
using cbm::algo::kfp::V0FinderCutter;
using cbm::kfp::V0FinderTask;
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderTask::ApplyConfiguration()
{
  namespace fs        = boost::filesystem;
  namespace yml       = cbm::algo::yaml;
  fs::path configPath = fsConfigName.Data();
  if (!fs::exists(configPath)) {
    std::stringstream msg;
    msg << fName << ": configuration file " << configPath.string() << " does not exist";
    throw std::runtime_error(msg.str());
  }
 
  LOG_(info, 1) << "applying configuration from " << configPath.string();
  auto config = yml::ReadFromFile<cbm::algo::kfp::V0FinderConfig>(configPath);
 
  LOG_(info, 1) << config.ToString();
 
 
  // -- Read the config
  if (config.reconstructPdg != 3122) {  // At the moment only Lambda analysis is possible
    std::stringstream msg;
    msg << fName << ": at the moment only lambda finding is possible. Provided PDG: " << config.reconstructPdg;
    throw std::runtime_error(msg.str());
  }
 
  // Check daughter particles:
  auto& particles = config.cuts.particles;
  int iPion       = -1;
  int iProton     = -1;
  for (int iPart = 0; iPart < int(particles.size()); ++iPart) {
    const auto& particle = particles[iPart];
    if (particle.pdg == -211) {
      if (iPion == -1) {
        iPion = iPart;
      }
      else {
        std::stringstream msg;
        msg << fName << ": pion entry is defined more then one time in the config.cuts.particles";
        throw std::runtime_error(msg.str());
      }
    }
    else if (particle.pdg == 2212) {
      if (iProton == -1) {
        iProton = iPart;
      }
      else {
        std::stringstream msg;
        msg << fName << ": proton entry is defined more then one time in the config.cuts.particles";
        throw std::runtime_error(msg.str());
      }
    }
  }
  if (iProton == -1 || iPion == -1) {
    std::stringstream msg;
    msg << fName << ": config cuts/particles: either pion or proton settings are not found";
    throw std::runtime_error(msg.str());
  }
 
  const auto& pion{particles[iPion]};
  const auto& proton{particles[iProton]};
 
  SetMinPionDca(pion.minDca);
  SetMinProtonDca(proton.minDca);
 
  SetTzeroOffset(config.tZeroOffset);
  SetQpAssignedUncertainty(config.qpAssignedUncertainty);
  AddDecayToReconstructionList(config.reconstructPdg);  // Lambda
 
  fPrimaryAssignedPdg = config.primaryAssignedPdg;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
template<int PdgCode>
void V0FinderTask::AssignMomentum(CbmGlobalTrack* pTrack, double beta)
{
  /* clang-format off */
  static constexpr double qm = ( PdgCode == -211 ?  -1. / kPionMass   : 
                               ( PdgCode == 2212 ?   1. / kProtonMass :  
                                 std::numeric_limits<double>::quiet_NaN()  ));
  /* clang-format on */
 
  // NOTE:  p = gamma * beta * m
  //        q/p  = q/m / (gamma * beta) = q/m * ((1/beta)^2 - 1)
  double inverseBeta{1. / beta};
  double qp{qm * std::sqrt(inverseBeta * inverseBeta - 1)};
  double qpVariance{qp * qp * fQpAssignedUncertainty * fQpAssignedUncertainty};
 
  // Update the track with the PID hypothesis and momentum
  AssignMomentum(pTrack, qp, qpVariance);
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderTask::AssignMomentum(CbmGlobalTrack* pTrack, double qp, double qpVar)
{
  // Update the track with the PID hypothesis and momentum
  FairTrackParam parFst(*pTrack->GetParamFirst());
  FairTrackParam parLst(*pTrack->GetParamLast());
 
  parFst.SetQp(qp);
  parLst.SetQp(qp);
  parFst.SetCovariance(4, 4, qpVar);
  parFst.SetCovariance(4, 4, qpVar);
 
  pTrack->SetParamFirst(&parFst);
  pTrack->SetParamLast(&parLst);
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
template<bool UseEvent>
bool V0FinderTask::ProcessEvent(const CbmEvent* pEvent)
{
  int nTracks = UseEvent ? pEvent->GetNofData(ECbmDataType::kGlobalTrack) : fpBrGlobalTracks->GetEntriesFast();
  fCounters[ECounter::TracksTotal] += nTracks;
 
 
  // Local event counters
  int nProtons{0};
  int nPions{0};
  int nTracksSelected{0};
 
  // Local containers (per event)
  std::vector<const CbmGlobalTrack*> vpSelectedTracks;  // Selected tracks [n selected tracks]
  vpSelectedTracks.reserve(nTracks);
  std::vector<int> vSelectedTrackIds;  // Indices of selected tracks [n selected tracks]
  vSelectedTrackIds.reserve(nTracks);
 
  // ----- Shift TOF hit times to t0
  if constexpr (UseEvent) {
    auto t0 = ShiftTofHitsToTzero(pEvent);
    if (std::isnan(t0)) {
      ++fCounters[ECounter::EventsWoTzero];
      return false;
    }
  }
 
  // ----- Select tracks
  for (int iTrkEvt{0}; iTrkEvt < nTracks; ++iTrkEvt) {
    const int iTrk = UseEvent ? pEvent->GetIndex(ECbmDataType::kGlobalTrack, iTrkEvt) : iTrkEvt;
    auto* pGlobalTrack{static_cast<CbmGlobalTrack*>(fpBrGlobalTracks->At(iTrk))};
    auto& trackInfo{fvTrackInfo[iTrk]};
 
    //* Define track extra information: beta, dca to origin and PID hypo
    trackInfo     = CollectTrackInfo(pGlobalTrack);
    auto trackPid = InferTrackPidTopo(trackInfo.dca);
 
    // NOTE: define PID hypothesis before applying cuts. It is important for QA
    pGlobalTrack->SetPidHypo(cbm::algo::kfp::ToPdg(trackPid));
    AssignMomentum(pGlobalTrack, std::numeric_limits<double>::quiet_NaN(), 0.);  // Undefined momentum for each track
 
    //* Select track
    if (!fpCutter->SelectTrack(
          Cutter_t::TrackProperty{.dcaRelToOrigin = trackInfo.dca, .beta = trackInfo.beta, .pid = trackPid})) {
      continue;
    }
    ++nTracksSelected;
 
 
    //* Assign momentum to the global track
    switch (pGlobalTrack->GetPidHypo()) {
      case -211:  // pions
        AssignMomentum<-211>(pGlobalTrack, trackInfo.beta);
        ++nPions;
        break;
      case 2212:  // protons
        AssignMomentum<2212>(pGlobalTrack, trackInfo.beta);
        ++nProtons;
        break;
    }
 
    vpSelectedTracks.push_back(pGlobalTrack);
    vSelectedTrackIds.push_back(iTrk);
  }  // END LOOP: tracks in event
 
 
  fCounters[ECounter::TracksSelected] += nTracksSelected;
  fCounters[ECounter::Pions] += nPions;
  fCounters[ECounter::Protons] += nProtons;
 
  if (nProtons < 1 || nPions < 1) {
    return false;  // Nothing to search for in the event
  }
 
  ++fCounters[ECounter::EventsLambdaCand];
 
  // NOTE: The variable is used to assign a KfParticle to a corresponding primary vertex:
  //        a) pvChi2 < 3 => primary particle
  //        b) pvChi2 > 3 => secondary particle
  // Here we assign all particles as secondary (TODO: test and try to separate primary and secondary particles)
  std::vector<float> vChi2ToPv;             // Chi-square to primary vertex (NOTE: not used)
  vChi2ToPv.resize(nTracksSelected, 999.);  // Assign all tracks as secondary
 
  KFPTrackVector kfpTracksFst = MakeKfpTrackVector(vpSelectedTracks, vSelectedTrackIds, vChi2ToPv, true);
  KFPTrackVector kfpTracksLst = MakeKfpTrackVector(vpSelectedTracks, vSelectedTrackIds, vChi2ToPv, false);
 
  // ----- Reconstruct topology in the event
  fpTopoReconstructorEvent->Clear();
  fpTopoReconstructorEvent->Init(kfpTracksFst, kfpTracksLst);
  if (EPvUsageMode::Target == fPvUsageMode) {
    float x{float(fpOrigin->GetX())};
    float y{float(fpOrigin->GetY())};
    float z{float(fpOrigin->GetZ())};
    // NOTE: the target size can be applied as a covariance matrix, now it has all nulls
    fpTopoReconstructorEvent->AddPV(MakeKfpPrimaryVertex(x, y, z), std::vector<int>{});
  }
  fpTopoReconstructorEvent->SortTracks();
  fpTopoReconstructorEvent->ReconstructParticles();
 
  // ----- Count number of found lambda-candidates
  int nLambda = 0;
  //std::count_if(particles.begin(), particles.end(), [](const auto& p) { return p.GetPDG() == 3122; });
  for (const auto& particle : fpTopoReconstructorEvent->GetParticles()) {
    if (particle.GetPDG() == 3122) {
      ++nLambda;
    }
  }
  fCounters[ECounter::KfpLambdaCandidates] += nLambda;
  if (nLambda > 0) {
    auto& triggers = (*fpBrEventTriggers)[pEvent ? pEvent->GetNumber() : 0];
    triggers.Set(CbmEventTriggers::ETrigger::Lambda);
    ++fCounters[ECounter::KfpEventsLambdaCand];
  }
 
  // ----- Store particles to the run topology
  StoreParticles(pEvent);
 
  return true;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
bool V0FinderTask::CheckTrackParam(const FairTrackParam* pParam)
{
  bool bOk{true};
  bOk &= std::isfinite(pParam->GetX());
  bOk &= std::isfinite(pParam->GetY());
  bOk &= std::isfinite(pParam->GetZ());
  bOk &= std::isfinite(pParam->GetTx());
  bOk &= std::isfinite(pParam->GetTy());
  bOk &= std::isfinite(pParam->GetQp());
  std::array<double, 15> covMatrix = {0.};
  if (bOk) {
    for (int i = 0, iCov = 0; i < 5; ++i) {
      for (int j = i; j < 5; j++, iCov++) {
        // NOTE: In the FairTrackParam::GetCovariance(i, j) one can exchange the i and j indices, but for more
        //       efficient calculation it is recommended to keep i < j.
        covMatrix[iCov] = pParam->GetCovariance(i, j);
        bOk &= std::isfinite(covMatrix[iCov]);
      }
    }
  }
  return bOk;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
V0FinderTask::TrackInfo V0FinderTask::CollectTrackInfo(const CbmGlobalTrack* pTrack) const
{
  // NOTE: No cuts are applied on this step
  // NOTE: isnan(res.dca)  <==> not enough STS hits attached to the track
  // NOTE: isnan(res.beta) <==> no TOF hits attached to the track
 
  TrackInfo res;
  //* Get distance of closest approach to origin
  const int iStsTrk = pTrack->GetStsTrackIndex();
  if (iStsTrk > -1) {
    const auto* pStsTrack{static_cast<const CbmStsTrack*>(fpBrStsTracks->At(iStsTrk))};
    res.dca = EstimateDcaToOrigin(pStsTrack).fAbs;
  }
 
  //* Get beta
  const int iTofTrk = pTrack->GetTofTrackIndex();
  if (iTofTrk > -1) {
    const auto* pTofTrack{static_cast<const CbmTofTrack*>(fpBrTofTracks->At(iTofTrk))};
    assert(pTofTrack);
 
    int nTofHits = pTofTrack->GetNofTofHits();
    if (nTofHits > 0) {
      int iLstTofHit = pTofTrack->GetTofHitIndex(nTofHits - 1);  // NOTE: maybe not guarantied?
      const auto* pLstTofHit{static_cast<CbmTofHit*>(fpBrTofHits->At(iLstTofHit))};
      assert(pLstTofHit);
      res.beta    = EstimateBeta(pLstTofHit);  // NOTE: can be < 0 or > 1 due to TOF/BMON miscalibration
      res.pTofHit = pLstTofHit;
    }
  }
 
  return res;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
V0FinderTask::DcaVector V0FinderTask::EstimateDcaToOrigin(const CbmStsTrack* pStsTrack) const
{
  DcaVector res;
  if (pStsTrack->GetNofStsHits() < 2) {
    // Too few STS hits
    return res;
  }
  const auto* pHitFst{static_cast<const CbmStsHit*>(fpBrStsHits->At(pStsTrack->GetStsHitIndex(0)))};  // first hit
  const auto* pHitSnd{static_cast<const CbmStsHit*>(fpBrStsHits->At(pStsTrack->GetStsHitIndex(1)))};  // second hit
  assert(pHitFst);
  assert(pHitSnd);
  double factor{(pHitFst->GetZ() - fpOrigin->GetZ()) / (pHitSnd->GetZ() - pHitFst->GetZ())};
  res.fX   = pHitFst->GetX() - fpOrigin->GetX() - factor * (pHitSnd->GetX() - pHitFst->GetX());
  res.fY   = pHitFst->GetY() - fpOrigin->GetY() - factor * (pHitSnd->GetY() - pHitFst->GetY());
  res.fAbs = std::sqrt(res.fX * res.fX + res.fY * res.fY);
  res.fX /= res.fAbs;
  res.fY /= res.fAbs;
  return res;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
double V0FinderTask::V0FinderTask::EstimateBeta(const CbmTofHit* pTofHit) const
{
  double x{pTofHit->GetX() - fpOrigin->GetX()};
  double y{pTofHit->GetY() - fpOrigin->GetY()};
  double z{pTofHit->GetZ() - fpOrigin->GetZ()};
  double x2{x * x};
  double y2{y * y};
  double z2{z * z};
  double r2{x2 + y2 + z2};
  return std::sqrt(r2) / (pTofHit->GetTime() * kSpeedOfLight);
 
  //* Uncertainty estimation: at the moment not used, so commented out
  //double r4{r2 * r2};
  //double t2{pTofHit->GetTime() * pTofHit->GetTime()};
  // FIXME: In general case, the x, y, z coordinates of the origin can have non-zero covariances, which should be
  //        accounted in the calculation.
  //double xVar{pTofHit->GetDx() * pTofHit->GetDx() + fpOrigin->GetCovariance(0, 0)};
  //double yVar{pTofHit->GetDy() * pTofHit->GetDy() + fpOrigin->GetCovariance(1, 1)};
  //double zVar{pTofHit->GetDz() * pTofHit->GetDz() + fpOrigin->GetCovariance(2, 2)};
  //double tVar{pTofHit->GetTimeError() * pTofHit->GetTimeError()};
  //double factor{(x2 * xVar + y2 * yVar + z2 * zVar) / r4 + tVar / t2};
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderTask::Exec(Option_t*)
{
  int nTracks = fpBrGlobalTracks->GetEntriesFast();
  fvTrackDca.clear();
  fvTrackDca.resize(nTracks);
  fvTrackInfo.clear();
  fvTrackInfo.resize(nTracks);
  fvParticleInfo.clear();
 
  // ----- Reset data containers per timeslice
  fpTopoReconstructorTs->Clear();
 
  // ----- Process timeslice
  if (EProcessingMode::EventBased == fProcessingMode) {
    ++fCounters[ECounter::EventsTotal];
    fpBrEventTriggers->clear();
    fpBrEventTriggers->resize(1);
    ProcessEvent</*UseEvent = */ false>(nullptr);
  }
  else {
    const int nEvents{fpBrRecoEvents->GetEntriesFast()};
 
    fCounters[ECounter::EventsTotal] += nEvents;
    fpBrEventTriggers->clear();
    fpBrEventTriggers->resize(nEvents);
    for (int iEvent = 0; iEvent < nEvents; ++iEvent) {
      const auto* pEvent = static_cast<CbmEvent*>(fpBrRecoEvents->At(iEvent));
      ProcessEvent</*UseEvent = */ true>(pEvent);
    }
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderTask::Finish()
{
  // ----- Counters summary log
  std::stringstream msg;
  msg << "*** Counters summary ***\n";
  msg << "\n  tracks total:                         " << fCounters[ECounter::TracksTotal];
  msg << "\n  tracks passed to KFP-finder:          " << fCounters[ECounter::TracksSelected];
  //msg << "\n  tracks w/ infinit parameters:         " << fCounters[ECounter::TracksInfiniteParam];
  //msg << "\n  tracks w/o TOF hits:                  " << fCounters[ECounter::TracksWoTofHits];
  //msg << "\n  tracks w/ last TOF hit having t<0:    " << fCounters[ECounter::TracksWNegativeTofHitTime];
  //msg << "\n  tracks w/o STS hits:                  " << fCounters[ECounter::TracksWoStsHits];
  //msg << "\n  tracks w/o PID:                       " << fCounters[ECounter::TracksWoPid];
  //msg << "\n  tracks w/o momentum:                  " << fCounters[ECounter::TracksWoMomentum];
  //msg << "\n  tracks w/ beta > 1                    " << fCounters[ECounter::TracksWithUnphysicalBeta];
  msg << "\n  pion candidates:                      " << fCounters[ECounter::Pions];
  msg << "\n  proton candidates:                    " << fCounters[ECounter::Protons];
  msg << "\n  events total:                         " << fCounters[ECounter::EventsTotal];
  msg << "\n  events w/o t-zero:                    " << fCounters[ECounter::EventsWoTzero];
  msg << "\n  events w/ potential Lambda-candidate: " << fCounters[ECounter::EventsLambdaCand];
  msg << "\n  events w/ Lambda-candidate from KFPF: " << fCounters[ECounter::KfpEventsLambdaCand];
  msg << "\n  Lambda-candidates:                    " << fCounters[ECounter::KfpLambdaCandidates];
  LOG_(info, 1) << msg.str();
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
cbm::algo::kfp::ETrackPid V0FinderTask::InferTrackPidTopo(double dca) const
{
  if (std::isnan(dca)) {
    return ETrackPid::Undef;  // Track PID cannot be infered
  }
  else if (dca > fMinPionDca) {
    return ETrackPid::Pim;
  }
  else if (dca > fMinProtonDca) {
    return ETrackPid::P;
  }
  else {
    return ETrackPid::Undef;  // Primary: PID is undefined
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
InitStatus V0FinderTask::Init()
{
  LOG_(info, 1) << "initializing the task ... ";
 
  LOG_(info, 1) << "using the following configuration:\n  processing mode: " << ToString(fProcessingMode)
                << "\n  PID approach: " << ToString(fPidApproach) << "\n  " << ToString(fPvUsageMode);
 
  // Temporary not implemented:
  if (EPidApproach::Mc == fPidApproach || EPvUsageMode::Mc == fPvUsageMode
      || EPvUsageMode::Reconstructed == fPvUsageMode || EPvUsageMode::ReconstructSingle == fPvUsageMode
      || EPvUsageMode::ReconstructMultiple == fPvUsageMode) {
    ERR_ << "Desirable configuration of PID or PV handling was not implemented yet";
    return kFATAL;
  }
 
  if (!fsConfigName.IsNull()) {
    try {
      ApplyConfiguration();
    }
    catch (const std::exception& err) {
      ERR_ << "configuration from a config was required, but failed. Reason: " << err.what();
      return kFATAL;
    }
  }
 
  if (fsCutsConfigName.IsNull()) {
    ERR_ << "cut configuration file was not defined";
    return kFATAL;
  }
  else {
    try {
      fpCutter->ReadCutsFromYaml(fsCutsConfigName.Data());
    }
    catch (const std::exception& err) {
      ERR_ << "cut configuration failed. Reason: " << err.what();
      return kFATAL;
    }
  }
 
  LOG_(info, 1) << fpCutter->ToString();
  fpCutter->ApplyCutsToKFParticleFinder(GetKFParticleFinder());
 
  // ----- Input data branches initialization
  const auto* pTarget = kf::Target::Instance();  // CBM target info
 
  auto* pFairManager = FairRootManager::Instance();
  if (!pFairManager) {
    ERR_ << "FairRootManager not found";
    return kERROR;
  }
 
  auto InitBranch = [&](TClonesArray*& branch, const char* name) -> bool {
    branch = dynamic_cast<TClonesArray*>(pFairManager->GetObject(name));
    if (!branch) {
      ERR_ << "branch \"" << name << "\" not found";
      return false;
    }
    return true;
  };
 
  bool bBranchesInitialized{true};
  if (EProcessingMode::TimeBased == fProcessingMode) {
    bBranchesInitialized = InitBranch(fpBrRecoEvents, "CbmEvent") && bBranchesInitialized;
  }
  bBranchesInitialized = InitBranch(fpBrGlobalTracks, "GlobalTrack") && bBranchesInitialized;
  bBranchesInitialized = InitBranch(fpBrStsTracks, "StsTrack") && bBranchesInitialized;
  bBranchesInitialized = InitBranch(fpBrTrdTracks, "TrdTrack") && bBranchesInitialized;
  bBranchesInitialized = InitBranch(fpBrTofTracks, "TofTrack") && bBranchesInitialized;
  bBranchesInitialized = InitBranch(fpBrStsHits, "StsHit") && bBranchesInitialized;
  bBranchesInitialized = InitBranch(fpBrTrdHits, "TrdHit") && bBranchesInitialized;
  bBranchesInitialized = InitBranch(fpBrTofHits, "TofHit") && bBranchesInitialized;
 
  if (EPvUsageMode::Reconstructed == fPvUsageMode) {
    fpBrPrimaryVertex = dynamic_cast<CbmVertex*>(pFairManager->GetObject("PrimaryVertex."));
    if (!fpBrPrimaryVertex) {
      ERR_ << "branch \"PrimaryVertex.\" not found";
      bBranchesInitialized = false;
    }
  }
  else if (EPvUsageMode::Target == fPvUsageMode) {
    double x{pTarget->GetX()};
    double y{pTarget->GetY()};
    double z{pTarget->GetZ()};
    LOG_(info, 1) << "using target center as origin: r = (" << x << ", " << y << ", " << z << ") [cm]";
    TMatrixFSym covMatrix(3);
    covMatrix(1, 0) = 0.;
    covMatrix(2, 0) = 0.;
    covMatrix(2, 1) = 0.;
    //double transverseVariance{pTarget->GetRmax() * pTarget->GetRmax() / 12.};
    //covMatrix(0, 0) = transverseVariance;
    //covMatrix(1, 1) = transverseVariance;
    //covMatrix(2, 2) = pTarget->GetDz() * pTarget->GetDz() / 3.;
    covMatrix(0, 0) = 0.;
    covMatrix(1, 1) = 0.;
    covMatrix(2, 2) = 0.;
    fpOrigin->SetVertex(x, y, z, 0., 1, 0, covMatrix);
  }
 
  if (!bBranchesInitialized) {
    return kERROR;
  }
 
  // ----- Output branches initialization
  fpBrEventTriggers = new std::vector<CbmEventTriggers>();
  if (pFairManager->GetObject("CbmEventTriggers")) {
    LOG(error) << "Branch \"CbmEventTriggers\" already exists!";
    return kFATAL;
  }
  pFairManager->RegisterAny("CbmEventTriggers", fpBrEventTriggers, true);
  LOG_(info, 1) << "registering branch \"CbmEventTriggers\"";
 
  // ----- Topology reconstructor initialization
  fpTopoReconstructorTs->SetTarget({float(pTarget->GetX()), float(pTarget->GetY()), float(pTarget->GetZ())});
  fpTopoReconstructorEvent->SetTarget(fpTopoReconstructorTs->GetTargetPosition());
  fpTopoReconstructorEvent->CopyCuts(fpTopoReconstructorTs.get());
  fpTopoReconstructorEvent->GetKFParticleFinder()->SetReconstructionList(
    fpTopoReconstructorTs->GetKFParticleFinder()->GetReconstructionList());
 
  // ----- Auxilary variables initialization
  std::fill(fCounters.begin(), fCounters.end(), 0);
 
  LOG_(info, 1) << "initializing the task ... done";
  return kSUCCESS;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
KFPTrackVector V0FinderTask::MakeKfpTrackVector(const std::vector<const CbmGlobalTrack*>& vpTracks,
                                                const std::vector<int>& vTrackIds, const std::vector<float>& vChi2ToPv,
                                                bool bAtFirstPoint) const
{
  // TODO: Cross-check with the CbmKFParticleFinder results
  // NOTE: Actually tracks themselves are not needed here, so one could fill and pass a vector of FairTrackParam*,
  //       put in KFPTrackVector there is an option to set number of Pixel hits, so we will keep the signature of
  //       the function as it is.
  KFPTrackVector trackVector;
  int nTracks = vpTracks.size();
  trackVector.Resize(nTracks);
  for (int iTrk = 0; iTrk < nTracks; ++iTrk) {
    const auto* pTrack{vpTracks[iTrk]};
    const auto* pParam{bAtFirstPoint ? pTrack->GetParamFirst() : pTrack->GetParamLast()};
 
    // ----- Parameters definition
    int pdg{pTrack->GetPidHypo()};
    const double& tx{pParam->GetTx()};
    const double& ty{pParam->GetTy()};
    const double& qp{pParam->GetQp()};
    int q{qp > 0. ? 1 : -1};
    if (std::fabs(pdg) == 10000020030 || std::fabs(pdg == 1000020040)) {  // He3 and He4
      // NOTE: at the moment not called, but never the less leave it here for the future pid procedures
      q *= 2;
    }
    double p{q / qp};
    double p2{p * p};
    double t2inv{1. / (1. + tx * tx + ty * ty)};
    double pz{std::sqrt(t2inv * p2)};
    double px{tx * pz};
    double py{ty * pz};
 
    trackVector.SetParameter(pParam->GetX(), 0, iTrk);
    trackVector.SetParameter(pParam->GetY(), 1, iTrk);
    trackVector.SetParameter(pParam->GetZ(), 2, iTrk);
    trackVector.SetParameter(px, 3, iTrk);
    trackVector.SetParameter(py, 4, iTrk);
    trackVector.SetParameter(pz, 5, iTrk);
 
    // Jacobian matrix for (tx, ty, qp) -> (px, py, pz)
    std::array<std::array<double, 3>, 3> Jp;
    Jp[2][0] = -t2inv * px;         // d(pz)/d(tx)
    Jp[2][1] = -t2inv * py;         // d(pz)/d(ty)
    Jp[2][2] = -pz / qp;            // d(pz)/d(qp)
    Jp[0][0] = tx * Jp[2][0] + pz;  // d(px)/d(tx)
    Jp[0][1] = tx * Jp[2][1];       // d(px)/d(ty)
    Jp[0][2] = tx * Jp[2][2];       // d(px)/d(qp)
    Jp[1][0] = ty * Jp[2][0];       // d(py)/d(tx)
    Jp[1][1] = ty * Jp[2][1] + pz;  // d(py)/d(ty)
    Jp[1][2] = ty * Jp[2][2];       // d(py)/d(qp)
 
    // ----- Covariance matrix definition
    // Position covariances
    trackVector.SetCovariance(pParam->GetCovariance(0, 0), 0, iTrk);  // var(x)
    trackVector.SetCovariance(pParam->GetCovariance(0, 1), 1, iTrk);  // cov(x, y)
    trackVector.SetCovariance(pParam->GetCovariance(1, 1), 2, iTrk);  // var(y)
 
    // Momentum-position covariances
    auto MomPosCovariance = [&](const int k, const int l) constexpr->double
    {
      double val{0.};
      const auto& JpA = Jp[k];
      for (int i = 0; i < 3; ++i) {
        val += JpA[i] * pParam->GetCovariance(i + 2, l);
      }
      return val;
    };
    trackVector.SetCovariance(MomPosCovariance(0, 0), 6, iTrk);   // cov(x, px)
    trackVector.SetCovariance(MomPosCovariance(0, 1), 7, iTrk);   // cov(y, px)
    trackVector.SetCovariance(MomPosCovariance(1, 0), 10, iTrk);  // cov(x, py)
    trackVector.SetCovariance(MomPosCovariance(1, 1), 11, iTrk);  // cov(y, py)
    trackVector.SetCovariance(MomPosCovariance(2, 0), 15, iTrk);  // cov(x, pz)
    trackVector.SetCovariance(MomPosCovariance(2, 1), 16, iTrk);  // cov(y, pz)
 
    // Momentum covariances
    auto MomentumCovariance = [&](const int k, const int l) constexpr->double
    {
      double val{0.};
      const auto& JpA = Jp[k];
      const auto& JpB = Jp[l];
      for (int i = 0; i < 3; ++i) {
        double factor{0.};
        for (int j = 0; j < 3; ++j) {
          factor += JpB[j] * pParam->GetCovariance(i + 2, j + 2);
        }
        val += JpA[i] * factor;
      }
      return val;
    };
    trackVector.SetCovariance(MomentumCovariance(0, 0), 9, iTrk);   // var(px)
    trackVector.SetCovariance(MomentumCovariance(1, 0), 13, iTrk);  // cov(px, py)
    trackVector.SetCovariance(MomentumCovariance(1, 1), 14, iTrk);  // var(py)
    trackVector.SetCovariance(MomentumCovariance(2, 0), 18, iTrk);  // cov(px, pz)
    trackVector.SetCovariance(MomentumCovariance(2, 1), 19, iTrk);  // cov(py, pz)
    trackVector.SetCovariance(MomentumCovariance(2, 2), 20, iTrk);  // var(pz)
 
    // Zero covariances
    // NOTE: from the tracking point of view a z-coordinate in known precisely, so all the corresponding covariances
    //       should be set to null.
    trackVector.SetCovariance(0.f, 3, iTrk);   // cov(x,z)
    trackVector.SetCovariance(0.f, 4, iTrk);   // cov(y,z)
    trackVector.SetCovariance(0.f, 5, iTrk);   // var(z)
    trackVector.SetCovariance(0.f, 8, iTrk);   // cov(z,px)
    trackVector.SetCovariance(0.f, 12, iTrk);  // cov(z,py)
    trackVector.SetCovariance(0.f, 17, iTrk);  // var(z,pz)
 
    // ----- Other parameters
    // Magnetic field
    // NOTE: In mCBM there is no magnetic field, so here we define the coefficients with zeros
    for (int iF = 0; iF < 10; ++iF) {
      trackVector.SetFieldCoefficient(0.f, iF, iTrk);
    }
 
    trackVector.SetId(vTrackIds[iTrk], iTrk);
    trackVector.SetPDG(pdg, iTrk);
    trackVector.SetQ(q, iTrk);
    trackVector.SetNPixelHits(0, iTrk);  // Number of MVD hits in track (used to define PV inside the KFPFinder)
 
    if (EPvUsageMode::Reconstructed == fPvUsageMode || EPvUsageMode::Mc == fPvUsageMode) {
      if (vChi2ToPv[iTrk] < kChi2PvPrimThrsh) {
        trackVector.SetPVIndex(0, iTrk);  // Primary track
      }
      else {
        trackVector.SetPVIndex(-1, iTrk);  // Secondary track
      }
    }
    else {
      trackVector.SetPVIndex(-1, iTrk);  // Secondary track
    }
  }
  return trackVector;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
KFVertex V0FinderTask::MakeKfpPrimaryVertex(float x, float y, float z) const
{
  KFVertex kfVertex;
  kfVertex.X() = x;
  kfVertex.Y() = y;
  kfVertex.Z() = z;
  for (int iC = 0; iC < 6; ++iC) {
    kfVertex.Covariance(iC) = 0.f;
  }
  kfVertex.Chi2() = -100.f;
  return kfVertex;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
double V0FinderTask::ShiftTofHitsToTzero(const CbmEvent* pEvent)
{
  // Define t0 from the Bmon hits
  double t0{std::numeric_limits<double>::signaling_NaN()};
  int nTofHits = pEvent->GetNofData(ECbmDataType::kTofHit);
  for (int iTofHitEvt{0}; iTofHitEvt < nTofHits; ++iTofHitEvt) {
    int iTofHit = pEvent->GetIndex(ECbmDataType::kTofHit, iTofHitEvt);
    const auto* pTofHit{static_cast<const CbmTofHit*>(fpBrTofHits->At(iTofHit))};
    //if (5 == CbmTofAddress::GetSmType(pTofHit->GetAddress())) {   // selection by the supermodule type
    if (pTofHit->GetZ() == 0) {  // Take some small z-coordinate to identify BMon (TODO: provide a flag by some task)
      t0 = pTofHit->GetTime();
    }
  }
  // NOTE: t0 must be defined for each event, since the Bmon digis are used to seed the digi event builder. Basically,
  //       the tZero must be defined as a field of CbmEvent, e.g. in a separate FairTask, or directly
  //       in CbmAlgoBuildRawEvent
  if (std::isnan(t0)) {
    return t0;
  }
 
  // Shift TOF times to found t0:
  for (int iTofHitEvt{0}; iTofHitEvt < nTofHits; ++iTofHitEvt) {
    int iTofHit = pEvent->GetIndex(ECbmDataType::kTofHit, iTofHitEvt);
    auto* pTofHit{static_cast<CbmTofHit*>(fpBrTofHits->At(iTofHit))};
    pTofHit->SetTime(pTofHit->GetTime() - t0 - fTzeroOffset);
  }
  return t0;
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
void V0FinderTask::StoreParticles(const CbmEvent* pEvent)
{
  int indexOffset = fpTopoReconstructorTs->GetParticles().size();
  for (int iPv = 0; iPv < fpTopoReconstructorEvent->NPrimaryVertices(); ++iPv) {
    std::vector<int> vPVTrackIndexArray = fpTopoReconstructorEvent->GetPVTrackIndexArray(iPv);
    std::for_each(vPVTrackIndexArray.begin(), vPVTrackIndexArray.end(), [&](auto& item) { item += indexOffset; });
    fpTopoReconstructorTs->AddPV(fpTopoReconstructorEvent->GetPrimKFVertex(iPv), vPVTrackIndexArray);
  }
 
  // Particle
  size_t nParticlesEvt{fpTopoReconstructorEvent->GetParticles().size()};
  std::vector<ParticleInfo> vParticleInfoEvt(nParticlesEvt);
  for (size_t iP = 0; iP < nParticlesEvt; ++iP) {
    const KFParticle& particleEvt{fpTopoReconstructorEvent->GetParticles()[iP]};
    KFParticle particle{particleEvt};
    particle.CleanDaughtersId();
    int nDaughters{particleEvt.NDaughters()};
    auto& particleInfo = vParticleInfoEvt[iP];
    if (nDaughters == 1) {  // Global track index
      int iTrk{particleEvt.DaughterIds()[0]};
      particle.AddDaughterId(iTrk);
      const auto& trkInfo  = fvTrackInfo[iTrk];
      particleInfo.beta    = trkInfo.beta;
      particleInfo.dca     = trkInfo.dca;
      particleInfo.pTofHit = trkInfo.pTofHit;
    }
    else if (nDaughters > 1) {
      for (int iD = 0; iD < particleEvt.NDaughters(); ++iD) {
        particle.AddDaughterId(particleEvt.DaughterIds()[iD] + indexOffset);
      }
    }
    particleInfo.eventID = pEvent->GetNumber();
    particle.SetId(particleEvt.Id() + indexOffset);
    fpTopoReconstructorTs->AddParticle(particle);
  }
  fvParticleInfo.insert(fvParticleInfo.end(), vParticleInfoEvt.begin(), vParticleInfoEvt.end());
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
std::string V0FinderTask::ToString(EProcessingMode mode)
{
  switch (mode) {
    case EProcessingMode::EventBased: return "event-based";
    case EProcessingMode::TimeBased: return "time-based";
    default: return "";
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
std::string V0FinderTask::ToString(EPidApproach pid)
{
  switch (pid) {
    case EPidApproach::Topo: return "track topology";
    case EPidApproach::Mc: return "MC-truth";
    default: return "";
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
std::string V0FinderTask::ToString(EPvUsageMode pvMode)
{
  switch (pvMode) {
    case EPvUsageMode::Reconstructed: return "reconstructed primary vertex is used";
    case EPvUsageMode::ReconstructSingle: return "a single primary vertex will be reconstructed";
    case EPvUsageMode::ReconstructMultiple: return "multiple primary vertices will be reconstructed";
    case EPvUsageMode::Target: return "target center is used for primary vertex";
    case EPvUsageMode::Mc: return "MC-true primary vertex is used";
    default: return "";
  }
}
 
// ---------------------------------------------------------------------------------------------------------------------
//
std::string V0FinderTask::ToString(const FairTrackParam* pParam)
{
  std::stringstream msg;
  msg << "Track parameters: r=(" << pParam->GetX() << ", " << pParam->GetY() << ", " << pParam->GetZ()
      << "), tx=" << pParam->GetTx() << ", ty=" << pParam->GetTy() << ", q/p=" << pParam->GetQp()
      << ", Covariance matrix:";
  for (int i = 0; i < 5; ++i) {
    msg << '\n';
    for (int j = 0; j <= i; ++j) {
      msg << std::setw(15) << pParam->GetCovariance(i, j) << ' ';
    }
  }
  return msg.str();
}