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]";
  }
  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;
            /* SimVertex parentvtx = sim_vertices[sim_track.parentVtxIdx()]; */
            /* SimVertex parentvtx = sim_vertices[sim_track.parentVtxIdx()]; */
            cout << "\t" << sim_track.pdgId() << " "
                 << sim_track.pt() << " "
                 << sim_track.eta() << " "
                 << sim_track.phi() << " "
                 << sim_track.parentVtxIdx() << " "
                 << 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
        }
    }
    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;
}