Call method for every object in RVec in a single line

Hi all,

I have a std::vector<XYZVector> stored in TTree.

Is there an easy way to get absolute momentum XYZVector.R() for each element of the array within RDataFrame?

The one I am currently using is:

ROOT.gInterpreter.Declare('''
using namespace ROOT::Math;
using namespace ROOT::VecOps;

RVec <double> getAbsOfVectors(const RVec<XYZVector>& vecs){
    auto getAbs = [&](const XYZVector& vec) { return vec.R(); };
    return Map(vecs, getAbs);
}
''')

    df = ROOT.RDataFrame("tree", "file.root")
    df = df.Define("abs_mom", "getAbsOfVectors(mom)")

not very convenient if I need to predefine separate function for evey getPhi(), getTheta(), getX(), getY(), getZ()… Then the advantage of using XYZVector rather than plain px, py, pz vanishes.

I have tried this one-liner, which I was expecting to work, but it doesn’t…

    df = df.Define("abs_mom", "Map(mom, [&](const XYZVector& vec) { return vec.R(); })")

Any advice?

ROOT version: 6.24/06

cheers,
Bohdan

Hello Bohdan,
very annoyingly, to make it work you need return Map instead of just Map, because there is a return in the expression so RDF does not insert another one. This should work:

.Define("absv", "return Map(v, [&](const XYZVector& vec) { return vec.R(); })")

Minimal repro that works for me:

import ROOT

ROOT.RDataFrame(2)
   .Define("v", "ROOT::RVec<ROOT::Math::XYZVector>(3)")
   .Define("absv", "return Map(v, [](ROOT::Math::XYZVector& vec) { return vec.R(); })")
   .Display().Print()

I agree that’s more characters than one would like to write, but at least you don’t need the extra Declare stuff. However if you do these operations over and over again you might want to just put some helper functions in a library that you load at the beginning. Write them once, use them many times.

For ROOT, we could:

  • have a bunch of RVec helper functions that work on RVecs of XYZVectors (or 4-vectors in general)
  • allow Map(vec, &XYZVector::R), to invoke a method on each element of the RVec. unfortunately C++ forces us to have that & in there

I’m open to better ideas :slight_smile:

Cheers,
Enrico

1 Like

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.