PredictionHandler.h

Go to the documentation of this file.

// PredictionHandler Class
// This is a refactoring of the Nus18PredictionHandler written by Mike Wallbank
// 
// PredictionHandler acts as a base class with FDPredictionHandler and NDPredictionHandler
// acting as derived classes
//
// author : Adam Lister
// email  : alister1@fnal.gov

#include "CAFAna/Core/Loaders.h"
#include "CAFAna/Core/SpectrumLoader.h"
#include "CAFAna/Core/Spectrum.h"
#include "CAFAna/Core/Binning.h"
#include "CAFAna/Core/ISyst.h"
#include "CAFAna/Core/SystShifts.h"
#include "CAFAna/Cuts/TimingCuts.h"
#include "CAFAna/Decomp/CheatDecomp.h"
#include "CAFAna/Decomp/ProportionalDecomp.h"
#include "CAFAna/Prediction/PredictionGenerator.h"
#include "CAFAna/Prediction/PredictionNoExtrap.h"
#include "CAFAna/Prediction/PredictionInterp.h"
#include "CAFAna/Vars/HistAxes.h"

#include "3FlavorAna/Cuts/NumuCuts2018.h"

#include "NuXAna/Cuts/NusCuts18.h"
#include "NuXAna/Prediction/NDPredictionSterile.h"
#include "NuXAna/Prediction/FDPredictionSterile.h"
#include "NuXAna/Prediction/PredictionSterile.h"
#include "NuXAna/Prediction/PredictionGeneratorNuX.h"

#include "OscLib/OscCalcSterile.h"

using namespace ana;

// base class
class PredictionHandler 
{
  public:
    PredictionHandler(std::string predOpt, bool isData = false, Loaders* loaders = nullptr);
    virtual ~PredictionHandler();

    virtual void AddVar(std::string name, HistAxis axis, Cut cuts, SystShifts shift, Var weight, bool cutdata = true) = 0; 

    virtual void AddVar(std::string name, std::string label, Var var, Binning binning, Cut cuts, SystShifts shift, Var weight, bool cutdata = true) = 0;
   
                virtual void AddVar(std::string name, 
                        std::string labelx, std::string labely,
                        Var varx, Var vary, 
                        Binning binningx, Binning binningy,
                        Cut cuts, SystShifts shift, Var weight, bool cutdata = true) = 0; 

                virtual void AddVar(std::string name, 
                        std::string labelx, std::string labely, std::string labelz,
                        Var varx, Var vary, Var varz,
                        Binning binningx, Binning binningy, Binning binningz,
                        Cut cuts, SystShifts shift, Var weight, bool cutdata = true) = 0; 


    virtual void Save(TFile* outFile) = 0; 

    void AddLoaders(Loaders* loaders);

    void AddData(SpectrumLoader* loadersData);
  
    void AddSystematics(std::vector<const ISyst*> systs, osc::OscCalcSterile* calc);

  protected:
    // Loaders                                        
    Loaders*        fLoaders;                                
    SpectrumLoader* fLoadersData;                     
    SpectrumLoader* fLoadersCos;                      

    // Systematics                                    
    std::vector<const ISyst*>  fSysts;                 
    osc::OscCalcSterile* fOscCalc;              

    // Data product containers                        
    std::map<std::string, IPrediction*> fPredictions; 
    std::map<std::string, Spectrum*>    fSpectra;        
    std::map<std::string, Spectrum*>    fCosSpec;        
    std::map<std::string, CheatDecomp*> fDecompositions;                            
    std::map<std::string, std::map<std::string, std::pair<Spectrum*, Spectrum*> > > fShiftedSpectra;                                                                  

    // options                                        
    std::string fPredOpt;                             
    bool        fData;
    bool        fRunSysts;
    bool        fCosmics;

};

PredictionHandler::PredictionHandler(std:: string predOpt,
                                     bool         isData ,
                                     Loaders*     loaders) :
  fPredOpt   (predOpt),
  fLoaders   (loaders),
  fData      (isData) ,
  fRunSysts  (false)  ,
  fCosmics   (false) 
{

  if (fPredOpt != "Extrap" && fPredOpt != "FDExtrap" && fPredOpt != "NoExtrap" && fPredOpt != "Decomp")
  {
    std::cout << "PredictionHandler not set up for " 
              << fPredOpt
              << "\nCurrent allowed options are:"
              << "\n- FD: "
              << "\n-- Extrap"
              << "\n-- FDExtrap"
              << "\n-- NoExtrap"
              << "\n-- Decomp"
              << std::endl;
    std::abort();
  }
}

PredictionHandler::~PredictionHandler() 
{
  for (const auto& pred : fPredictions)
    delete pred.second;                
  for (const auto& cos : fCosSpec)     
    delete cos.second;                 
  for (const auto& spectrum : fSpectra)
    delete spectrum.second;            
  for (const auto& decomp : fDecompositions) 
    delete decomp.second;                    
  for (const auto& shift : fShiftedSpectra) {
    for (const auto& syst : shift.second) {  
      delete syst.second.first;              
      delete syst.second.second;             
    }                                        
  }                                          
}

void PredictionHandler::AddLoaders(Loaders* loaders){
  fLoaders = loaders;
}

void PredictionHandler::AddData(SpectrumLoader* loadersData){
  fData = true;
  fLoadersData = loadersData;
}

void PredictionHandler::AddSystematics(std::vector<const ISyst*> systs, 
                                       osc::OscCalcSterile* calc) 
{
  fRunSysts = true;                                                                                         
  fSysts = systs;                                                                                           
  fOscCalc = calc;                                                                                          
}                                                                                                           

// ND Prediction handler class
class NDPredictionHandler : public PredictionHandler
{
  public:
    NDPredictionHandler(std::string predOpt, bool isData = false, Loaders* loaders = nullptr)
      : PredictionHandler(predOpt, isData, loaders)
    {
    }

    ~NDPredictionHandler()
    {
    }

    void AddVar(std::string name, HistAxis axis, Cut cuts, SystShifts shift, Var weight, bool cutdata = true); 
    void AddVar(std::string name, std::string label, Var var, Binning binning, Cut cuts, SystShifts shift, Var weight, bool cutdata = true);
                void AddVar(std::string name, 
                std::string labelx, std::string labely,
                Var varx, Var vary, 
                Binning binningx, Binning binningy,
                Cut cuts, SystShifts shift, Var weight, bool cutdata = true); 

                void AddVar(std::string name, 
                std::string labelx, std::string labely, std::string labelz,
                Var varx, Var vary, Var varz,
                Binning binningx, Binning binningy, Binning binningz,
                Cut cuts, SystShifts shift, Var weight, bool cutdata = true); 

    void Save(TFile* outFile); 
};

void NDPredictionHandler::AddVar(std::string name, 
                                 std::string label, 
                                 Var var, 
                                 Binning binning, 
                                 Cut cuts, 
                                 SystShifts shift, 
                                 Var weight,
                                                                                                                                 bool cutdata) 
{

  const HistAxis axis(label, binning, var);

  return AddVar(name, axis, cuts, shift, weight, cutdata);

}

void NDPredictionHandler::AddVar(std::string name, 
                                 std::string labelx,
                                                                                                                                 std::string labely,
                                 Var varx,
                                                                                                                                 Var vary, 
                                 Binning binningx,
                                                                                                                                 Binning binningy,
                                 Cut cuts, 
                                 SystShifts shift, 
                                 Var weight,
                                                                                                                                 bool cutdata) 
{

  const HistAxis axis(labelx, binningx, varx, labely, binningy, vary);

  return AddVar(name, axis, cuts, shift, weight, cutdata);

}

void NDPredictionHandler::AddVar(std::string name, 
                                 std::string labelx, std::string labely, std::string labelz,
                                 Var varx, Var vary, Var varz,
                                 Binning binningx, Binning binningy, Binning binningz,
                                 Cut cuts, SystShifts shift, Var weight, bool cutdata)
{
        const HistAxis axis ({labelx, labely, labelz}, {binningx, binningy, binningz}, {varx, vary, varz});

        return AddVar(name, axis, cuts, shift, weight, cutdata);
}

void NDPredictionHandler::AddVar(std::string name, 
                                 HistAxis axis, 
                                 Cut cuts, 
                                 SystShifts shift, 
                                 Var weight,
                                                                                                                                 bool cutdata) 
{

  CheatDecomp* decomp = nullptr;
  IPrediction* pred   = nullptr;
  Spectrum* spec      = nullptr;

  // Decomposition
  if (fPredOpt == "Decomp")
    decomp = new CheatDecomp(fLoaders->GetLoader(caf::kNEARDET, Loaders::kMC),
                                               axis, cuts, shift, weight);

  // Predictions
  if (fPredOpt == "NDExtrap") {
    const NDPredictionGenerator genND(axis, cuts, shift, weight);
    pred = genND.Generate(*fLoaders).release();
  }

  // Systematics
  std::map<std::string, std::pair<Spectrum*, Spectrum*> > shifts;
  if (fRunSysts) {
    for (const auto& shift : fSysts) {
      Spectrum* shiftUp   = new Spectrum(fLoaders->GetLoader(caf::kNEARDET, Loaders::kMC),
                                                                 axis, cuts, SystShifts(shift, +1), weight);
      Spectrum* shiftDown = new Spectrum(fLoaders->GetLoader(caf::kNEARDET, Loaders::kMC),
                                                                       axis, cuts, SystShifts(shift, -1), weight);
      shifts[shift->ShortName()] = std::make_pair(shiftUp, shiftDown);
    }
  }

  // Make data spectrum
  if (fData && cutdata)
    spec = new Spectrum(*fLoadersData, axis, cuts);

  if (decomp)
    fDecompositions[name] = decomp;
  if (pred)
    fPredictions[name] = pred;
  if (shifts.size())
    fShiftedSpectra[name] = shifts;
  if (spec)
    fSpectra[name] = spec;

  return;

}

void NDPredictionHandler::Save(TFile* outFile) 
{

  std::cout << "saving file..."  << std::endl;

  // Save MC decompositions to file
  for (std::map<std::string, CheatDecomp*>::const_iterator decompIt = fDecompositions.begin(); decompIt != fDecompositions.end(); ++decompIt)
    decompIt->second->SaveTo(outFile->mkdir(Form("ND%sDecomp", decompIt->first.c_str())));

  // Save systematics to file
  for (std::map<std::string, std::map<std::string, std::pair<Spectrum*, Spectrum*> > >::const_iterator shiftIt = fShiftedSpectra.begin();
       shiftIt != fShiftedSpectra.end(); ++shiftIt) {
    TString subdir = Form("ND%sDecompSysts", shiftIt->first.c_str());
    outFile->cd();
    outFile->mkdir(subdir);
    for (std::map<std::string, std::pair<Spectrum*, Spectrum*> >::const_iterator systIt = shiftIt->second.begin();
         systIt != shiftIt->second.end(); ++systIt) {
      TString subsubdir = systIt->first.c_str();
      outFile->cd(subdir);
      gDirectory->mkdir(subsubdir);
      gDirectory->cd(subsubdir);
      systIt->second.first->SaveTo(gDirectory->mkdir("ShiftUp"));
      systIt->second.second->SaveTo(gDirectory->mkdir("ShiftDown"));
    }
  }    

  // Save data spectra to file
  if (fData){
    for (std::map<std::string, Spectrum*>::const_iterator specIt = fSpectra.begin(); specIt != fSpectra.end(); ++specIt){
      specIt->second->SaveTo(outFile->mkdir(Form("ND%sSpec", specIt->first.c_str())));
    }
  }


  return;

}

// FD Prediction handler class
class FDPredictionHandler : public PredictionHandler
{
  public:
    FDPredictionHandler(std::string predOpt, bool isData = false, Loaders* loaders = nullptr)
      : PredictionHandler(predOpt, isData, loaders)
    {
    }

    ~FDPredictionHandler()
    {
    }

    void AddVar(std::string name, HistAxis axis, Cut cuts, SystShifts shift, Var weight, bool cutdata = true); 
    void AddVar(std::string name, std::string label, Var var, Binning binning, Cut cuts, SystShifts shift, Var weight, bool cutdata = true);
                void AddVar(std::string name, 
                std::string labelx, std::string labely,
                Var varx, Var vary,
                Binning binningx, Binning binningy,
                Cut cuts, SystShifts shift, Var weight, bool cutdata = true); 

                void AddVar(std::string name, 
                std::string labelx, std::string labely, std::string labelz,
                Var varx, Var vary, Var varz,
                Binning binningx, Binning binningy, Binning binningz,
                Cut cuts, SystShifts shift, Var weight, bool cutdata = true); 

    void Save(TFile* outFile); 
    void AddCosmics(SpectrumLoader* loadersCosmic);
};

void FDPredictionHandler::AddVar(std::string name, std::string label, Var var, Binning binning, Cut cuts, SystShifts shift, Var weight, bool cutdata) {

  const HistAxis axis(label, binning, var);

  return AddVar(name, axis, cuts, shift, weight, cutdata);

}

void FDPredictionHandler::AddVar(std::string name, 
                                 std::string labelx,
                                                                                                                                 std::string labely,
                                 Var varx,
                                                                                                                                 Var vary, 
                                 Binning binningx,
                                                                                                                                 Binning binningy,
                                 Cut cuts, 
                                 SystShifts shift, 
                                 Var weight,
                                                                                                                                 bool cutdata) 
{

  const HistAxis axis(labelx, binningx, varx, labely, binningy, vary);

  return AddVar(name, axis, cuts, shift, weight, cutdata);

}

void FDPredictionHandler::AddVar(std::string name, 
                                 std::string labelx, std::string labely, std::string labelz,
                                 Var varx, Var vary, Var varz,
                                 Binning binningx, Binning binningy, Binning binningz,
                                 Cut cuts, SystShifts shift, Var weight, bool cutdata)
{
        const HistAxis axis ({labelx, labely, labelz}, {binningx, binningy, binningz}, {varx, vary, varz});

        return AddVar(name, axis, cuts, shift, weight, cutdata);
}

void FDPredictionHandler::AddVar(std::string name, HistAxis axis, Cut cuts, SystShifts shift, Var weight, bool cutdata) {

  if (!fLoaders) {
    std::cout << "Error: Loaders must be set before variables can be added." << std::endl;
    abort();
  }

  Spectrum* spec = nullptr;
  Spectrum* cos = nullptr;
  IPrediction* pred = nullptr;

  // Make the prediction generator
  IPredictionGenerator* pred_gen = nullptr;
  if (fPredOpt == "Extrap") {
    SterileGenerator* genExtrap = new SterileGenerator(axis, kNus18BinsNumuCCAxis,
                                                                                           cuts, 
                                                       kNus18ND, kNumu2020ND, 
                                                       shift, weight);
    pred_gen = genExtrap;
  }
  else if (fPredOpt == "FDExtrap") {
    FDPredictionGenerator* genFD = new FDPredictionGenerator(axis, cuts, shift, weight);
    pred_gen = genFD;
  }
  else if (fPredOpt == "NoExtrap") {
    NoExtrapGenerator* genNoExtrap = new NoExtrapGenerator(axis, cuts, weight);
    pred_gen = genNoExtrap;
  }

  // Get the prediction
  if (fRunSysts) {
    PredictionInterp* pred_interp = new PredictionInterp(fSysts, fOscCalc, *pred_gen, *fLoaders);
    pred = pred_interp;
  }
  else
    pred = pred_gen->Generate(*fLoaders).release();

  // Make data spectrum
  if (fData && cutdata)
    spec = new Spectrum(*fLoadersData, axis, cuts && kInBeamSpill);

  // Make cosmic spectrum
  if (fCosmics && cutdata)
    cos = new Spectrum(*fLoadersCos, axis, cuts && kInCosmicTimingWindow);

  // Save the data products
  if (pred)
    fPredictions[name] = pred;
  if (spec)
    fSpectra[name] = spec;
  if (cos)
    fCosSpec[name] = cos;

  return;

}

void FDPredictionHandler::Save(TFile* outFile) {

  // Save MC predictions to file
  for (std::map<std::string, IPrediction*>::const_iterator predIt = fPredictions.begin();
       predIt != fPredictions.end(); ++predIt)
    predIt->second->SaveTo(outFile->mkdir(Form("FD%sPred", predIt->first.c_str())));

  // Save data spectra to file
  if (fData)
    for (std::map<std::string, Spectrum*>::const_iterator specIt = fSpectra.begin();
         specIt != fSpectra.end(); ++specIt)
      specIt->second->SaveTo(outFile->mkdir(Form("FD%sSpec", specIt->first.c_str())));

  // Save cosmics to file
  if (fCosmics)
    for (std::map<std::string, Spectrum*>::const_iterator cosIt = fCosSpec.begin();
         cosIt != fCosSpec.end(); ++cosIt)
      cosIt->second->SaveTo(outFile->mkdir(Form("FD%sCosSpec", cosIt->first.c_str())));

  return;

}

void FDPredictionHandler::AddCosmics(SpectrumLoader* loadersCosmic){
    fCosmics = true;
    fLoadersCos = loadersCosmic;
}