Variable as argument of function in RDataFrame Define

Hello,

I am currently trying to write a python script using a RDataFrame in which I would like to create several columns from a user defined function with different arguments for each column as in :

rdf = rdf.Define('genDauInCone_02','''IsgenDauInCone(0.2,genLepton_daughters,HPStau_p4)''')

where rdf is my RDataFrame and ‘IsgenDauInCone’ is a C++ function defined with gInterpreter. Here 0.2 is the value I would like to vary but instead of writing as many lines as needed values I would like to write something like :

for dRmax in np.arange(0,1,0.1):
 a = str(int(dRmax*10))
 rdf=rdf.Define('genDauInCone_'+a,'''IsgenDauInCone(dRmax,genLepton_daughters,HPStau_p4)''')

but this gets me an error ‘dRmax’ undefined’. Is there a way to perform this ?

Thank you very much.
Mario

ROOT Version: 6.24/00
Platform: linuxx8664gcc


Hi Mario,
and welcome to the ROOT forum.

The simplest solution I can think of:

rdf.Define('genDauInCone_'+a, f"IsgenDauInCone({dRmax},genLepton_daughters,HPStau_p4)")

The most versatile solution is probably using a functor class:

gInterpreter.Declare("""
struct IsgenDauInConeFunctor {
  double _dRmax;

  IsgenDauInConeFunctor(double dRmax) : _dRmax(dRmax) {}
 
  double operator()(double genLepton_daughters, double HPStau_p4) {
     return IsgenDauInCone(_dRmax, genLepton_daughters, HPStau_p4);
  }
};
""")

for dRmax in np.arange(0,1,0.1):
 a = str(int(dRmax*10))
 f = ROOT.IsgenDauInConeFunctor(dRmax);
 rdf = rdf.Define('genDauInCone_'+a, f, ["genLepton_daughters","HPStau_p4"])

Cheers,
Enrico

Thank you for your prompt reply, this is working.

Best,
Mario

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