BmnFunctionAna.cxx

Go to the documentation of this file.

#include "BmnFunctionAna.h"
 
#include "BmnFieldMapSym3.h"
#include "BmnGlobalTrack.h"
#include "BmnMCInfoDst.h"
#include "BmnNewFieldMap.h"
#include "CbmMCTrack.h"
#include "CbmTrackMatch.h"
#include "CbmVertex.h"
#include "DstRunHeader.h"
#include "FairLogger.h"
#include "FairTrackParam.h"
#include "TChain.h"
#include "TClonesArray.h"
#include "TFile.h"
#include "TMath.h"
#include "TStopwatch.h"
#include "TSystem.h"
 
#include <ROOT/RDFHelpers.hxx>
#include <ROOT/RDataFrame.hxx>
#include <fstream>
#include <iostream>
#include <list>
#include <string>
using namespace std;
 
// get magnetic field and run information from a specified DST file
// fileName: DST file path [in]
// fieldScale: magnetic field scale [out]; fieldType: magnetic field type [out]
// beam, target: names of the beam and target, e.g. 'Xe' and 'CsI (2%)' [out, optional]
// beam_energy: beam kinetic energy per nucleon in AGeV [out, optional]
// beam_a, beam_z, target_a, target_z: mass and charge numbers of the beam and target nuclei [out, optional]
// NB: the names are empty and A and Z are equal to -1 if the nucleus was not defined at the reconstruction stage,
// target A is equal to 0 if the empty target
// returns magnetic field, nullptr - if some errors occured
FairField* BmnFunctionAna::GetDSTField(TString dstFileName,
                                       Double_t& fieldScale,
                                       Int_t& fieldType,
                                       TString* beam,
                                       TString* target,
                                       Double_t* beam_energy,
                                       Int_t* beam_a,
                                       Int_t* beam_z,
                                       Int_t* target_a,
                                       Int_t* target_z)
{
    gSystem->ExpandPathName(dstFileName);
    if (gSystem->AccessPathName(dstFileName.Data()) == true) {
        LOG(error) << "GetDSTField: no specified file: " << dstFileName;
        return nullptr;
    }
 
    TFile* dst_file = TFile::Open(dstFileName.Data());
    if (dst_file->IsZombie()) {
        LOG(error) << "GetDSTField: opening the input file failed: " << dstFileName;
        return nullptr;
    }
 
    // Get a pointer to "DstRunHeader"
    DstRunHeader* pRunHeader = dynamic_cast<DstRunHeader*>(dst_file->Get("DstRunHeader"));
    if (!pRunHeader) {
        LOG(error) << "GetDSTField: No 'DstRunHeader' object found in the input file: " << dstFileName;
        dst_file->Close();
        return nullptr;
    }
 
    // get the beam and target information from the 'DstRunHeader' object
    if (beam)
        *beam = pRunHeader->GetBeam();
    if (target)
        *target = pRunHeader->GetTarget();
    if (beam_a)
        *beam_a = pRunHeader->GetBeamA();
    if (beam_z)
        *beam_z = pRunHeader->GetBeamZ();
    if (beam_energy)
        *beam_energy = pRunHeader->GetBeamEnergy();
    if (target_a)
        *target_a = pRunHeader->GetTargetA();
    if (target_z)
        *target_z = pRunHeader->GetTargetZ();
 
    // get field type and scale from the 'DstRunHeader' object
    fieldType = pRunHeader->GetFieldType();
    fieldScale = pRunHeader->GetFieldScale();
    TString fieldName = pRunHeader->GetFieldName();
    FairField* magField = nullptr;
    switch (fieldType) {
        case 1: {
            magField = new BmnNewFieldMap(fieldName.Data());
            (dynamic_cast<BmnNewFieldMap*>(magField))->SetScale(fieldScale);
            magField->Init();
            break;
        }
        case 3: {
            magField = new BmnFieldMapSym3(fieldName.Data());
            (dynamic_cast<BmnFieldMapSym3*>(magField))->SetScale(fieldScale);
            magField->Init();
            break;
        }
        default:
            LOG(error) << "GetDSTField: This type of the magnetic field is not supported: " << fieldType;
    }
 
    dst_file->Close();
    return magField;
}
 
// write the reconstructed global tracks from the simulated DST files to a text file
int BmnFunctionAna::ConvertRecoToTxt(TString input_list_file, TString output_file, TString configFileName)
{
    TStopwatch timer;
    timer.Start();
 
    gSystem->ExpandPathName(input_list_file);
 
    list<string> file_list;
    ifstream list_file(input_list_file.Data());
    string s;
    while (getline(list_file, s))
        file_list.push_back(s);
    if (file_list.empty()) {
        cout << "WARNING: There are no files in the list of the input text file: " << input_list_file << endl;
        return 0;
    }
 
    TChain chain("bmndata");
    for (auto const& cur_file_path : file_list) {
        TString strFilePath(cur_file_path);
        gSystem->ExpandPathName(strFilePath);
 
        chain.Add(strFilePath);
    }
 
    /*
    // read configuration variables
    configFileName = gSystem->ExpandPathName("$VMCWORKDIR/config/" + configFileName);
    ifstream configFile(configFileName.Data());
    if (!configFile.is_open()) {
        cout << "[ERROR] Cannot open config file " << configFileName << endl;
        return -2;
    }
    json config;
    try {
        configFile >> config;
    }
    catch (const json::parse_error& ex) {
        cout << "[ERROR] Wrong config file format: " << ex.what() << endl;
        return -3;
    }
    Int_t minNTracks = config.value("minNTracks", 2);
    Double_t minZ = config.value("minZ", -TMath::Infinity());
    Double_t maxZ = config.value("maxZ", TMath::Infinity());
    vector<string> excludedBranches = config.value("excludedBranches", vector<string>());
 
    if (outFormat != "root") {
        cout << "[ERROR] Only \"root\" format is supported at the moment" << endl;
        return -4;
    }
    */
    /*
        TFile* inFile = TFile::Open(inFilePath.Data());
        if (inFile == nullptr) {
            cout << "[ERROR] Cannot open input DST file " << inFilePath << endl;
            return -5;
        }
        TTree* inTree = (TTree*) inFile->Get("bmndata");
        if (!inTree) {
            cout << "[ERROR] Input file " << inFilePath << " does not contain \"bmndata\" tree" << endl;
            return -6;
        }
    *//*
    // exclude branches
    vector<string> branches;
    for (TObject* branch: *inTree->GetListOfBranches()) {
        string branchName = ((TBranch*)branch)->GetName();
        if (find(excludedBranches.begin(), excludedBranches.end(), branchName) == excludedBranches.end())
            branches.push_back(branchName);
    }
*/
    ofstream file(output_file.Data());
    file << "event# firstVx firstVy firstVz lastVx lastVy lastVz firstTx firstTy lastTx lastTy firstQp lastQp "
            "firstCov[0] firstCov[1] firstCov[2] firstCov[3] firstCov[4] firstCov[5] firstCov[6] firstCov[7] "
            "firstCov[8] firstCov[9] firstCov[10] firstCov[11] firstCov[12] firstCov[13] firstCov[14] "
            "lastCov[0] lastCov[1] lastCov[2] lastCov[3] lastCov[4] lastCov[5] lastCov[6] lastCov[7] "
            "lastCov[8] lastCov[9] lastCov[10] lastCov[11] lastCov[12] lastCov[13] lastCov[14] "
            "Chi2 NDF Length Nhits Beta400 Beta700 fDCAInVertex PDGCode"
         << endl;
 
    // velocity measured by ToF-400 (tof_id = 1) or ToF-700 (tof_id = 2). The global track does not store it since the
    // change of the BmnGlobalTrack format, but it stores the squared mass computed from it by BmnGlobalTracking as
    // m^2 = p^2 * (1/beta^2 - 1) with p = 1/(q/p) at the first track point, so the velocity is restored from the mass.
    // -1000 is returned if the track has no ToF measurement, as the former BmnGlobalTrack::GetBeta() did
    auto get_beta = [](BmnGlobalTrack* track, Int_t tof_id) -> Double_t {
        Float_t mass2 = track->GetMass2(tof_id);
        Double_t qp = track->GetParamFirst()->GetQp();
        if ((mass2 == -1000.0) || (qp == 0))
            return -1000.0;
        Double_t denominator = 1. + mass2 * qp * qp;
        return (denominator > 0) ? 1. / TMath::Sqrt(denominator) : -1000.0;
    };
 
    auto save_to_txt = [&file, &get_beta](const TClonesArray& global_tracks, const TClonesArray& sts_tracks,
                                          const TClonesArray& sts_match, const BmnMCInfoDst& mc_info, int col_iter) {
        for (int i = 0; i < global_tracks.GetEntries(); i++) {
            BmnGlobalTrack* cur_global_track = static_cast<BmnGlobalTrack*>(global_tracks.UncheckedAt(i));
 
            // check that there is information on Beta from TOF-400 or TOF-700
            // (an invalid Beta is defined by the range check, consistently with BmnGlobalTrack::GetMass2)
            Double_t beta400 = get_beta(cur_global_track, 1), beta700 = get_beta(cur_global_track, 2);
            if ((beta400 < -999.0) and (beta700 < -999.0))
                continue;
 
            // get PDG code
            CbmTrackMatch* sts_match_track =
                static_cast<CbmTrackMatch*>(sts_match.UncheckedAt(cur_global_track->GetGemTrackIndex()));
            if (sts_match_track->GetMCTrackId() < 0)
                continue;
            CbmMCTrack* mc_track =
                static_cast<CbmMCTrack*>(mc_info.GetMCTracks()->UncheckedAt(sts_match_track->GetMCTrackId()));
            // cout<<"PDG code = "<<mc_track->GetPdgCode()<<endl;
 
            FairTrackParam* param_first = cur_global_track->GetParamFirst();
            FairTrackParam* param_last = cur_global_track->GetParamLast();
 
            file << col_iter << " " << param_first->GetX() << " " << param_first->GetY() << " " << param_first->GetZ()
                 << " " << param_last->GetX() << " " << param_last->GetY() << " " << param_last->GetZ() << " "
                 << param_first->GetTx() << " " << param_first->GetTy() << " " << param_last->GetTx() << " "
                 << param_last->GetTy() << " " << param_first->GetQp() << " " << param_last->GetQp() << " ";
            Double_t cov_matrix1[15];
            param_first->CovMatrix(cov_matrix1);
            for (int j = 0; j < 15; j++)
                file << cov_matrix1[j] << " ";
            // the covariance matrix at the last point of the track ('lastCov' columns)
            Double_t cov_matrix2[15];
            param_last->CovMatrix(cov_matrix2);
            for (int j = 0; j < 15; j++)
                file << cov_matrix2[j] << " ";
            file << cur_global_track->GetChi2() << " " << cur_global_track->GetNDF() << " "
                 << cur_global_track->GetLength() << " " << cur_global_track->GetNHits() << " " << beta400 << " "
                 << beta700 << " " << cur_global_track->GetDCAInVertex() << " " << mc_track->GetPdgCode() << endl;
        }
        return;
    };
 
    // function to cut events without vertex
    auto cutVertex = [](const CbmVertex& vtx) { return (vtx.GetNTracks() > 1); };
 
    int col_iter = -1;
    ROOT::RDataFrame data(chain);
    auto w_iterator =
        data.Filter(cutVertex, {"PrimaryVertex."})
            //.Define("pdg_code", definePDG, {"BmnGlobalTrack", "StsVector", "StsTrackMatch", "BmnMCInfo."})
            .Define("col_iter", [&col_iter] { return ++col_iter; });
 
    ROOT::RDF::Experimental::AddProgressBar(data);
    w_iterator.Foreach(save_to_txt, {"BmnGlobalTrack", "StsVector", "StsTrackMatch", "BmnMCInfo.", "col_iter"});
 
    file.close();
    timer.Stop();
    Double_t rtime = timer.RealTime(), ctime = timer.CpuTime();
    printf("RealTime=%f seconds, CpuTime=%f seconds\n", rtime, ctime);
    return 0;
}