EGamma_ElectronTrackingValidation/analysis/tracking_validation.cpp

#include <iostream>
#include <vector>
#include <utility>
#include <numeric>
#include <limits>
#include <cmath>

#include <TSystem.h>

#include "filval/filval.hpp"
#include "filval/root/filval.hpp"

#include <boost/range/combine.hpp>
#include <boost/format.hpp>

#include "TrackingNtuple.h"
#include "analysis/obj_types.cpp"

using namespace std;
using namespace fv;
using namespace fv::root;

using namespace std;

#define HIT_TYPE_PIXEL 0
#define HIT_TYPE_GLUED 1
#define HIT_TYPE_STRIP 2

vector<string> seedTypes =
    {"initialStepSeeds",
     "highPtTripletStepSeeds",
     "mixedTripletStepSeeds",
     "pixelLessStepSeeds",
     "tripletElectronSeeds",
     "pixelPairElectronSeeds",
     "stripPairElectronSeeds"};
int nSeedTypes = seedTypes.size();

Value<vector<Seed>>* seeds;
Value<vector<PixRecHit>>* pixrec_hits;
Value<vector<Track>>* tracks;
Value<vector<SimHit>>* sim_hits;
Value<vector<SimTrack>>* sim_tracks;

void register_objects(TrackingDataSet& tds){
    seeds = register_seeds(tds);

    pixrec_hits = register_pixrec_hits(tds);
    tracks = register_tracks(tds);

    sim_hits = register_sim_hits(tds);
    sim_tracks = register_sim_tracks(tds);
}

void setup_matched_tracks(TrackingDataSet& tds){
    // Finds pairs of (track,sim_track) where the cycle
    // track->seed->rec_hit->sim_hit->sim_track->track
    // is closed
    typedef std::tuple<Seed, PixRecHit, Track, SimHit, SimTrack> MatchedTrack;

    auto& find_matched_tracks =
    func<vector<MatchedTrack>(vector<Seed>,
                              vector<PixRecHit>,
                              vector<Track>,
                              vector<SimHit>,
                              vector<SimTrack>)>("find_matched_tracks",
        FUNC(([](const vector<Seed>& seeds,
                 const vector<PixRecHit> pixrec_hits,
                 const vector<Track>& tracks,
                 const vector<SimHit> sim_hits,
                 const vector<SimTrack> sim_tracks){
            INFO("New event");
            INFO(boost::format("Number of tracks is %d, number of seeds is %d, number of hits %d, number of simhits %d")
                    % tracks.size() % seeds.size() % pixrec_hits.size() % sim_hits.size());
            vector<MatchedTrack> matched_tracks;
            for(const Track &track : tracks){
                const Seed& seed = seeds[track.seedIdx];
                if( seed.algoOriginal < 0 || seed.algoOriginal >= nSeedTypes) continue;
                INFO(boost::format("       seed indx= %d   algorithm= %d   type %s\n")
                        % track.seedIdx % seed.algoOriginal % seedTypes[seed.algoOriginal]);
                for(auto tup : boost::combine(seed.hitIdx, seed.hitType)){ //for hits in track's seed
                    int hitIdx, hitType;
                    boost::tie(hitIdx, hitType) = tup;
                    if(hitType != HIT_TYPE_PIXEL) continue; // take only pixel hits for now
                    const PixRecHit &rec_hit = pixrec_hits[hitIdx];

                    TVector3 recoHitPoint;
                    recoHitPoint.SetXYZ( rec_hit.x, rec_hit.y, rec_hit.z);
                    float Rreco = recoHitPoint.Perp();
                    INFO(boost::format("            hit= %d is at radius= %f\n") % hitIdx % Rreco);

                    for(const int& simHitIdx : rec_hit.simHitIdx){ // for sim-hits matching rec-hit
                        const SimHit &sim_hit = sim_hits[simHitIdx];
                        const SimTrack &sim_track = sim_tracks[sim_hit.simTrkIdx];
                        for(const int &trkIdx : sim_track.trkIdx){ // for tracks matched to this sim_track
                            if(trkIdx == track.idx){  // sim_track matches with our original track
                                matched_tracks.push_back({seed, rec_hit, track, sim_hit, sim_track});
                                break;
                            }
                        }
                    }
                }
            }
            return matched_tracks;
        })));

    auto matched_tracks =
        fv::apply(find_matched_tracks, fv::tuple(seeds, pixrec_hits, tracks, sim_hits, sim_tracks), "matched_tracks");

    auto delta_pt = fv::map(func<float(MatchedTrack)>("matched_track_delta_pt",
        FUNC(([](const MatchedTrack& matched_track){
            const Track& track = std::get<Track>(matched_track);
            const SimTrack& sim_track = std::get<SimTrack>(matched_track);
            return track.pt - sim_track.pt;
        }))), matched_tracks, "matched_track_delta_pt");

    tds.register_container<ContainerTH1Many<float>>("matched_track_delta_pt", "Matched Track delta P_t",
                                                      delta_pt, 20, -10, 10, "P_t");
    /* auto dphi_vs_eta_B1 = */
    /*     fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("1", 1)), "dphi_vs_eta_B1"); */
    /* auto dphi_vs_eta_B2 = */
    /*     fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("2", 2)), "dphi_vs_eta_B2"); */
    /* auto dphi_vs_eta_B3 = */
    /*     fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("3", 3)), "dphi_vs_eta_B3"); */
    /* auto dphi_vs_eta_B4 = */
    /*     fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("4", 4)), "dphi_vs_eta_B4"); */


    /* tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B1", "dphi vs eta (BPIX L1)", */
    /*                                                   dphi_vs_eta_B1, 20, -3, 3, 100, -0.002, 0.002, */
    /*                                                   "eta",  "dphi"); */
    /* tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B2", "dphi vs eta (BPIX L2)", */
    /*                                                   dphi_vs_eta_B2, 20, -3, 3, 100, -0.002, 0.002, */
    /*                                                   "eta",  "dphi"); */
    /* tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B3", "dphi vs eta (BPIX L3)", */
    /*                                                   dphi_vs_eta_B3, 20, -3, 3, 100, -0.002, 0.002, */
    /*                                                   "eta",  "dphi"); */
    /* tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B4", "dphi vs eta (BPIX L4)", */
    /*                                                   dphi_vs_eta_B4, 20, -3, 3, 100, -0.002, 0.002, */
    /*                                                   "eta",  "dphi"); */

}


/* void setup_dz_vs_eta(TrackingDataSet& tds){ */
/*     auto& calc_dz_vs_eta = */
/*     func<std::pair<vector<float>,vector<float>>(vector<Seed>, */
/*                                                 vector<PixRecHit>, */
/*                                                 vector<SimHit>, int)>("calc_dz_vs_eta", */
/*             FUNC(([](const vector<Seed>& seeds, */
/*                      const vector<PixRecHit> pixrec_hits, */
/*                      const vector<SimHit> sim_hits, */
/*                      int layer){ */
/*                 vector<float> dz; */
/*                 vector<float> eta; */
/*                 for(const Seed &seed : seeds){ */
/*                     int hitIdx, hitType; */
/*                     for(auto tup : boost::combine(seed.hitIdx, seed.hitType)){ */
/*                         boost::tie(hitIdx, hitType) = tup; */
/*                         if(hitType == HIT_TYPE_PIXEL){ */
/*                             const PixRecHit &hit = pixrec_hits[hitIdx]; */
/*                             if(layer && hit.lay != layer) continue; */
/*                             for(const int& simHitIdx : hit.simHitIdx){ */
/*                                 dz.push_back(hit.z-sim_hits[simHitIdx].z); */
/*                                 eta.push_back(seed.eta); */
/*                             } */
/*                         } */
/*                     } */
/*                 } */
/*                 return std::make_pair(eta,dz); */
/*             }))); */

/*     auto dz_vs_eta_B1 = */
/*     fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("1", 1)), "dz_vs_eta_B1"); */
/*     auto dz_vs_eta_B2 = */
/*     fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("2", 2)), "dz_vs_eta_B2"); */
/*     auto dz_vs_eta_B3 = */
/*     fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("3", 3)), "dz_vs_eta_B3"); */
/*     auto dz_vs_eta_B4 = */
/*     fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("4", 4)), "dz_vs_eta_B4"); */


/*     tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B1", "dz vs eta (BPIX L1)", */
/*                                                       dz_vs_eta_B1, 20, -3, 3, 20, -.01, .01, */
/*                                                       "eta",  "dz"); */
/*     tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B2", "dz vs eta (BPIX L2)", */
/*                                                       dz_vs_eta_B2, 20, -3, 3, 20, -.01, .01, */
/*                                                       "eta",  "dz"); */
/*     tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B3", "dz vs eta (BPIX L3)", */
/*                                                       dz_vs_eta_B3, 20, -3, 3, 20, -.01, .01, */
/*                                                       "eta",  "dz"); */
/*     tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B4", "dz vs eta (BPIX L4)", */
/*                                                       dz_vs_eta_B4, 20, -3, 3, 20, -.01, .01, */
/*                                                       "eta",  "dz"); */

/* } */

void run_analysis(const vector<fv::util::DataFileDescriptor>& dfds, const string output_filename, bool silent){
    gSystem->Load("libfilval.so");
    auto replace_suffix = [](const std::string& input, const std::string& new_suffix){
        return input.substr(0, input.find_last_of(".")) + new_suffix;
    };

    string log_filename = replace_suffix(output_filename, ".log");
    fv::util::Log::init_logger(log_filename, fv::util::LogPriority::kLogDebug);

    TrackingDataSet tds(output_filename, dfds, "trackingNtuple/tree");
    /* tds.set_max_events(10); */
    register_objects(tds);

    /* setup_dz_vs_eta(tds); */
    /* setup_dphi_vs_eta(tds); */
    setup_matched_tracks(tds);


    tds.process(silent);
    tds.save_all();
}

int main(int argc, char * argv[])
{
    fv::util::ArgParser args(argc, argv);
    bool silent = args.cmdOptionExists("-s");
    string output_filename = args.cmdOptionExists("-o") ? args.getCmdOption("-o") : "output.root";
    if(args.cmdOptionExists("-F")){
        auto file_list = fv::util::read_input_list(args.getCmdOption("-F"));
        run_analysis(file_list, output_filename, silent);
    }else if(args.cmdOptionExists("-f")) {
        string input_filename = args.getCmdOption("-f");
        string data_label = args.cmdOptionExists("-l") ? args.getCmdOption("-l") : "";
        fv::util::DataFileDescriptor dfd(input_filename, data_label);
        run_analysis(std::vector<fv::util::DataFileDescriptor>({dfd}), output_filename, silent);
    }else {
        cout << "Usage: ./tracking_validation (-s) {-o output_filename} -F datafiles.txt" << endl;
        cout << "       ./tracking_validation (-s) {-l DATA_LABEL} {-o output_filename} -f treefile.root" << endl;
    }
}