diff --git a/test/test_rdf.cc b/test/test_rdf.cc index bcf9a4e47..af48bc20e 100644 --- a/test/test_rdf.cc +++ b/test/test_rdf.cc @@ -22,14 +22,17 @@ int main(int argc, char* argv[]) { .Define("MCParticles_cosTheta", edm4hep::utils::cos_theta, {"MCParticles"}) .Define("SimTrackerHits_r", edm4hep::utils::r, {"SimTrackerHits"}) .Define("SimTrackerHit_pt", edm4hep::utils::pt, {"SimTrackerHits"}) - .Define("TrackerHits_r", edm4hep::utils::r, {"TrackerHitPlanes"}); + .Define("TrackerHits_r", edm4hep::utils::r, {"TrackerHitPlanes"}) + .Define("MCParticle_p4", edm4hep::utils::p4M, {"MCParticles"}) + + .Define("MCParticle_energy", edm4hep::utils::E, {"MCParticle_p4"}); std::string outfilename = "edm4hep_events_rdf.root"; std::cout << "Writing snapshot to disk ... \t" << outfilename << std::endl; df2.Snapshot("events", outfilename, {"MCParticles_pt", "MCParticles_eta", "MCParticles_cosTheta", "SimTrackerHits_r", "SimTrackerHit_pt", - "TrackerHits_r"}); + "TrackerHits_r", "MCParticle_energy"}); return 0; } diff --git a/test/test_rdf.py b/test/test_rdf.py index 05fc50974..0a38e5346 100644 --- a/test/test_rdf.py +++ b/test/test_rdf.py @@ -20,6 +20,8 @@ .Define('SimTrackerHits_r', 'edm4hep::utils::r(SimTrackerHits)') .Define('SimTrackerHits_pt', 'edm4hep::utils::pt(SimTrackerHits)') .Define('TrackerHits_r', 'edm4hep::utils::r(TrackerHitPlanes)') + .Define('MCParticle_p4', 'edm4hep::utils::p4M(MCParticles)') + .Define('MCParticle_energy', 'edm4hep::utils::E(MCParticle_p4)') ) filename = 'edm4hep_events_py_rdf.root' @@ -31,4 +33,5 @@ 'MCParticles_cosTheta', 'SimTrackerHits_r', 'SimTrackerHits_pt', - 'TrackerHits_r']) + 'TrackerHits_r', + 'MCParticle_energy']) diff --git a/test/utils/CMakeLists.txt b/test/utils/CMakeLists.txt index fde786392..edd86737e 100644 --- a/test/utils/CMakeLists.txt +++ b/test/utils/CMakeLists.txt @@ -24,4 +24,5 @@ add_test(NAME pyunittests COMMAND python -m unittest discover -s ${CMAKE_CURRENT set_property(TEST pyunittests PROPERTY ENVIRONMENT LD_LIBRARY_PATH=$:$ENV{LD_LIBRARY_PATH} + ROOT_INCLUDE_PATH=${PROJECT_SOURCE_DIR}/utils/include ) diff --git a/test/utils/test_kinematics.py b/test/utils/test_kinematics.py index 3616d1a4e..df94d329f 100644 --- a/test/utils/test_kinematics.py +++ b/test/utils/test_kinematics.py @@ -11,7 +11,7 @@ print('Cannot load edm4hep dictionary for tests') sys.exit(1) -ROOT.gInterpreter.LoadFile(os.path.dirname(__file__) + '/../../utils/include/edm4hep/utils/kinematics.h') +ROOT.gInterpreter.LoadFile('edm4hep/utils/kinematics.h') from ROOT import edm4hep diff --git a/utils/include/edm4hep/utils/common.h b/utils/include/edm4hep/utils/common.h new file mode 100644 index 000000000..e7f6883c5 --- /dev/null +++ b/utils/include/edm4hep/utils/common.h @@ -0,0 +1,19 @@ +#ifndef EDM4HEP_UTILS_COMMON_H +#define EDM4HEP_UTILS_COMMON_H + +#include "Math/Vector4D.h" + +namespace edm4hep { +/** + * A LorentzVector with (px, py, pz) and M + */ +using LorentzVectorM = ROOT::Math::PxPyPzMVector; + +/** + * A LorentzVector with (px, py, pz) and E + */ +using LorentzVectorE = ROOT::Math::PxPyPzEVector; + +} // namespace edm4hep + +#endif diff --git a/utils/include/edm4hep/utils/dataframe.h b/utils/include/edm4hep/utils/dataframe.h index 554c65aca..372930928 100644 --- a/utils/include/edm4hep/utils/dataframe.h +++ b/utils/include/edm4hep/utils/dataframe.h @@ -1,6 +1,8 @@ #ifndef EDM4HEP_UTILS_DATAFRAME_H #define EDM4HEP_UTILS_DATAFRAME_H +#include "edm4hep/utils/common.h" + #include "ROOT/RVec.hxx" namespace edm4hep::utils { @@ -21,6 +23,22 @@ ROOT::VecOps::RVec cos_theta(ROOT::VecOps::RVec const& in); template ROOT::VecOps::RVec r(ROOT::VecOps::RVec const& in); +/// Get the 4 momentum of the passed particle using the momentum and the mass +template +ROOT::VecOps::RVec p4M(ROOT::VecOps::RVec const& in); + +/// Get the 4 momentum of the passed particle using the momentum and the energy +template +ROOT::VecOps::RVec p4E(ROOT::VecOps::RVec const& in); + +/// Get the energy from a four momentum vector +template +ROOT::VecOps::RVec E(ROOT::VecOps::RVec const& fourMom); + +/// Get the mass from a four momentum vector +template +ROOT::VecOps::RVec M(ROOT::VecOps::RVec const& fourMom); + } // namespace edm4hep::utils #endif // EDM4HEP_UTILS_DATAFRAME_H diff --git a/utils/include/edm4hep/utils/kinematics.h b/utils/include/edm4hep/utils/kinematics.h index 8a6c04de9..eef362361 100644 --- a/utils/include/edm4hep/utils/kinematics.h +++ b/utils/include/edm4hep/utils/kinematics.h @@ -1,21 +1,11 @@ #ifndef EDM4HEP_UTILS_KINEMATICS_H #define EDM4HEP_UTILS_KINEMATICS_H -#include "Math/Vector4D.h" +#include "edm4hep/utils/common.h" #include namespace edm4hep { -/** - * A LorentzVector with (px, py, pz) and M - */ -using LorentzVectorM = ROOT::Math::PxPyPzMVector; - -/** - * A LorentzVector with (px, py, pz) and E - */ -using LorentzVectorE = ROOT::Math::PxPyPzEVector; - namespace utils { /** diff --git a/utils/src/dataframe.cc b/utils/src/dataframe.cc index ef12f0342..685e3ee7a 100644 --- a/utils/src/dataframe.cc +++ b/utils/src/dataframe.cc @@ -1,4 +1,5 @@ #include "edm4hep/utils/dataframe.h" +#include "edm4hep/utils/common.h" #include "edm4hep/CalorimeterHitData.h" #include "edm4hep/ClusterData.h" @@ -16,44 +17,55 @@ namespace edm4hep::utils { template ROOT::VecOps::RVec pt(ROOT::VecOps::RVec const& in) { - ROOT::VecOps::RVec result; - result.reserve(in.size()); - 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; + return ROOT::VecOps::Map( + in, [](const auto& p) { return std::sqrt(p.momentum.x * p.momentum.x + p.momentum.y + p.momentum.y); }); } template ROOT::VecOps::RVec eta(ROOT::VecOps::RVec const& in) { - ROOT::VecOps::RVec result; - result.reserve(in.size()); - for (size_t i = 0; i < in.size(); ++i) { - ROOT::Math::XYZVector lv{in[i].momentum.x, in[i].momentum.y, in[i].momentum.z}; - result.push_back(lv.Eta()); - } - return result; + return ROOT::VecOps::Map(in, [](const auto& p) { + ROOT::Math::XYZVector lv{p.momentum.x, p.momentum.y, p.momentum.z}; + return lv.Eta(); + }); } template ROOT::VecOps::RVec cos_theta(ROOT::VecOps::RVec const& in) { - ROOT::VecOps::RVec result; - result.reserve(in.size()); - for (size_t i = 0; i < in.size(); ++i) { - ROOT::Math::XYZVector lv{in[i].momentum.x, in[i].momentum.y, in[i].momentum.z}; - result.push_back(cos(lv.Theta())); - } - return result; + return ROOT::VecOps::Map(in, [](const auto& p) { + ROOT::Math::XYZVector lv{p.momentum.x, p.momentum.y, p.momentum.z}; + return std::cos(lv.Theta()); + }); } template ROOT::VecOps::RVec r(ROOT::VecOps::RVec const& in) { - ROOT::VecOps::RVec result; - result.reserve(in.size()); - for (size_t i = 0; i < in.size(); ++i) { - result.push_back(std::sqrt(in[i].position.x * in[i].position.x + in[i].position.y * in[i].position.y)); - } - return result; + return ROOT::VecOps::Map(in, [](const auto& p) { + return std::sqrt(p.position.x * p.position.x + p.position.y * p.position.y + p.position.z + p.position.z); + }); +} + +template +ROOT::VecOps::RVec p4M(ROOT::VecOps::RVec const& in) { + return ROOT::VecOps::Map(in, [](const auto& p) { + return edm4hep::LorentzVectorM{p.momentum.x, p.momentum.y, p.momentum.z, p.mass}; + }); +} + +template +ROOT::VecOps::RVec p4E(ROOT::VecOps::RVec const& in) { + return ROOT::VecOps::Map(in, [](const auto& p) { + return edm4hep::LorentzVectorE{p.momentum.x, p.momentum.y, p.momentum.z, p.energy}; + }); +} + +template +ROOT::VecOps::RVec E(ROOT::VecOps::RVec const& fourMom) { + return ROOT::VecOps::Map(fourMom, [](const auto& p) { return p.E(); }); +} + +template +ROOT::VecOps::RVec M(ROOT::VecOps::RVec const& fourMom) { + return ROOT::VecOps::Map(fourMom, [](const auto& p) { return p.M(); }); } // Explicitly instantiate the template functions here to have them available in @@ -69,13 +81,15 @@ ROOT::VecOps::RVec r(ROOT::VecOps::RVec const& in) { INST_DATA_TO_FLOAT_VEC_FUNC(eta, DATATYPE); \ INST_DATA_TO_FLOAT_VEC_FUNC(cos_theta, DATATYPE) -// Macro to instantiate all position related functions for a datatype -#define INST_POSITION_FUNCS(DATATYPE) INST_DATA_TO_FLOAT_VEC_FUNC(r, DATATYPE) - INST_MOMENTUM_FUNCS(edm4hep::MCParticleData); INST_MOMENTUM_FUNCS(edm4hep::ReconstructedParticleData); INST_MOMENTUM_FUNCS(edm4hep::SimTrackerHitData); +#undef INST_MOMENTUM_FUNCS + +// Macro to instantiate all position related functions for a datatype +#define INST_POSITION_FUNCS(DATATYPE) INST_DATA_TO_FLOAT_VEC_FUNC(r, DATATYPE) + INST_POSITION_FUNCS(edm4hep::SimTrackerHitData); INST_POSITION_FUNCS(edm4hep::TrackerHitData); INST_POSITION_FUNCS(edm4hep::TrackerHitPlaneData); @@ -84,8 +98,26 @@ INST_POSITION_FUNCS(edm4hep::CalorimeterHitData); INST_POSITION_FUNCS(edm4hep::ClusterData); INST_POSITION_FUNCS(edm4hep::VertexData); -#undef INST_DATA_TO_FLOAT_VEC #undef INST_POSITION_FUNCS -#undef INST_MOMENTUM_FUNCS + +INST_DATA_TO_FLOAT_VEC_FUNC(E, edm4hep::LorentzVectorE); +INST_DATA_TO_FLOAT_VEC_FUNC(E, edm4hep::LorentzVectorM); +INST_DATA_TO_FLOAT_VEC_FUNC(M, edm4hep::LorentzVectorE); +INST_DATA_TO_FLOAT_VEC_FUNC(M, edm4hep::LorentzVectorM); + +#undef INST_DATA_TO_FLOAT_VEC + +#define INST_4MOM_MASS_FUNCS(DATATYPE) \ + template ROOT::VecOps::RVec p4M(ROOT::VecOps::RVec const&) + +#define INST_4MOM_ENERGY_FUNCS(DATATYPE) \ + template ROOT::VecOps::RVec p4E(ROOT::VecOps::RVec const&) + +INST_4MOM_ENERGY_FUNCS(edm4hep::ReconstructedParticleData); +INST_4MOM_MASS_FUNCS(edm4hep::ReconstructedParticleData); +INST_4MOM_MASS_FUNCS(edm4hep::MCParticleData); + +#undef INST_4MOM_ENERGY_FUNCS +#undef INST_4MOM_MASS_FUNCS } // namespace edm4hep::utils