EGamma_ElectronTrackingValidation/analysis/tracking_eff.cpp

#include <iostream>
#include <vector>
#include <map>
#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 "analysis/TrackingNtuple.h"
#include "analysis/TrackingNtupleObjs.hpp"
#include "analysis/common.hpp"

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

vector<string> seedTypes =
    {"initialStepSeeds",
     "highPtTripletStepSeeds",
     "mixedTripletStepSeeds",
     "pixelLessStepSeeds",
     "tripletElectronSeeds",
     "pixelPairElectronSeeds",
     "stripPairElectronSeeds"};

enum class HitType {
  Pixel = 0,
  Strip = 1,
  Glued = 2,
  Invalid = 3,
  Phase2OT = 4,
  Unknown = 99
};

SeedCollection seeds;
SimTrackCollection sim_tracks;
SimHitCollection sim_hits;
SimVertexCollection sim_vertices;
TrackCollection gsf_tracks;
PixRecHitCollection pixrec_hits;

void register_objects(TrackingDataSet& tds){
    seeds.init(tds);
    sim_tracks.init(tds);
    sim_hits.init(tds);
    sim_vertices.init(tds);
    gsf_tracks.init(tds);
    pixrec_hits.init(tds);
}

float rad(const float& x, const float& y) {
    return sqrt(x*x + y*y);
}

struct {
    Value<float>* x;
    Value<float>* y;
    Value<float>* z;
    Value<float>* sigma_x;
    Value<float>* sigma_y;
    Value<float>* sigma_z;
} bs;

bool in_lum_region(const SimVertex& vertex) {
    /* Here approximate the luminous region as a cylinder with length 5*\sigma_z and radius
     * sqrt(\simga_x^2 + \sigma_y^2) centered as beamspot position.
     */
    float dx = vertex.x() - bs.x->get_value();
    float dy = vertex.y() - bs.y->get_value();
    float dz = vertex.z() - bs.z->get_value();
    float radius = 5*sqrt(pow(bs.sigma_x->get_value(), 2) + pow(bs.sigma_y->get_value(), 2));
    float half_len = 5*bs.sigma_z->get_value();
    return sqrt(dx*dx + dy*dy) < radius and abs(dz) < half_len;
};

bool is_good_sim(size_t simTrkIdx) {
    const auto& track = sim_tracks[simTrkIdx];
    const auto& vertex = sim_vertices[track.parentVtxIdx()];
    return abs(track.pdgId()) == 11 and in_lum_region(vertex);
};

void run_analysis(const vector<fv::util::DataFileDescriptor>& dfds, const string output_filename, bool silent){
    TrackingDataSet tds(output_filename, dfds, "trackingNtuple/tree");
    register_objects(tds);

    bs = {
        tds.track_branch<float>("bsp_x"),
        tds.track_branch<float>("bsp_y"),
        tds.track_branch<float>("bsp_z"),
        tds.track_branch<float>("bsp_sigmax"),
        tds.track_branch<float>("bsp_sigmay"),
        tds.track_branch<float>("bsp_sigmaz")
    };

    // Indices of prompt-ish sim electrons
    auto good_sim_electrons = func_value<vector<size_t>>("good_sim_electrons",
        FUNC(([](){
            vector<size_t> idxs;
            for(const auto& sim_track : sim_tracks) {
                if (is_good_sim(sim_track.idx)) idxs.push_back(sim_track.idx);
            }
            return idxs;
        })));

    // Indices of prompt-ish sim electrons that have a matched seed
    function<bool(size_t)> with_seed = [](const size_t& e_idx) {
        return seeds.size() != 0 and sim_tracks[e_idx].seedIdx().size() != 0;
    };

    // Indices of prompt-ish sim electrons that have a matched gsftrack
    function<bool(size_t)> with_track = [](const size_t& e_idx) {
        return sim_tracks[e_idx].trkIdx().size() > 0;
    };

    function<float(size_t)> sim_pt = [](const size_t& sim_idx) {
        return sim_tracks[sim_idx].pt();
    };

    function<float(size_t)> sim_eta = [](const size_t& sim_idx) {
        return sim_tracks[sim_idx].eta();
    };

    function<float(size_t)> sim_phi = [](const size_t& sim_idx) {
        return sim_tracks[sim_idx].phi();
    };

    float pt_cut = 20; // GeV
    float eta_cut = 2.4;
    function<bool(size_t)> sim_in_pt = [pt_cut](size_t sim_idx) {
        return sim_tracks[sim_idx].pt() > pt_cut;
    };

    function<bool(size_t)> sim_in_eta = [eta_cut](size_t sim_idx) {
        return abs(sim_tracks[sim_idx].eta()) < eta_cut;
    };

    function<bool(size_t)> sim_in_eta_and_pt = [eta_cut, pt_cut](size_t sim_idx) {
        return abs(sim_tracks[sim_idx].eta()) < eta_cut and
               sim_tracks[sim_idx].pt() > pt_cut;
    };

    tds.register_container<EfficiencyContainer<size_t>>("seed_eff_v_pt",
        good_sim_electrons, TH1Params::lookup("eff_v_pt"), &sim_in_eta, &with_seed, &sim_pt);

    tds.register_container<EfficiencyContainer<size_t>>("seed_eff_v_eta",
        good_sim_electrons, TH1Params::lookup("eff_v_eta"), &sim_in_pt, &with_seed, &sim_eta);

    tds.register_container<EfficiencyContainer<size_t>>("seed_eff_v_phi",
        good_sim_electrons, TH1Params::lookup("eff_v_phi"), &sim_in_eta_and_pt, &with_seed, &sim_phi);

    tds.register_container<EfficiencyContainer<size_t>>("track_eff_v_pt",
        good_sim_electrons, TH1Params::lookup("eff_v_pt"), &sim_in_eta, &with_track, &sim_pt);

    tds.register_container<EfficiencyContainer<size_t>>("track_eff_v_eta",
        good_sim_electrons, TH1Params::lookup("eff_v_eta"), &sim_in_pt, &with_track, &sim_eta);

    tds.register_container<EfficiencyContainer<size_t>>("track_eff_v_phi",
        good_sim_electrons, TH1Params::lookup("eff_v_phi"), &sim_in_eta_and_pt, &with_track, &sim_phi);

    // Indices of ecal-driven seeds
    auto ecal_seeds = func_value<vector<size_t>>("ecal_seeds",
        FUNC(([](){
            vector<size_t> idxs;
            for (const auto& seed : seeds) {
                if (seed.ecalDriven())
                    idxs.push_back(seed.idx);
            }
            return idxs;
        })));

    function<bool(size_t)> seed_with_sim_electron = [](const size_t& seed_idx) {
        for (const int& simTrkIdx : seeds[seed_idx].simTrkIdx()) {
            if (is_good_sim(simTrkIdx)) return true;
        }
        return false;
    };

    function<float(size_t)> seed_pt = [](const size_t& seed_idx) {
        return seeds[seed_idx].pt();
    };

    function<float(size_t)> seed_eta = [](const size_t& seed_idx) {
        return seeds[seed_idx].eta();
    };

    function<float(size_t)> seed_phi = [](const size_t& seed_idx) {
        return seeds[seed_idx].phi();
    };

    function<bool(size_t)> seed_in_pt = [pt_cut](size_t seed_idx) {
        return seeds[seed_idx].pt() > pt_cut;
    };

    function<bool(size_t)> seed_in_eta = [eta_cut](size_t seed_idx) {
        return abs(seeds[seed_idx].eta()) < eta_cut;
    };

    function<bool(size_t)> seed_in_eta_and_pt = [eta_cut, pt_cut](size_t seed_idx) {
        return abs(seeds[seed_idx].eta()) < eta_cut and
               seeds[seed_idx].pt() > pt_cut;
    };

    tds.register_container<EfficiencyContainer<size_t>>("seed_pur_v_pt",
        ecal_seeds, TH1Params::lookup("pur_v_pt"), &seed_in_eta, &seed_with_sim_electron, &seed_pt);

    tds.register_container<EfficiencyContainer<size_t>>("seed_pur_v_eta",
        ecal_seeds, TH1Params::lookup("pur_v_eta"), &seed_in_pt, &seed_with_sim_electron, &seed_eta);

    tds.register_container<EfficiencyContainer<size_t>>("seed_pur_v_phi",
        ecal_seeds, TH1Params::lookup("pur_v_phi"), &seed_in_eta_and_pt, &seed_with_sim_electron, &seed_phi);


    // Indices of ecal-driven gsf-tracks
    auto ecal_tracks = func_value<vector<size_t>>("ecal_tracks",
        FUNC(([](){
            vector<size_t> idxs;
            for (const auto& track : gsf_tracks) {
                const auto& seed = seeds[track.seedIdx()];
                if (seed.ecalDriven())
                    idxs.push_back(track.idx);
            }
            return idxs;
        })));

    function<bool(size_t)> gsf_track_with_sim_electron = [](const size_t& track_idx) {
        for (const int& simTrkIdx : gsf_tracks[track_idx].simTrkIdx()) {
            if (is_good_sim(simTrkIdx)) return true;
        }
        return false;
    };

    function<float(size_t)> track_pt = [](const size_t& track_idx) {
        return gsf_tracks[track_idx].pt();
    };

    function<float(size_t)> track_eta = [](const size_t& track_idx) {
        return gsf_tracks[track_idx].eta();
    };

    function<float(size_t)> track_phi = [](const size_t& track_idx) {
        return gsf_tracks[track_idx].phi();
    };

    function<bool(size_t)> track_in_pt = [pt_cut](size_t track_idx) {
        return gsf_tracks[track_idx].pt() > pt_cut;
    };

    function<bool(size_t)> track_in_eta = [eta_cut](size_t track_idx) {
        return abs(gsf_tracks[track_idx].eta()) < eta_cut;
    };

    function<bool(size_t)> track_in_eta_and_pt = [eta_cut, pt_cut](size_t track_idx) {
        return abs(gsf_tracks[track_idx].eta()) < eta_cut and
               gsf_tracks[track_idx].pt() > pt_cut;
    };

    tds.register_container<EfficiencyContainer<size_t>>("track_pur_v_pt",
        ecal_tracks, TH1Params::lookup("pur_v_pt"), &track_in_eta, &gsf_track_with_sim_electron, &track_pt);

    tds.register_container<EfficiencyContainer<size_t>>("track_pur_v_eta",
        ecal_tracks, TH1Params::lookup("pur_v_eta"), &track_in_pt, &gsf_track_with_sim_electron, &track_eta);

    tds.register_container<EfficiencyContainer<size_t>>("track_pur_v_phi",
        ecal_tracks, TH1Params::lookup("pur_v_phi"), &track_in_eta_and_pt, &gsf_track_with_sim_electron, &track_phi);


    auto simpixlay_v_eta = func_value<std::pair<vector<float>,vector<float>>>("sim_pixlay_v_eta",
        FUNC(([](){
            vector<float> etas;
            vector<float> nLays;
            for(const auto& sim_track : sim_tracks) {
                if(sim_track.pt() > 20 and is_good_sim(sim_track.idx)) {
                    bool hasLHit[] = {false, false, false, false,  // <- 4 barrel layers
                                      false, false, false};        // <- 3 fwd layers
                    for(const auto& sim_hit_idx : sim_track.simHitIdx()) {
                        const auto& sim_hit = sim_hits[sim_hit_idx];
                        if (sim_hit.subdet() == 1) hasLHit[sim_hit.layer()-1] = true;
                        if (sim_hit.subdet() == 2) hasLHit[sim_hit.layer()+3] = true;
                    }
                    int layers = 0;
                    for(size_t i=0; i<7; i++) {
                        if (hasLHit[i]) layers++;
                    }
                    etas.push_back(sim_track.eta());
                    nLays.push_back(layers);
                }
            }
            return std::make_pair(etas, nLays);
        })));

    auto recpixlay_v_eta = func_value<std::pair<vector<float>,vector<float>>>("recpixlay_v_eta",
        FUNC(([](){
            vector<float> etas;
            vector<float> nLays;
            for(const auto& gsf_track : gsf_tracks) {
                if(gsf_track.pt() > 20) {
                    bool hasLHit[] = {false, false, false, false,  // <- 4 barrel layers
                                      false, false, false};        // <- 3 fwd layers
                    const auto& hitIdx = gsf_track.hitIdx();
                    const auto& hitType = gsf_track.hitType();
                    size_t nHits = hitIdx.size();
                    for(size_t i=0; i<nHits; i++) {
                        if (HitType(hitType[i]) != HitType::Pixel) continue;
                        const auto& pixrec_hit = pixrec_hits[hitIdx[i]];
                        if (pixrec_hit.subdet() == 1) hasLHit[pixrec_hit.layer()-1] = true;
                        if (pixrec_hit.subdet() == 2) hasLHit[pixrec_hit.layer()+3] = true;
                    }
                    int layers = 0;
                    for(size_t i=0; i<7; i++) {
                        if (hasLHit[i]) layers++;
                    }
                    etas.push_back(gsf_track.eta());
                    nLays.push_back(layers);
                }
            }
            return std::make_pair(etas, nLays);
        })));


    tds.register_container<ContainerTH2Many<float>>("simpixlay_v_eta", simpixlay_v_eta, "simpixlay_v_eta",
        TH2Params::lookup("n_hit_v_eta"));
    tds.register_container<ContainerTH2Many<float>>("recpixlay_v_eta", recpixlay_v_eta, "recpixlay_v_eta",
        TH2Params::lookup("n_hit_v_eta"));

//    auto simtrack_hits = func_value<int>("simtrack_hits",
//        FUNC(([](){
//            vector<float> etas;
//            vector<float> nhits;
//            for(const auto& sim_track : sim_tracks) {
//                if(sim_track.pt() < 20 or !is_good_sim(sim_track.idx)) continue;
//                cout << "New Track: nPixel=" << sim_track.nPixel() << endl;
//                for(const auto& sim_hit_idx : sim_track.simHitIdx()) {
//                    const auto& sim_hit = sim_hits[sim_hit_idx];
//                    if (sim_hit.subdet() > 2) continue;
//                    cout << "  Hit: " << sim_hit_idx << endl
//                         << "    pos: " << sim_hit.x() << " " << sim_hit.y() << " " << sim_hit.z() << endl
//                         << "    subdet: " << sim_hit.subdet()  << endl
//                         << "    layler: " << sim_hit.layer() << endl;
//                }
//            }
//            return 0;
//        })));
//
//    ofstream out("blah.txt");
//    tds.register_container<Printer<int>>("blah", simtrack_hits, &out);

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

int main(int argc, char * argv[]){
    using namespace fv::util;
    ArgParser args(argc, argv);
    bool silent = args.cmdOptionExists("-s");
    if(args.cmdOptionExists("-c")) {
        init_config(args.getCmdOption("-c"));
        auto output_filename = the_config->get_output_filename();
        if (output_filename == "") return -1;

        init_log(fv::util::LogPriority::kLogInfo);
        gSystem->Load("libfilval.so");

        auto file_list = the_config->get_source_files();
        run_analysis(file_list, output_filename, silent);
    } else {
        cout << "Usage: ./tracking_eff (-s) -c config_file.yaml" << endl;
    }
    return 0;
}