BmnRootIterator.cxx

Go to the documentation of this file.

/*
    BM@N alignment routine
    BM@N experiment at NICA complex, JINR, 2025
 
    Department: Math & Soft Group of HEP lab
    Author:     Igor Polev, polev@jinr.int
 
    BmnRootIterator class implementation
*/
// TODO: check if we can use std::move() across the project to optimize a little
#include "BmnRootIterator.h"
 
#include "BmnHit.h"
#include "BmnWeightedLine.h"
#include "fairlogger/Logger.h"
 
#include <utility>
 
namespace
{
constexpr Int_t SIBT_STATIONS{3};   // the beam line needs one hit in each, numbered from 0 along the beam
}   // anonymous namespace
 
BmnRootIterator::BmnRootIterator(const char* filesPathList,
                                 const BmnDetectorModel& detModel,
                                 TrackFilter_t trackFilter,
                                 HitFilter_t hitFilter,
                                 BmnSimpleProgressBar* pb,
                                 Bool_t verbose,
                                 Int_t maxTracks)
    : BmnDataIterator(pb)
    , fpDetModel(&detModel)
    , fTrackFilter(std::move(trackFilter))
    , fHitFilter(std::move(hitFilter))
    , fpData(new BmnDataReader(filesPathList, verbose))
    , fpReader(&fpData->Reader())
    , fTrackLimit(maxTracks)
{
    if (!fTrackFilter)
        fTrackFilter = [](const Track_t&, Int_t) { return kTRUE; };
    if (!fHitFilter)
        fHitFilter = [](const Hit_t&) { return kTRUE; };
 
    // detect whether the source data has a DstEventHeader branch
    Bool_t hasEventHeader{fpData->Chain().GetBranch(BMN_EVENT_BRANCH_NAME) != nullptr};
 
    const Bool_t withVertex{detModel.TargetID() >= 0};
    if (withVertex) {
        if (detModel.Target().zEnd <= detModel.Target().zStart) {
            LOG(error) << " <!> BmnRootIterator : the detector model has a target element but no"
                          " target geometry - pass the target .geo file to its Init()";
            return;
        }
        if (!fpData->Chain().GetBranch(BMN_BEAM_BRANCH_NAME)) {
            LOG(error) << " <!> BmnRootIterator : the detector model has a target element but the data"
                          " has no " BMN_BEAM_BRANCH_NAME " branch to build the vertex hit from";
            return;
        }
    }
 
    // turn off unneeded branches
    // notice: turning off branches works weirdly and does not give any speedup
    // but let it be for now - maybe it will give something later
    fpData->Chain().SetBranchStatus("*", kFALSE);
    fpData->Chain().SetBranchStatus(BMN_TRACKS_BRANCH_NAME "*", kTRUE);
    fpData->Chain().SetBranchStatus(BMN_HITS_BRANCH_NAME "*", kTRUE);
    if (hasEventHeader)
        fpData->Chain().SetBranchStatus(BMN_EVENT_BRANCH_NAME "*", kTRUE);
    if (withVertex)
        fpData->Chain().SetBranchStatus(BMN_BEAM_BRANCH_NAME "*", kTRUE);
 
    // Tryed to use TTreeReaderArray but failed to make it work
    fpRootTracks = new TTreeReaderValue<TClonesArray>{*fpReader, BMN_TRACKS_BRANCH_NAME};
    fpRootHits = new TTreeReaderValue<TClonesArray>{*fpReader, BMN_HITS_BRANCH_NAME};
    if (hasEventHeader)
        fpRootEventHeader = new TTreeReaderValue<DstEventHeader>{*fpReader, BMN_EVENT_BRANCH_NAME};
    if (withVertex)
        fpRootBeamHits = new TTreeReaderValue<TClonesArray>{*fpReader, BMN_BEAM_BRANCH_NAME};
 
    fProgressByTracks = kFALSE;                   // progress by events until tracks are counted
    fTotalEvents = fpReader->GetEntries(kTRUE);   // without kTRUE result is not guaranteed
    SetTickProgress(1.0f / fTotalEvents);
 
    fInitialized = kTRUE;
}
 
BmnRootIterator::~BmnRootIterator()
{
    delete fpRootTracks;
    delete fpRootHits;
    delete fpRootBeamHits;
    delete fpRootEventHeader;
    delete fpData;
}
 
void BmnRootIterator::CountElements()
{
    // Count all tracks and hits that matches minimum number
    // of hits per track and maximum allowed chi2 per hit,
    // Filtration is built in ResetAll() and NextTrack() directly.
    LOG_IF(warning, fTrackLimit) << "  ! track limit is set to " << fTrackLimit << " tracks";
    LOG(info) << "    counting suitable tracks...";
 
    SetTickProgress(1.0f / fTotalEvents);
    fProgressByTracks = kFALSE;   // progress by events
    ProgressReset();
    fTotalTracks = 0;
    fTotalHits = 0;
    ResetAll();
    while (!EndOfTracks()) {
        fTotalTracks++;
        fTotalHits += HitsInTrack();
        NextTrack();
    }
    ProgressClear();
    fProgressByTracks = kTRUE;
    SetTickProgress(1.0f / fTotalTracks);
 
    LOG(info) << "    " << fTotalTracks << " total suitable tracks found";
    LOG(info) << "    " << fTotalHits << " total hits counted";
    LOG(info) << "    " << fTotalEvents << " events processed";
}
 
void BmnRootIterator::ResetAll()
{
    if (!fInitialized)
        return;
 
    if (fpRootBeamHits)
        fVertexZ = 0.5 * (fpDetModel->Target().zStart + fpDetModel->Target().zEnd);
    UnsetProperties();
    fpReader->Restart();
    Bool_t event = fpReader->Next();
    if (!event)
        return;   // no events at all
    if (!fProgressByTracks)
        ProgressTick();
 
    // search for event with tracks
    fTracksInEvent = (*fpRootTracks)->GetEntries();
    while (event && fTracksInEvent <= 0) {
        event = fpReader->Next();
        if (!fProgressByTracks)
            ProgressTick();
        fTracksInEvent = (*fpRootTracks)->GetEntries();
    }
    if (!event)
        return;   // no events with tracks
    if (fpRootEventHeader)
        fEventID = (*fpRootEventHeader)->GetEventId();
    ScanEvent();
 
    fpCurrentTrack = (Track_t*)(*fpRootTracks)->At(0);
    fTrackIdx = 0;
    fEndOfTracks = kFALSE;
    // CountUsableHits() first: the track predicate is told how many hits will
    // actually be read, not how many the track claims
    if (!fTrackFilter(*fpCurrentTrack, CountUsableHits()))
        NextTrack();
    else {
        fTrackCount = 1;
        fEventTracks = 1;
        ResetHits();
    }
}
 
void BmnRootIterator::NextTrack()
{
    if (!fInitialized || fEndOfTracks)
        return;
    if (fTrackLimit && fTrackCount == fTrackLimit) {
        UnsetProperties();
        return;
    }
    do {
        if (++fTrackIdx >= fTracksInEvent) {
            Bool_t event{fpReader->Next()};
            if (!fProgressByTracks)
                ProgressTick();
            if (event) {
                fTracksInEvent = (*fpRootTracks)->GetEntries();
                while (event && fTracksInEvent <= 0) {
                    event = fpReader->Next();
                    if (!fProgressByTracks)
                        ProgressTick();
                    fTracksInEvent = (*fpRootTracks)->GetEntries();
                }
            }
            // (!event) must be checked againg, don't use else
            if (!event) {   // no more events - no more tracks
                UnsetProperties();
                return;
            }
            if (fpRootEventHeader)
                fEventID = (*fpRootEventHeader)->GetEventId();
            fTrackIdx = 0;
            fEventTracks = 0;
            ScanEvent();
        }
        fpCurrentTrack = (Track_t*)(*fpRootTracks)->At(fTrackIdx);
    } while (!fTrackFilter(*fpCurrentTrack, CountUsableHits()));
    fTrackCount++;
    fEventTracks++;
    ResetHits();
    if (fProgressByTracks)
        ProgressTick();
}
 
/*
 *  Builds the vertex hit of the event just read. The beam line is fitted
 *  through the SiBT hits and taken at the target centre. Every track that
 *  passes the filters takes that hit; BmnAligner::Align() later disables the
 *  ones whose residual shows the track does not come from the vertex. The beam
 *  gives one measurement per event, so its weight is shared out among those N
 *  tracks instead of being counted N times.
 */
void BmnRootIterator::ScanEvent()
{
    fVertexTrack.clear();
    if (!fpRootBeamHits)
        return;
 
    const TClonesArray& beamHits{**fpRootBeamHits};
    if (SIBT_STATIONS != beamHits.GetEntriesFast())
        return;
    BmnWeightedLine lineX, lineY;
    UInt_t stations{0};
    for (Int_t i{0}; i < SIBT_STATIONS; i++) {
        const BmnHit& hit{*static_cast<BmnHit*>(beamHits.At(i))};
        if (hit.GetStation() >= 0 && hit.GetStation() < SIBT_STATIONS)
            stations |= 1u << hit.GetStation();
        lineX.Add(hit.GetZ(), hit.GetX(), BmnHitRecord::Weight(hit.GetDx()));
        lineY.Add(hit.GetZ(), hit.GetY(), BmnHitRecord::Weight(hit.GetDy()));
    }
    if (stations != (1u << SIBT_STATIONS) - 1)
        return;
 
    fVertexTrack.assign(fTracksInEvent, kFALSE);
    Int_t vertexTracks{0};
    for (Int_t i{0}; i < fTracksInEvent; i++) {
        const Track_t& track{*static_cast<Track_t*>((*fpRootTracks)->At(i))};
        Int_t usableHits{0};
        for (Int_t h{0}; h < track.GetNStsHits(); h++)
            if (fHitFilter(*static_cast<Hit_t*>((*fpRootHits)->At(track.GetStsHitIndex(h)))))
                usableHits++;
        if (!fTrackFilter(track, usableHits))
            continue;
        fVertexTrack[i] = kTRUE;
        vertexTracks++;
    }
    if (!vertexTracks)
        return;
 
    fVertexX = lineX.At(fVertexZ);
    fVertexY = lineY.At(fVertexZ);
    fVertexVarX = lineX.Var(fVertexZ);
    fVertexVarY = lineY.Var(fVertexZ);
    fVertexWx = 1.0 / (fVertexVarX * vertexTracks);
    fVertexWy = 1.0 / (fVertexVarY * vertexTracks);
}