Running RDataFrame's Define In For Loop

ROOT
3

Hi,
here’s one way to do it in v6.16:

auto latestDF = std::make_unique<RNode>(df);                                          
for (auto i = 0u; i < nDefines; ++i)                                                                             
   latestDF = std::make_unique<RNode>(latestDF->Define(names[i], expressions[i]));

The other option is using a recursive function:

auto ApplyDefines(RNode df, const std::vector<std::string> &names, const std::vector<std::string> &exprs,
                  unsigned int i = 0)
{
   if (i == names.size())
      return df;

   return ApplyDefines(df.Define(names[i], exprs[i]), names, exprs, i + 1);
}

A full, working example with method 1:

#include <ROOT/RDataFrame.hxx>
#include <TApplication.h>

using ROOT::RDF::RNode;

int main()
{
   TApplication app("app", nullptr, nullptr);

   ROOT::RDataFrame df(100);

   const auto names = std::vector<std::string>({"a", "b"});
   const auto exprs = std::vector<std::string>({"rdfentry_", "rdfentry_*rdfentry_"});

   auto latestDF = std::make_unique<RNode>(df);                                          
   for (auto i = 0u; i < names.size(); ++i)                                                                             
      latestDF = std::make_unique<RNode>(latestDF->Define(names[i], exprs[i]));                                     

   auto g = latestDF->Graph("a", "b");
   g->Draw();
   app.Run();
   return 0;
}

And with method 2 (I like it better, less pointers is always nice):

#include <ROOT/RDataFrame.hxx>
#include <TApplication.h>

using ROOT::RDF::RNode;

auto ApplyDefines(RNode df, const std::vector<std::string> &names, const std::vector<std::string> &exprs,
                  unsigned int i = 0)
{
   if (i == names.size())
      return df;

   return ApplyDefines(df.Define(names[i], exprs[i]), names, exprs, i + 1);
}

int main()
{
   TApplication app("app", nullptr, nullptr);

   ROOT::RDataFrame df(100);

   const auto names = std::vector<std::string>({"a", "b"});
   const auto expressions = std::vector<std::string>({"rdfentry_", "rdfentry_*rdfentry_"});

   auto dfWithDefines = ApplyDefines(df, names, expressions);

   auto g = dfWithDefines.Graph("a", "b");
   g->Draw();
   app.Run();
   return 0;
}

Hope this helps!
Enrico

2 Likes