CovMxManager.cxx

Go to the documentation of this file.

// CAFAna includes
#include "CAFAna/Core/LoadFromFile.h"
#include "CAFAna/Core/Progress.h"
#include "CAFAna/Systs/CovMxSysts.h"
#include "CAFAna/Prediction/PredictionInterp.h"
#include "CAFAna/Analysis/CovMxManager.h"
#include "CAFAna/Analysis/Exposures.h"

#include "OscLib/OscCalcSterile.h"

// ROOT includes
#include "TKey.h"
#include "TMatrixD.h"

#include <Eigen/Dense>

// C++ includes
#include <iostream>
#include <numeric>

namespace ana {

  using std::back_inserter;
  using std::begin;
  using std::cout;
  using std::end;
  using std::endl;
  using std::get;
  using std::make_pair;
  using std::make_unique;
  using std::ostringstream;
  using std::runtime_error;
  using std::stoi;
  using std::string;
  using std::transform;
  using std::unique_ptr;
  using std::vector;

  using Eigen::MatrixXd;

  using covmx::CovarianceMatrix;
  using covmx::Patch;
  using covmx::Sample;
  using covmx::SamplePair;

  //----------------------------------------------------------------------
  CovMxManager::CovMxManager(TDirectory* dir, std::vector<Sample> samples)
    : fTopDir(dir) {

    transform(begin(samples), end(samples), back_inserter(fSampleIDs),
      [](Sample s) -> unsigned int { return s.GetID(); });
    if (!dir) throw runtime_error("Directory passed is not valid!");
    TObjString tag("CovMxManager");
    dir->WriteTObject(&tag);
  }

  //----------------------------------------------------------------------
  void CovMxManager::AddSystematic(vector<Sample>& samples,
    const ISyst* syst, size_t patch) {

    // Check samples are consistent
    if (true) { // Let objects in this code block fall out of scope
      ostringstream err;
      err << "Vector of samples passed to AddSystematic is inconsistent "
        << "with the one passed to the CovMxManager constructor.";
      if (samples.size() != fSampleIDs.size()) {
        throw runtime_error(err.str());
      }
      for (size_t i = 0; i < samples.size(); ++i) {
        if (samples[i].GetID() != fSampleIDs[i]) {
          throw runtime_error(err.str());
        }
      } // for sample i
    } // if true

    // Trivial oscillation calculator
    osc::IOscCalcAdjustable* calc = new osc::OscCalcSterileTrivial();

    // Making a matrix across many samples can be computationally intensive,
    // so we want to be able to break it up into smaller chunks. However,
    // because it's a symmetric matrix and we only actually generate in a
    // triangular manner and then mirror when necessary, splitting it up
    // evenly isn't entirely trivial.

    // Overall we will be generating n factorial matrices, where n is the
    // number of samples. That's what the "total" variable is
    size_t total = 0;
    for (size_t i = 1; i <= samples.size(); ++i) {
      total += i;
    }

    size_t nPatches = (samples.size() * (samples.size()+1)) / 2;
    if (patch >= nPatches) {
      ostringstream err;
      err << "Requested patch number " << patch+1 << ", but there are only"
        << nPatches << " patches in this matrix.";
      throw runtime_error(err.str());
    }

    // Loop over samples
    size_t counter = 0;

    for (size_t iS = 0; iS < samples.size(); ++iS) {
      for (size_t jS = iS; jS < samples.size(); ++jS) {

        // Skip matrices until we get to the first one we want
        if (counter < patch) {
    ++counter;
    continue;
        }

        IPrediction* iP = samples[iS].GetPrediction();
        IPrediction* jP = samples[jS].GetPrediction();

        // Number of bins & components for each of the current samples
        size_t iNBins = samples[iS].GetBinning().NBins();
        size_t jNBins = samples[jS].GetBinning().NBins();
        auto iComps = covmx::GetComponents(samples[iS]);
        auto jComps = covmx::GetComponents(samples[jS]);
        size_t iNComps = iComps.size();
        size_t jNComps = jComps.size();

        // Get exposures
        double iPOT = samples[iS].GetPOT();
        double jPOT = samples[jS].GetPOT();

        ostringstream msg;
        msg << "Producing " << syst->ShortName() << " matrix for samples "
    << samples[iS].GetName() << " and " << samples[jS].GetName();
        Progress p(msg.str());
        size_t nUniv = 1e5;

        // Nominal spectra
        vector< vector<double> > iNom(iNComps, vector<double>(iNBins));
        vector< vector<double> > jNom(jNComps, vector<double>(jNBins));

        // Populate nominal spectra for sample i
        for (size_t iC = 0; iC < iNComps; ++iC) {
    TH1D* iH = iP->PredictComponent(calc, get<0>(iComps[iC]),
      get<1>(iComps[iC]), get<2>(iComps[iC])).ToTH1(iPOT);
    for (size_t iB = 0; iB < iNBins; ++iB) {
      iNom[iC][iB] = iH->GetBinContent(iB+1);
    }
    delete iH;
        } // for component i

        // Populate nominal spectra for sample j
        for (size_t jC = 0; jC < jNComps; ++jC) {
    TH1D* jH = jP->PredictComponent(calc, get<0>(jComps[jC]),
      get<1>(jComps[jC]), get<2>(jComps[jC])).ToTH1(jPOT);
    for (size_t jB = 0; jB < jNBins; ++jB) {
      jNom[jC][jB] = jH->GetBinContent(jB+1);
    }
    delete jH;
        } // for component i

        // Prepare key and matrix vector inputs
        SamplePair key = make_pair(samples[iS], samples[jS]);
        Patch matrix(iNComps, vector< TMatrixD >(jNComps, TMatrixD(iNBins, jNBins)));

        TRandom3 rnd(0);
        
        for (size_t u = 0; u < nUniv; ++u) {

    p.SetProgress((double(u)+1.)/(double(nUniv)+1.));

    // Throw the systematic
    SystShifts shift;
    double val = rnd.Gaus();
    const KeySyst* tmp = dynamic_cast<const KeySyst*>(syst);
    if (tmp && tmp->IsOneSided()) {
      cout << "Key syst " << tmp->ShortName() << " is one-sided! Throwing positive shift." << endl;
      val = fabs(val);
    }
    shift.SetShift(syst, val);

    // Get systematic shifts
    vector< vector<double> > iShift(iNComps, vector<double>(iNBins));
    vector< vector<double> > jShift(jNComps, vector<double>(jNBins));

    // Populate shifted spectra for sample i
    for (size_t iC = 0; iC < iNComps; ++iC) {
      TH1D* iH = iP->PredictComponentSyst(calc, samples[iS].GetSystShifts(shift),
        get<0>(iComps[iC]), get<1>(iComps[iC]), get<2>(iComps[iC])).ToTH1(iPOT);
      for (size_t iB = 0; iB < iNBins; ++iB) {
        iShift[iC][iB] = iH->GetBinContent(iB+1);
      }
      delete iH;
    } // for component i

    // Populate shifted spectra for sample j
    for (size_t jC = 0; jC < jNComps; ++jC) {
      TH1D* jH = jP->PredictComponentSyst(calc, samples[jS].GetSystShifts(shift), 
        get<0>(jComps[jC]), get<1>(jComps[jC]), get<2>(jComps[jC])).ToTH1(jPOT);
      for (size_t jB = 0; jB < jNBins; ++jB) {
        jShift[jC][jB] = jH->GetBinContent(jB+1);
      }
      delete jH;
    } // for component j

    // Get fractional shifts
    vector< vector<double> > iFracShift(iNComps, vector<double>(iNBins));
    vector< vector<double> > jFracShift(jNComps, vector<double>(jNBins));

    // Populate fractional shift for sample i
    for (size_t iC = 0; iC < iNComps; ++iC) {
      for (size_t iB = 0; iB < iNBins; ++iB) {
        if (iNom[iC][iB] > 0) {
    iFracShift[iC][iB] = ((iShift[iC][iB]+0.01)/(iNom[iC][iB]+0.01)) - 1;
        }
        else {
    iFracShift[iC][iB] = 0;
        }
      } // for bin i
    } // for component i

    // Populate fractional shift for sample j
    for (size_t jC = 0; jC < jNComps; ++jC) {
      for (size_t jB = 0; jB < jNBins; ++jB) {
        if (jNom[jC][jB] > 0) {
    jFracShift[jC][jB] = ((jShift[jC][jB]+0.01)/(jNom[jC][jB]+0.01)) - 1;
        }
        else {
    jFracShift[jC][jB] = 0;
        }
      } // for bin j
    } // for component j

    // Now populate matrix
    for (size_t iC = 0; iC < iNComps; ++iC) {
      for (size_t jC = 0; jC < jNComps; ++jC) {
        for (size_t iB = 0; iB < iNBins; ++iB) {
    for(size_t jB = 0; jB < jNBins; ++jB) {
      matrix[iC][jC](iB, jB) += (iFracShift[iC][iB] * jFracShift[jC][jB]) * (1.0 / (float)nUniv);
    } // for bin j
        } // for bin i
      } // for component j
    } // for component i
        } // for universe u

        p.Done();
        fCovMx.insert(make_pair(key, matrix));

        // Save the matrices and remove them from memory
        SaveSyst(syst->ShortName());
        fCovMx.clear();

        return; // Done processing this patch

      } // for sample j
    } // for sample i

  } // function CovMxManager::AddSystematic

  //----------------------------------------------------------------------
  /// Constructor for LoadFrom implementation
  CovMxManager::CovMxManager(TDirectory* dir)
    : fTopDir(dir) {}

  //----------------------------------------------------------------------
  /// Destructor
  CovMxManager::~CovMxManager() {}

  //----------------------------------------------------------------------
  unique_ptr<CovarianceMatrix> CovMxManager::GetCovarianceMatrix(
    vector<Sample> samples, vector<const ISyst*> systs)
  {
    // Quick sanity check
    if (samples.empty()) throw runtime_error("Sample vector is empty!");
    if (systs.empty()) throw runtime_error("Systematics vector is empty!");

    // Temporary matrix just to figure out binning
    CovarianceMatrix mxTmp(samples);
    size_t nBinsFull = mxTmp.GetFullBinning().NBins();
    MatrixXd mxFull(nBinsFull, nBinsFull);
    mxFull.setZero();

    cout << endl;
    cout << "Making covariance matrix for systematics:" << endl;
    for (const ISyst* syst : systs) {
      cout << "  " << syst->ShortName() << endl;
    }
    cout << endl;

    Progress p("Producing covariance matrix");

    // Loop over each systematic
    int counter = 0;
    int max = systs.size() * samples.size() * samples.size();
    for (size_t iSyst = 0; iSyst < systs.size(); ++iSyst) {

      LoadSyst(systs[iSyst]->ShortName());

      // Loop over samples in i
      for (size_t iS = 0; iS < samples.size(); ++iS) {
        size_t iNBins = samples[iS].GetBinning().NBins();
        auto iComps = covmx::GetComponents(samples[iS]);
        size_t iNComps = iComps.size();
        // Figure out offset in i
        size_t iOffset(0);
        for (size_t i = 0; i < iS; ++i) {
    // Offset the full matrix index by the total number of matrix elements in preceding samples
    iOffset += samples[i].GetBinning().NBins() * covmx::GetComponents(samples[i]).size();
        }
        // Loop over samples in j
        for (size_t jS = 0; jS < samples.size(); ++jS) {
    p.SetProgress((double)counter/(double)max);
    ++counter;
    size_t jNBins = samples[jS].GetBinning().NBins();
    auto jComps = covmx::GetComponents(samples[jS]);
    size_t jNComps = jComps.size();
    // Figure out offset in j
    size_t jOffset(0);
    for (size_t i = 0; i < jS; ++i) {
      // Offset the full matrix index by the total number of matrix elements in preceding samples
      jOffset += samples[i].GetBinning().NBins() * covmx::GetComponents(samples[i]).size();
    }
    SamplePair key = make_pair(samples[iS], samples[jS]);
    Patch mx = fCovMx[key];
    if (mx.size() == 0) {
      ostringstream err;
      err << "Manager not set up for samples " << samples[iS].GetName()
          << " and " << samples[jS].GetName() << endl;
      throw runtime_error(err.str());
    }
    // Loop over individual components and bins within samples
    for (size_t iC = 0; iC < iNComps; ++iC) {
      for (size_t jC = 0; jC < jNComps; ++jC) {
        for (size_t iB = 0; iB < iNBins; ++iB) {
    size_t iBin = iOffset + (iC * iNBins) + iB;
    for (size_t jB = 0; jB < jNBins; ++jB) {
      size_t jBin = jOffset + (jC * jNBins) + jB;
      mxFull(iBin,jBin) += mx[iC][jC](iB,jB);
    } // for bin j
        } // for bin i
      } // for component j
    } // for component i
        } // for sample j
      } // for sample i
    } // for syst

    p.Done();

    return make_unique<CovarianceMatrix>(mxFull, samples);

  } // function CovMxManager::GetCovarianceMatrix

  //----------------------------------------------------------------------
  /// Standard  implementation
  unique_ptr<CovMxManager> CovMxManager::LoadFrom(
    TDirectory* dir)
  {
    unique_ptr<CovMxManager> ret(new CovMxManager(dir));

    return move(ret);

  } // function CovMxManager::LoadFrom

  //----------------------------------------------------------------------
  /// Function to save matrix for a given systematic
  void CovMxManager::SaveSyst(string mxName)
  {
    // Create directory for this systematic
    TDirectory* systsDir = fTopDir->GetDirectory("systs");
    if (!systsDir) systsDir = fTopDir->mkdir("systs");
    TDirectory* dir = systsDir->GetDirectory(mxName.c_str());
    if (!dir) dir = systsDir->mkdir(mxName.c_str());

    for (auto mx : fCovMx) {

      // Write a key for this pair of samples
      ostringstream samples;
      samples << mx.first.first.GetID() << "_" << mx.first.second.GetID();

      // Create a directory
      TDirectory* patchDir = dir->GetDirectory(samples.str().c_str());
      if (!patchDir) patchDir = dir->mkdir(samples.str().c_str());

      // Save number of components & bins for each sample pair
      int c1 = covmx::GetComponents(mx.first.first).size();
      int b1 = mx.first.first.GetBinning().NBins();
      int c2 = covmx::GetComponents(mx.first.second).size();
      int b2 = mx.first.second.GetBinning().NBins();
      vector<int> bins = { c1, b1, c2, b2 };
      patchDir->WriteObject(&bins, "bins");
      // Save matrix patch for each component
      for (int iC = 0; iC < c1; ++iC) {
        for (int jC = 0; jC < c2; ++jC) {
    // Save each matrix with the corresponding key
    ostringstream key;
    key << iC << "_" << jC;
    patchDir->WriteTObject(&mx.second[iC][jC], key.str().c_str());
        } // for component j
      } // for component i
    } // for sample pair mx

  } // function CovMxManager::SaveSyst

  //----------------------------------------------------------------------
  /// Function to load matrix for a given systematic
  void CovMxManager::LoadSyst(std::string mxName)
  {
    DontAddDirectory guard;
    fCovMx.clear();

    // We're loading the covariance matrix for a given systematic.
    // This means looping over all the samples in this systematic's
    // directory, and loading the corresponding matrix patch.
    if (!fTopDir) {
      throw runtime_error("CovMxManager directory not valid!");
    }
    TDirectory* dir = fTopDir->GetDirectory("systs")->GetDirectory(mxName.c_str());
    if (!dir) {
      ostringstream err;
      err << "CovMxManager directory for systematic " << mxName
          << " not valid!";
      throw runtime_error(err.str());
    }
    TIter next(dir->GetListOfKeys());
    TKey* key;
    while ((key = (TKey*)next())) {

      TDirectory* patchDir = (TDirectory*)key->ReadObj();
      string name(key->GetName());

      // Parse key string to get sample information
      size_t start = 0;
      size_t end = name.find("_", start);
      int id = stoi(name.substr(start, end));
      Sample s1(id);

      start = end + 1;
      id = stoi(name.substr(start));
      Sample s2(id);

      // Now we make a sample pair, and create the patch of covariance matrix
      // corresponding to this pair of samples
      SamplePair pair = make_pair(s1, s2);
      vector<int>* bins;
      patchDir->GetObject("bins", bins);
      int c1 = bins->at(0);
      int b1 = bins->at(1);
      int c2 = bins->at(2);
      int b2 = bins->at(3);

      Patch matrix(c1, vector< TMatrixD >(c2, TMatrixD(b1, b2)));

      // Now we loop over and load all the matrices for the various beam
      // components which make up this patch of covariance matrix
      for (int iC = 0; iC < c1; ++iC) {
        for (int jC = 0; jC < c2; ++jC) {
    ostringstream mxname;
    mxname << iC << "_" << jC;
    TMatrixD* mxtemp = (TMatrixD*)patchDir->Get(mxname.str().c_str());
    if (!mxtemp) throw runtime_error("Failed to load matrix with name "+mxname.str());
    matrix[iC][jC] = *mxtemp;
    delete mxtemp;
        } // for component j
      } // for component i

      // Add to map
      fCovMx[pair] = matrix;

      // If we're on an off-diagonal patch, then there's a corresponding
      // mirrored patch on the other side of the diagonal. It's exactly
      // equivalent to this one, so rather than storing redundant information
      // we just mirror this patch into that one.
      if (s1 != s2) {
        SamplePair pairTrans = make_pair(s2, s1);
        Patch matrixTrans(c2, vector< TMatrixD >(c1, TMatrixD(b2, b1)));
        for (int iC = 0; iC < c1; ++iC) {
    for (int jC = 0; jC < c2; ++jC) {
      matrixTrans[jC][iC].Transpose(matrix[iC][jC]);
    } // for component j
        } // for component i
        fCovMx[pairTrans] = matrixTrans;
      } // if transposing

      delete patchDir;

    } // for sample pair key in directory

    delete dir;

  } // function CovMxManager::LoadSyst

} // namespace ana