Skip to content
Snippets Groups Projects
rdf_test.cxx 1.83 KiB
Newer Older
  • Learn to ignore specific revisions
  • #include "ROOT/RDataFrame.hxx"
    #include <iostream>
    
    #include "dd4pod/Geant4ParticleCollection.h"
    
    //namespace edm4hep {
    //
    //std::vector<float> pt (std::vector<MCParticleData> const& in){
    //  std::vector<float> result;
    //  for (size_t i = 0; i < in.size(); ++i) {
    //    result.push_back(std::sqrt(in[i].momentum.x * in[i].momentum.x + in[i].momentum.y * in[i].momentum.y));
    //  }
    //  return result;
    //}
    //
    //std::vector<float> eta(std::vector<MCParticleData> const& in){
    //  std::vector<float> result;
    //  ROOT::Math::PxPyPzMVector lv;
    //  for (size_t i = 0; i < in.size(); ++i) {
    //    lv.SetCoordinates(in[i].momentum.x, in[i].momentum.y, in[i].momentum.z, in[i].mass);
    //    result.push_back(lv.Eta());
    //  }
    //  return result;
    //}
    //
    //std::vector<float> cos_theta(std::vector<MCParticleData> const& in){
    //  std::vector<float> result;
    //  ROOT::Math::PxPyPzMVector lv;
    //  for (size_t i = 0; i < in.size(); ++i) {
    //    lv.SetCoordinates(in[i].momentum.x, in[i].momentum.y, in[i].momentum.z, in[i].mass);
    //    result.push_back(cos(lv.Theta()));
    //  }
    //  return result;
    //}
    //
    //}
    
    
    std::vector<float> pt (std::vector<dd4pod::Geant4ParticleData> const& in){
      std::vector<float> result;
      for (size_t i = 0; i < in.size(); ++i) {
        result.push_back(std::sqrt(in[i].psx * in[i].psx + in[i].psy * in[i].psy));
      }
      return result;
    }
    
    int rdf_test() {
    
      ROOT::EnableImplicitMT();
    
      ROOT::RDataFrame df("events", "sim_barrel_clusters.root");
    
      auto df2 =
          df.Define("MCParticles_pt", pt, {"mcparticles"});
              //.Define("MCParticles_eta", edm4hep::eta, {"mcparticles"})
              //.Define("MCParticles_cosTheta", edm4hep::cos_theta, {"mcparticles"});
    
      std::string outfilename = "rdf_test.root";
    
      df2.Snapshot("events", outfilename,
                   {
                       "MCParticles_pt",
                       "mcparticles"
                   });
      return 0;
    }