misaligned_friends.cpp

#include <ROOT/RDataFrame.hxx>
#include <TTreeReader.h>
#include <TCanvas.h>
#include <TTreeReaderValue.h>
#include <TChain.h>
#include <TFile.h>
#include <TTree.h>
#include <TH1D.h>
#include <iostream>

void CompareWithTreeFriend(const std::string &frname, bool withIndex = false) {
   TFile fmain("main.root");
   auto tmain = fmain.Get<TTree>("main");
   TFile ffriend(frname.c_str());
   auto tfriend = ffriend.Get<TTree>("friend");

   if (withIndex)
      tfriend->BuildIndex("idx");

   tmain->AddFriend(tfriend);

   int counter = 0;
   TTreeReader r(tmain);
   while (r.Next())
      ++counter;
   std::cout << "\tTTreeReader count " << counter << '\n';

   TH1D h1("h1", "h1", 100, 0, 0);
   tmain->Draw("x>>h1");
   TH1D h2("h2", "h2", 100, 0, 0);
   tmain->Draw("friend.y>>h2");
   std::cout << "\tTTree::Draw count " << h1.GetEntries() << '\n';
   std::cout << "\tTTree::Draw count " << h2.GetEntries() << '\n';
}

void CompareWithChainFriend(const std::string &frname, bool withIndex = false) {
   TChain cmain("main");
   cmain.Add("main.root");
   TChain cfriend("friend");
   cfriend.Add(frname.c_str());

   if (withIndex)
      cfriend.BuildIndex("idx");

   cmain.AddFriend(&cfriend);

   int counter = 0;
   TTreeReader r(&cmain);
   while (r.Next())
      ++counter;
   std::cout << "\tTTreeReader count " << counter << '\n';

   TH1D h1("h1", "h1", 100, 0, 0);
   cmain.Draw("x>>h1");
   TH1D h2("h2", "h2", 100, 0, 0);
   cmain.Draw("friend.y>>h2");
   std::cout << "\tTTree::Draw count " << h1.GetEntries() << '\n';
   std::cout << "\tTTree::Draw count " << h2.GetEntries() << '\n';
}

int main() {
   ROOT::RDataFrame(10)
      .Define("idx", [] { return 1; })
      .Define("x", [] { return 42; })
      .Snapshot<int, int>("main", "main.root", {"x", "idx"});
   ROOT::RDataFrame(5).Define("y", [] { return 1; })
      .Define("idx", [] { return 1; })
      .Snapshot<int, int>("friend", "smaller_friend.root", {"y", "idx"});
   ROOT::RDataFrame(15)
      .Define("y", [] { return 1; })
      .Define("idx", [] { return 1; })
      .Snapshot<int, int>("friend", "larger_friend.root", {"y", "idx"});

   TCanvas c; // just to silence ROOT's Info messages

   std::cout << "Tree, smaller friend\n";
   CompareWithTreeFriend("smaller_friend.root");
   std::cout << "Tree, larger friend\n";
   CompareWithTreeFriend("larger_friend.root");

   std::cout << "Chain, smaller friend\n";
   CompareWithChainFriend("smaller_friend.root");
   std::cout << "Chain, larger friend\n";
   CompareWithChainFriend("larger_friend.root");

   std::cout << "Tree, smaller indexed friend\n";
   CompareWithTreeFriend("smaller_friend.root", /*withIndex=*/true);
   std::cout << "Tree, larger indexed friend\n";
   CompareWithTreeFriend("larger_friend.root", /*withIndex=*/true);

   std::cout << "Chain, smaller indexed friend\n";
   CompareWithChainFriend("smaller_friend.root", /*withIndex=*/true);
   std::cout << "Chain, larger indexed friend\n";
   CompareWithChainFriend("larger_friend.root", /*withIndex=*/true);

   return 0;
}