[df] RedefinePerSample

[gpetruc]

Include RedefinePerSample

What and why

Currently RDataFrame has Define and Redefine for ordinary columns, but only DefinePerSample; this PR adds RedefinePerSample.

One usecase is to redefine some per-sample quantities (e.g. cross section) that may were saved as columns in the input sample, if those original values were incorrect.

How

The implementation is easy, as just like for the ordinary Define and Redefine the two only differ in the precondition checks, the way the column is eventually defined is the same.

Similarly to DefineandRedefine, the new implementation of DefinePerSampleandRedefinePerSample` just calls into private implementation functions with an additional argument specifying whether it's a define or redefine and applying the proper preconditions. The final implementation

A few unit tests are added

• Check that RedefinePerSample throws if the column doesn't exist already
• Check that RedefinePerSample works for a column already present in the input file
• For completeness, check that DefinePerSample would instead throw an exception for a column from the input file

Longer example.

#include <cstdio>
#include <ROOT/RDataFrame.hxx>
#include <ROOT/RDFHelpers.hxx>
#include <ROOT/RDF/RDatasetSpec.hxx>
#include <ROOT/RDF/RSample.hxx>
#include <ROOT/RSnapshotOptions.hxx>
#include <ROOT/RHist.hxx>
#include <ROOT/Hist/ConvertToTH1.hxx>
#include <TH1.h>

void gen(const std::string & outfile, float xsec, float xoffset) {
  ROOT::RDataFrame df(10);
  ROOT::RDF::RSnapshotOptions options;
  options.fMode ="RECREATE";
  auto dfX = df.Define("x", [xoffset](ULong64_t e) -> Float_t { return e + xoffset; }, {"rdfentry_"});
  if (xsec) {
    dfX.DefinePerSample("xsec", [xsec](unsigned int slot, const ROOT::RDF::RSampleInfo &id) -> Float_t { return xsec; })
    .Snapshot("Events", outfile, {"x", "xsec"}, options);
    printf("Wrote %s with xsec=%.1f\n", outfile.c_str(), xsec);
  } else {
    dfX.Snapshot("Events", outfile, {"x"}, options);
    printf("Wrote %s without xsec\n", outfile.c_str());
  }
}

bool analyze(const std::string & infile1, float xsec1, const std::string & infile2, float xsec2, const std::string & xsecName = "xsec", bool usePerSample = true, bool redefinePerSample = false) {
   ROOT::RDF::Experimental::RMetaData meta1, meta2;
   if (usePerSample) {
    meta1.Add("xsec", xsec1);
    meta2.Add("xsec", xsec2);
   }
   auto spec = ROOT::RDF::Experimental::RDatasetSpec();
   spec.AddSample(ROOT::RDF::Experimental::RSample("sample1", "Events", infile1, meta1));
   spec.AddSample(ROOT::RDF::Experimental::RSample("sample2", "Events", infile2, meta2));
   auto df = ROOT::RDataFrame(spec);
   ROOT::RDF::RResultPtr<TH1D> result;
   if (redefinePerSample) {
        printf("Using RedefinePerSample with name %s\n", xsecName.c_str());
        result = df.RedefinePerSample(xsecName, [](unsigned int slot, const ROOT::RDF::RSampleInfo &id) -> Float_t { return id.GetD("xsec"); })
               .Define("weight", [](Float_t xs) -> Float_t { return xs * 10.0f; }, {xsecName})
               .Histo1D({"hx", "hx", 20, 0.0, 20.0}, "x", "weight");
   } else if (usePerSample) {
        printf("Using DefinePerSample with name %s\n", xsecName.c_str());
        result = df.DefinePerSample(xsecName, [](unsigned int slot, const ROOT::RDF::RSampleInfo &id) -> Float_t { return id.GetD("xsec"); })
               .Define("weight", [](Float_t xs) -> Float_t { return xs * 10.0f; }, {xsecName})
               .Histo1D({"hx", "hx", 20, 0.0, 20.0}, "x", "weight");   } else {
        printf("Using existing xsec with name %s\n", xsecName.c_str());
        result = df.Define("weight", [](Float_t xs) -> Float_t { return xs * 10.0f; }, {xsecName})
                .Histo1D({"hx", "hx", 20, 0.0, 20.0}, "x", "weight");
   }
   auto hist = *result; // trigger the computation of the histogram
   bool success = true;
   for (unsigned int i = 1; i < 20; ++i) {
        float obs = hist.GetBinContent(i);
        float exp1 = (i >= 2 && i <= 11) ? xsec1 * 10.0f : 0.0f;
        float exp2 = (i >= 7 && i <= 16) ? xsec2 * 10.0f : 0.0f;
        bool ok = std::abs(obs - (exp1 + exp2)) < 1e-5;
        printf("bin %2d: seen %6.1f  exp %6.1f  diff %6.2f  %s\n", i, obs, exp1 + exp2, obs - (exp1 + exp2), ok ? "OK" : "FAIL");
        success = success && ok;
   }
   return success;
}


int main(int argc, char** argv) {
    float xsec1 = 0.2, xsec2 = 0.3;
    printf("\nTest 1: using per-sample metadata using DefinePerSample in the second step\n");
    gen("file1.root", 0, 1.5);
    gen("file2.root", 0, 6.5);
    analyze("file1.root", xsec1, "file2.root", xsec2, "xsec", true);
    
    printf("\nTest 2: using xsec embedded from first step\n");
    gen("file1.root", xsec1, 1.5);
    gen("file2.root", xsec2, 6.5);
    analyze("file1.root", xsec1, "file2.root", xsec2, "xsec", false);
    
    printf("\nTest 3: DefinePerSample overriding wrong xsec from step1 using a new name\n");
    gen("file1.root", xsec1*3, 1.5);
    gen("file2.root", xsec2, 6.5);
    analyze("file1.root", xsec1, "file2.root", xsec2, "xsecFixed", true);
    
    printf("\nTest 4: DefinePerSample overriding wrong xsec from step1, keeping same name\n");
    gen("file1.root", xsec1*3, 1.5);
    gen("file2.root", xsec2, 6.5);
    analyze("file1.root", xsec1, "file2.root", xsec2, "xsec", true, true);    
}