plot_3flavor_withsysts.C

Go to the documentation of this file.

#include "TCanvas.h"
#include "TColor.h"
#include "TGraph.h"
#include "TH1D.h"
#include "TH2D.h"
#include "THashList.h"
#include "TLegend.h"
#include "TLine.h"
#include "TMarker.h"
#include "TMath.h"
#include "TProfile.h"

#include "CAFAna/Analysis/Plots.h"
#include "CAFAna/Fit/Bayesian1DMarginal.h"
#include "CAFAna/Fit/BayesianSurface.h"
#include "CAFAna/Fit/MCMCSamples.h"

#include "OscLib/OscCalcAnalytic.h"

#include "Utilities/func/MathUtil.h"

#include "3FlavorAna/MCMC/VarsWithPriors.h"
#include "3FlavorAna/MCMC/MCMC3FShared.h"

#include "CAFAna/3flavor/Ana2020/joint_fit_2020_loader_tools.h"

namespace mcmc
{
  const std::map<const ana::IFitVar*, std::pair<double, double>> fitVarDrawRanges
  {
    {&ana::fitSsqTh23_UniformPriorSsqTh23,       {.3, .7}},
    {&ana::fitDmSq32Scaled_UniformPrior,         {2.2, 2.8}},
    {&ana::fitDeltaInPiUnits_UniformPriordCP,    {0,   2}},
  };

  int GetRedHeatPalette()
  {
    const int NRGBs = 9;
    const int n_color_contours = 999;
    static bool initialized=false;
    static int* colors=new int[n_color_contours];
    static int colmin = 0;

    if(!initialized){
      // White -> red
      Double_t stops[NRGBs] = { 0.00, 0.125, 0.250, 0.375, 0.500, 0.625, 0.750, 0.875, 1.000};
      Double_t red[NRGBs]   = { 1.00, 1.00, 0.99, 0.99, 0.98, 0.94, 0.80, 0.65, 0.40 };
      Double_t green[NRGBs] = { 0.96, 0.88, 0.73, 0.57, 0.42, 0.23, 0.09, 0.06, 0.00 };
      Double_t blue[NRGBs]  = { 0.94, 0.82, 0.63, 0.45, 0.29, 0.17, 0.11, 0.08, 0.05 };
      colmin=TColor::CreateGradientColorTable(NRGBs, stops, red, green, blue, n_color_contours);
      for(uint i=0; i<n_color_contours; ++i) colors[i]=colmin+i;

      initialized=true;
    }

    return colmin;
  }

}

void plot_3flavor_withsysts(const std::string & samplesFilename,
                            const std::string & fakeDataFilename,
                            const std::string& outDirPrefix,
                            bool drawFakeDataPredComparison=true,
                            bool drawSurfaces=true,
                            bool draw1Dmarginals=true,
                            bool draw2Dmarginals=true,
                            bool drawPulls=true)
{
  const auto TFText = [](bool yes) { return yes ? "YES" : "NO"; };
  std::cout << "Drawing best fit - fake data comparison: " << TFText(drawFakeDataPredComparison) << std::endl;
  std::cout << "Drawing osc param surfaces: " << TFText(drawSurfaces) << std::endl;
  std::cout << "Drawing 1D marginals: " << TFText(draw1Dmarginals) << std::endl;
  std::cout << "Drawing 2D marginals: " << TFText(draw2Dmarginals) << std::endl;
  std::cout << "Drawing pulls: " << TFText(drawPulls) << std::endl;

  // load the systs so that they're in memory when we try to reload the systTruePulls below
  ana::GetJointFitSystematicList(true);

  mcmc::FileContents fileContents = mcmc::LoadFromFile(samplesFilename);
  std::unique_ptr<ana::MCMCSamples> warmup = std::move(fileContents.warmup);
  std::unique_ptr<ana::MCMCSamples> samples = std::move(fileContents.samples);

  std::map<std::string, ana::Spectrum> fakeDataSpectra = mcmc::LoadFakeData(fakeDataFilename, "spectra");

  fileContents = mcmc::LoadFromFile(fakeDataFilename);
  std::unique_ptr<ana::SystShifts> systTruePulls = std::move(fileContents.trueSystPulls);
  std::unique_ptr<osc::IOscCalcAdjustable> calcTruth = std::move(fileContents.trueOscParams);

  std::cout << "samples: " << samples.get() << std::endl;
  std::cout << "syst truth: " << systTruePulls.get() << std::endl;
  std::cout << "calc truth: " << calcTruth.get() << std::endl;
  assert (samples && systTruePulls && calcTruth);

  // todo: what to do here?  really want the maximum of the joint marginal distribution over all the variables,
  //                         but the single point with the best LL isn't a foolproof estimator of that.
  auto bestFitIdx = samples->BestFitSampleIdx();
  auto calc = std::make_unique<osc::OscCalcAnalytic>();
  mcmc::ResetCalculator(*calc);
  for (const auto & var : samples->Vars())
    var->SetValue(calc.get(), samples->SampleValue(var, bestFitIdx));
  auto shifts = std::make_unique<ana::SystShifts>();
  for (const auto & syst : samples->Systs())
    shifts->SetShift(syst, samples->SampleValue(syst, bestFitIdx));

  //auto bfDir = dynamic_cast<TDirectory*>(inf.Get("bestfit"));
  //calc.reset(dynamic_cast<osc::OscCalcPMNSOpt*>(LoadFrom<osc::IOscCalc>(dynamic_cast<TDirectory*>(bfDir->Get("osc"))).release()));
  //shifts = LoadFrom<SystShifts>(dynamic_cast<TDirectory*>(bfDir->Get("systs")));
  //assert(shifts);

  if (!fakeDataSpectra.empty() && drawFakeDataPredComparison)
  {
    // draw a comparison of the fake data and the best fit prediction
    auto preds = mcmc::LoadPreds();
    TCanvas c;
    for (const auto & predBundle : preds)
    {
      c.Clear();

      // calc we passed to the surface now contains best fit params
      ana::DataMCComparison(fakeDataSpectra.at(predBundle.name),
                            predBundle.pred,  // missing the cosmics for the moment...
                            calc.get(),
                            *shifts,
                            kBinDensity);
      c.SaveAs(mcmc::FullFilename(outDirPrefix, "bestfitpred." + predBundle.name + ".systs.png").c_str());
    }
  }

  mcmc::GetRedHeatPalette();
  TCanvas c;

  if (drawSurfaces)
  {
    // draw the contours in oscillation parameter space overlaid on the LL surface.
    // first, atmospheric
    ana::BayesianSurface surf = samples->MarginalizeTo(&ana::fitSsqTh23_UniformPriorSsqTh23, 100, .3, .7,
                                                       &ana::fitDmSq32Scaled_UniformPrior, 100, 2.2, 2.8);
    surf.Draw();
    surf.DrawContour(surf.QuantileSurface(ana::Quantile::kGaussian1Sigma), 7, kGreen + 2); // dashed
    surf.DrawContour(surf.QuantileSurface(ana::Quantile::kGaussian2Sigma), 7, kBlack); // dashed
    surf.DrawBestFit(kGray);
    TMarker marker(util::sqr(sin(calcTruth->GetTh23())), calcTruth->GetDmsq32() * 1e3, kFullStar);
    marker.SetMarkerColor(kBlue);
    marker.SetMarkerSize(3);
    marker.Draw();
    c.SaveAs(mcmc::FullFilename(outDirPrefix, "surface_contour_ssth23-dm32.png").c_str());

    // now the 'reactor' contour
    c.Clear();
    surf = samples->MarginalizeTo(&ana::fitDeltaInPiUnits_UniformPriordCP, 75, 0, 2,
                                  &ana::fitSsqTh23_UniformPriorSsqTh23, 75, 0.3, 0.7);
    surf.Draw();
    surf.DrawContour(surf.QuantileSurface(ana::Quantile::kGaussian1Sigma), 7, kGreen + 2); // dashed
    surf.DrawContour(surf.QuantileSurface(ana::Quantile::kGaussian2Sigma), 7, kBlack); // dashed
    surf.DrawBestFit(kGray);
    marker = TMarker(calcTruth->GetdCP() / TMath::Pi(), util::sqr(sin(calcTruth->GetTh23())), kFullStar);
    marker.SetMarkerColor(kBlue);
    marker.SetMarkerSize(3);
    marker.Draw();
    c.SaveAs(mcmc::FullFilename(outDirPrefix, "surface_contour_ssth23-dcp.png").c_str());
  }

  if (draw1Dmarginals)
  {
    // draw the marginals and the pulls
    for (const auto &fitVarPtr : samples->Vars())
    {
      c.Clear();
      ana::Bayesian1DMarginal marg = samples->MarginalizeTo(fitVarPtr);
      auto h = marg.ToTH1(ana::Binning::Simple(50, samples->MinValue(fitVarPtr), samples->MaxValue(fitVarPtr)));
      h.Draw("hist");

      double truthVal = fitVarPtr->GetValue(calcTruth.get());
      TLine l(truthVal, h.GetMinimum(), truthVal, h.GetMaximum());
      l.SetLineColor(kRed);
      l.SetLineWidth(2);
      l.SetLineStyle(kDashed);
      l.Draw();

      c.SaveAs(Form(mcmc::FullFilename(outDirPrefix, "marg_%s.systs.png").c_str(), fitVarPtr->ShortName().c_str()));
    }
    for (const auto &systPtr : shifts->ActiveSysts())
    {
      c.Clear();
      ana::Bayesian1DMarginal marg = samples->MarginalizeTo(systPtr);
      auto h = marg.ToTH1(ana::Binning::Simple(50, samples->MinValue(systPtr), samples->MaxValue(systPtr)));
      h.Draw("hist");

      double truthVal = systTruePulls->GetShift(systPtr);
      TLine l(truthVal, h.GetMinimum(), truthVal, h.GetMaximum());
      l.SetLineColor(kRed);
      l.SetLineWidth(2);
      l.SetLineStyle(kDashed);
      l.Draw();

      c.SaveAs(Form(mcmc::FullFilename(outDirPrefix, "marg_%s.systs.png").c_str(), systPtr->ShortName().c_str()));
    }
  }

  if (draw2Dmarginals)
  {
    // pairs of fit vars
    for (const auto fitVar1 : samples->Vars())
    {
      for (auto itVar2 = samples->Vars().rbegin(); itVar2 != samples->Vars().rend(); itVar2++)
      {
        const ana::IFitVar *fitVar2 = *itVar2;
        if (fitVar2 == fitVar1)
          break;

        c.Clear();
        std::pair<double, double> drawRange1 = mcmc::fitVarDrawRanges.at(fitVar1);
        std::pair<double, double> drawRange2 = mcmc::fitVarDrawRanges.at(fitVar2);
        ana::BayesianSurface surf = samples->MarginalizeTo(fitVar1, 100, drawRange1.first, drawRange1.second,
                                                           fitVar2, 100, drawRange2.first, drawRange2.second);
        surf.Draw();
        surf.DrawContour(surf.QuantileSurface(ana::Quantile::kGaussian1Sigma), 7, kGreen + 2); // dashed
        surf.DrawBestFit(kGray);
        TMarker marker(fitVar1->GetValue(calcTruth.get()), fitVar2->GetValue(calcTruth.get()), kFullStar);
        marker.SetMarkerColor(kMagenta);
        marker.SetMarkerSize(3);
        marker.Draw();
        c.SaveAs(Form(mcmc::FullFilename(outDirPrefix, "surface_contour_%s-%s.systs.png").c_str(),
                      fitVar1->ShortName().c_str(),
                      fitVar2->ShortName().c_str()));
      }
    }

    // 2D syst space
    auto systs = shifts->ActiveSysts();
    std::set<std::string> systNames;
    std::transform(systs.begin(), systs.end(),
                   std::inserter(systNames, systNames.end()), [](const ana::ISyst *s)
                   { return s->ShortName(); });
    assert(systNames.size() == systs.size());
    for (const auto &systName1 : systNames)
    {
      auto systPtr1 = *std::find_if(systs.begin(), systs.end(),
                                    [&systName1](const ana::ISyst *s)
                                    { return s->ShortName() == systName1; });
      for (auto itSystName2 = systNames.rbegin(); itSystName2 != systNames.rend(); itSystName2++)
      {
        if (*itSystName2 == systName1)
          break;

        const ana::ISyst *systPtr2 = *std::find_if(systs.begin(), systs.end(),
                                                   [&itSystName2](const ana::ISyst *s)
                                                   { return s->ShortName() == *itSystName2; });

        c.Clear();
        ana::BayesianSurface marg = samples->MarginalizeTo(systPtr1, 30, samples->MinValue(systPtr1),
                                                           samples->MaxValue(systPtr1),
                                                           systPtr2, 30, samples->MinValue(systPtr2),
                                                           samples->MaxValue(systPtr2));
        auto h2 = marg.ToTH2();
        h2->Draw("colz");
        marg.DrawContour(marg.QuantileSurface(ana::Quantile::kGaussian1Sigma), 7, kGreen + 2); // dashed

        TMarker marker(systTruePulls->GetShift(systPtr1), systTruePulls->GetShift(systPtr2), kFullStar);
        marker.SetMarkerColor(kBlack);
        marker.SetMarkerSize(3);
        marker.Draw();

        std::string outf = Form(mcmc::FullFilename(outDirPrefix, "marg_%s-%s.allSysts.png").c_str(),
                                systPtr1->ShortName().c_str(),
                                systPtr2->ShortName().c_str());
        c.SaveAs(outf.c_str());
      }
    }
  }

  if (drawPulls && !samples->Systs().empty())
  {  // show that the pulls were what we expect
    const ana::Binning margBins = ana::Binning::Simple(100, -2.95,
                                                       2.95);  // then labels don't collide in the ratio plot
    TH2D h_fittedPulls("fitted_pulls", ";Systematic;Pull (#sigma)",
                       shifts->ActiveSysts().size(), 0, shifts->ActiveSysts().size(),
                       margBins.NBins(), margBins.Min(), margBins.Max());
    TH1D h_truePulls("true_pulls", ";Systematic;Pull (#sigma)", shifts->ActiveSysts().size(), 0,
                     shifts->ActiveSysts().size());
    std::size_t systIdx = 0;
    double maxShift = 0;
    auto systsSorted = shifts->ActiveSysts();
    std::sort(systsSorted.begin(),
              systsSorted.end(),
              [](const auto syst1, const auto syst2)
              {
                return syst1->ShortName() < syst2->ShortName();
              });
    for (const auto &syst : systsSorted)
    {
      ++systIdx;
      auto systMarginal = samples->MarginalizeTo(syst);
      TH1D h_systMarginal = systMarginal.ToTH1(margBins);
      for (int binIdx = 0; binIdx <= h_systMarginal.GetNbinsX(); binIdx++)
        h_fittedPulls.SetBinContent(systIdx, binIdx, h_systMarginal.GetBinContent(binIdx));
      maxShift = std::max(maxShift, std::abs(shifts->GetShift(syst)));
      h_truePulls.SetBinContent(systIdx, systTruePulls->GetShift(syst));
      maxShift = std::max(maxShift, std::abs(systTruePulls->GetShift(syst)));
      for (auto &h : {dynamic_cast<TH1 *>(&h_fittedPulls), dynamic_cast<TH1 *>(&h_truePulls)})
        h->GetXaxis()->SetBinLabel(systIdx, syst->ShortName().c_str());
    }
    maxShift = std::max(3., 1.2 * maxShift);
    h_truePulls.GetYaxis()->SetRangeUser(-maxShift, maxShift);
    h_fittedPulls.GetYaxis()->SetRangeUser(-maxShift, maxShift);

    std::unique_ptr<TProfile> h_fittedPulls_prof(h_fittedPulls.ProfileX(ana::UniqueName().c_str(), 1, -1, "s"));
    h_fittedPulls_prof->SetLineColor(kBlue);

    c.Clear();
    c.SetCanvasSize(900, 600);
    TAxis ax(*h_fittedPulls.GetXaxis());
    h_fittedPulls.GetXaxis()->GetLabels()->Clear();

    TPad upper("", "", 0, 0, 1.0, 1);
    upper.cd();
    h_fittedPulls.SetMaximum(0.25);  // so that more of the color winds up red
    mcmc::GetRedHeatPalette();
    h_fittedPulls.SetFillColor(kRed + 1);  // for the legend
    h_fittedPulls.GetZaxis()->SetTitle("Fraction of MCMC samples");
    h_fittedPulls.GetZaxis()->RotateTitle();
    h_fittedPulls.Draw("colz");
    h_truePulls.Draw("hist x same");
    h_fittedPulls_prof->Draw("pe same");
    upper.SetLeftMargin(0.15);
    upper.SetRightMargin(0.15);
    upper.SetBottomMargin(0.4);

    TPad lower("", "", 0, 0, 1.0, 1);
    lower.SetFillStyle(0);
    lower.cd();

    // a "pull" plot
    TGraph fitPulls;
    for (int bin = 1; bin <= h_fittedPulls_prof->GetNbinsX(); bin++)
    {
      double pull = 0;
      if (h_fittedPulls_prof->GetBinError(bin) != 0)
        pull = (h_fittedPulls_prof->GetBinContent(bin) - h_truePulls.GetBinContent(bin))
               / h_fittedPulls_prof->GetBinError(bin);
      fitPulls.SetPoint(fitPulls.GetN(), double(bin) - 0.5, pull);
    }
    fitPulls.Draw("al");
    *fitPulls.GetXaxis() = ax;
    fitPulls.SetTitle(";;#frac{True - Fit mean}{Fit RMS}");
    fitPulls.GetYaxis()->SetRangeUser(-0.95, 0.95);
    lower.SetLeftMargin(0.15);
    lower.SetRightMargin(0.15);
    lower.SetTopMargin(0.6);
    lower.SetBottomMargin(0.3);

    c.cd();
    upper.Draw();
    lower.Draw();
    TLegend leg(0.2, 0.9, 0.8, 1.0);
    leg.SetFillStyle(0);
    leg.SetBorderSize(0);
    leg.SetNColumns(3);
    leg.AddEntry(&h_truePulls, "True", "l");
    leg.AddEntry(&h_fittedPulls, "MCMC samples", "f");
    leg.AddEntry(h_fittedPulls_prof.get(), "MCMC (mean+RMS)", "lpe");
    leg.SetColumnSeparation(0.15);
    leg.Draw();

//  c.SetBottomMargin(0.5);
    c.SaveAs(mcmc::FullFilename(outDirPrefix, "pulls.systs.png").c_str());
  }
}