TrackingChain.cxx

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

/// \file   TrackingChain.cxx
/// \date   14.09.2023
/// \brief  A chain class to execute CA tracking algorithm in online reconstruction (implementation)
/// \author S.Zharko <s.zharko@gsi.de>

#include "TrackingChain.h"

#include "CaConstants.h"
#include "CaHit.h"
#include "CaInitManager.h"
#include "CaParameters.h"
#include "tof/Config.h"
#include "yaml/Yaml.h"

#include <boost/archive/text_oarchive.hpp>

#include <fstream>
#include <unordered_map>

#include <xpu/host.h>

using namespace cbm::algo;

using cbm::algo::TrackingChain;
using cbm::algo::ca::EDetectorID;
using cbm::algo::ca::Framework;
using cbm::algo::ca::HitTypes_t;
using cbm::algo::ca::InitManager;
using cbm::algo::ca::InputQa;
using cbm::algo::ca::OutputQa;
using cbm::algo::ca::Parameters;
using cbm::algo::ca::Track;
using cbm::algo::ca::constants::clrs::CL;   // clear text
using cbm::algo::ca::constants::clrs::GNb;  // grin bald text

// ---------------------------------------------------------------------------------------------------------------------
//
TrackingChain::TrackingChain(std::shared_ptr<HistogramSender> histoSender)
  : fInputQa(InputQa(histoSender))
  , fOutputQa(OutputQa(histoSender))
{
}

// ---------------------------------------------------------------------------------------------------------------------
//
void TrackingChain::Init()
{
  // ------ Read tracking chain parameters from the config
  fConfig = yaml::ReadFromFile<TrackingChainConfig>(Opts().ParamsDir() / "CaConfig.yaml");

  // ------ Read parameters from binary
  auto paramFile = Opts().ParamsDir() / fConfig.fsParName;
  L_(info) << "Tracking Chain: reading CA parameters file " << GNb << paramFile.string() << CL << '\n';
  auto manager = InitManager{};
  manager.ReadParametersObject(paramFile.string());
  auto parameters = manager.TakeParameters();
  L_(info) << "Tracking Chain: parameters object: \n" << parameters.ToString(1) << '\n';

  // ------ Initialize CA framework
  fCaMonitor.Reset();
  fCaFramework.SetNofThreads(fConfig.fNofThreads);
  fCaFramework.ReceiveParameters(std::move(parameters));
  fCaFramework.Init(ca::Framework::TrackingMode::kMcbm);

  // ------ Initialize QA modules
  if (fInputQa.IsSenderDefined()) {
    fInputQa.RegisterParameters(&fCaFramework.GetParameters());
    fInputQa.Init();
  }

  if (fOutputQa.IsSenderDefined()) {
    fOutputQa.RegisterParameters(&fCaFramework.GetParameters());
    fOutputQa.Init();
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
TrackingChain::Output_t TrackingChain::Run(Input_t recoResults)
{
  xpu::scoped_timer t_("CA");  // TODO: pass timings to monitoring for throughput?
  fCaMonitorData.Reset();
  fCaMonitorData.StartTimer(ca::ETimer::TrackingChain);

  // ----- Init input data ---------------------------------------------------------------------------------------------
  fCaMonitorData.StartTimer(ca::ETimer::PrepareInputData);
  this->PrepareInput(recoResults);
  fCaMonitorData.StopTimer(ca::ETimer::PrepareInputData);

  // ----- Run reconstruction ------------------------------------------------------------------------------------------
  fCaFramework.SetMonitorData(fCaMonitorData);
  fCaFramework.fTrackFinder.FindTracks();
  fCaMonitorData = fCaFramework.GetMonitorData();

  // ----- Init output data --------------------------------------------------------------------------------------------
  return PrepareOutput();
}

// ---------------------------------------------------------------------------------------------------------------------
//
void TrackingChain::Finalize()
{
  L_(info) << fCaMonitor.ToString();
  if (fConfig.fbStoreMonitor) {
    auto fileName = "./" + fConfig.fsMoniOutName;
    std::ofstream ofs(fileName);
    boost::archive::text_oarchive oa(ofs);
    oa << fCaMonitor;
  }
}

// ---------------------------------------------------------------------------------------------------------------------
//
void TrackingChain::PrepareInput(Input_t recoResults)
{
  fNofHitKeys  = 0;
  int nHitsTot = recoResults.stsHits.NElements() + recoResults.tofHits.NElements();
  L_(info) << "Tracking chain: input has " << nHitsTot << " hits";
  fCaDataManager.ResetInputData(nHitsTot);
  faHitExternalIndices.clear();
  faHitExternalIndices.reserve(nHitsTot);
  ReadHits<EDetectorID::Sts>(recoResults.stsHits);
  ReadHits<EDetectorID::Tof>(recoResults.tofHits);
  faHitExternalIndices.shrink_to_fit();
  fCaDataManager.SetNhitKeys(fNofHitKeys);
  L_(info) << "Tracking chain:" << fCaDataManager.GetNofHits() << " hits will be passed to the ca::Framework";
  fCaFramework.ReceiveInputData(fCaDataManager.TakeInputData());
}

// ---------------------------------------------------------------------------------------------------------------------
//
TrackingChain::Output_t TrackingChain::PrepareOutput()
{
  Output_t output;
  output.tracks = std::move(fCaFramework.fRecoTracks);
  int nTracks   = output.tracks.size();

  output.stsHitIndices.reset(nTracks);
  output.tofHitIndices.reset(nTracks);

  int trackFirstHit = 0;
  for (int iTrk = 0; iTrk < nTracks; ++iTrk) {
    output.stsHitIndices[iTrk].clear();
    output.tofHitIndices[iTrk].clear();
    int nHits = output.tracks[iTrk].fNofHits;
    for (int iHit = 0; iHit < nHits; ++iHit) {
      int iHitInternal = fCaFramework.GetInputData().GetHit(fCaFramework.fRecoHits[trackFirstHit + iHit]).Id();
      const auto [detID, iPartition, iPartHit] = faHitExternalIndices[iHitInternal];
      switch (detID) {
        case ca::EDetectorID::Sts: output.stsHitIndices[iTrk].push_back(std::make_pair(iPartition, iPartHit)); break;
        case ca::EDetectorID::Tof: output.tofHitIndices[iTrk].push_back(std::make_pair(iPartition, iPartHit)); break;
        default: break;
      }
    }
    fCaMonitorData.IncrementCounter(ca::ECounter::RecoStsHit, output.stsHitIndices[iTrk].size());
    fCaMonitorData.IncrementCounter(ca::ECounter::RecoTofHit, output.tofHitIndices[iTrk].size());
    trackFirstHit += nHits;
  }

  L_(info) << "TrackingChain: Timeslice contains " << fCaMonitorData.GetCounterValue(ca::ECounter::RecoTrack)
           << " tracks, with " << fCaMonitorData.GetCounterValue(ca::ECounter::RecoStsHit) << " sts hits, "
           << fCaMonitorData.GetCounterValue(ca::ECounter::RecoTofHit) << " tof hits; the FindTracks routine ran "
           << fCaMonitorData.GetTimer(ca::ETimer::FindTracks).GetTotal() << " s";


  // QA
  if (fInputQa.IsSenderDefined()) {
    fCaMonitorData.StartTimer(ca::ETimer::InputQa);
    fInputQa.RegisterInputData(&fCaFramework.GetInputData());
    fInputQa.RegisterTracks(&output.tracks);
    fInputQa.RegisterRecoHitIndices(&fCaFramework.fRecoHits);
    fInputQa.Exec();
    fCaMonitorData.StopTimer(ca::ETimer::InputQa);
  }

  if (fOutputQa.IsSenderDefined()) {
    fCaMonitorData.StartTimer(ca::ETimer::OutputQa);
    fOutputQa.RegisterInputData(&fCaFramework.GetInputData());
    fOutputQa.RegisterTracks(&output.tracks);
    fOutputQa.RegisterRecoHitIndices(&fCaFramework.fRecoHits);
    fOutputQa.Exec();
    fCaMonitorData.StopTimer(ca::ETimer::OutputQa);
  }

  fCaMonitorData.StopTimer(ca::ETimer::TrackingChain);
  fCaMonitor.AddMonitorData(fCaMonitorData);
  output.monitorData = fCaMonitorData;

  return output;
}

// ---------------------------------------------------------------------------------------------------------------------
//
template<EDetectorID DetID>
void TrackingChain::ReadHits(PartitionedSpan<const ca::HitTypes_t::at<DetID>> hits)
{
  std::ofstream out;
  if constexpr (kDEBUG) {
    out.open(std::string("./Debug_hits_") + ca::kDetName[DetID] + ".txt");
  }

  int nSt = fCaFramework.GetParameters().GetNstationsActive();

  using Hit_t           = ca::HitTypes_t::at<DetID>;
  constexpr bool IsMvd  = (DetID == EDetectorID::Mvd);
  constexpr bool IsSts  = (DetID == EDetectorID::Sts);
  constexpr bool IsMuch = (DetID == EDetectorID::Much);
  constexpr bool IsTrd  = (DetID == EDetectorID::Trd);
  constexpr bool IsTof  = (DetID == EDetectorID::Tof);

  xpu::t_add_bytes(hits.NElements() * sizeof(Hit_t));  // Assumes call from Run, for existence of timer!

  int64_t dataStreamDet         = static_cast<int64_t>(DetID) << 60;  // detector part of the data stream
  int64_t dataStream            = 0;
  for (size_t iPartition = 0; iPartition < hits.NPartitions(); ++iPartition, ++dataStream) {
    const auto& [vHits, extHitAddress] = hits.Partition(iPartition);
    // ---- Define data stream and station index
    //int64_t dataStream = dataStreamDet | extHitAddress;
    int iStLocal = -1;
    // FIXME: This definition of the station index works only for STS, and there is no any guaranty, that it will
    //        work for other mCBM setups.
    if constexpr (IsSts) {
      iStLocal = (extHitAddress >> 4) & 0xF;
    }
    if constexpr (IsTof) {
      iStLocal = cbm::algo::tof::Config::GetTofTrackingStation(extHitAddress);
      if (tof::Config::IsBmon(extHitAddress)) {
        continue;
      }  // skip hits from Bmon
    }

    int iStActive  = (iStLocal != -1) ? fCaFramework.GetParameters().GetStationIndexActive(iStLocal, DetID) : -1;
    //size_t iOffset = hits.Offsets()[iPartition];
    if (iStActive < 0) {
      continue;  // legit
    }
    if (iStActive >= nSt) {
      L_(error) << "TrackingChain: found hit with wrong active station index above the upper limit: " << iStActive
                << ", detector: " << ca::kDetName[DetID];
      continue;
    }
    double lastTime = -1e9;
    double prevTime = -1;
    ca::HitKeyIndex_t firstHitKey = fNofHitKeys;
    for (size_t iPartHit = 0; iPartHit < vHits.size(); ++iPartHit) {
      const auto& hit = vHits[iPartHit];
      //int iHitExt     = iOffset + iPartHit;
      // ---- Fill ca::Hit
      ca::Hit caHit;
      if constexpr (IsSts) {
        caHit.SetFrontKey(firstHitKey + hit.fFrontClusterId);
        caHit.SetBackKey(firstHitKey + hit.fBackClusterId);
        //L_(info) << ", hit="  << iHitExt << ", f=" << hit.fFrontClusterId << ", b=" << hit.fBackClusterId;
      }
      else {
        caHit.SetFrontKey(firstHitKey + iPartHit);
        caHit.SetBackKey(caHit.FrontKey());
      }
      caHit.SetX(hit.X());
      caHit.SetY(hit.Y());
      caHit.SetZ(hit.Z());
      caHit.SetT(hit.Time());
      caHit.SetDx2(hit.Dx() * hit.Dx() + fCaFramework.GetParameters().GetMisalignmentXsq(DetID));
      caHit.SetDy2(hit.Dy() * hit.Dy() + fCaFramework.GetParameters().GetMisalignmentYsq(DetID));
      if constexpr (IsSts) caHit.SetDxy(hit.fDxy);
      caHit.SetDt2(hit.TimeError() * hit.TimeError() + fCaFramework.GetParameters().GetMisalignmentTsq(DetID));
      /// FIXME: Define ranges from the hit, when will be available
      caHit.SetRangeX(3.5 * hit.Dx());
      caHit.SetRangeY(3.5 * hit.Dy());
      caHit.SetRangeT(3.5 * hit.TimeError());
      //L_(info) << ">>>>>>>>>>>> " << iStActive;
      caHit.SetStation(iStActive);
      caHit.SetId(fCaDataManager.GetNofHits());
      if (caHit.Check()) {
        if ((caHit.T() < lastTime - 1000.) && (dataStream < 100000)) {
          dataStream++;
        }
        lastTime = caHit.T();
        fCaDataManager.PushBackHit(caHit, dataStreamDet | dataStream);
        faHitExternalIndices.push_back(std::make_tuple(DetID, iPartition, iPartHit));
        if constexpr (kDEBUG) {
          out << (dataStreamDet | dataStream) << " ----- " << caHit.ToString() << '\n';
          if (prevTime > caHit.T()) {
            out << "TIME IS UNSORTED\n";
          }
        }

        if (fNofHitKeys <= caHit.FrontKey()) {
          fNofHitKeys = caHit.FrontKey() + 1;
        }
        if (fNofHitKeys <= caHit.BackKey()) {
          fNofHitKeys = caHit.BackKey() + 1;
        }
      }
      else {
        if constexpr (IsMvd) {
          fCaMonitorData.IncrementCounter(ca::ECounter::UndefinedMvdHit);
        }
        if constexpr (IsSts) {
          fCaMonitorData.IncrementCounter(ca::ECounter::UndefinedStsHit);
        }
        if constexpr (IsMuch) {
          fCaMonitorData.IncrementCounter(ca::ECounter::UndefinedMuchHit);
        }
        if constexpr (IsTrd) {
          fCaMonitorData.IncrementCounter(ca::ECounter::UndefinedTrdHit);
        }
        if constexpr (IsTof) {
          fCaMonitorData.IncrementCounter(ca::ECounter::UndefinedTofHit);
        }
      }
      prevTime = caHit.T();
      // ---- Update number of hit keys
    }  // iPartHit
  }    // iPartition
}

// template void TrackingChain::ReadHits<EDetectorID::Sts>(const PartitionedPODVector<HitTypes_t::at<EDetectorID::Sts>>&);