expdataNN.C

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.ru
 
    ROOT data export macro for NN alignment of each detector element.
    Exports hits for each detector element into individual CSV files.
    For each track computes best-fit straight line through all its hits
    and stores the line coefficients alongside hit data.
 
    Line fit model (independent fits in xz and yz planes):
        x = xA + xB * z
        y = yA + yB * z
    Ordinary least squares, minimizing residuals in x and y.
*/
 
#ifndef __CLING__
#include "BmnRootIterator.h"
#include "BmnRun8cgDetModel.h"
#include "BmnSimpleProgressBar.h"
#include "TStopwatch.h"
 
#include <fstream>
#include <map>
#include <string>
#include <vector>
#endif
 
#define BMN_IN_FILELIST "../cfg/dst7466noa-full.txt"
#define BMN_EXPORT_PREFIX "/home/igor/DATA/bmn/nn/data_det"
 
int expdataNN(int argc = 0, char** argv = nullptr)
{
    TStopwatch timer;
 
    // prepare ROOT file readers
    BmnRun8cgDetModel detModel;
    detModel.Init();
    BmnSimpleProgressBar pb;
    BmnRootIterator data(BMN_IN_FILELIST, detModel, &pb);
 
    std::cout << "Exporting data defined in " << BMN_IN_FILELIST << std::endl;
    std::cout << data.GetTotalEvents() << " total events to process" << std::endl;
    std::cout << detModel.Alignables() << " detector elements expected" << std::endl;
 
    // per-detectorID output files and hit counters (opened on first use)
    std::map<Int_t, std::ofstream*> files;
    std::map<Int_t, Int_t> detHitCnt;
 
    auto GetFile = [&](Int_t detID) -> std::ofstream& {
        auto it = files.find(detID);
        if (it != files.end())
            return *(it->second);
        std::string name = std::string(BMN_EXPORT_PREFIX) + std::to_string(detID) + ".csv";
        auto* f = new std::ofstream(name, std::ofstream::out);
        *f << "x,y,z,xA,xB,yA,yB\n";
        files[detID] = f;
        detHitCnt[detID] = 0;
        return *f;
    };
 
    // process ROOT data
    Int_t trackID{0}, hitCnt{0};
    struct HitRec
    {
        Int_t detID;
        Double_t x, y, z;
    };
 
    pb.Init();
    data.ResetAll();
    while (!data.EndOfTracks()) {
 
        // accumulate regression sums and collect hits in one pass
        Double_t xSum{0.0}, ySum{0.0}, zSum{0.0}, z2Sum{0.0}, zxSum{0.0}, zySum{0.0};
        std::vector<HitRec> hits;
        hits.reserve(data.HitsInTrack());
 
        while (!data.EndOfHits()) {
            hits.push_back({data.HitDetectorID(), data.HitX(), data.HitY(), data.HitZ()});
            xSum += data.HitX();
            ySum += data.HitY();
            zSum += data.HitZ();
            z2Sum += data.HitZ() * data.HitZ();
            zxSum += data.HitZ() * data.HitX();
            zySum += data.HitZ() * data.HitY();
            data.NextHit();
        }
 
        // linear regression fit
        Double_t n{static_cast<Double_t>(data.HitsInTrack())}, factor{1.0 / (n * z2Sum - zSum * zSum)},
            xA{(xSum * z2Sum - zSum * zxSum) * factor}, yA{(ySum * z2Sum - zSum * zySum) * factor},
            xB{(n * zxSum - zSum * xSum) * factor}, yB{(n * zySum - zSum * ySum) * factor};
 
        // write each hit to its detector file
        for (const auto& h : hits) {
            std::ofstream& f = GetFile(h.detID);
            f << h.x << "," << h.y << "," << h.z << "," << xA << "," << xB << "," << yA << "," << yB << "\n";
            detHitCnt[h.detID]++;
            hitCnt++;
        }
 
        trackID++;
        data.NextTrack();
    }
    pb.Clear();
 
    // close all output files
    for (auto& kv : files) {
        kv.second->close();
        delete kv.second;
    }
 
    std::cout << "Tracks exported: " << trackID << std::endl;
    std::cout << "Hits exported:   " << hitCnt << std::endl;
    std::cout << "Output files:    " << files.size() << std::endl;
 
    if (files.size() != detModel.Alignables())
        std::cout << "WARNING! Files count does not match detector count:"
                  << " some data may be missing!" << std::endl;
    for (const auto& kv : detHitCnt)
        std::cout << "  det ID " << std::setw(10) << kv.first << " : " << std::setw(10) << kv.second << " hits\n";
 
    // no memory cleanup in ROOT macro mode
    std::cout << "Done in " << timer.RealTime() << "s" << std::endl;
    return 0;
}