SpectrumHandler.cxx

Go to the documentation of this file.

#include "CAFAna/Analysis/SpectrumHandler.h"

// Standard includes
#include <iostream>
#include <exception>
#include <string>
#include <vector>

// CAFE includes
#include "CAFAna/Core/Binning.h"
#include "CAFAna/Core/Var.h"
#include "CAFAna/Core/MultiVar.h"
#include "CAFAna/Core/Cut.h"
#include "CAFAna/Core/Spectrum.h"
#include "CAFAna/Core/SpectrumLoader.h"
#include "CAFAna/Core/SystShifts.h"

// ROOT includes
#include <TROOT.h>
#include <TFile.h>

// This constructor creates a new SpectrumHandler using existing loaders. This is greatly beneficial when 
ana::SpectrumHandler::SpectrumHandler(const std::vector<SpectrumLoaderBase*> loaders, const std::vector<std::string> dataset_labels){
  // this->SetLoaders(loaders, dataset_labels);
  fNLoaders = loaders.size();
  fLoaders = loaders;
  fLoader_labels = dataset_labels;

  fStatus = kSpecHelpLoaders;
}

bool ana::SpectrumHandler::SetLoaders(const std::vector<std::string>& datasets, const std::vector<std::string>& dataset_labels){
  if (fStatus != kSpecHelpUnloaded)
    throw std::runtime_error("Attempt to set loaders while SpecHelper is not unloaded.");

  if (datasets.size() != dataset_labels.size()){
    std::cerr << "Number of datasets not equal to number of dataset labels. Please fix that." << std::endl;
    return false;
  }

  fNLoaders = datasets.size();
  for (unsigned int idataset = 0; idataset < fNLoaders; idataset++){
    SpectrumLoaderBase * loader;
    loader = new SpectrumLoader(datasets[idataset]);
    fLoaders.push_back(loader);
    fLoader_labels.push_back(dataset_labels[idataset]);
  }
  fStatus = kSpecHelpLoaders;
  return true;
}

bool ana::SpectrumHandler::SetLoaders(const std::vector<SpectrumLoaderBase*> loaders, const std::vector<std::string> loader_labels){
  if (fStatus != kSpecHelpUnloaded)
    throw std::runtime_error("Attempt to set loaders while SpecHelper is not unloaded.");

  if (loaders.size() != loader_labels.size()){
    std::cerr << "Number of loaders not equal to number of loader labels." << std::endl;
    return false;
  }

  fNLoaders = loaders.size();
  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++){
    SpectrumLoaderBase * loader = loaders[iloader];
    if (!loader || loader == NULL)
      throw std::runtime_error("Attempt to use a non-existent loader: " + loader_labels[iloader]);
    fLoaders.push_back(loader);
    fLoader_labels.push_back(loader_labels[iloader]);
  }
  fStatus = kSpecHelpLoaders;
  return true;
}

bool ana::SpectrumHandler::SetSpillCuts(const std::vector<SpillCut> spillcuts){
  if (fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set spill cuts sometime NOT right after setting loaders. Only allow setting spill cuts IMMEDIATELY after setting loaders.");

  if (fNLoaders != spillcuts.size()){
    std::cerr << "Warning: Number of spill cuts ( " << std::to_string(spillcuts.size()) << " ) not equal to number of loaders ( " << std::to_string(fNLoaders) << " ).\n\tFailed to set spill cuts." << std::endl;
    return false;
  }

  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++)
    fLoaders[iloader]->SetSpillCut(spillcuts[iloader]);

  return true;
}

bool ana::SpectrumHandler::SetSpillCuts(const SpillCut cut){
  if (fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set spill cuts sometime NOT right after setting loaders. Only allow setting spill cuts IMMEDIATELY after setting loaders.");

  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++)
    fLoaders[iloader]->SetSpillCut(cut);

  return true;
}

ana::SpectrumLoaderBase * ana::SpectrumHandler::GetLoader(const std::string label){
  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++)
    if (fLoader_labels[iloader] == label)
      return fLoaders[iloader];
  return NULL;
}

std::vector<ana::SpectrumLoaderBase*> ana::SpectrumHandler::GetLoaders(const std::vector<std::string> loader_labels){
  std::vector<SpectrumLoaderBase*> loaderList;
  for (unsigned int ilabel = 0; ilabel < loader_labels.size(); ilabel++){
    SpectrumLoaderBase * loader = GetLoader(loader_labels[ilabel]);
    if (loader == NULL)
      throw std::runtime_error("Could not find loader with label \"" + loader_labels[ilabel] + "\". Skipping.");
    loaderList.push_back(loader);
  }
  return loaderList;
}

bool ana::SpectrumHandler::SetVars(const std::vector<Var> vars, const std::vector<std::string> var_labels, const std::vector<std::string> var_titles, const std::vector<Binning> var_bins){
  if (fStatus != kSpecHelpUnloaded && fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set the vars too late. The spectrums are either ready, or already had go() called on them.");

  if (vars.size() != var_labels.size() || vars.size() != var_titles.size() || vars.size() != var_bins.size()){
    std::cerr << "Number of vars not equal to number of var labels / titles. Fix that." << std::endl;
    return false;
  }

  fVars = vars;
  fVar_labels = var_labels;
  fVar_titles = var_titles;
  fVar_bins   = var_bins;

  fNVars = fVars.size();

  return true;
}

bool ana::SpectrumHandler::SetCuts(const std::vector<Cut> cuts, const std::vector<std::string> cut_labels){
  if (fStatus != kSpecHelpUnloaded && fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set the vars too late. The spectrums are either ready, or already had go() called on them.");

  if (cuts.size() != cut_labels.size()){
    std::cerr << "Number of cuts not equal to number of cut labels. Fix that." << std::endl;
    return false;
  }

  fCuts = cuts;
  fCut_labels = cut_labels;

  fNCuts = fCuts.size();

  return true;
}

bool ana::SpectrumHandler::SetWeights(const std::vector<Var> weights, const std::vector<std::string> weight_labels){
  if (fStatus != kSpecHelpUnloaded && fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set the vars too late. The spectrums are either ready, or already had go() called on them.");

  if (weights.size() != weight_labels.size()){
    std::cerr << "Number of Weights not equal to number of weight labels. Fix that." << std::endl;
    return false;
  }

  fWeights = weights;
  fWeight_labels = weight_labels;

  fNWeights = fWeights.size();

  return true;
}

bool ana::SpectrumHandler::SetShifts(const std::vector<SystShifts> shifts, const std::vector<std::string> shift_labels){
  if (fStatus != kSpecHelpUnloaded && fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set the vars too late. The spectrums are either ready, or already had go() called on them.");

  if (shifts.size() != shift_labels.size()){
    std::cerr << "Number of shifts not equal to number of shift labels. Ignoring shifts." << std::endl;
    return false;
  }

  fShifts = shifts;
  fShift_labels = shift_labels;

  fNShifts = fShifts.size();

  return true;
}

bool ana::SpectrumHandler::SetHistAxes(const std::vector<HistAxis> haxes, const std::vector<std::string> haxis_labels){
  if (fStatus != kSpecHelpUnloaded && fStatus != kSpecHelpLoaders)
    throw std::runtime_error("Attempt to set the vars too late. The spectrums are either ready, or already had go() called on them.");

  if (haxes.size() != haxis_labels.size()){
    std::cerr << "Number of HistAxis not equal to number of HistAxis labels. Ignoring HistAxes." << std::endl;
    return false;
  }

  fHistAxes = haxes;
  fHistAxis_labels = haxis_labels;

  fNHistAxes = fHistAxes.size();

  return true;
}

bool ana::SpectrumHandler::CheckReady(){
  if (fStatus != kSpecHelpLoaders)
    return false;
  bool loadersgood = fLoaders.size() && fLoaders.size() == fLoader_labels.size();
  bool varsgood = fVars.size() > 0 && fVars.size() == fVar_labels.size() && fVars.size() == fVar_titles.size() && fVars.size() == fVar_bins.size();
  bool histaxesgood = fHistAxes.size() > 0 && fHistAxes.size() == fHistAxis_labels.size();
  bool cutsgood = fCuts.size() > 0 && fCuts.size() == fCut_labels.size();
  bool weightsgood   = fWeights.size() > 0   && fWeight_labels.size() == fWeights.size();
  bool shiftsgood   = fShifts.size() > 0   && fShift_labels.size() == fShifts.size();

  bool allready = loadersgood && (varsgood || histaxesgood) && cutsgood && weightsgood && shiftsgood;
  if (allready)
    fStatus = kSpecHelpReady;
  return allready;
}

bool ana::SpectrumHandler::CreateSpectrums(){
  if (fStatus != kSpecHelpReady){
    if (fStatus != kSpecHelpLoaders || !this->CheckReady()){
      std::cerr << "Not ready to create spectrums. Doing nothing..." << std::endl;
      return false;
    } // We just checked that we had a loader, and CheckReady() returned true. Can continue.
    assert(fStatus == kSpecHelpReady); // Now we KNOW it is ready... right?
  }

  unsigned long numSpec = fNLoaders * (fNVars + fNHistAxes) * fNCuts * fNWeights * fNShifts;
  std::cout << "Spectrum Helper: Creating all spectrum definitions.\n\tThere are a total of " << std::to_string(numSpec) << " Spectrums from:" << 
    "\n\tLoaders:  " << std::to_string(fNLoaders) <<
    "\n\tVars:     " << std::to_string(fNVars) <<
    "\n\tHistAxes: " << std::to_string(fNHistAxes) <<
    "\n\tCuts:     " << std::to_string(fNCuts) <<
    "\n\tWeights:  " << std::to_string(fNWeights) <<
    "\n\tShifts:   " << std::to_string(fNShifts) << std::endl;

  CreateVarSpectrums();
  CreateHistAxisSpectrums();

  fStatus = kSpecHelpReady;

  return true;
}

void ana::SpectrumHandler::CreateVarSpectrums(){
  // Loop through loaders
  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++){
    std::vector<std::vector<std::vector<std::vector<Spectrum* > > > > lspecs;
    // Loop through vars
    for (unsigned int ivar = 0; ivar < fNVars; ivar++){
      std::vector<std::vector<std::vector<Spectrum* > > > vspecs;
      // Loop through cuts
      for (unsigned int icut = 0; icut < fNCuts; icut++){
        std::vector<std::vector<Spectrum* > > cspecs;
        // Now some actual logic.
        for (unsigned int iweight = 0; iweight < fNWeights; iweight++){
          std::vector<Spectrum* > wspecs;
          // Loop through shifts
          for (unsigned int ishift = 0; ishift < fNShifts; ishift++)
            wspecs.push_back(new Spectrum(
              fVar_titles[ivar],
              fVar_bins[ivar],
              *fLoaders[iloader],
              fVars[ivar],
              fCuts[icut],
              fShifts[ishift],
              fWeights[iweight] ));
          cspecs.push_back(wspecs);
        }
        vspecs.push_back(cspecs);
      }
      lspecs.push_back(vspecs);
    }
    fSpecs.push_back(lspecs); // Put into the actual spec holder (5D std::vector: loader -> var -> cut -> weight -> shift)
  }
}

void ana::SpectrumHandler::CreateHistAxisSpectrums(){
  // Loop through loaders
  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++){
    std::vector<std::vector<std::vector<std::vector<Spectrum* > > > > lspecs;
    // Loop through vars
    for (unsigned int ihaxis = 0; ihaxis < fNHistAxes; ihaxis++){
      std::vector<std::vector<std::vector<Spectrum* > > > hspecs;
      // Loop through cuts
      for (unsigned int icut = 0; icut < fNCuts; icut++){
        std::vector<std::vector<Spectrum* > > cspecs;
        // Now some actual logic.
        for (unsigned int iweight = 0; iweight < fNWeights; iweight++){
          std::vector<Spectrum* > wspecs;
          // Loop through shifts
          for (unsigned int ishift = 0; ishift < fNShifts; ishift++)
            wspecs.push_back(new Spectrum(
              *fLoaders[iloader],
              fHistAxes[ihaxis],
              fCuts[icut],
              fShifts[ishift],
              fWeights[iweight] ));
          cspecs.push_back(wspecs);
        }
        hspecs.push_back(cspecs);
      }
      lspecs.push_back(hspecs);
    }
    fHistAxisSpecs.push_back(lspecs); // Put into the actual spec holder (5D std::vector: loader -> var -> cut -> weight -> shift)
  }
}

bool ana::SpectrumHandler::Go(){
  if (fStatus != kSpecHelpReady){
    std::cerr << "You didn't create any spectrums yet." << std::endl;
    return false;
  }
  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++){
    SpectrumLoaderBase * loader = fLoaders[iloader];
    if (!loader || loader == NULL)
      throw std::runtime_error("Trying to call Go() on Loader[" + std::to_string(iloader) + "], which was not defined.");
    else if (loader->Gone()){
      std::cout << "Loader: \"" + fLoader_labels[iloader] + "\" has already Gone() (likely by another instance of SpectrumHandler)." << std::endl;
      continue;
    }
    fLoaders[iloader]->Go();
  }

  fStatus = kSpecHelpWent;
  return true;
}

bool ana::SpectrumHandler::Gone(){
  if (fStatus == kSpecHelpUnloaded){
    std::cerr << "No loaders have been set yet." << std::endl;
    return false;
  }
  // Loop throuh all loaders
  for (SpectrumLoaderBase * l : fLoaders)
    // If any loader has had Go() already called, return true
    if (l->Gone())
      return true;

  // No loaders had Go() called yet.
  return false;
}

bool ana::SpectrumHandler::SaveSpectrums(TFile * fout){
  if (fStatus != kSpecHelpWent){
    std::cerr << "You didn't create the spectrums and fill them yet!!! Create them, and call Go() on loaders." << std::endl;
    return false;
  }
  // Check file exists, is writable here?
  if (!fout || fout == NULL || fout->IsZombie()){
    std::cerr << "File not writable. Aborting save." << std::endl;
    return false;
  }

  // Only loop through loaders once to prevent trying to make directories twice
  for (unsigned int iloader = 0; iloader < fNLoaders; iloader++){
    // Switch to the top-level directory of the output file.
    fout->cd();
    TDirectory* loaderDir = fout->GetDirectory("/");

    // If multiple loaders, each gets its own directory for simplicity.
    if (fLoader_labels.size() > iloader && !fLoader_labels[iloader].empty())
      loaderDir = fout->mkdir(fLoader_labels[iloader].c_str());

    SaveVarSpectrums(loaderDir, iloader);
    SaveHistAxisSpectrums(loaderDir, iloader);

  } // End loop through loaders

  return true;
}

void ana::SpectrumHandler::SaveVarSpectrums(TDirectory * loaderDir, unsigned int iloader){
  // Loop through vars
  for (unsigned int ivar = 0; ivar < fNVars; ivar++){
    std::string vlabel = fVar_labels[ivar];

    // Loop through cuts
    for (unsigned int icut = 0; icut < fNCuts; icut++){
      std::string clabel = fCut_labels[icut];

      // Loop through weights
      for (unsigned int iweight = 0; iweight < fNWeights; iweight++){
        std::string wlabel = fWeight_labels[iweight];

        // Loop through shifts
        for (unsigned int ishift = 0; ishift < fNShifts; ishift++){
          std::string slabel = fShift_labels[ishift];

          std::string specname = vlabel + "_" + clabel;
          if (!wlabel.empty())
            specname = specname + "_" + wlabel;
          if (!slabel.empty())
            specname = specname + "_" + slabel;

          fSpecs[iloader][ivar][icut][iweight][ishift]->SaveTo(loaderDir, specname);

        }

      }
    }
  }
}

void ana::SpectrumHandler::SaveHistAxisSpectrums(TDirectory * loaderDir, unsigned int iloader){
  // Loop through hist axes
  for (unsigned int ihaxis = 0; ihaxis < fNHistAxes; ihaxis++){
    std::string hlabel = fHistAxis_labels[ihaxis];

    // Loop through cuts
    for (unsigned int icut = 0; icut < fNCuts; icut++){
      std::string clabel = fCut_labels[icut];

      // Loop through weights
      for (unsigned int iweight = 0; iweight < fNWeights; iweight++){
        std::string wlabel = fWeight_labels[iweight];

        // Loop through shifts
        for (unsigned int ishift = 0; ishift < fNShifts; ishift++){
          std::string slabel = fShift_labels[ishift];

          std::string specname = hlabel + "_" + clabel;
          if (!wlabel.empty())
            specname = specname + "_" + wlabel;
          if (!slabel.empty())
            specname = specname + "_" + slabel;

          fHistAxisSpecs[iloader][ihaxis][icut][iweight][ishift]->SaveTo(loaderDir, specname);

        }

      }
    }
  }
}