EGamma_ElectronTrackingValidation/looper/analysis/dump_good_sim.cpp

#include <iostream>
#include <vector>
#include <cmath>
#include <numeric>
#include <TSystem.h>
#include <Math/VectorUtil.h>

#include "filval.hpp"
#include "root_filval.hpp"
#include "common.hpp"

#include "analysis/TrackingNtupleObjs.hpp"

typedef ROOT::Math::LorentzVector<ROOT::Math::PxPyPzE4D<float>> Vec4;

SeedCollection seeds;
SimTrackCollection sim_tracks;
SimVertexCollection sim_vertices;
TrackCollection gsf_tracks;
SuperClusterCollection scls;

void register_objects(TrackingDataSet& tds){
    /* seeds.init(tds); */
    sim_tracks.init(&tds);
    sim_vertices.init(&tds);
    /* gsf_tracks.init(tds); */
    /* scls.init(tds); */
}

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

template <typename T>
struct KinColl {
    T* pt;
    T* eta;
    T* phi;
};


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();
    float dy = vertex.y() - bs.y->get();
    float dz = vertex.z() - bs.z->get();
    auto radius = static_cast<float>(5 * sqrt(pow(bs.sigma_x->get(), 2) + pow(bs.sigma_y->get(), 2)));
    float half_len = 5*bs.sigma_z->get();
    return sqrt(dx*dx + dy*dy) < radius and abs(dz) < half_len;
};

bool is_good_sim(const SimTrack& sim_track) {
    const auto& vertex = sim_vertices[sim_track.parentVtxIdx()];
    return abs(sim_track.pdgId()) == 11 and in_lum_region(vertex);
};


template <typename T>
std::ostream& operator<< (std::ostream& out, const std::vector<T>& v) {
  if ( !v.empty() ) {
    out << '[';
    std::copy(v.begin(), v.end(), std::ostream_iterator<T>(out, ", "));
    out << "\b\b]";
  } else {
    cout << "[]";
  }
  return out;
}

void run(){
    using namespace std;
    using namespace fv;
    using namespace fv_root;
    using namespace fv_util;
    auto file_list = the_config->get_source_files();
    string output_filename = the_config->get_output_filename();
    TrackingDataSet tds(output_filename, file_list, "trackingNtuple/tree");
    register_objects(tds);

    // Declare TTree

    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")
    };

    while (tds.next()) {

        cout << "Event: " << tds.get_current_event() << endl;
        /* for (const auto& sim_track : sim_tracks) { */
        /*     if (!is_good_sim(sim_track)) continue; */
        /*     /1* SimVertex parentvtx = sim_vertices[sim_track.parentVtxIdx()]; *1/ */
        /*     /1* SimVertex parentvtx = sim_vertices[sim_track.parentVtxIdx()]; *1/ */
        /*     cout << "\t" << sim_track.pdgId() << " " */
        /*          << sim_track.pt() << " " */
        /*          << sim_track.eta() << " " */
        /*          << sim_track.phi() << " " */
        /*          << sim_track.parentVtxIdx() << "(" */
        /*          << sim_vertices[sim_track.parentVtxIdx()].processType() << ") " */
        /*          << sim_track.decayVtxIdx() << " "; */
        /*     if (sim_track.decayVtxIdx().size() > 0) { */
        /*         cout << sim_tracks[sim_vertices[sim_track.decayVtxIdx()[0]].daughterSimIdx()[0]].pdgId(); */
        /*     } */
        /*     cout << endl; */
        /*     // Log Sim Track */
        /* } */

        int n_pvx = 0;
        cout << setprecision(3);
        for (const auto& sim_vtx : sim_vertices) {
            if (sim_vtx.sourceSimIdx().size() > 0) continue;
            n_pvx++;
            cout << n_pvx << " ("
                 << sim_vtx.x() <<  ", "
                 << sim_vtx.y() <<  ", "
                 << sim_vtx.z() <<  ") "
                 << sim_vtx.processType() << endl;
            for (const int& trk_idx : sim_vtx.daughterSimIdx()) {
                const auto& trk = sim_tracks[trk_idx];
                cout << "\t"
                     << setw(5) << trk.pdgId()
                     << "("
                     << setw(5) << trk.pt()  << ", "
                     << setw(5) << trk.eta() << ", "
                     << setw(5) << trk.phi() << ")"
                     << endl;

            }
            /* cout << sim_vtx.sourceSimIdx() << endl; */
        }
    }
    tds.save_all();
}

int main(int argc, char * argv[]){
    using namespace fv_util;
    ArgParser args(argc, argv);
    if(args.cmd_option_exists("-c")) {
        init_config(args.get_cmd_option("-c"));
        args.update_config();
        if (args.cmd_option_exists("-s")) {
            the_config->update_key("silent", true);
        }
        if (args.cmd_option_exists("-b")) {
            the_config->update_key("batch", true);
        }
        init_log(LogPriority::kLogInfo);
//        gSystem->Load("libfilval.so");
        run();
    } else {
        cout << "Usage: ./" << argv[0] << " (-s) (-b) -c config_file.yaml" << endl;
    }
    return 0;
}