iguana_ex_cpp_dataframes.cc

Go to the documentation of this file.

 
#include <hipo4/RIguanaDS.hxx>
#include <hipo4/ThreadedAlgo.hxx>
 
#include <iguana/algorithms/clas12/EventBuilderFilter/Algorithm.h>
#include <iguana/algorithms/clas12/SectorFinder/Algorithm.h>
#include <iguana/algorithms/clas12/rga/MomentumCorrection/Algorithm.h>
 
#include <TApplication.h>
#include <TCanvas.h>

int main(int argc, char** argv)
{
 
  // parse arguments
  if(argc <= 1) {
    fmt::println(stderr, "USAGE: {} [HIPO_FILE] <NUM_INSPECT> <INTERACTIVE_MODE>", argv[0]);
    fmt::println(stderr, "  NUM_INSPECT is optional; by default a few events are inspected");
    fmt::println(stderr, "  INTERACTIVE_MODE=\"true\" will draw the plots interactively; default is non-interactive");
    return 2;
  }
  char const* in_file         = argc > 1 ? argv[1] : "data.hipo";
  int const num_events        = argc > 2 ? std::stoi(argv[2]) : 100;
  bool const interactive_mode = argc > 3 ? std::string(argv[3]) == "true" : false;
 
  // enable interactive mode
  auto app = interactive_mode ? new TApplication("app", &argc, argv) : nullptr;
 
  // enable multi-threading
  // ROOT::EnableImplicitMT();
 
  // open the HIPO data file, using `RIguanaDS`, a data source (`RDataSource`) subclass; we will apply Iguana
  // algorithms to this object, then transform it to an `RDataFrame` afterward
  auto ds = std::make_unique<RIguanaDS>(in_file, num_events);
 
  // clone the `REC::Particle` bank, since we want to compare its values to those after Iguana processing
  // FIXME: we could use "ORIG::REC::Particle" and every `::` will be replaced by `_`, but currently (HIPO v4.5.0)
  // only the first `::` gets replaced;
  // see https://code.jlab.org/hallb/clas12/hipo-cpp/-/work_items/12
  ds->CloneBankTo("REC::Particle", "ORIG_REC_Particle");
 
  // define iguana algorithms
  // - for dataframe usage, they must be `ThreadedAlgo` types
  iguana::ThreadedAlgo<iguana::clas12::EventBuilderFilter> algo_eventbuilder_filter; // filter by Event Builder PID (a filter algorithm)
  iguana::ThreadedAlgo<iguana::clas12::rga::MomentumCorrection> algo_momentum_correction; // momentum corrections (a transformer algorithm)
  // - creator-type algorithms must be `ThreadedCreatorAlgo` types, and they require knowledge of the `RIguanaDS`
  iguana::ThreadedCreatorAlgo<iguana::clas12::SectorFinder> algo_sector_finder(ds.get()); // get the sector for each particle (a creator algorithm)
 
  // configure iguana algorithms
  // FIXME: `ThreadAlgo` defines one algorithm instance per thread ("slot"), but currenlty (HIPO v4.5.0)
  // does not provide a way to loop over the instances, so we have to re-define the number of slots here in order
  // to be able to call a function on each algorithm instance;
  // see https://code.jlab.org/hallb/clas12/hipo-cpp/-/work_items/11
  unsigned int nSlots = ROOT::IsImplicitMTEnabled() ? ROOT::GetThreadPoolSize() : 1;
  // now we can loop over all the instances, calling `SetConfigFile` on each
  for(unsigned int i = 0; i < nSlots; i++) {
    algo_eventbuilder_filter[i].SetConfigFile("examples/config_for_examples.yaml");
    algo_momentum_correction[i].SetConfigFile("examples/config_for_examples.yaml");
    algo_sector_finder[i].SetConfigFile("examples/config_for_examples.yaml");
  }
 
  // start the algorithms
  algo_eventbuilder_filter.Start();
  algo_momentum_correction.Start();
  algo_sector_finder.Start();
 
  // filtered banks will need a synthetic boolean column, a "mask";
  // - `algo_eventbuilder_filter` will filter `REC::Particle`
  // - NOTE: the column name will be uses as is, that is, `::` will not be replaced by `_`
  ds->AddIguanaMask("REC::Particle", "MASK_REC_Particle");
 
  // define bank indices
  // - banks read from the input HIPO file
  int i_particle     = ds->GetBankIndex("REC::Particle");
  int i_config       = ds->GetBankIndex("RUN::config");
  int i_track        = ds->GetBankIndex("REC::Track");
  int i_calorimeter  = ds->GetBankIndex("REC::Calorimeter");
  int i_scintillator = ds->GetBankIndex("REC::Scintillator");
  // - banks created by iguana
  int i_sector = ds->GetBankIndex("REC::Particle::Sector");
 
  // describe how to call iguana algorithms' `Run` functions with a callback function definition
  auto iguana_run_callback = [&](unsigned int slot, hipo::banklist& banks) -> bool {
    auto& bank_particle     = banks[i_particle];
    auto& bank_config       = banks[i_config];
    auto& bank_track        = banks[i_track];
    auto& bank_calorimeter  = banks[i_calorimeter];
    auto& bank_scintillator = banks[i_scintillator];
    auto& bank_sector       = banks[i_sector];
    if(!algo_eventbuilder_filter[slot].Run(bank_particle))
      return false;
    if(!algo_sector_finder[slot].Run(bank_particle, bank_track, bank_calorimeter, bank_scintillator, bank_sector))
      return false;
    if(!algo_momentum_correction[slot].Run(bank_particle, bank_sector, bank_config))
      return false;
    return true;
  };
  ds->SetEventCallback(iguana_run_callback);
 
  // build the dataframe (`RDataFrame`) from the data source object (`RIguanaDS`)
  auto df = ROOT::RDataFrame(std::move(ds));
 
  // print the column names
  auto const column_names = df.GetColumnNames();
  fmt::println("\nDATAFRAME COLUMNS:");
  for(decltype(column_names)::size_type i = 0; i < column_names.size(); i++) {
    fmt::print("{:<30}", column_names[i]);
    if((i + 1) % 3 == 0 || i + 1 == column_names.size())
      fmt::print("\n");
  }
  fmt::print("\n");
 
  // the column `MASK_REC_Particle` is composed of `RVec<int>` objects, where each element of the `RVec`
  // indicates whether (`1`) or not (`0`) the corresponding row of `REC::Particle` passed the Iguana filter algorithms;
  // to filter, use the expression `"<col_name>[MASK_REC_Particle == 1]"`, where `<col_name>` should be replaced by the column name
  // to be filtered; let's do this for the PDG distribution
  int pdg_max  = 300;
  int pdg_min  = -pdg_max;
  int pdg_bins = pdg_max - pdg_min;
  // define fitered distribution:
  auto hist_pdg_filtered = df
                               .Define("pid", "REC_Particle_pid[MASK_REC_Particle == 1]")
                               .Histo1D({"pid_filtered", "PDG filtered", pdg_bins, static_cast<double>(pdg_min), static_cast<double>(pdg_max)}, "pid");
  // and let's compare it to the unfiltered PDG distribution:
  auto hist_pdg_unfiltered = df.Histo1D({"pid_unfiltered", "PDG unfiltered", pdg_bins, static_cast<double>(pdg_min), static_cast<double>(pdg_max)}, "REC_Particle_pid");
 
  // the created `REC::Particle::Sector` bank has the same rows as `REC::Particle`, so we can apply the filter the same way, making
  // a sector distribution; let's also restrict it to electrons only
  // FIXME: we have to make an alias, to workaround https://code.jlab.org/hallb/clas12/hipo-cpp/-/work_items/12
  auto hist_ele_sector = df
                             .Alias("REC_Particle_Sector_sector", "REC_Particle::Sector_sector")
                             .Define("sector", "REC_Particle_Sector_sector[MASK_REC_Particle == 1 && REC_Particle_pid == 11]")
                             .Histo1D({"ele_sector", "Electron Sector", 6, 0, 6}, "sector");
 
  // if you want to filter _all_ the `REC::Particle` columns, you could use an `RNode`, but note that may be slower
  // than just filtering just the columns that you need; here we define filtered `REC::Particle` columns, named
  // `MASK_REC_Particle_*`, by applying the mask to each;
  std::vector<std::string> rec_particle_vars{"pid", "px", "py", "pz", "vx", "vy", "vz", "vt", "charge", "beta", "chi2pid", "status"};
  ROOT::RDF::RNode df_filtered = df;
  for(auto const& var : rec_particle_vars)
    df_filtered = df_filtered.Define(
        fmt::format("MASK_REC_Particle_{}", var),
        fmt::format("REC_Particle_{}[MASK_REC_Particle == 1]", var));
  // let's also do this for the `ORIG_REC_Particle` columns, calling them `MASK_ORIG_REC_Particle_*`
  for(auto const& var : rec_particle_vars)
    df_filtered = df_filtered.Define(
        fmt::format("MASK_ORIG_REC_Particle_{}", var),
        fmt::format("ORIG_REC_Particle_{}[MASK_REC_Particle == 1]", var));
 
  // now let's plot the effect of momentum corrections on all electrons
  // NOTE: this is _all_ electrons that pass the Iguana filter, _i.e._, not just the scattered electron
  auto df_mom_corr = df_filtered
                         .Alias("px_orig", "MASK_ORIG_REC_Particle_px")
                         .Alias("py_orig", "MASK_ORIG_REC_Particle_py")
                         .Alias("pz_orig", "MASK_ORIG_REC_Particle_pz")
                         .Alias("px", "MASK_REC_Particle_px")
                         .Alias("py", "MASK_REC_Particle_py")
                         .Alias("pz", "MASK_REC_Particle_pz")
                         .Define("p_orig", "sqrt(px_orig*px_orig + py_orig*py_orig + pz_orig*pz_orig)")
                         .Define("p", "sqrt(px*px + py*py + pz*pz)")
                         .Define("delta_p", "p - p_orig");
  // make the 2D histogram and profile:
  auto hist_mom_corr = df_mom_corr.Histo2D({"mom_corr", "Momentum Correction;p [GeV];#Delta p [GeV]", 10, 0, 12, 100, -0.2, 0.2}, "p_orig", "delta_p");
  auto prof_mom_corr = df_mom_corr.Profile1D({"prof_mom_corr", "", 10, 0, 12}, "p_orig", "delta_p");
 
  // draw or write canvases
  auto canv = new TCanvas("canv", "canv", 1600, 1200);
  canv->SetGrid(1, 1);
  canv->Divide(2, 2);
  for(int p = 1; p <= 4; p++) {
    auto pad = canv->GetPad(p);
    pad->cd();
    switch(p) {
    case 1:
      pad->SetLogy();
      hist_pdg_unfiltered->Draw();
      break;
    case 2:
      pad->SetLogy();
      hist_pdg_filtered->Draw();
      break;
    case 3:
      hist_ele_sector->Draw();
      break;
    case 4:
      hist_mom_corr->Draw("colz");
      prof_mom_corr->SetLineColor(kBlack);
      prof_mom_corr->SetLineWidth(5);
      prof_mom_corr->SetErrorOption("s");
      prof_mom_corr->Draw("same");
      break;
    }
  }
  if(interactive_mode) {
    fmt::print("\n\nShowing plots interactively;\npress ^C to exit.\n\n");
    app->Run();
  }
  else
    canv->SaveAs("out-iguana-dataframe-example.png");
 
  return 0;
}