BmnGlobalTracking.cxx

Go to the documentation of this file.

/** BmnGlobalTracking.cxx
 * \author Sergey Merts <Sergey.Merts@gmail.com>
 * @since 2014
 * @version 1.0
 **/
 
#include "BmnGlobalTracking.h"
 
#include "BmnEventHeader.h"
#include "BmnFieldMap.h"
#include "BmnGemStripHit.h"
#include "TFile.h"
#include "TVectorD.h"
 
#include <TKey.h>
#include <algorithm>
#include <iterator>
#include <map>
#include <vector>
 
using namespace TMath;
 
const Int_t numChambers = 59;
const std::vector<std::string> PARTICLES = {"pi", "K", "p", "d", "t", "He3"};
const std::vector<std::string> ZONES = {"clean", "overlap", "accepted"};
 
BmnGlobalTracking::BmnGlobalTracking()
    : fCscHits(nullptr)
    , fTof1Hits(nullptr)
    , fTof2Hits(nullptr)
    , fStsTracks(nullptr)
    , fGlobalTracks(nullptr)
    , fScWall(nullptr)
    , fNdet(nullptr)
    , fInnerTrackBranchName("StsVector")
    , fPeriod(8)
    , fEventNo(0)
    , fVertex(nullptr)
    , fKalman(nullptr)
    , fTime(0.0)
    , fNMatchedTof400(0)
    , fNMatchedTof700(0)
    , fNInnerTracks(0)
    , fNGoodInnerTracks(0)
{}
 
BmnGlobalTracking::~BmnGlobalTracking()
{
    if (fGlobalTracks) {
        fGlobalTracks->Delete();
        delete fGlobalTracks;
    }
}
 
InitStatus BmnGlobalTracking::Init()
{
    if (fVerbose > 1)
        cout << "BmnGlobalTracking::Init started\n";
 
    fKalman = new BmnKalmanFilter();
 
    FairRootManager* ioman = FairRootManager::Instance();
    if (!ioman)
        Fatal("Init", "FairRootManager is not instantiated");
 
    fCscHits = (TClonesArray*)ioman->GetObject("BmnCSCHit");
    fTof1Hits = (TClonesArray*)ioman->GetObject("BmnTof400Hit");
    fTof2Hits = (TClonesArray*)ioman->GetObject("BmnTof700Hit");
 
    fVertex = (CbmVertex*)ioman->GetObject("MpdVertex.");
    fScWall = (BmnScWallEvent*)ioman->GetObject("ScWallEvent");
    fNdet = (BmnNdetEvent*)ioman->GetObject("NdetEvent");
 
    fStsTracks = (TClonesArray*)ioman->GetObject(fInnerTrackBranchName);
    if (fStsTracks) {
        printf("BmnGlobalTracking::Init(): branch %s was found!\n", fInnerTrackBranchName.Data());
        fGlobalTracks = new TClonesArray("BmnGlobalTrack", 100);   // out
        ioman->Register("BmnGlobalTrack", "GLOBAL", fGlobalTracks, kTRUE);
    } else {
        printf("BmnGlobalTracking::Init(): branch with inner tracks was not found! Task will be deactivated \n");
        SetActive(kFALSE);
        return kERROR;
    }
 
    if (fPeriod == 8) {
        TString workdir = getenv("VMCWORKDIR");
 
        LoadAllAkimaCorrections(workdir + "/input/OutputStages_FitResults_combined_tof400_akima_v2_it16.root", kTRUE);
        LoadAllAkimaCorrections(workdir + "/input/OutputStages_FitResults_combined_tof700_akima_it12.root", kFALSE);
    }
 
    if (fVerbose > 1)
        cout << "BmnGlobalTracking::Init finished\n";
    return kSUCCESS;
}
 
void BmnGlobalTracking::Exec(Option_t* opt)
{
    TStopwatch sw;
    sw.Start();
 
    if (!IsActive())
        return;
 
    fEventNo++;
    fGlobalTracks->Delete();
 
    if (fVerbose > 1)
        cout << "\n======================== Global tracking exec started =====================\n" << endl;
 
    if (!fVertex)
        return;
    if (fVertex->GetNTracks() < 2)
        return;
 
    for (Int_t i = 0; i < fStsTracks->GetEntriesFast(); ++i) {
        CbmStsTrack* cbmTrack = (CbmStsTrack*)fStsTracks->At(i);
 
        fNInnerTracks++;
 
        if (cbmTrack->GetNStsHits() < 4)
            continue;
 
        BmnGlobalTrack globTr;
 
        globTr.SetGemTrackIndex(i);
        globTr.SetParamFirst(*(cbmTrack->GetParamFirst()));
        globTr.SetParamLast(*(cbmTrack->GetParamLast()));
        globTr.SetNHits(cbmTrack->GetNStsHits());
        globTr.SetNDF(cbmTrack->GetNDF());
        globTr.SetChi2(cbmTrack->GetChi2());
        FairTrackParam par(*(globTr.GetParamLast()));
        Int_t pdg = (globTr.GetP() > 0) ? 2212 : -211;
        Double_t len = 0.0;
        const Float_t tZ = -0.01942;   // z-coordinate of target for RUN-8
        if (fKalman->TGeoTrackPropagate(&par, tZ, pdg, nullptr, &len) == kBMNERROR)
            continue;
        globTr.SetLength(len);
 
        FairTrackParam parF(*(globTr.GetParamFirst()));
        if (fKalman->TGeoTrackPropagate(&parF, fVertex->GetZ(), pdg, nullptr, nullptr) == kBMNERROR)
            continue;
        Float_t dca = Sqrt(Sq(fVertex->GetX() - parF.GetX()) + Sq(fVertex->GetY() - parF.GetY()));
        globTr.SetParamFirst(parF);
        globTr.SetDCAInVertex(dca);
 
        BmnHit virtualVertexHit;
        virtualVertexHit.SetXYZ(fVertex->GetX(), fVertex->GetY(), fVertex->GetZ());
        virtualVertexHit.SetDxyz(Sqrt(fVertex->GetCovariance(0, 0)), Sqrt(fVertex->GetCovariance(1, 1)),
                                 Sqrt(fVertex->GetCovariance(2, 2)));
 
        Double_t chi = 0;
        fKalman->Update(&parF, &virtualVertexHit, chi);
        globTr.SetChi2InVertex(chi);
 
        fNGoodInnerTracks++;
        new ((*fGlobalTracks)[fGlobalTracks->GetEntriesFast()]) BmnGlobalTrack(globTr);
    }
 
    MatchingTOF(kTRUE);    // tof400
    MatchingTOF(kFALSE);   // tof700
    // MatchingScWall();
    MatchingNdet();
 
    sw.Stop();
    fTime += sw.RealTime();
 
    if (fVerbose > 1)
        cout << "\n======================== Global tracking exec finished ====================\n" << endl;
}
 
BmnStatus BmnGlobalTracking::MatchingTOF(Bool_t tof400)
{
    Int_t tofID = (tof400) ? 1 : 2;
    TString tofName = (tof400) ? "TOF-400" : "TOF-700";
    if (fVerbose)
        cout << "Matching of " << tofName << " started" << endl;
    TClonesArray* tofHits = (tof400) ? fTof1Hits : fTof2Hits;
    if (!tofHits)
        return kBMNERROR;
 
    for (Int_t i = 0; i < fGlobalTracks->GetEntriesFast(); i++)
        ((BmnGlobalTrack*)fGlobalTracks->At(i))->SetUsing(kFALSE);
 
    map<Float_t, pair<Int_t, Int_t>> distPairs = FindPairsTof(fGlobalTracks, tofHits, tof400);
 
    for (auto const& distPair : distPairs) {
        // Float_t dist = distPair.first;
        pair<Int_t, Int_t> idxes = distPair.second;
        BmnGlobalTrack* track = (BmnGlobalTrack*)fGlobalTracks->At(idxes.first);
        if (track->IsUsed())
            continue;
        BmnHit* hit = (BmnHit*)tofHits->At(idxes.second);
        if (hit->IsUsed())
            continue;
 
        track->SetUsing(kTRUE);
        hit->SetUsing(kTRUE);
        track->SetNHits(track->GetNHits() + 1);
        if (tof400)
            fNMatchedTof400++;
        else
            fNMatchedTof700++;
 
        Double_t len = 0.0;
        FairTrackParam par(*(track->GetParamLast()));
        Int_t pdg = (track->GetP() > 0) ? 2212 : -211;
        fKalman->TGeoTrackPropagate(&par, hit->GetZ(), pdg, nullptr, &len);
        len += track->GetLength();
        hit->SetLength(len);
        Float_t mom = track->GetP();
        Float_t beta = len / hit->GetTimeStamp() / (TMath::C() * 1e-7);
        Float_t mass2 = mom * mom * (1.0 / beta / beta - 1.0);
        Float_t xExtr = par.GetX();
        Float_t yExtr = par.GetY();
        Int_t chamber = ((hit->GetDetectorID() & 0x0000FF00) >> 8) - 1;
        Float_t xExtr_stage2 = xExtr + GetAkimaCorrection(chamber, mom, kTRUE, "stage2", tof400);
        Float_t yExtr_stage2 = yExtr + GetAkimaCorrection(chamber, mom, kFALSE, "stage2", tof400);
        Float_t sigmaX_stage2 = GetAkimaSigma(chamber, mom, kTRUE, "stage2", tof400);
        Float_t sigmaY_stage2 = GetAkimaSigma(chamber, mom, kFALSE, "stage2", tof400);
 
        // for kCommon
        Float_t xDist = xExtr_stage2 - hit->GetX();
        Float_t yDist = yExtr_stage2 - hit->GetY();
        bool isWithinSigma = (Abs(xDist) < 3 * sigmaX_stage2) && (Abs(yDist) < 3 * sigmaY_stage2);
        if (isWithinSigma) {
            if (tof400)
                track->SetTof1HitIndex(idxes.second, static_cast<BmnGlobalTrack::ParticleType>(0));
            else
                track->SetTof2HitIndex(idxes.second, static_cast<BmnGlobalTrack::ParticleType>(0));
            track->SetMass2(tofID, static_cast<BmnGlobalTrack::ParticleType>(0), mass2);
            track->SetResX(tofID, static_cast<BmnGlobalTrack::ParticleType>(0), xDist);
            track->SetResY(tofID, static_cast<BmnGlobalTrack::ParticleType>(0), yDist);
        }
 
        std::string zone = "accepted";
        for (Int_t hypId = 1; hypId < BmnGlobalTrack::kNumParticleTypes; ++hypId) {   // all hypothises except kCommon
            TString particle = track->GetParticleName(static_cast<BmnGlobalTrack::ParticleType>(hypId));
            Float_t xCorr_stage3 = GetAkimaCorrection(chamber, mom, kTRUE, "stage3", tof400, particle.Data(), zone);
            Float_t yCorr_stage3 = GetAkimaCorrection(chamber, mom, kFALSE, "stage3", tof400, particle.Data(), zone);
            Float_t sigmaX_stage3 = GetAkimaSigma(chamber, mom, kTRUE, "stage3", tof400, particle.Data(), zone);
            Float_t sigmaY_stage3 = GetAkimaSigma(chamber, mom, kFALSE, "stage3", tof400, particle.Data(), zone);
 
            xDist = xExtr_stage2 + xCorr_stage3 - hit->GetX();
            yDist = yExtr_stage2 + yCorr_stage3 - hit->GetY();
            isWithinSigma = (Abs(xDist) < 3 * sigmaX_stage3) && (Abs(yDist) < 3 * sigmaY_stage3);
            if (isWithinSigma) {
                if (tof400)
                    track->SetTof1HitIndex(idxes.second, static_cast<BmnGlobalTrack::ParticleType>(hypId));
                else
                    track->SetTof2HitIndex(idxes.second, static_cast<BmnGlobalTrack::ParticleType>(hypId));
                track->SetMass2(tofID, static_cast<BmnGlobalTrack::ParticleType>(hypId), mass2);
                track->SetResX(tofID, static_cast<BmnGlobalTrack::ParticleType>(hypId), xDist);
                track->SetResY(tofID, static_cast<BmnGlobalTrack::ParticleType>(hypId), yDist);
            }
        }
    }
    return kBMNSUCCESS;
}
 
BmnStatus BmnGlobalTracking::MatchingScWall()
{
    if (fVerbose)
        cout << "Matching of ScWall started " << fScWall << endl;
    if (!fScWall)
        return kBMNERROR;
 
    for (Int_t trIdx = 0; trIdx < fGlobalTracks->GetEntriesFast(); ++trIdx) {
        BmnGlobalTrack* tr = (BmnGlobalTrack*)fGlobalTracks->At(trIdx);
        FairTrackParam par(*(tr->GetParamLast()));
        Double_t minDX = DBL_MAX;
        Double_t minDY = DBL_MAX;
        Float_t xCut = 20.0;   // cm
        Float_t yCut = 20.0;   // cm
        Int_t minIdx = -1;
 
        Float_t z = 741.5;
        Int_t nCells = 174;   // BmnScWallEvent::fgkMaxCells; //174
        Int_t pdg = (tr->GetP() > 0) ? 2212 : -211;
 
        if (fKalman->TGeoTrackPropagate(&par, z, pdg, nullptr, nullptr) == kBMNERROR)
            continue;
 
        for (Int_t i = 1; i <= nCells; ++i) {
            BmnScWallCell* cell = fScWall->GetCell(i);
            if (cell->GetSignal() < 0.001)
                continue;   // skip empty cells
 
            Float_t dX = par.GetX() - cell->GetX();
            Float_t dY = par.GetY() - cell->GetY();
 
            if (Abs(dX) < xCut && Abs(dY) < yCut && Abs(dX) < Abs(minDX) && Abs(dY) < Abs(minDY)) {
                minDX = dX;
                minDY = dY;
                minIdx = i;
            }
        }
        // cout << minDX << "   " << minDY << endl;
 
        if (minIdx == -1)
            continue;
        tr->SetScWallCellId(minIdx);
        tr->SetScWallSignal(fScWall->GetCell(minIdx)->GetSignal());
    }
    return kBMNSUCCESS;
}
 
BmnStatus BmnGlobalTracking::MatchingNdet()
{
    if (fVerbose)
        cout << "Matching of Ndet started " << fNdet << endl;
    if (!fNdet)
        return kBMNERROR;
    auto* cells = fNdet->GetCells();
    if (cells->GetEntriesFast() == 0)
        return kBMNSUCCESS;
 
    for (Int_t trIdx = 0; trIdx < fGlobalTracks->GetEntriesFast(); ++trIdx) {
        BmnGlobalTrack* tr = (BmnGlobalTrack*)fGlobalTracks->At(trIdx);
        FairTrackParam par(*(tr->GetParamLast()));
 
        for (int ice = 0, nce = cells->GetEntriesFast(); ice < nce; ++ice) {
            auto* cell = static_cast<BmnNdetCell*>(cells->At(ice));
            if (!cell)
                continue;
            Int_t pdg = (tr->GetP() > 0) ? 2212 : -211;
 
            if (fKalman->TGeoTrackPropagate(&par, cell->GetZ(), pdg, nullptr, nullptr) == kBMNERROR)
                continue;
 
            Float_t dX = par.GetX() - cell->GetX();
            Float_t dY = par.GetY() - cell->GetY();
 
            if (Abs(dX) < 5 * cell->GetDx() && Abs(dY) < 5 * cell->GetDy()) {   // TODO remove 5* (after alignment)
                tr->SetNdetHitIndex(ice);
                // break;
            }
        }
    }
    return kBMNSUCCESS;
}
 
map<Float_t, pair<Int_t, Int_t>> BmnGlobalTracking::FindPairsTof(TClonesArray* tracks,
                                                                 TClonesArray* hits,
                                                                 Bool_t tof400)
{
 
    map<Float_t, pair<Int_t, Int_t>> distancesPairs;
 
    Double_t minZ = 10000.0;
    for (Int_t hitIdx = 0; hitIdx < hits->GetEntriesFast(); ++hitIdx) {
        BmnHit* hit = (BmnHit*)hits->At(hitIdx);
        if (hit->GetZ() < minZ)
            minZ = hit->GetZ();
    }
 
    for (Int_t trIdx = 0; trIdx < tracks->GetEntriesFast(); ++trIdx) {
        BmnGlobalTrack* tr = (BmnGlobalTrack*)tracks->At(trIdx);
        FairTrackParam parMinZ(*(tr->GetParamLast()));
        Int_t pdg = (tr->GetP() > 0) ? 2212 : -211;
        if (fKalman->TGeoTrackPropagate(&parMinZ, minZ, pdg, nullptr, nullptr) == kBMNERROR)
            continue;
        for (Int_t hitIdx = 0; hitIdx < hits->GetEntriesFast(); ++hitIdx) {
 
            BmnHit* hit = (BmnHit*)hits->At(hitIdx);
            FairTrackParam param = parMinZ;
            if (fKalman->TGeoTrackPropagate(&param, hit->GetZ(), pdg, nullptr, nullptr) == kBMNERROR)
                continue;
            Float_t mom = 1.0 / param.GetQp();
            Int_t chamber = ((hit->GetDetectorID() & 0x0000FF00) >> 8) - 1;
            Float_t xMomCorr = GetAkimaCorrection(chamber, mom, kTRUE, "stage1", tof400);
            Float_t yMomCorr = GetAkimaCorrection(chamber, mom, kFALSE, "stage1", tof400);
 
            Float_t dX = param.GetX() + xMomCorr - hit->GetX();
            Float_t dY = param.GetY() + yMomCorr - hit->GetY();
            Float_t dist = Sqrt(dX * dX + dY * dY);
 
            distancesPairs.insert(pair<Float_t, pair<Int_t, Int_t>>(dist, pair<Int_t, Int_t>(trIdx, hitIdx)));
        }
    }
 
    return distancesPairs;
}
 
BmnStatus BmnGlobalTracking::UpdateMomentum(BmnGlobalTrack* tr)
{
    FairTrackParam par = *(tr->GetParamFirst());
    Int_t pdg = (par.GetQp() > 0) ? 2212 : -211;
    fKalman->TGeoTrackPropagate(&par, -200, pdg, nullptr, nullptr);
    Double_t Alpha_in = ATan(par.GetTx());
    // Double_t Bdl = 2.856;     //Abs(MagFieldIntegral(par, -200.0, 1000.0, 1.0) * 0.001);
    Double_t Bdl = Abs(MagFieldIntegral(par, -200.0, 600.0, 1.0) * 0.001);
    fKalman->TGeoTrackPropagate(&par, 600, pdg, nullptr, nullptr);
    Double_t Alpha_out = ATan(par.GetTx());
    // if (tr->GetUpstreamTrackIndex() != -1) {
    //     BmnTrack* upTrack = (BmnTrack*)fUpstreamTracks->At(tr->GetUpstreamTrackIndex());
    //     if (upTrack != nullptr)
    //         Alpha_in = ATan(upTrack->GetParamLast()->GetTx());
    // }
    Double_t momInt = 0.3 * Bdl / (Sin(Alpha_out) - Sin(Alpha_in));
    // Double_t momInt = 0.3 * Bdl / (Alpha_out - Alpha_in);
    if (!IsNaN(momInt) && Abs(momInt) > 0.01 && Abs(momInt) < 100) {
        tr->GetParamFirst()->SetQp(1.0 / momInt);
    }
    // cout << Alpha_in << " " << Alpha_out << " " << Bdl << " " << momInt << endl;
 
    return kBMNSUCCESS;
}
 
void BmnGlobalTracking::Finish()
{
    delete fKalman;
    printf("Work time of BmnGlobalTracking: %4.2f sec.\n", fTime);
    PrintStatistics();
}
 
Double_t BmnGlobalTracking::MagFieldIntegral(FairTrackParam& par, Double_t zMin, Double_t zMax, Double_t step)
{
    /*
    field in kG
    step in cm
    */
    Int_t pdg = (par.GetQp() > 0.) ? 2212 : -211;
    FairField* field = FairRunAna::Instance()->GetField();
    fKalman->TGeoTrackPropagate(&par, zMin, pdg, NULL, NULL, kTRUE);
    Double_t z = par.GetZ();
    Double_t integral = 0.0;
    while (z < zMax) {
        z += step;
        fKalman->TGeoTrackPropagate(&par, z, pdg, NULL, NULL, kTRUE);
        integral += field->GetBy(par.GetX(), par.GetY(), par.GetZ());
    }
    integral *= step;
    return integral;
}
 
void BmnGlobalTracking::PrintStatistics()
{
    printf("================================================================================\n");
    printf("=======  Statistics of Global tracking\n");
    printf("=======  Number of input inner tracks: %ld\n", fNInnerTracks);
    printf("=======  Number of input inner tracks with 4+ hits extrapolated to the Vertex: %ld (%4.1f%%)\n",
           fNGoodInnerTracks, fNGoodInnerTracks * 100.0 / fNInnerTracks);
    printf("=======  Number of tracks matched with TOF-400: %ld\n", fNMatchedTof400);
    printf("=======  Number of tracks matched with TOF-700: %ld\n", fNMatchedTof700);
    printf("================================================================================\n");
}
 
void BmnGlobalTracking::OnlineWrite(const std::unique_ptr<TTree>& resultTree)
{
    if (!IsActive())
        return;
 
    if (fStsTracks) {
        resultTree->Branch("BmnGlobalTrack", &fGlobalTracks);
    }
    resultTree->Fill();
}
 
void BmnGlobalTracking::LoadAllAkimaCorrections(const TString& filename, Bool_t isTOF400)
{
    std::cout << "\n=== Loading Akima corrections for " << (isTOF400 ? "TOF400" : "TOF700") << " ===" << std::endl;
    std::cout << "=== File: " << filename << " ===" << std::endl;
 
    AkimaContainer& cont = isTOF400 ? fAkima400 : fAkima700;
 
    TFile* file = TFile::Open(filename);
    if (!file || file->IsZombie()) {
        std::cerr << "ERROR: Cannot open file: " << filename << std::endl;
        return;
    }
 
    for (Int_t chamber = 0; chamber < numChambers; ++chamber) {
        TString chamberPath = Form("chamber_%d", chamber);
        TDirectory* chamberDir = dynamic_cast<TDirectory*>(file->Get(chamberPath));
        if (!chamberDir)
            continue;
 
        TDirectory* stage1Dir = dynamic_cast<TDirectory*>(chamberDir->Get("stage1"));
        if (stage1Dir) {
            TDirectory* dxPosDir = dynamic_cast<TDirectory*>(stage1Dir->Get("dx_positive"));
            if (dxPosDir) {
                TIter next(dxPosDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        cont.dx_stage1_pos[chamber] = (AkimaFunction*)obj;
                        cont.loaded_count++;
                        break;
                    }
                }
            }
 
            TDirectory* dxNegDir = dynamic_cast<TDirectory*>(stage1Dir->Get("dx_negative"));
            if (dxNegDir) {
                TIter next(dxNegDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        cont.dx_stage1_neg[chamber] = (AkimaFunction*)obj;
                        cont.loaded_count++;
                        break;
                    }
                }
            }
 
            TDirectory* dyPosDir = dynamic_cast<TDirectory*>(stage1Dir->Get("dy_positive"));
            if (dyPosDir) {
                TIter next(dyPosDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        cont.dy_stage1_pos[chamber] = (AkimaFunction*)obj;
                        cont.loaded_count++;
                        break;
                    }
                }
            }
 
            TDirectory* dyNegDir = dynamic_cast<TDirectory*>(stage1Dir->Get("dy_negative"));
            if (dyNegDir) {
                TIter next(dyNegDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        cont.dy_stage1_neg[chamber] = (AkimaFunction*)obj;
                        cont.loaded_count++;
                        break;
                    }
                }
            }
        }
 
        TDirectory* stage2Dir = dynamic_cast<TDirectory*>(chamberDir->Get("stage2"));
        if (stage2Dir) {
            TDirectory* dxPosDir = dynamic_cast<TDirectory*>(stage2Dir->Get("dx_positive"));
            if (dxPosDir) {
                TIter next(dxPosDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        AkimaFunction* func = (AkimaFunction*)obj;
                        TString name = func->GetName();
                        if (name.Contains("mean")) {
                            cont.dx_stage2_mean_pos[chamber] = func;
                            cont.loaded_count++;
                        } else if (name.Contains("sigma")) {
                            cont.dx_stage2_sigma_pos[chamber] = func;
                            cont.loaded_count++;
                        }
                    }
                }
            }
 
            TDirectory* dxNegDir = dynamic_cast<TDirectory*>(stage2Dir->Get("dx_negative"));
            if (dxNegDir) {
                TIter next(dxNegDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        AkimaFunction* func = (AkimaFunction*)obj;
                        TString name = func->GetName();
                        if (name.Contains("mean")) {
                            cont.dx_stage2_mean_neg[chamber] = func;
                            cont.loaded_count++;
                        } else if (name.Contains("sigma")) {
                            cont.dx_stage2_sigma_neg[chamber] = func;
                            cont.loaded_count++;
                        }
                    }
                }
            }
 
            TDirectory* dyPosDir = dynamic_cast<TDirectory*>(stage2Dir->Get("dy_positive"));
            if (dyPosDir) {
                TIter next(dyPosDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        AkimaFunction* func = (AkimaFunction*)obj;
                        TString name = func->GetName();
                        if (name.Contains("mean")) {
                            cont.dy_stage2_mean_pos[chamber] = func;
                            cont.loaded_count++;
                        } else if (name.Contains("sigma")) {
                            cont.dy_stage2_sigma_pos[chamber] = func;
                            cont.loaded_count++;
                        }
                    }
                }
            }
 
            TDirectory* dyNegDir = dynamic_cast<TDirectory*>(stage2Dir->Get("dy_negative"));
            if (dyNegDir) {
                TIter next(dyNegDir->GetListOfKeys());
                TKey* key;
                while ((key = (TKey*)next())) {
                    TObject* obj = key->ReadObj();
                    if (obj->InheritsFrom("AkimaFunction")) {
                        AkimaFunction* func = (AkimaFunction*)obj;
                        TString name = func->GetName();
                        if (name.Contains("mean")) {
                            cont.dy_stage2_mean_neg[chamber] = func;
                            cont.loaded_count++;
                        } else if (name.Contains("sigma")) {
                            cont.dy_stage2_sigma_neg[chamber] = func;
                            cont.loaded_count++;
                        }
                    }
                }
            }
        }
 
        TDirectory* stage3Dir = dynamic_cast<TDirectory*>(chamberDir->Get("stage3"));
        if (stage3Dir) {
            for (const auto& particle : PARTICLES) {
                TDirectory* particleDir = dynamic_cast<TDirectory*>(stage3Dir->Get(particle.c_str()));
                if (!particleDir)
                    continue;
 
                for (const auto& zone : ZONES) {
                    TString dxDirName = Form("dx_%s_%s_positive", particle.c_str(), zone.c_str());
                    TDirectory* dxDir = dynamic_cast<TDirectory*>(particleDir->Get(dxDirName));
                    if (dxDir) {
                        TIter next(dxDir->GetListOfKeys());
                        TKey* key;
                        while ((key = (TKey*)next())) {
                            TObject* obj = key->ReadObj();
                            if (obj->InheritsFrom("AkimaFunction")) {
                                AkimaFunction* func = (AkimaFunction*)obj;
                                TString name = func->GetName();
                                Stage3Key k = {chamber, particle, zone};
                                if (name.Contains("mean")) {
                                    cont.dx_stage3_mean[k] = func;
                                    cont.loaded_count++;
                                } else if (name.Contains("sigma")) {
                                    cont.dx_stage3_sigma[k] = func;
                                    cont.loaded_count++;
                                }
                            }
                        }
                    }
 
                    TString dyDirName = Form("dy_%s_%s_positive", particle.c_str(), zone.c_str());
                    TDirectory* dyDir = dynamic_cast<TDirectory*>(particleDir->Get(dyDirName));
                    if (dyDir) {
                        TIter next(dyDir->GetListOfKeys());
                        TKey* key;
                        while ((key = (TKey*)next())) {
                            TObject* obj = key->ReadObj();
                            if (obj->InheritsFrom("AkimaFunction")) {
                                AkimaFunction* func = (AkimaFunction*)obj;
                                TString name = func->GetName();
                                Stage3Key k = {chamber, particle, zone};
                                if (name.Contains("mean")) {
                                    cont.dy_stage3_mean[k] = func;
                                    cont.loaded_count++;
                                } else if (name.Contains("sigma")) {
                                    cont.dy_stage3_sigma[k] = func;
                                    cont.loaded_count++;
                                }
                            }
                        }
                    }
                }
            }
        }
    }
 
    file->Close();
    std::cout << "Loaded " << cont.loaded_count << " AkimaFunction objects for " << (isTOF400 ? "TOF400" : "TOF700")
              << std::endl;
}
 
Float_t BmnGlobalTracking::GetAkimaCorrection(Int_t chamber,
                                              Float_t qp,
                                              Bool_t isDx,
                                              const TString& stage,
                                              Bool_t isTOF400,
                                              const TString& particle,
                                              const TString& zone)
{
    if (chamber < 0 || chamber >= numChambers)
        return 0.0;
 
    AkimaContainer& cont = isTOF400 ? fAkima400 : fAkima700;
    qp = AdjustQP(qp, isTOF400);
    std::string region = GetRegionFromQP(qp);
 
    if (stage == "stage1") {
        if (isDx) {
            if (region == "negative") {
                auto it = cont.dx_stage1_neg.find(chamber);
                if (it != cont.dx_stage1_neg.end() && it->second)
                    return -it->second->Eval(qp);
                it = cont.dx_stage1_pos.find(chamber);
                if (it != cont.dx_stage1_pos.end() && it->second)
                    return -it->second->Eval(-qp);
            } else {
                auto it = cont.dx_stage1_pos.find(chamber);
                if (it != cont.dx_stage1_pos.end() && it->second)
                    return -it->second->Eval(qp);
            }
        } else {
            if (region == "negative") {
                auto it = cont.dy_stage1_neg.find(chamber);
                if (it != cont.dy_stage1_neg.end() && it->second)
                    return -it->second->Eval(qp);
                auto it_pos = cont.dy_stage1_pos.find(chamber);
                if (it_pos != cont.dy_stage1_pos.end() && it_pos->second)
                    return -it_pos->second->Eval(-qp);
            } else {
                auto it = cont.dy_stage1_pos.find(chamber);
                if (it != cont.dy_stage1_pos.end() && it->second)
                    return -it->second->Eval(qp);
            }
        }
    } else if (stage == "stage2") {
        if (isDx) {
            if (region == "negative") {
                auto it = cont.dx_stage2_mean_neg.find(chamber);
                if (it != cont.dx_stage2_mean_neg.end() && it->second)
                    return -it->second->Eval(qp);
                auto it_pos = cont.dx_stage2_mean_pos.find(chamber);
                if (it_pos != cont.dx_stage2_mean_pos.end() && it_pos->second)
                    return -it_pos->second->Eval(-qp);
            } else {
                auto it = cont.dx_stage2_mean_pos.find(chamber);
                if (it != cont.dx_stage2_mean_pos.end() && it->second)
                    return -it->second->Eval(qp);
            }
        } else {
            if (region == "negative") {
                auto it = cont.dy_stage2_mean_neg.find(chamber);
                if (it != cont.dy_stage2_mean_neg.end() && it->second)
                    return -it->second->Eval(qp);
                auto it_pos = cont.dy_stage2_mean_pos.find(chamber);
                if (it_pos != cont.dy_stage2_mean_pos.end() && it_pos->second)
                    return -it_pos->second->Eval(-qp);
            } else {
                auto it = cont.dy_stage2_mean_pos.find(chamber);
                if (it != cont.dy_stage2_mean_pos.end() && it->second)
                    return -it->second->Eval(qp);
            }
        }
    } else if (stage == "stage3") {
        if (qp < 0)
            return GetAkimaCorrection(chamber, -qp, isDx, "stage2", isTOF400, "", "");
        if (particle == "" || zone == "")
            return GetAkimaCorrection(chamber, qp, isDx, "stage2", isTOF400);
 
        Stage3Key key = {chamber, particle.Data(), zone.Data()};
 
        if (isDx) {
            auto it = cont.dx_stage3_mean.find(key);
            if (it != cont.dx_stage3_mean.end() && it->second)
                return -it->second->Eval(qp);
        } else {
            auto it = cont.dy_stage3_mean.find(key);
            if (it != cont.dy_stage3_mean.end() && it->second)
                return -it->second->Eval(qp);
        }
        return GetAkimaCorrection(chamber, qp, isDx, "stage2", isTOF400);
    }
 
    return 0.0;
}
 
Float_t BmnGlobalTracking::GetAkimaSigma(Int_t chamber,
                                         Float_t qp,
                                         Bool_t isDx,
                                         const TString& stage,
                                         Bool_t isTOF400,
                                         const TString& particle,
                                         const TString& zone)
{
    if (chamber < 0 || chamber >= numChambers)
        return 1.0;
 
    AkimaContainer& cont = isTOF400 ? fAkima400 : fAkima700;
    qp = AdjustQP(qp, isTOF400);
    std::string region = GetRegionFromQP(qp);
 
    if (stage == "stage2") {
        if (isDx) {
            if (region == "negative") {
                auto it = cont.dx_stage2_sigma_neg.find(chamber);
                if (it != cont.dx_stage2_sigma_neg.end() && it->second) {
                    Float_t sigma = it->second->Eval(qp);
                    return (sigma > 0) ? sigma : 1.0;
                }
                auto it_pos = cont.dx_stage2_sigma_pos.find(chamber);
                if (it_pos != cont.dx_stage2_sigma_pos.end() && it_pos->second) {
                    Float_t sigma = it_pos->second->Eval(-qp);
                    return (sigma > 0) ? sigma : 1.0;
                }
            } else {
                auto it = cont.dx_stage2_sigma_pos.find(chamber);
                if (it != cont.dx_stage2_sigma_pos.end() && it->second) {
                    Float_t sigma = it->second->Eval(qp);
                    return (sigma > 0) ? sigma : 1.0;
                }
            }
        } else {
            if (region == "negative") {
                auto it = cont.dy_stage2_sigma_neg.find(chamber);
                if (it != cont.dy_stage2_sigma_neg.end() && it->second) {
                    Float_t sigma = it->second->Eval(qp);
                    return (sigma > 0) ? sigma : 1.0;
                }
                auto it_pos = cont.dy_stage2_sigma_pos.find(chamber);
                if (it_pos != cont.dy_stage2_sigma_pos.end() && it_pos->second) {
                    Float_t sigma = it_pos->second->Eval(-qp);
                    return (sigma > 0) ? sigma : 1.0;
                }
            } else {
                auto it = cont.dy_stage2_sigma_pos.find(chamber);
                if (it != cont.dy_stage2_sigma_pos.end() && it->second) {
                    Float_t sigma = it->second->Eval(qp);
                    return (sigma > 0) ? sigma : 1.0;
                }
            }
        }
        return 1.0;
    } else if (stage == "stage3") {
        if (qp < 0)
            return GetAkimaSigma(chamber, -qp, isDx, "stage2", isTOF400, "", "");
        if (particle == "" || zone == "")
            return GetAkimaSigma(chamber, qp, isDx, "stage2", isTOF400);
 
        Stage3Key key = {chamber, particle.Data(), zone.Data()};
 
        if (isDx) {
            auto it = cont.dx_stage3_sigma.find(key);
            if (it != cont.dx_stage3_sigma.end() && it->second) {
                Float_t sigma = it->second->Eval(qp);
                return (sigma > 0) ? sigma : 1.0;
            }
        } else {
            auto it = cont.dy_stage3_sigma.find(key);
            if (it != cont.dy_stage3_sigma.end() && it->second) {
                Float_t sigma = it->second->Eval(qp);
                return (sigma > 0) ? sigma : 1.0;
            }
        }
        return GetAkimaSigma(chamber, qp, isDx, "stage2", isTOF400);
    }
 
    return 1.0;
}
 
std::string BmnGlobalTracking::GetRegionFromQP(Float_t qp)
{
    return (qp < 0) ? "negative" : "positive";
}
 
Float_t BmnGlobalTracking::AdjustQP(Float_t qp, Bool_t isTOF400)
{
    if (isTOF400) {
        if (qp < 0 && qp > -0.5)
            return -0.5;
        if (qp > 0 && qp < 0.5)
            return 0.5;
    } else {
        if (qp < 0 && qp > -0.7)
            return -0.7;
        if (qp > 0 && qp < 0.7)
            return 0.7;
    }
    return qp;
}