CbmKfFitTracksTask.cxx

Go to the documentation of this file.

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

 
 
#include "CbmKfFitTracksTask.h"
 
#include "CbmGlobalTrack.h"
#include "CbmKfUtil.h"
#include "CbmMuchTrack.h"
#include "CbmMvdHit.h"
#include "CbmStsHit.h"
#include "CbmStsSetup.h"
#include "CbmStsTrack.h"
#include "CbmTofTrack.h"
#include "CbmTrdTrack.h"
#include "FairRootManager.h"
#include "TClonesArray.h"
 
#include <cmath>
#include <iostream>
#include <thread>
 
using namespace cbm::algo;
 
 
CbmKfFitTracksTask::CbmKfFitTracksTask(FitMode mode, Int_t iVerbose)
  : FairTask("CbmKfFitTracksTask", iVerbose)
  , fFitMode(mode)
{
}
 
CbmKfFitTracksTask::~CbmKfFitTracksTask() {}
 
 
InitStatus CbmKfFitTracksTask::Init()
{
 
  if (fNthreads < 1) {
 
    // check if the number of threads is set in the invironment variable
    const char* env = std::getenv("OMP_NUM_THREADS");
    if (env) {
      fNthreads = TString(env).Atoi();
      LOG(info) << "  found environment variable OMP_NUM_THREADS = \"" << env << "\", read as integer: " << fNthreads;
    }
 
    // ensure that at least one thread is set
    if (fNthreads < 1) {
      fNthreads = 1;
    }
    // fNthreads = 1; // enforce n threads to one
  }
 
  fFitter.Init();
 
  //Get ROOT Manager
  FairRootManager* ioman = FairRootManager::Instance();
 
  if (!ioman) {
    LOG(error) << "Init: RootManager not instantiated!";
    return kERROR;
  }
 
  // Get global tracks
 
  fGlobalTracks = dynamic_cast<TClonesArray*>(ioman->GetObject("GlobalTrack"));
 
  // Get detector tracks
  fStsTracks  = dynamic_cast<TClonesArray*>(ioman->GetObject("StsTrack"));
  fMuchTracks = dynamic_cast<TClonesArray*>(ioman->GetObject("MuchTrack"));
  fTrdTracks  = dynamic_cast<TClonesArray*>(ioman->GetObject("TrdTrack"));
  fTofTracks  = dynamic_cast<TClonesArray*>(ioman->GetObject("TofTrack"));
 
  if (FitMode::kSts != fFitMode && !fGlobalTracks) {
    LOG(error) << "Init: Global track array not found!";
    return kERROR;
  }
 
  if (FitMode::kSts == fFitMode && !fStsTracks) {
    LOG(error) << "Init: Sts track array not found!";
    return kERROR;
  }
 
  return kSUCCESS;
}
 
void CbmKfFitTracksTask::Exec(Option_t* /*opt*/)
{
 
  LOG(info) << "exec event N " << fNeventsProcessed++ << " with " << fNthreads << " threads";
 
  int statNfittedTracks = 0;
  int statNtracks       = 0;
 
  fFitter.SetDefaultMomentumForMs(0.1);  // 0.1 GeV/c
  // TODO: automatically fix the momentum for Ms when it is not measured - for mCBM or Trd-only tracks
 
  // fix the momentum to default for the Multiple Scattering calculation in mCBM
  fFitter.SetEnforceDefaultMomentumForMs(FitMode::kMcbm == fFitMode);
 
  if (FitMode::kSts == fFitMode && fStsTracks) {
    statNtracks = fStsTracks->GetEntriesFast();
  }
  else if ((FitMode::kGlobal == fFitMode || FitMode::kMcbm == fFitMode) && fGlobalTracks) {
    statNtracks = fGlobalTracks->GetEntriesFast();
  }
 
  std::vector<int> statNfittedThreadTracks(fNthreads, 0);
 
  auto fitThread = [&](int iThread) {
    CbmKfTrackFitter<> threadFitter(fFitter);
 
    for (int iTr = iThread; iTr < statNtracks; iTr += fNthreads) {
 
      CbmKfTrackFitter<>::Trajectory t;
      if (FitMode::kSts == fFitMode) {
        if (!threadFitter.CreateFromMvdStsTrack(t, iTr, false)) {
          LOG(fatal) << "can not create the sts track for the fit! ";
          return;
        }
      }
      else if (FitMode::kMcbm == fFitMode || FitMode::kGlobal == fFitMode) {
        if (!threadFitter.CreateFromGlobalTrack(t, iTr, false)) {
          LOG(fatal) << "can not create the global track for the fit! ";
          return;
        }
      }
      else {
        LOG(fatal) << "unknown fit mode!";
        return;
      }
 
      if (t.GetNofMeasurements() == 0) {
        LOG(fatal) << "the track has no nodes with measurements!";
      }
 
      threadFitter.SetNofIterations(1);
      // first time, fit with the smoother to get the linearisation at each node
      if (!threadFitter.FitTrajectory(t)) {
        continue;
      }
      // second time, fit without the smoother to get the final results at the first and at the last nodes
      // using the linearisation from the previous fit
 
      kf::TrackParamD parUp, parDn;
      if (!threadFitter.FitTrajectoryUpstream(t, parUp)) {
        continue;
      }
 
      if (!threadFitter.FitTrajectoryDownstream(t, parDn)) {
        continue;
      }
 
      statNfittedThreadTracks[iThread]++;
 
      FairTrackParam trackFirst = cbm::kf::ConvertTrackParam(parUp);
      FairTrackParam trackLast  = cbm::kf::ConvertTrackParam(parDn);
 
      if (FitMode::kSts == fFitMode) {
        CbmStsTrack* track = dynamic_cast<CbmStsTrack*>(fStsTracks->At(iTr));
        if (!track) {
          LOG(fatal) << "null pointer to the sts track!";
          return;
        }
        track->SetParamFirst(&trackFirst);
        track->SetParamLast(&trackLast);
      }
      else if (FitMode::kMcbm == fFitMode || FitMode::kGlobal == fFitMode) {
        CbmGlobalTrack* globalTrack = dynamic_cast<CbmGlobalTrack*>(fGlobalTracks->At(iTr));
        if (!globalTrack) {
          LOG(fatal) << "null pointer to the global track!";
          return;
        }
        globalTrack->SetParamFirst(&trackFirst);
        globalTrack->SetParamLast(&trackLast);
      }
    }  // end of loop over tracks in the thread
  };   // end of fitThread lambda
 
  std::vector<std::thread> threads(fNthreads);
 
  // run n threads
  for (int i = 0; i < fNthreads; i++) {
    threads[i] = std::thread(fitThread, i);
  }
 
  // wait for the threads to finish
  for (auto& th : threads) {
    th.join();
  }
 
  // collect statistics
  for (auto& n : statNfittedThreadTracks) {
    statNfittedTracks += n;
  }
 
  fStatNtracks += statNtracks;
  fStatNfittedTracks += statNfittedTracks;
 
  if (statNfittedTracks != statNtracks) {
    LOG(error) << "can not fit " << 100. * (statNtracks - statNfittedTracks) / statNtracks << "\% of " << statNtracks
               << " tracks! " << statNtracks - statNfittedTracks << " tracks are not fitted! ";
  }
}
 
 
void CbmKfFitTracksTask::Finish()
{
  if (fStatNfittedTracks != fStatNtracks) {
    LOG(error) << "can not fit " << 100. * (fStatNtracks - fStatNfittedTracks) / fStatNtracks << "\% of "
               << fStatNtracks << " tracks! " << fStatNtracks - fStatNfittedTracks << " tracks are not fitted! ";
  }
  else {
    LOG(info) << "fitted " << fStatNfittedTracks << " tracks out of " << fStatNtracks
              << " tracks! All tracks are fitted! ";
  }
}