KinemDistributionTutorial.cpp

// MaCh3 includes
#include "Fitters/MaCh3Factory.h"
#include "SamplesTutorial/SampleHandlerTutorial.h"

struct PlotKinematicCut {
  std::string ParamToCutOn = "";
  double LowerBound = M3::_BAD_DOUBLE_;
  double UpperBound = M3::_BAD_DOUBLE_;
};

struct Plot {
  std::string Name;
  std::vector<std::string> VarStrings;
  std::vector<PlotKinematicCut> SelectionCuts = {};
};

int main(int argc, char **argv) {
  //JM OscillationChannel kinematic variable does not exist, so only supports looping modes
  int Selection = 0; //0 to loop modes, 1 to loop osc channels

  // Initialise manger responsible for config handling
  auto FitManager = MaCh3ManagerFactory(argc, argv);

  // Initialise covariance class reasonable for Systematics
  auto xsec = MaCh3CovarianceFactory<ParameterHandlerGeneric>(FitManager.get(), "Xsec");

  // Initialise samplePDF
  auto SampleConfig = Get<std::vector<std::string>>(FitManager->raw()["General"]["TutorialSamples"], __FILE__ , __LINE__);
  auto mySamples = MaCh3SampleHandlerFactory<SampleHandlerTutorial>(SampleConfig, xsec.get());

  // Output file is named after output filename in config plus plot variable(s)
  std::string configOutName = FitManager->raw()["General"]["OutputFile"].as<std::string>();
  std::stringstream OutName;
  OutName<<configOutName.substr(0,configOutName.find(".root"));

  TString OutputName = TString(OutName.str()) + "_KinemPlot" + ".pdf";

  std::vector<std::string> vecParams = {"TrueNeutrinoEnergy", "TrueQ2", "RecoNeutrinoEnergy"};
  
  //JM Read in kinematic distribution plots from config
  std::vector<Plot> PlotsToDraw = {};
  auto ConfigPlots = FitManager->raw()["KinematicDistributionPlots"];

  for (const auto& ConfigPlot : ConfigPlots) {
    Plot PlotToDraw;
    PlotToDraw.Name = Get<std::string>(ConfigPlot["Name"],  __FILE__, __LINE__);
    PlotToDraw.VarStrings = Get<std::vector<std::string>>(ConfigPlot["VarStrings"], __FILE__, __LINE__);
    if (PlotToDraw.VarStrings.size() != 1 && PlotToDraw.VarStrings.size() != 2) {
      MACH3LOG_ERROR("Error in yaml file: In KinemDistribtuion Plot {}, VarStrings is of size {}. VarString should be of size 1 or 2 (higher dimensional histogram plotting is not yet supported)");
      throw MaCh3Exception(__FILE__,__LINE__);
    }

    for (const auto& Cut : ConfigPlot["KinematicCuts"]) {
      PlotKinematicCut SelectionCut;
      SelectionCut.ParamToCutOn = Cut["VarString"].as<std::string>();
      std::vector<double> range = Cut["Range"].as<std::vector<double>>();

      if (range.size() != 2) {
        MACH3LOG_ERROR("Error in yaml file: In KinemDistribution Plot {}, KinematicCut {} has range of size {}. Range should be of size 2.", PlotToDraw.Name, SelectionCut.ParamToCutOn, range.size());
        throw MaCh3Exception(__FILE__,__LINE__);
      }
      SelectionCut.LowerBound = range[0];
      SelectionCut.UpperBound = range[1];
      PlotToDraw.SelectionCuts.push_back(SelectionCut);
    }
    PlotsToDraw.push_back(PlotToDraw);
  }

  TCanvas* Canv = new TCanvas("Canv","");
  Canv->Divide(1,2);
  Canv->Print(OutputName+"[");

  for (size_t iParam = 0; iParam < vecParams.size(); iParam++) {
    TString ParamDirName = vecParams[iParam];

    for (size_t iPDF=0; iPDF < mySamples.size(); iPDF++) {
      MACH3LOG_INFO("Number of samples: {}", mySamples[iPDF]->GetNSamples());
       for(int iSample = 0; iSample < mySamples[iPDF]->GetNSamples(); iSample++){
        THStack* Stack = new THStack(*mySamples[iPDF]->ReturnStackedHistBySelection1D(iSample, vecParams[iParam], Selection));
        TLegend* Legend = new TLegend(*mySamples[iPDF]->ReturnStackHistLegend());

        Canv->cd(1);
        Stack->Draw("HIST");
        //Due to crappy TStack design, you need to draw THStack first then assign axis titles
        Stack->SetTitle(mySamples[iPDF]->GetSampleTitle(iSample).c_str());
        Stack->GetXaxis()->SetTitle((vecParams[iParam]).c_str());
        Canv->cd(2);
        Legend->Draw();

        Canv->Update();
        Canv->Print(OutputName);
        delete Stack; Stack = nullptr;
        delete Legend; Legend = nullptr;
      }
    }
  }

  delete Canv;
  Canv = new TCanvas("Canv","");
  int WeightStyle = 1;
  for (size_t iHist=0; iHist<PlotsToDraw.size(); iHist++) {
    MACH3LOG_INFO("Plotting kinematic distributions in config: {} / {}", iHist+1, PlotsToDraw.size());  

    std::vector<std::string> PlotVar_Str = PlotsToDraw[iHist].VarStrings;
    int histdim = PlotVar_Str.size();

    for (size_t iPDF = 0;iPDF < mySamples.size(); iPDF++) {
      for(int iSample = 0; iSample < mySamples[iPDF]->GetNSamples(); iSample++){
        std::vector<KinematicCut> EventSelectionVector = {};
        std::vector<KinematicCut> SubEventSelectionVector = {};

        for (size_t iCut=0; iCut<PlotsToDraw[iHist].SelectionCuts.size(); iCut++) {
          KinematicCut Selection;
          Selection.LowerBound = PlotsToDraw[iHist].SelectionCuts[iCut].LowerBound;
          Selection.UpperBound = PlotsToDraw[iHist].SelectionCuts[iCut].UpperBound;

          if (mySamples[iPDF]->IsSubEventVarString(PlotsToDraw[iHist].SelectionCuts[iCut].ParamToCutOn)) {
            Selection.ParamToCutOnIt = mySamples[iPDF]->ReturnKinematicVectorFromString(PlotsToDraw[iHist].SelectionCuts[iCut].ParamToCutOn);
            SubEventSelectionVector.push_back(Selection);
          }
          else {
            Selection.ParamToCutOnIt = mySamples[iPDF]->ReturnKinematicParameterFromString(PlotsToDraw[iHist].SelectionCuts[iCut].ParamToCutOn);
            EventSelectionVector.push_back(Selection);
          }
        }
        std::unique_ptr<TH1> Hist;
        if (histdim == 1) {
          Hist = mySamples[iPDF]->Get1DVarHist(iSample, PlotVar_Str[0], EventSelectionVector, WeightStyle, SubEventSelectionVector);
          Hist->GetYaxis()->SetTitle("Events");
        }
        else {
          Hist = mySamples[iPDF]->Get2DVarHist(iSample, PlotVar_Str[0], PlotVar_Str[1], EventSelectionVector, WeightStyle, SubEventSelectionVector);
          Hist->GetYaxis()->SetTitle(PlotVar_Str[1].c_str());
        }
        Canv->cd(1);
        Hist->SetTitle(PlotsToDraw[iHist].Name.c_str());
        Hist->GetXaxis()->SetTitle(PlotVar_Str[0].c_str());
        Hist->SetStats(false);
        Hist->Draw("COLZ");
        Canv->Print(OutputName);
      }
    }
  }
  Canv->Print(OutputName+"]");
  delete Canv;
  return 0;
}