RDataFrame apply method on vector of LorentzVectors

Newbie
,
4

Right,

The following snippet should then be closer to your case:

import ROOT

if __name__ == "__main__":
    df = (
        ROOT.RDataFrame(5)
            .Define("photon_pt", "ROOT::RVecD(3, 1.)")
            .Define("photon_eta", "ROOT::RVecD(3, 2.)")
            .Define("photon_phi", "ROOT::RVecD(3, 3.)")
            .Define("photon_E", "ROOT::RVecD(3, 4.)")
            .Define("photons", "ROOT::VecOps::Construct<ROOT::Math::PtEtaPhiEVector>(photon_pt, photon_eta, photon_phi, photon_E)")
            .Define("ph_pt", "return ROOT::VecOps::Map(photons, [](const ROOT::Math::PtEtaPhiEVector &v){return v.Pt();}); ")
    )
    df.Display(("photons",)).Print()

First off, with respect to your original example which you said doesn’t work, I don’t think photons.Pt() is doing what you think. The docs of the Map function state that you should pass a callable that will be applied to all the elements of the collection. In your case, the collection is photons (i.e. an RVec<ROOT::Math::PtEtaPhiEVector>). The other argument you are passing, i.e. photons.Pt(), is not a callable. Logically, that would mean the result of the method Pt, which is a double. Not even photons.Pt would work since that is a member function of the ROOT::Math::PtEtaPhiEVector class, and not the RVec class which is the type of photons. So, practically, you have to pass a callable that takes a ROOT::Math::PtEtaPhiEVector as input argument and calls its Pt method (that is what you see in my example above).

That being said, pay extra attention in this particular case: passing a C++ lambda to the Map call inside the string hits a non-ideal behaviour which is documented in the docs (incidentally with exactly your use case as an example) and has also been explained in the forum in other posts (here for example). Make sure to always explicitly use return in the strings you pass to Define and other operations whenever another return keyword is already present in the string.

Cheers,
Vincenzo