BmnAlignerTest.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.ru
 
    BmnAlignerTest class implementation
*/
 
#include "BmnAlignerTest.h"
 
#include "BmnAlignDefines.h"
#include "Math/Functions.h"
#include "TGraph.h"
#include "TStopwatch.h"
#include "fairlogger/Logger.h"
 
#include <fstream>
#include <nlohmann/json.hpp>
 
BmnAlignerTest::BmnAlignerTest(const BmnMeasureModel* measureModel,
                               const BmnATestDetector* modelDetector,
                               Int_t nTracks,
                               const char* constraintsPath,
                               Int_t nThreads)
    : BmnAligner(measureModel, modelDetector, constraintsPath, nThreads)
    , fpDetectors(modelDetector)
{
    LOG(state) << " > BmnAlignerTest : initialization started";
 
    // BmnAligner protected fields initialization
    fTotalTracks = nTracks;
    fTracksUsableCnt = nTracks;
    // each track penetrates every (infinite) detector plane
    fTotalHits = fTotalTracks * BMN_MODULE_COUNT;
    LOG(detail) << "  total hits: " << fTotalHits << ", " << sizeof(BmnATestHit) << " bytes per hit";
    auto ram_size = fTotalHits * sizeof(BmnATestHit);
    LOG(info) << "    RAM required:    " << ram_size / 1073741824.0 << " GB, ";
    if (ram_size > BMN_MAX_HIT_RAM_SIZE) {
        LOG(error) << "<!> required RAM size is too large for test:" << std::endl
                   << "<!> decrease number of tracks or detector planes" << std::endl
                   << "<!> max RAM allowed: " << BMN_MAX_HIT_RAM_SIZE / 1073741824.0 << " GB";
        return;
    }
    // all dynamic memory allocated here is released in BmnAligner destructor
    fpFirstHits = new Int_t[fTotalTracks + 1];   // one extra element for iterator convenience
    fpTrackUsable = new std::vector<Bool_t>(fTracksUsableCnt, kTRUE);
    fpHitsData = new Byte_t[ram_size];
    fpHits = reinterpret_cast<BmnATestHit*>(fpHitsData);
    fpRamIterator = new BmnRamIterator(fpHits, fpFirstHits, fTotalTracks);
    fpIteratorMain = fpRamIterator;
 
    // prepare set of iterators for parallelization by data
    InitIterators();
 
    fInitialized = kTRUE;
    LOG(state) << " < BmnAlignerTest : initialization completed";
}
 
BmnAlignerTest::~BmnAlignerTest()   // virtual override
{
    if (fpTestTracksData)
        delete[] fpTestTracksData;
    // the rest of dynamic memory is released in BmnAligner destructor
}
 
Bool_t BmnAlignerTest::PrepareData(const char* tracksPath)   // virtual override
{
    LOG(state) << " > BmnAlignerTest : data preparation started";
    if (!fInitialized) {
        LOG(error) << "  <!> BmnAlignerTest is not initialized";
        return kFALSE;
    }
    TStopwatch Timer;
    Bool_t loadTracks{tracksPath && std::filesystem::exists(tracksPath)};
 
    // tracks data preparation
    if (loadTracks) {
        if (!LoadTracks(tracksPath)) {
            LOG(error) << "  <!> failed to load tracks from file " << tracksPath;
            return kFALSE;
        }
    } else {
        LOG(info) << "    Generating random tracks";
        fpTestTracksData = new Byte_t[fTotalTracks * sizeof(BmnATestTrack)];
        fpTestTracks = reinterpret_cast<BmnATestTrack*>(fpTestTracksData);
        for (auto iTrack = 0; iTrack < fTotalTracks; iTrack++)
            new (fpTestTracks + iTrack)
                BmnATestTrack(BMN_TEST_ALPHA1_SIGMA, BMN_TEST_ALPHA2_MIN, BMN_TEST_ALPHA2_MAX,
                              BMN_MODULE_COUNT   // in test procedure each track penetrates every detector plane
                );
        if (tracksPath && !SaveTracks(tracksPath))
            LOG(warning) << "  failed to save tracks to file " << tracksPath;
    }
 
    // calculate hits as intersections of track and detector planes
    SVect3 HitPoint, HitPointLocal;
    Double_t Z;
    Int_t iTrack, iHit{0};
    for (iTrack = 0; iTrack < fTotalTracks; iTrack++) {
        fpFirstHits[iTrack] = iHit;
        BmnATestTrack& Track = fpTestTracks[iTrack];
        for (Int_t iDet = 0; iDet < BMN_MODULE_COUNT; iDet++, iHit++) {
            const BmnATestDetectorPlane& Detector = fpDetectors->GetDetectors()[iDet];
            // generate true hit in lab coordinates
            Z = (-Detector.GetConst() - ROOT::Math::Dot(Detector.GetNormal(), Track.GetVertex()))
                / ROOT::Math::Dot(Detector.GetNormal(), Track.GetDirection());
            HitPoint = {Track.X(Z), Track.Y(Z), Z};
            // recalculate hit to local coordinates of detector
            HitPointLocal = Detector.GetCrdInvRotation() * (HitPoint - Detector.GetCrdShift());
            LOG_IF(error, HitPointLocal(2) > BMN_CLOSE_TO_ZERO)
                << "    <!> non-zero local z coordinate of test hit: " << HitPointLocal(2);
            // store hit as data measured by detector
            auto pHit = new (fpHits + iHit)
                BmnATestHit(HitPointLocal(0) + BMN_RND.Gaus(0.0, BMN_TEST_AVERAGE_DX),   // x measurement + noise
                            HitPointLocal(1) + BMN_RND.Gaus(0.0, BMN_TEST_AVERAGE_DY),   // y measurement + noise
                            Detector.GetPosition(),   // z = nominal position of detector
                            BMN_TEST_AVERAGE_DX, BMN_TEST_AVERAGE_DY, &Detector);
            Track.AddHit(*pHit);   // tracks should remember its hits for visualization purposes
        }   // for iDet
 
    }   // for iTrack
    fpFirstHits[iTrack] = iHit;   // last syntetic element for iterator convenience
 
    // calculate zero approximation of linear tracks
    std::vector<SVectLC> arrAlpha(fTotalTracks);
    for (iTrack = 0; iTrack < fTotalTracks; iTrack++) {
        BmnATestHit &firstHit = fpHits[fpFirstHits[iTrack]], &lastHit = fpHits[fpFirstHits[iTrack + 1] - 1];
        Double_t dX{lastHit.GetX() - firstHit.GetX()}, dY{lastHit.GetY() - firstHit.GetY()},
            dZ{lastHit.GetZ() - firstHit.GetZ()}, AlphaX1{dX / dZ}, AlphaY1{dY / dZ};
        arrAlpha[iTrack] = SVectLC(firstHit.GetX() - AlphaX1 * firstHit.GetZ(), AlphaX1,
                                   firstHit.GetY() - AlphaY1 * firstHit.GetZ(), AlphaY1);
    }
 
    // MSE scale factors calculation
    fOmegaScaleFactor = (BMN_TEST_AVERAGE_DX + BMN_TEST_AVERAGE_DY) * 0.5;
    fOmegaScaleFactor *= fOmegaScaleFactor;
 
    StoreZeroSolution(arrAlpha);
 
    /*
    // DEBUG : Test on precise solution
    for (Int_t i = 0; i < BMN_MODULE_COUNT; i++)
        fpResultCurrent->pA[i] = fpDetectors->GetDetectors()[i].GetAlignShift();
    for (Int_t i = 0; i < fTotalTracks; i++)
        fpResultCurrent->pAlpha[i] = fpTestTracks[i].GetAlpha();
    CalculateMSE(kTRUE);
    */
 
    Timer.Stop();
    LOG(info) << "    " << fTotalTracks << " tracks ready, total hits: " << iHit << ", time spent: " << Timer.CpuTime()
              << " s (CPU)";
    LOG(state) << " < BmnAlignerTest : data preparation completed";
    return LoadConstraints();
}
 
Bool_t BmnAlignerTest::IterateAlignment()   // virtual override
{
    Bool_t success = BmnAligner::IterateAlignment();
    if (!success)
        return kFALSE;
 
    // monitor local parameters convergence
    SVectLC delta;
    for (auto i = 0; i < TotalTracks(); i++) {
        const SVectLC& Alpha = GetResult()->Alpha(i);
        delta += ROOT::Math::fabs(Alpha - GetTestTrack(i).GetAlpha());
    }
    delta /= TotalTracks();
    LOG(info) << "      === average local parameters convergence:";
    LOG(info) << "        error by X = (" << delta(iX0) << ", " << delta(iX1) << ")";
    LOG(info) << "        error by Y = (" << delta(iY0) << ", " << delta(iY1) << ")";
 
    return kTRUE;
}
 
Bool_t BmnAlignerTest::SaveTracks(const char* path) const
{
    LOG(info) << "    Saving tracks data to " << path;
    if (!fpTestTracks || fTotalTracks <= 0)
        return kFALSE;
    std::ofstream ofs(path);
    if (!ofs.good()) {
        LOG(error) << "  <!> failed to open JSON file " << path;
        return kFALSE;
    }
 
    nlohmann::json j;
    j["tracks"] = fTotalTracks;
    j["alphas"] = nlohmann::json::array();
    for (Int_t i = 0; i < fTotalTracks; i++) {
        const auto& tr = fpTestTracks[i];
        const auto& Alpha = tr.GetAlpha();
        j["alphas"].push_back({Alpha(iX0), Alpha(iX1), Alpha(iY0), Alpha(iY1)});
    }
 
    ofs << j.dump(2);
    return kTRUE;
}
 
Bool_t BmnAlignerTest::LoadTracks(const char* path)
{
    LOG(info) << "    Loading tracks data from " << path;
    std::ifstream ifs(path);
    if (!ifs.good()) {
        LOG(error) << "  <!> failed to open JSON file " << path;
        return kFALSE;
    }
 
    nlohmann::json j;
    ifs >> j;
 
    // re-allocate new buffer and construct records
    fTotalTracks = j.value("tracks", 0);
    if (fpTestTracksData)
        delete[] fpTestTracksData;
    fpTestTracksData = new Byte_t[fTotalTracks * sizeof(BmnATestTrack)];
    fpTestTracks = reinterpret_cast<BmnATestTrack*>(fpTestTracksData);
 
    const auto& arr = j["alphas"];
    for (Int_t i = 0; i < fTotalTracks; i++) {
        const auto& a = arr[i];
        SVectLC Alpha(a[iX0].get<double>(), a[iX1].get<double>(), a[iY0].get<double>(), a[iY1].get<double>());
        new (fpTestTracks + i) BmnATestTrack(Alpha, BMN_MODULE_COUNT);
    }
    return kTRUE;
}
 
// visualization works in ROOT macro mode only
void BmnAlignerTest::Draw()
{
    PrepareDrawing();
 
    // Draw tracks
    fpPad3D->cd();
    fpPad3D->SetTitle("Some random tracks");
    const Double_t zMax = fpDetectors->GetMaxZ() * 1.1;
    for (Int_t i = 0; i < BMN_MODULE_COUNT; i++)
        fpDetectors->GetDetectors()[i].Draw();
    TRandom trueRnd{0};
    std::vector<Int_t> drawn_tracks(BMN_DRAW_TRACKS_CNT, -1);
    for (Int_t i = 0, j; i < BMN_DRAW_TRACKS_CNT; i++) {
        j = trueRnd.Uniform(0, TotalTracks());
        if (drawn_tracks.end() == std::find(drawn_tracks.begin(), drawn_tracks.end(), j) && fpTrackUsable->at(i)) {
            GetTestTrack(j).Draw(kTRUE, zMax);
            BmnATestTrack trk(GetResult()->Alpha(j), BMN_MODULE_COUNT);
            trk.Draw(kFALSE, zMax);
            drawn_tracks[i] = j;
        } else
            i--;
    }
 
    double xValues[BMN_GLPARAMS_TOTAL];
    for (auto i = 0; i < BMN_GLPARAMS_TOTAL; i++)
        xValues[i] = static_cast<double>(i + 1);
 
    // Draw solution
    Double_t A[BMN_GLPARAMS_TOTAL];
    for (Int_t i = 0, k = 0; i < BMN_MODULE_COUNT; i++)
        for (Int_t j = 0; j < BMN_GLOBAL_PARAMS_PD; j++, k++)
            A[k] = GetResult()->A(i)(j);
    fpPadG1->cd();
    fpPadG1->SetTitle("Solution");
    TGraph* graphSolution = new TGraph(BMN_GLPARAMS_TOTAL, xValues, A);
    graphSolution->SetTitle("Alignment values");
    graphSolution->SetFillColor(kRed);
    graphSolution->SetMinimum(-0.25);   // DEBUG
    graphSolution->SetMaximum(0.35);    // DEBUG
    graphSolution->GetXaxis()->SetRangeUser(xValues[0], xValues[BMN_GLPARAMS_TOTAL - 1]);
    graphSolution->Draw("AB");
    Double_t T[BMN_GLPARAMS_TOTAL];
    for (Int_t i = 0, k = 0; i < BMN_MODULE_COUNT; i++)
        for (Int_t j = 0; j < BMN_GLOBAL_PARAMS_PD; j++, k++)
            T[k] = fpDetectors->GetTestA()[i](j);
    TGraph* graphTest = new TGraph(BMN_GLPARAMS_TOTAL, xValues, T);
    graphTest->SetLineColor(kBlack);
    graphTest->Draw("*SAME");
 
    /*
    // Draw SVD singularity values
    TVectorD SVDsig {GetResult()->SVDSig()};
    fpPadG4->cd();
    fpPadG4->SetTitle("SVD singularity values");
    fpPadG4->SetLogy();
    TGraph* graphSVD = new TGraph(
        BMN_GLPARAMS_TOTAL,
        xValues,
        SVDsig.GetMatrixArray()
    );
    // narrow drawing range to interesting part
    Int_t SVDdrawIdx;
    for (SVDdrawIdx = 0; SVDdrawIdx < SVDsig.GetNrows(); SVDdrawIdx++)
        if (SVDsig(SVDdrawIdx) <= BMN_SVD_LOW) break;
    SVDdrawIdx = std::max(
        0,
        std::min(
            SVDdrawIdx,
            BMN_GLPARAMS_TOTAL - 10
        )
    );
    graphSVD->GetXaxis()->SetRangeUser(
        static_cast<Double_t>(SVDdrawIdx),
        static_cast<Double_t>(BMN_GLPARAMS_TOTAL)
    );
    graphSVD->SetTitle("SVD singularity values (tale)");
    graphSVD->SetFillColor(kGray);
    graphSVD->Draw("AB");
    */
 
    DrawResiduals();
}
 
void BmnAlignerTest::ReportResults() const
{
    BmnAligner<BmnATestHit>::ReportResults();
    const SVectGL* Atest{fpDetectors->GetTestA()};
    BmnAlignResult *pBeforeResult{GetResult(kFALSE)}, *pAfterResult{GetResult(kTRUE)};
    Double_t beforeValue{0.0}, afterValue{0.0};
    for (Int_t i = 0; i < BMN_MODULE_COUNT; i++) {
        beforeValue += ROOT::Math::Mag2(Atest[i] - pBeforeResult->A(i));
        afterValue += ROOT::Math::Mag2(Atest[i] - pAfterResult->A(i));
    }
    LOG(info) << " === distance to test A before/after alignment: " << std::sqrt(beforeValue) << " / "
              << std::sqrt(afterValue);
}
 
Double_t BmnAlignerTest::CalculateMSE(Bool_t withResiduals)
{
    fpResultCurrent->ResetValueMSE();
    if (withResiduals) {
        fpResultCurrent->ResidualsX().Reset();
        fpResultCurrent->ResidualsY().Reset();
    }
    for (Int_t iTrack = 0; iTrack < fTotalTracks; iTrack++) {
        if (!fpTrackUsable->at(iTrack))
            continue;
        fpTestTracks[iTrack].AddMSE(this, withResiduals);
    }
    return fpResultCurrent->GetValueMSE();
}