OscProcessor.cpp

Go to the documentation of this file.

#include "OscProcessor.h"
 
_MaCh3_Safe_Include_Start_ //{
// ROOT includes
#include "TArrow.h"
#include "TEllipse.h"
#include "TLatex.h"
#include "TVector2.h"
#include "TComplex.h"
_MaCh3_Safe_Include_End_ //}
 
#pragma GCC diagnostic ignored "-Wfloat-conversion"
 
// ****************************
OscProcessor::OscProcessor(const std::string &InputFile) : MCMCProcessor(InputFile) {
// ****************************
  //KS: WARNING this only work when you project from Chain, will nor work when you try SetBranchAddress etc. Turn it on only if you know how to use it
  PlotJarlskog = false;
 
  Sin2Theta13Index = M3::_BAD_INT_;
  Sin2Theta12Index = M3::_BAD_INT_;
  Sin2Theta23Index = M3::_BAD_INT_;
  DeltaCPIndex = M3::_BAD_INT_;
  DeltaM2_23Index = M3::_BAD_INT_;
}
 
// ****************************
// The destructor
OscProcessor::~OscProcessor() {
// ****************************
}
 
// ***************
// Read the Osc cov file and get the input central values and errors
void OscProcessor::LoadAdditionalInfo() {
// ***************
  // KS: Check if OscParams were enabled, in future we will also get
  for(size_t i = 0; i < ParameterGroup.size(); i++) {
    if(ParameterGroup[i] == "Osc"){
      OscEnabled = true;
      break;
    }
  }
 
  if(OscEnabled)
  {
    for (int i = 0; i < nDraw; ++i)
    {
      //Those keep which parameter type we run currently and relative number
      const int ParamEnum = ParamType[i];
      const int ParamNo = i - ParamTypeStartPos[ParameterEnum(ParamEnum)];
      const std::string CurrentName = ParamNames[ParamEnum][ParamNo].Data();
 
      if (CurrentName == "sin2th_13") {
        Sin2Theta13Index = i;
        Sin2Theta13Name  = CurrentName;
      } else if (CurrentName == "sin2th_12") {
        Sin2Theta12Index = i;
        Sin2Theta12Name  = CurrentName;
      } else if (CurrentName == "sin2th_23") {
        Sin2Theta23Index = i;
        Sin2Theta23Name  = CurrentName;
      } else if (CurrentName == "delta_cp") {
        DeltaCPIndex     = i;
        DeltaCPName      = CurrentName;
      } else if (CurrentName == "delm2_23") {
        DeltaM2_23Index  = i;
        DeltaM2_23Name   = CurrentName;
      }
    }
  } else{
    MACH3LOG_WARN("Didn't find oscillation parameters");
  }
 
  if(PlotJarlskog && OscEnabled)
  {
    Chain->SetAlias("J_cp", "TMath::Sqrt(sin2th_13)*TMath::Sqrt(1.-sin2th_13)*TMath::Sqrt(1.-sin2th_13)*TMath::Sqrt(sin2th_12)*TMath::Sqrt(1.-sin2th_12)*TMath::Sqrt(sin2th_23)*TMath::Sqrt(1.-sin2th_23)*TMath::Sin(delta_cp)");
    BranchNames.push_back("J_cp");
    ParamType.push_back(kXSecPar);
    nParam[kXSecPar]++;
    nDraw++;
 
    ParamCentral[kXSecPar].push_back( 0. );
    ParamErrors[kXSecPar].push_back( 1. );
    // Push back the name
    ParamNames[kXSecPar].push_back("J_cp");
    ParamFlat[kXSecPar].push_back( false );
  } else if(PlotJarlskog && !OscEnabled) {
    MACH3LOG_ERROR("Trying to enable Jarlskog without oscillations");
    throw MaCh3Exception(__FILE__,__LINE__);
  }
}
 
// ***************
// Calculate Jarlskog Invariant using oscillation parameters
double OscProcessor::CalcJarlskog(const double s2th13, const double s2th23, const double s2th12, const double dcp) const {
// ***************
  const double s13  = std::sqrt(s2th13);
  const double s23  = std::sqrt(s2th23);
  const double s12  = std::sqrt(s2th12);
  const double sdcp = std::sin(dcp);
  const double c13  = std::sqrt(1.-s2th13);
  const double c12  = std::sqrt(1.-s2th12);
  const double c23  = std::sqrt(1.-s2th23);
 
  const double j = s13*c13*c13*s12*c12*s23*c23*sdcp;
 
  return j;
}
 
// ***************
double OscProcessor::SamplePriorForParam(const int paramIndex, const std::unique_ptr<TRandom3>& randGen, const std::vector<double>& FlatBounds) const {
// ***************
  TString Title = "";
  double Prior = 1.0, PriorError = 1.0;
  bool FlatPrior = false;
 
  // Get info for this parameter
  GetNthParameter(paramIndex, Prior, PriorError, Title);
 
  ParameterEnum ParType = ParamType[paramIndex];
  int ParamTemp = paramIndex - ParamTypeStartPos[ParType];
  FlatPrior = ParamFlat[ParType][ParamTemp];
 
  if (FlatPrior) {
    return randGen->Uniform(FlatBounds[0], FlatBounds[1]);
  } else {
    // Gaussian prior centered at Prior with width PriorError
    return randGen->Gaus(Prior, PriorError);
  }
}
 
// ***************
void OscProcessor::Get1DReactorConstraintInfo(std::pair<double, double>& Sin13_NewPrior, bool& DoReweight) const {
// ***************
  if(ReweightNames.size() != 1){
    MACH3LOG_INFO("Found more than 1 weight for Jarlskog analysis");
    return;
  }
  // Now read the MCMC file
  TFile *TempFile = M3::Open((MCMCFile + ".root"), "open", __FILE__, __LINE__);
 
  // Get the settings for the MCMC
  TMacro *Config = TempFile->Get<TMacro>("Reweight_Config");
 
  DoReweight = false;
  if (Config != nullptr) {
    MACH3LOG_INFO("Found Reweight_Config in chain");
    // Print the reweight configuration for user info
    YAML::Node Settings = TMacroToYAML(*Config);
    // Simple check: only enable DoReweight if it's a 1D sin2th_13 Gaussian reweight since Savage Dickey process later on generates values from the Gaussian
    if(CheckNodeExists(Settings, ReweightNames.back())) {
      YAML::Node firstReweight = Settings[ReweightNames.back()];
      int dimension = Get<int>(firstReweight["ReweightDim"], __FILE__ , __LINE__);
      std::string reweightType = Get<std::string>(firstReweight["ReweightType"],__FILE__ , __LINE__);
      auto paramNames = Get<std::vector<std::string>>(firstReweight["ReweightVar"], __FILE__ , __LINE__);
      if (dimension == 1 && reweightType == "Gaussian" && paramNames.size() == 1){
         auto Priors = Get<std::vector<std::pair<double, double>>>(firstReweight["ReweightPrior"], __FILE__,__LINE__);
         Sin13_NewPrior = Priors[0];
        DoReweight = true;
      } else {
        MACH3LOG_INFO("No valid reweighting configuration (1D Gaussian on sin2th_13 only) found for Jarlskog analysis");
      }
    } else {
      MACH3LOG_INFO("No reweighting configuration found for Jarlskog analysis");
    }
  }
 
  TempFile->Close();
  delete TempFile;
}
 
// ***************
// Perform Several Jarlskog Plotting
void OscProcessor::PerformJarlskogAnalysis() {
// ***************
  if(!OscEnabled ||
    Sin2Theta13Index == M3::_BAD_INT_ ||
    Sin2Theta12Index == M3::_BAD_INT_ ||
    Sin2Theta23Index == M3::_BAD_INT_ ||
    DeltaCPIndex == M3::_BAD_INT_||
    DeltaM2_23Index == M3::_BAD_INT_)
  {
    MACH3LOG_WARN("Will not {}, as oscillation parameters are missing", __func__);
    return;
  }
  MACH3LOG_INFO("Starting {}", __func__);
 
  double s2th13, s2th23, s2th12, dcp, dm2 = M3::_BAD_DOUBLE_;
  double weight = 1.0;
 
  bool DoReweight = false;
  std::pair<double, double> Sin13_NewPrior;
  Get1DReactorConstraintInfo(Sin13_NewPrior, DoReweight);
 
  TDirectory *JarlskogDir = OutputFile->mkdir("Jarlskog");
  JarlskogDir->cd();
 
  unsigned int step = 0;
  Chain->SetBranchStatus("*", false);
 
  Chain->SetBranchStatus(Sin2Theta13Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta13Name.c_str(), &s2th13);
 
  Chain->SetBranchStatus(Sin2Theta23Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta23Name.c_str(), &s2th23);
 
  Chain->SetBranchStatus(Sin2Theta12Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta12Name.c_str(), &s2th12);
 
  Chain->SetBranchStatus(DeltaCPName.c_str(), true);
  Chain->SetBranchAddress(DeltaCPName.c_str(), &dcp);
 
  Chain->SetBranchStatus(DeltaM2_23Name.c_str(), true);
  Chain->SetBranchAddress(DeltaM2_23Name.c_str(), &dm2);
 
  Chain->SetBranchStatus("step", true);
  Chain->SetBranchAddress("step", &step);
 
  if(DoReweight) {
    Chain->SetBranchStatus(ReweightNames.back().c_str(), true);
    Chain->SetBranchAddress(ReweightNames.back().c_str(), &weight);
  } else {
    MACH3LOG_WARN("Not applying reweighting weight");
    weight = 1.0;
  }
 
  // Original histograms
  auto jarl = std::make_unique<TH1D>("jarl", "jarl", 1000, -0.05, 0.05);
  jarl->SetDirectory(nullptr);
  auto jarl_th23 = std::make_unique<TH2D>("jarl_th23", "jarl_th23", 500, -0.05, 0.05, 500, 0.3, 0.7);
  jarl_th23->SetDirectory(nullptr);
  auto jarl_dcp = std::make_unique<TH2D>("jarl_dcp", "jarl_dcp", 500, -0.05, 0.05, 500, -1. * TMath::Pi(), TMath::Pi());
  jarl_dcp->SetDirectory(nullptr);
 
  jarl->SetTitle("Jarlskog Invariant;J #equiv s_{13}c_{13}^{2}s_{12}c_{12}s_{23}c_{23}sin#delta;Posterior probability");
  jarl_th23->SetTitle("Jarlskog Invariant;J #equiv s_{13}c_{13}^{2}s_{12}c_{12}s_{23}c_{23}sin#delta;Posterior probability");
 
  // Clones
  auto jarl_IH  = M3::Clone(jarl.get(), "jarl_IH");
  auto jarl_NH  = M3::Clone(jarl.get(), "jarl_NH");
 
  auto jarl_th23_IH  = M3::Clone(jarl_th23.get(), "jarl_th23_IH");
  auto jarl_th23_NH  = M3::Clone(jarl_th23.get(), "jarl_th23_NH");
 
  auto jarl_dcp_IH  = M3::Clone(jarl_dcp.get(), "jarl_dcp_IH");
  auto jarl_dcp_NH  = M3::Clone(jarl_dcp.get(), "jarl_dcp_NH");
 
  auto jarl_flatsindcp          = M3::Clone(jarl.get(), "jarl_flatsindcp");
  auto jarl_IH_flatsindcp       = M3::Clone(jarl.get(), "jarl_IH_flatsindcp");
  auto jarl_NH_flatsindcp       = M3::Clone(jarl.get(), "jarl_NH_flatsindcp");
 
  auto jarl_th23_flatsindcp     = M3::Clone(jarl_th23.get(), "jarl_th23_flatsindcp");
  auto jarl_th23_IH_flatsindcp  = M3::Clone(jarl_th23.get(), "jarl_th23_IH_flatsindcp");
  auto jarl_th23_NH_flatsindcp  = M3::Clone(jarl_th23.get(), "jarl_th23_NH_flatsindcp");
 
  auto jarl_prior               = M3::Clone(jarl.get(), "jarl_prior");
  auto jarl_prior_flatsindcp    = M3::Clone(jarl.get(), "jarl_prior_flatsindcp");
  std::unique_ptr<TH1D> jarl_wRC_prior, jarl_wRC_prior_flatsindcp, jarl_wRC_prior_t2kth23;
  // Only use this if chain has reweigh weight [mostly coming from Reactor Constrains]
  if(DoReweight){
    jarl_wRC_prior            = M3::Clone(jarl.get(), "jarl_wRC_prior");
    jarl_wRC_prior_flatsindcp = M3::Clone(jarl.get(), "jarl_wRC_prior_flatsindcp");
    jarl_wRC_prior_t2kth23    = M3::Clone(jarl.get(), "jarl_wRC_prior_flatsindcp");
  }
 
  // to apply a prior that is flat in sin(dcp) intead of dcp
  auto prior3 = std::make_unique<TF1>("prior3", "TMath::Abs(TMath::Cos(x))");
 
  // T2K prior is flat (and uncorrelated) in dcp, sin^2(th13), sin^2(th23)
  auto randGen = std::make_unique<TRandom3>(0);
  const Long64_t countwidth = nEntries/5;
 
  for(int i = 0; i < nEntries; ++i) {
    if (i % countwidth == 0) {
      M3::Utils::PrintProgressBar(i, nEntries);
      M3::Utils::EstimateDataTransferRate(Chain, i);
    } else {
      Chain->GetEntry(i);
    }
 
    if(step < BurnInCut) continue; // burn-in cut
 
    const double j = CalcJarlskog(s2th13, s2th23, s2th12, dcp);
    const double prior_weight = prior3->Eval(dcp);
 
    jarl->Fill(j, weight);
    jarl_th23->Fill(j, s2th23, weight);
    jarl_dcp->Fill(j, dcp, weight);
 
    jarl_flatsindcp->Fill(j, prior_weight*weight);
    jarl_th23_flatsindcp->Fill(j, s2th23, prior_weight*weight);
 
    const double prior_s2th13 = SamplePriorForParam(Sin2Theta13Index, randGen, {0.,1.});
    const double prior_s2th23 = SamplePriorForParam(Sin2Theta23Index, randGen, {0.,1.});
    const double prior_s2th12 = SamplePriorForParam(Sin2Theta12Index, randGen, {0.,1.});
    const double prior_dcp = SamplePriorForParam(DeltaCPIndex, randGen, {-1.*TMath::Pi(),TMath::Pi()});
    // KS: This is hardcoded but we always assume flat in delta CP so probably fine
    const double prior_sindcp = randGen->Uniform(-1., 1.);
 
    const double prior_s13          = std::sqrt(prior_s2th13);
    const double prior_s23          = std::sqrt(prior_s2th23);
    const double prior_s12          = std::sqrt(prior_s2th12);
    const double prior_sdcp         = std::sin(prior_dcp);
    const double prior_c13          = std::sqrt(1.-prior_s2th13);
    const double prior_c12          = std::sqrt(1.-prior_s2th12);
    const double prior_c23          = std::sqrt(1.-prior_s2th23);
    const double prior_j            = prior_s13*prior_c13*prior_c13*prior_s12*prior_c12*prior_s23*prior_c23*prior_sdcp;
    const double prior_flatsindcp_j = prior_s13*prior_c13*prior_c13*prior_s12*prior_c12*prior_s23*prior_c23*prior_sindcp;
 
    jarl_prior->Fill(prior_j);
    jarl_prior_flatsindcp->Fill(prior_flatsindcp_j);
 
    if(DoReweight) {
      const double prior_wRC_s2th13       = randGen->Gaus(Sin13_NewPrior.first, Sin13_NewPrior.second);
      const double prior_wRC_s13          = std::sqrt(prior_wRC_s2th13);
      const double prior_wRC_c13          = std::sqrt(1.-prior_wRC_s2th13);
      const double prior_wRC_j            = prior_wRC_s13*prior_wRC_c13*prior_wRC_c13*prior_s12*prior_c12*prior_s23*prior_c23*prior_sdcp;
      const double prior_wRC_flatsindcp_j = prior_wRC_s13*prior_wRC_c13*prior_wRC_c13*prior_s12*prior_c12*prior_s23*prior_c23*prior_sindcp;
      const double s23                    = std::sqrt(s2th23);
      const double c23                    = std::sqrt(1.-s2th23);
 
      jarl_wRC_prior->Fill(prior_wRC_j);
      jarl_wRC_prior_flatsindcp->Fill(prior_wRC_flatsindcp_j);
      jarl_wRC_prior_t2kth23->Fill(prior_wRC_s13*prior_wRC_c13*prior_wRC_c13*prior_s12*prior_c12*s23*c23*prior_sdcp);
    }
 
    if(dm2 > 0.) {
      jarl_NH->Fill(j, weight);
      jarl_th23_NH->Fill(j, s2th23, weight);
      jarl_dcp_NH->Fill(j, dcp, weight);
      jarl_NH_flatsindcp->Fill(j, prior_weight*weight);
      jarl_th23_NH_flatsindcp->Fill(j, s2th23, prior_weight*weight);
    }
    else if(dm2 < 0.) {
      jarl_IH->Fill(j, weight);
      jarl_th23_IH->Fill(j, s2th23, weight);
      jarl_dcp_IH->Fill(j, dcp, weight);
      jarl_IH_flatsindcp->Fill(j, prior_weight*weight);
      jarl_th23_IH_flatsindcp->Fill(j, s2th23, prior_weight*weight);
    }
  }
 
  jarl->Write("jarlskog_both");
  jarl_NH->Write("jarlskog_NH");
  jarl_IH->Write("jarlskog_IH");
  jarl_th23->Write("jarlskog_th23_both");
  jarl_th23_NH->Write("jarlskog_th23_NH");
  jarl_th23_IH->Write("jarlskog_th23_IH");
 
  jarl_dcp->Write("jarlskog_dcp_both");
  jarl_dcp_NH->Write("jarlskog_dcp_NH");
  jarl_dcp_IH->Write("jarlskog_dcp_IH");
 
 
  jarl_flatsindcp->Write("jarlskog_both_flatsindcp");
  jarl_NH_flatsindcp->Write("jarlskog_NH_flatsindcp");
  jarl_IH_flatsindcp->Write("jarlskog_IH_flatsindcp");
  jarl_th23_flatsindcp->Write("jarlskog_th23_both_flatsindcp");
  jarl_th23_NH_flatsindcp->Write("jarlskog_th23_NH_flatsindcp");
  jarl_th23_IH_flatsindcp->Write("jarlskog_th23_IH_flatsindcp");
 
  jarl_prior->Write("jarl_prior");
  jarl_prior_flatsindcp->Write("jarl_prior_flatsindcp");
  if(DoReweight) {
    jarl_wRC_prior->Write("jarl_wRC_prior");
    jarl_wRC_prior_flatsindcp->Write("jarl_wRC_prior_flatsindcp");
    jarl_wRC_prior_t2kth23->Write("jarl_wRC_prior_t2kth23");
  }
 
  MakeJarlskogPlot(jarl, jarl_flatsindcp,
                   jarl_NH, jarl_NH_flatsindcp,
                   jarl_IH, jarl_IH_flatsindcp);
 
  // Perform Savage Dickey analysis
  if(DoReweight) {
    SavageDickeyPlot(jarl, jarl_wRC_prior, "Jarlskog flat #delta_{CP}", 0);
    SavageDickeyPlot(jarl_flatsindcp, jarl_wRC_prior_flatsindcp, "Jarlskog flat sin#delta_{CP}", 0);
  } else {
    SavageDickeyPlot(jarl, jarl_prior, "Jarlskog flat #delta_{CP}", 0);
    SavageDickeyPlot(jarl_flatsindcp, jarl_prior_flatsindcp, "Jarlskog flat sin#delta_{CP}", 0);
  }
 
  JarlskogDir->Close();
  delete JarlskogDir;
 
  Chain->SetBranchStatus("*", true);
  OutputFile->cd();
}
 
 
// ***************
void OscProcessor::MakeJarlskogPlot(const std::unique_ptr<TH1D>& jarl,
                                    const std::unique_ptr<TH1D>& jarl_flatsindcp,
                                    const std::unique_ptr<TH1D>& jarl_NH,
                                    const std::unique_ptr<TH1D>& jarl_NH_flatsindcp,
                                    const std::unique_ptr<TH1D>& jarl_IH,
                                    const std::unique_ptr<TH1D>& jarl_IH_flatsindcp) {
// ***************
  MACH3LOG_INFO("Starting {}", __func__);
  int originalErrorLevel = gErrorIgnoreLevel;
  gErrorIgnoreLevel = kFatal;
 
  // 1-->NH, 0-->both, -1-->IH
  for(int hierarchy = -1; hierarchy <= 1; hierarchy++)
  {
    std::unique_ptr<TH1D> j_hist;
    std::unique_ptr<TH1D> j_hist_sdcp;
    if(hierarchy == 1) {
      j_hist = M3::Clone(jarl_NH.get(), "");
      j_hist_sdcp = M3::Clone(jarl_NH_flatsindcp.get(), "");
      j_hist->SetTitle(";J_{CP} #equiv s_{13}c^{2}_{13}s_{12}c_{12}s_{23}c_{23}sin#delta_{CP};Posterior probability");
    } else if(hierarchy == 0) {
      j_hist = M3::Clone(jarl.get(), "");
      j_hist_sdcp = M3::Clone(jarl_flatsindcp.get(), "");
      j_hist->SetTitle(";J_{CP} #equiv s_{13}c^{2}_{13}s_{12}c_{12}s_{23}c_{23}sin#delta_{CP};Posterior probability");
    } else if(hierarchy == -1) {
      j_hist = M3::Clone(jarl_IH.get(), "");
      j_hist_sdcp = M3::Clone(jarl_IH_flatsindcp.get(), "");
      j_hist->SetTitle(";J_{CP} #equiv s_{13}c^{2}_{13}s_{12}c_{12}s_{23}c_{23}sin#delta_{CP};Posterior probability");
    } else {
      MACH3LOG_ERROR("Invalid hierarchy option. 1 for NH, 0 for both, -1 for IH");
      throw MaCh3Exception(__FILE__ , __LINE__ );
    }
 
    j_hist->Rebin(7);
    j_hist_sdcp->Rebin(7);
 
    j_hist->SetLineColor(kAzure-2);
    j_hist_sdcp->SetLineColor(kOrange+1);
    j_hist->SetLineWidth(2);
    j_hist_sdcp->SetLineWidth(2);
 
    auto StyleAxis = [](TH1* h) {
      auto xAxis = h->GetXaxis();
      auto yAxis = h->GetYaxis();
 
      xAxis->SetLabelSize(0.04);
      xAxis->SetLabelFont(132);
      xAxis->SetTitleSize(0.04);
      xAxis->SetTitleOffset(0.80);
      xAxis->SetTitleFont(132);
      xAxis->SetNdivisions(505);
      xAxis->SetTickSize(0.04);
 
      yAxis->SetLabelSize(0.04);
      yAxis->SetLabelFont(132);
      yAxis->SetTitleSize(0.04);
      yAxis->SetTitleOffset(1.2);
      yAxis->SetTitleFont(132);
      yAxis->SetNdivisions(505);
      yAxis->SetTickSize(0.04);
    };
 
    StyleAxis(j_hist.get());
 
    j_hist->GetXaxis()->SetRangeUser(-0.04,0.04);
    j_hist->Scale(1./j_hist->Integral());
    j_hist_sdcp->Scale(1./j_hist_sdcp->Integral());
 
    std::unique_ptr<TH1D> j_hist_copy = M3::Clone(j_hist.get(), "j_hist_copy");
    std::unique_ptr<TH1D> j_hist_1sig = M3::Clone(j_hist.get(), "j_hist_1sig");
    std::unique_ptr<TH1D> j_hist_2sig = M3::Clone(j_hist.get(), "j_hist_2sig");
    std::unique_ptr<TH1D> j_hist_3sig = M3::Clone(j_hist.get(), "j_hist_3sig");
 
    //upper and lower edges
    double j_bf = j_hist_copy->GetXaxis()->GetBinCenter(j_hist_copy->GetMaximumBin());
    double j_1sig_low = 9999999.;
    double j_1sig_up  = -9999999.;
    double j_2sig_low = 9999999.;;
    double j_2sig_up  = -9999999.;
    double j_3sig_low = 9999999.;;
    double j_3sig_up  = -9999999.;
 
 
    std::unique_ptr<TH1D> j_hist_sdcp_copy = M3::Clone(j_hist_sdcp.get(), "j_hist_sdcp_copy");
    std::unique_ptr<TH1D> j_hist_sdcp_1sig = M3::Clone(j_hist_sdcp.get(), "j_hist_sdcp_1sig");
    std::unique_ptr<TH1D> j_hist_sdcp_2sig = M3::Clone(j_hist_sdcp.get(), "j_hist_sdcp_2sig");
    std::unique_ptr<TH1D> j_hist_sdcp_3sig = M3::Clone(j_hist_sdcp.get(), "j_hist_sdcp_3sig");
 
    //upper and lower edges
    double j_sdcp_1sig_low = 9999999.;
    double j_sdcp_1sig_up  = -9999999.;
    double j_sdcp_2sig_low = 9999999.;;
    double j_sdcp_2sig_up  = -9999999.;
    double j_sdcp_3sig_low = 9999999.;;
    double j_sdcp_3sig_up  = -9999999.;
 
    double contlevel1 = 0.68;
    double contlevel2 = 0.90;
    double contlevel4 = 0.99;
    double contlevel5 = 0.9973;
    double integral, tsum = 0.;
 
    integral = j_hist_copy->Integral();
 
    while((tsum/integral)<contlevel5) {
      double tmax = j_hist_copy->GetMaximum();
      int bin = j_hist_copy->GetMaximumBin();
      double xval = j_hist_copy->GetXaxis()->GetBinCenter(bin);
      double xwidth = j_hist_copy->GetXaxis()->GetBinWidth(bin);
      if((tsum/integral)<contlevel1) {
        j_hist_copy->SetBinContent(bin,-1.0);
        j_hist_1sig->SetBinContent(bin,0.);
        j_hist_2sig->SetBinContent(bin,0.);
        j_hist_3sig->SetBinContent(bin,0.);
        if(xval<j_1sig_low && xval<j_bf) j_1sig_low = xval - xwidth/2.;
        if(xval>j_1sig_up && xval>j_bf) j_1sig_up = xval + xwidth/2.;
      }
      if((tsum/integral)<contlevel2  && (tsum / integral > contlevel1) ) {
        j_hist_copy->SetBinContent(bin,-5.0);
        j_hist_2sig->SetBinContent(bin,0.);
        j_hist_3sig->SetBinContent(bin,0.);
        if(xval<j_2sig_low && xval<j_bf) j_2sig_low = xval - xwidth/2.;
        if(xval>j_2sig_up && xval>j_bf) j_2sig_up = xval + xwidth/2.;
      }
      if((tsum/integral)<contlevel4  && (tsum / integral > contlevel1) ) {
        j_hist_copy->SetBinContent(bin,-9.0);
        j_hist_3sig->SetBinContent(bin,0.);
        if(xval < j_3sig_low && xval <j_bf) j_3sig_low = xval - xwidth/2.;
        if(xval > j_3sig_up && xval > j_bf) j_3sig_up = xval + xwidth/2.;
      }
      tsum+=tmax;
    }
 
    integral = j_hist_sdcp_copy->Integral();
    tsum = 0.;
 
    while((tsum/integral)<contlevel5) {
      double tmax = j_hist_sdcp_copy->GetMaximum();
      int bin = j_hist_sdcp_copy->GetMaximumBin();
      double xval = j_hist_sdcp_copy->GetXaxis()->GetBinCenter(bin);
      double xwidth = j_hist_sdcp_copy->GetXaxis()->GetBinWidth(bin);
      if((tsum/integral)<contlevel1) {
        j_hist_sdcp_copy->SetBinContent(bin,-1.0);
        j_hist_sdcp_1sig->SetBinContent(bin,0.);
        j_hist_sdcp_2sig->SetBinContent(bin,0.);
        j_hist_sdcp_3sig->SetBinContent(bin,0.);
        if(xval<j_sdcp_1sig_low && xval<j_bf) j_sdcp_1sig_low = xval - xwidth/2.;
        if(xval>j_sdcp_1sig_up && xval>j_bf) j_sdcp_1sig_up = xval + xwidth/2.;
      }
      if((tsum/integral)<contlevel2  && (tsum / integral > contlevel1) ) {
        j_hist_sdcp_copy->SetBinContent(bin,-5.0);
        j_hist_sdcp_2sig->SetBinContent(bin,0.);
        j_hist_sdcp_3sig->SetBinContent(bin,0.);
        if(xval<j_sdcp_2sig_low && xval<j_bf) j_sdcp_2sig_low = xval - xwidth/2.;
        if(xval>j_sdcp_2sig_up && xval>j_bf) j_sdcp_2sig_up = xval + xwidth/2.;
      }
      if((tsum/integral)<contlevel4  && (tsum / integral > contlevel1) ) {
        j_hist_sdcp_copy->SetBinContent(bin,-9.0);
        j_hist_sdcp_3sig->SetBinContent(bin,0.);
        if(xval<j_sdcp_3sig_low && xval<j_bf) j_sdcp_3sig_low = xval - xwidth/2.;
        if(xval>j_sdcp_3sig_up && xval>j_bf) j_sdcp_3sig_up = xval + xwidth/2.;
      }
      tsum+=tmax;
    }
 
    j_hist_1sig->SetLineStyle(9);
    j_hist_sdcp_1sig->SetLineStyle(9);
    j_hist_2sig->SetLineStyle(7);
    j_hist_sdcp_2sig->SetLineStyle(7);
    j_hist_3sig->SetLineStyle(2);
    j_hist_sdcp_3sig->SetLineStyle(2);
 
    auto ldash = std::make_unique<TH1D>("ldash", "solid",  10, -0.04, 0.04);
    auto sdash = std::make_unique<TH1D>("sdash", "dashed", 10, -0.04, 0.04);
    auto fdash = std::make_unique<TH1D>("fdash", "fdashed",10, -0.04, 0.04);
    ldash->SetLineColor(kBlack);
    sdash->SetLineColor(kBlack);
    fdash->SetLineColor(kBlack);
    ldash->SetLineWidth(2);
    sdash->SetLineWidth(2);
    fdash->SetLineWidth(2);
    ldash->SetLineStyle(9);
    sdash->SetLineStyle(7);
    fdash->SetLineStyle(2);
 
    double vertUp = 0.5 * j_hist->GetMaximum();
    auto jline_1sig_low = std::make_unique<TLine>(j_1sig_low, 0., j_1sig_low, vertUp);
    auto jline_2sig_low = std::make_unique<TLine>(j_2sig_low, 0., j_2sig_low, vertUp);
    auto jline_3sig_low = std::make_unique<TLine>(j_3sig_low, 0., j_3sig_low, vertUp);
 
    auto jline_1sig_up = std::make_unique<TLine>(j_1sig_up, 0., j_1sig_up,vertUp);
    auto jline_2sig_up = std::make_unique<TLine>(j_2sig_up, 0., j_2sig_up,vertUp);
    auto jline_3sig_up = std::make_unique<TLine>(j_3sig_up, 0., j_3sig_up,vertUp);
 
    auto jline_sdcp_1sig_low = std::make_unique<TLine>(j_sdcp_1sig_low, 0., j_sdcp_1sig_low, vertUp);
    auto jline_sdcp_2sig_low = std::make_unique<TLine>(j_sdcp_2sig_low, 0., j_sdcp_2sig_low, vertUp);
    auto jline_sdcp_3sig_low = std::make_unique<TLine>(j_sdcp_3sig_low, 0., j_sdcp_3sig_low, vertUp);
 
    auto jline_sdcp_1sig_up = std::make_unique<TLine>(j_sdcp_1sig_up, 0., j_sdcp_1sig_up, vertUp);
    auto jline_sdcp_2sig_up = std::make_unique<TLine>(j_sdcp_2sig_up, 0., j_sdcp_2sig_up, vertUp);
    auto jline_sdcp_3sig_up = std::make_unique<TLine>(j_sdcp_3sig_up, 0., j_sdcp_3sig_up, vertUp);
 
    double arrowLength = 0.003;
    double arrowHeight = vertUp;
 
    auto MakeArrow = [&](double x, Color_t color, Width_t width) -> std::unique_ptr<TArrow> {
      auto arrow = std::make_unique<TArrow>(x, arrowHeight, x - arrowLength, arrowHeight, 0.02, ">");
      arrow->SetLineColor(color);
      arrow->SetLineWidth(width);
      return arrow;
    };
 
    auto j_arrow_1sig_up = MakeArrow(j_1sig_up, j_hist_1sig->GetLineColor(), j_hist_1sig->GetLineWidth());
    auto j_arrow_2sig_up = MakeArrow(j_2sig_up, j_hist_2sig->GetLineColor(), j_hist_2sig->GetLineWidth());
    auto j_arrow_3sig_up = MakeArrow(j_3sig_up, j_hist_3sig->GetLineColor(), j_hist_3sig->GetLineWidth());
 
    auto j_sdcp_arrow_1sig_up = MakeArrow(j_sdcp_1sig_up, j_hist_sdcp_1sig->GetLineColor(), j_hist_sdcp_1sig->GetLineWidth());
    auto j_sdcp_arrow_2sig_up = MakeArrow(j_sdcp_2sig_up, j_hist_sdcp_2sig->GetLineColor(), j_hist_sdcp_2sig->GetLineWidth());
    auto j_sdcp_arrow_3sig_up = MakeArrow(j_sdcp_3sig_up, j_hist_sdcp_3sig->GetLineColor(), j_hist_sdcp_3sig->GetLineWidth());
 
    MACH3LOG_DEBUG("j_1sig_low = {:.4f}, j_2sig_low = {:.4f}, j_3sig_low = {:.4f}", j_1sig_low, j_2sig_low, j_3sig_low);
    MACH3LOG_DEBUG("j_1sig_up = {:.4f}, j_2sig_up = {:.4f}, j_3sig_up = {:.4f}", j_1sig_up, j_2sig_up, j_3sig_up);
 
    auto CopyLineStyle = [](const TH1D* src, TLine* dst) {
      dst->SetLineColor(src->GetLineColor());
      dst->SetLineStyle(src->GetLineStyle());
      dst->SetLineWidth(src->GetLineWidth());
    };
 
    CopyLineStyle(j_hist_1sig.get(), jline_1sig_low.get());
    CopyLineStyle(j_hist_1sig.get(), jline_1sig_up.get());
    CopyLineStyle(j_hist_2sig.get(), jline_2sig_low.get());
    CopyLineStyle(j_hist_2sig.get(), jline_2sig_up.get());
    CopyLineStyle(j_hist_3sig.get(), jline_3sig_low.get());
    CopyLineStyle(j_hist_3sig.get(), jline_3sig_up.get());
 
    CopyLineStyle(j_hist_sdcp_1sig.get(), jline_sdcp_1sig_low.get());
    CopyLineStyle(j_hist_sdcp_1sig.get(), jline_sdcp_1sig_up.get());
    CopyLineStyle(j_hist_sdcp_2sig.get(), jline_sdcp_2sig_low.get());
    CopyLineStyle(j_hist_sdcp_2sig.get(), jline_sdcp_2sig_up.get());
    CopyLineStyle(j_hist_sdcp_3sig.get(), jline_sdcp_3sig_low.get());
    CopyLineStyle(j_hist_sdcp_3sig.get(), jline_sdcp_3sig_up.get());
 
    auto leg = std::make_unique<TLegend>(0.45, 0.60, 0.75, 0.90);
    leg->SetTextSize(0.05);
    leg->SetFillStyle(0);
    leg->SetNColumns(1);
    leg->SetTextFont(132);
    leg->SetBorderSize(0);
 
    leg->AddEntry(j_hist.get(), "Prior flat in #delta_{CP}", "l");
    leg->AddEntry(j_hist_sdcp.get(), "Prior flat in sin#delta_{CP}", "l");
    leg->AddEntry(ldash.get(), "68% CI", "l");
    leg->AddEntry(sdash.get(), "90% CI", "l");
    leg->AddEntry(fdash.get(), "99% CI", "l");
 
    j_hist->GetYaxis()->SetRangeUser(0., j_hist->GetMaximum()*1.15);
    j_hist->Draw("h");
    j_hist_sdcp->Draw("same h");
 
    jline_sdcp_1sig_up->Draw("same");
    jline_sdcp_2sig_up->Draw("same");
    jline_sdcp_3sig_up->Draw("same");
    jline_1sig_up->Draw("same");
    jline_2sig_up->Draw("same");
    jline_3sig_up->Draw("same");
 
    j_arrow_1sig_up->Draw();
    j_arrow_2sig_up->Draw();
    j_arrow_3sig_up->Draw();
    j_sdcp_arrow_1sig_up->Draw();
    j_sdcp_arrow_2sig_up->Draw();
    j_sdcp_arrow_3sig_up->Draw();
    leg->Draw("same");
 
    auto ttext = std::make_unique<TText>();
    ttext->SetNDC(); // Use normalized device coordinates
    ttext->SetTextSize(0.03); // Adjust size as needed
    ttext->SetTextAlign(13); // Align left-top
 
    if (hierarchy == 1) ttext->DrawText(0.15, 0.85, "Normal Ordering");
    else if (hierarchy == 0) ttext->DrawText(0.15, 0.85, "Both Orderings");
    else if (hierarchy == -1) ttext->DrawText(0.15, 0.85, "Inverted Ordering");
 
    gPad->RedrawAxis();
    Posterior->Update();
    gPad->Update();
 
    Posterior->Print(CanvasName);
 
    if(hierarchy == 1)  Posterior->Write("jarl1D_NH_comp");
    else if(hierarchy == 0)  Posterior->Write("jarl1D_both_comp");
    else if(hierarchy == -1)  Posterior->Write("jarl1D_IH_comp");
  }
 
  gErrorIgnoreLevel = originalErrorLevel;
}
 
// ***************
void OscProcessor::MakePiePlot() {
// ***************
  if(!OscEnabled || DeltaCPIndex == M3::_BAD_INT_)
  {
    MACH3LOG_WARN("Will not {}, as oscillation parameters are missing", __func__);
    return;
  }
  MACH3LOG_INFO("Starting {}", __func__);
 
  // get best fit for delta CP
  const double best_fit = (*Means_HPD)(DeltaCPIndex);
 
  const double sigma_p = (*Errors_HPD_Positive)(DeltaCPIndex);
  const double sigma_n = (*Errors_HPD_Negative)(DeltaCPIndex);
  // make sure result is between -pi and pi
  auto wrap_pi = [](double x) {
    while (x > TMath::Pi())  x -= 2*TMath::Pi();
    while (x < -TMath::Pi()) x += 2*TMath::Pi();
    return x;
  };
 
  std::array<double, 6> bds;
  bds[0] = wrap_pi(best_fit - 3.0 * sigma_n); // -3σ
  bds[1] = wrap_pi(best_fit - 2.0 * sigma_n); // -2σ
  bds[2] = wrap_pi(best_fit - 1.0 * sigma_n); // -1σ
  bds[3] = wrap_pi(best_fit + 1.0 * sigma_p); // +1σ
  bds[4] = wrap_pi(best_fit + 2.0 * sigma_p); // +2σ
  bds[5] = wrap_pi(best_fit + 3.0 * sigma_p); // +3σ
 
  constexpr double radius = 0.4;
  constexpr double rad_to_deg = 180.0 / TMath::Pi();
 
  // ROOT expects TEllipse angles in degrees, counterclockwise from the x-axis.
  // If phimax < phimin, ROOT draws counterclockwise across the full circle, causing overlaps.
  // This ensures threesigA slice stays within the intended range.
  auto normalize_angle = [](double rad) {
    // If rad is negative, add 2*pi to wrap into [0, 2*pi)
    if (rad < 0) rad += 2.0 * TMath::Pi();
    return rad;
  };
 
  TEllipse onesig   (0.5, 0.5, radius, radius, bds[2] * rad_to_deg, bds[4] * rad_to_deg);
  TEllipse twosigA  (0.5, 0.5, radius, radius, bds[1] * rad_to_deg, bds[2] * rad_to_deg);
  TEllipse twosigB  (0.5, 0.5, radius, radius, bds[3] * rad_to_deg, bds[4] * rad_to_deg);
 
  // three sigma slices
  TEllipse threesigA(0.5, 0.5, radius, radius, bds[0] * rad_to_deg, normalize_angle(bds[1]) * rad_to_deg);
  TEllipse threesigB(0.5, 0.5, radius, radius, bds[4] * rad_to_deg, bds[5] * rad_to_deg);
 
  // Remaining slices
  TEllipse rest(0.5, 0.5, radius, radius, bds[5]*rad_to_deg, bds[0]*rad_to_deg);
  TEllipse restA(0.5, 0.5, radius, radius, bds[5]*rad_to_deg, 180.0);
  TEllipse restB(0.5, 0.5, radius, radius, -180.0, bds[0]*rad_to_deg);
 
  onesig.SetFillColor(13);
  twosigA.SetFillColor(12);
  twosigB.SetFillColor(12);
  threesigA.SetFillColor(11);
  threesigB.SetFillColor(11);
  TLine line1(0.5 - radius, 0.5, 0.5 + radius, 0.5);
  line1.SetLineWidth(3);
 
  TLine line2(0.5, 0.5 - radius, 0.5, 0.5 + radius);
  line2.SetLineWidth(3);
 
  TArrow bf(0.5, 0.5, 0.5 + radius * cos(best_fit),0.5 + radius * sin(best_fit),0.04, "|>");
  bf.SetLineWidth(3);
  bf.SetLineColor(kRed);
  bf.SetFillColor(kRed);
 
  TCanvas canvas("canvas", "canvas", 0, 0, 1000, 1000);
  onesig.Draw();
  twosigA.Draw();
  twosigB.Draw();
  threesigA.Draw();
  threesigB.Draw();
 
  // Check if the rest wraps around the circle
  if (bds[5] > 0) {
    // Single rest slice
    rest.Draw();
  } else {
    // Split rest into two slices
    restA.Draw();
    restB.Draw();
  }
 
  line1.Draw();
  line2.Draw();
  bf.Draw();
 
  TLegend leg(0.0, 0.8, 0.23, 0.95);
  leg.AddEntry(&bf,        "Best Fit", "L");
  leg.AddEntry(&onesig,    "1#sigma",  "F");
  leg.AddEntry(&twosigA,   "2#sigma",  "F");
  leg.AddEntry(&threesigA, "3#sigma",  "F");
  leg.Draw();
 
  // KS: Simple lambda to avoid copy-pasting
  auto draw_text = [](auto& txt, Color_t color = kBlack) {
    txt.SetTextAlign(22);
    txt.SetTextColor(color);
    txt.SetTextFont(43);
    txt.SetTextSize(40);
    txt.SetTextAngle(0);
    txt.Draw();
  };
 
  //KS: If best fit point is somehow very close text we simply not plot it
  // Define a threshold for "too close"
  constexpr double too_close_threshold = 0.1;
 
  // Position of tbf
  const double tbf_x = 0.5 + (radius + 0.02) * cos(best_fit);
  const double tbf_y = 0.5 + (radius + 0.02) * sin(best_fit);
 
  // Function to calculate distance between two points
  auto distance = [](double x1, double y1, double x2, double y2) {
    return std::sqrt((x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2));
  };
 
  // Check and draw right 0
  constexpr double t0_x = 0.5 + radius + 0.02;
  constexpr double t0_y = 0.5;
  TText t0(t0_x, t0_y, "0");
  if (distance(tbf_x, tbf_y, t0_x, t0_y) > too_close_threshold) {
    draw_text(t0);
  }
 
  // Check and draw left pi
  constexpr double tp_x = 0.5 - radius - 0.02;
  constexpr double tp_y = 0.5;
  TLatex tp(tp_x, tp_y, "#pi");
  if (distance(tbf_x, tbf_y, tp_x, tp_y) > too_close_threshold) {
    draw_text(tp);
  }
 
  // Check and draw top pi/2
  constexpr double tp2_x = 0.5;
  constexpr double tp2_y = 0.5 + radius + 0.04;
  TLatex tp2(tp2_x, tp2_y, "#frac{#pi}{2}");
  if (distance(tbf_x, tbf_y, tp2_x, tp2_y) > too_close_threshold) {
    draw_text(tp2);
  }
 
  // Check and draw bottom -pi/2
  constexpr double tmp2_x = 0.5;
  constexpr double tmp2_y = 0.5 - radius - 0.04;
  TLatex tmp2(tmp2_x, tmp2_y, "-#frac{#pi}{2}");
  if (distance(tbf_x, tbf_y, tmp2_x, tmp2_y) > too_close_threshold) {
    draw_text(tmp2);
  }
 
  TLatex tbf(0.5 + (radius + 0.02) * cos(best_fit),
             0.5 + (radius + 0.02) * sin(best_fit),
             fmt::format("{:.2f}", best_fit).c_str());
  draw_text(tbf, kRed);
 
  canvas.Print(CanvasName);
}
 
// ***************
// MP: Calculate PMNS matrix elements
std::array<std::array<TComplex, 3>, 3> OscProcessor::CalculatePMNSElements(const double s2th13, const double s2th23, const double s2th12, const double dcp) const {
// ***************
  const double s13 = std::sqrt(s2th13);
  const double s23 = std::sqrt(s2th23);
  const double s12 = std::sqrt(s2th12);
 
  const double sdcp = std::sin(dcp);
  const double cdcp = std::cos(dcp);
  
  const double c13 = std::sqrt(1.-s2th13);
  const double c12 = std::sqrt(1.-s2th12);
  const double c23 = std::sqrt(1.-s2th23);
 
  double real_ue[3];
  double imag_ue[3];
  real_ue[0] = c12*c13;
  imag_ue[0] = 0.;
  real_ue[1] = s12*c13;
  imag_ue[1] = 0.;
  real_ue[2] = s13*cdcp;
  imag_ue[2] = -s13*sdcp;
 
  double real_umu[3];
  double imag_umu[3];
  real_umu[0] = -s12*c23 - c12*s23*s13*cdcp;
  imag_umu[0] = -c12*s23*s13*sdcp;
  real_umu[1] =  c12*c23 - s12*s23*s13*cdcp;
  imag_umu[1] = -s12*s23*s13*sdcp;
  real_umu[2] =  s23*c13;
  imag_umu[2] = 0.;
 
  double real_utau[3];
  double imag_utau[3];
  real_utau[0] =  s12*s23 - c12*c23*s13*cdcp;
  imag_utau[0] = -c12*c23*s13*sdcp;
  real_utau[1] = -c12*s23 - s12*c23*s13*cdcp;
  imag_utau[1] = -s12*c23*s13*sdcp;
  real_utau[2] =  c23*c13;
  imag_utau[2] = 0.;
 
  TComplex U[3][3];
  for (int i = 0; i < 3; i++) {
    U[0][i] = TComplex(real_ue[i], imag_ue[i]);
    U[1][i] = TComplex(real_umu[i], imag_umu[i]);
    U[2][i] = TComplex(real_utau[i], imag_utau[i]);
  }
 
  return {{{U[0][0], U[0][1], U[0][2]}, {U[1][0], U[1][1], U[1][2]}, {U[2][0], U[2][1], U[2][2]}}};
}
 
// ***************
// MP: Produce PMNS matrix elements
void OscProcessor::ProducePMNSElements() {
// ***************
  if(!OscEnabled ||
    Sin2Theta13Index == M3::_BAD_INT_ ||
    Sin2Theta12Index == M3::_BAD_INT_ ||
    Sin2Theta23Index == M3::_BAD_INT_ ||
    DeltaCPIndex == M3::_BAD_INT_)
  {
    MACH3LOG_WARN("Will not {}, as oscillation parameters are missing", __func__);
    return;
  }
  MACH3LOG_INFO("Starting {}", __func__);
 
  double s2th13, s2th23, s2th12, dcp = M3::_BAD_DOUBLE_;
 
 
  TDirectory *PMNSElementsDir = OutputFile->mkdir("PMNSElements");
  PMNSElementsDir->cd();
 
  unsigned int step = 0;
  Chain->SetBranchStatus("*", false);
 
  Chain->SetBranchStatus(Sin2Theta13Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta13Name.c_str(), &s2th13);
 
  Chain->SetBranchStatus(Sin2Theta23Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta23Name.c_str(), &s2th23);
 
  Chain->SetBranchStatus(Sin2Theta12Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta12Name.c_str(), &s2th12);
 
  Chain->SetBranchStatus(DeltaCPName.c_str(), true);
  Chain->SetBranchAddress(DeltaCPName.c_str(), &dcp);
 
  Chain->SetBranchStatus("step", true);
  Chain->SetBranchAddress("step", &step);
 
  double weight = 1.0;
  std::vector<double> reweights(ReweightNames.size(), 1.0);
  bool anyMissing = false;
  for (size_t i = 0; i < ReweightNames.size(); ++i) {
    const auto& name = ReweightNames[i];
    if (Chain->GetBranch(name.c_str())) {
      Chain->SetBranchStatus(name.c_str(), true);
      Chain->SetBranchAddress(name.c_str(), &reweights[i]);
    } else {
      anyMissing = true;
    }
  }
  if (anyMissing) {
    MACH3LOG_WARN("Some reweight branches were missing — using partial weight");
    reweights.clear();
  }
  constexpr int n_bins = 1000;
 
  std::unique_ptr<TH1D> h_ue[3];
  std::unique_ptr<TH1D> h_umu[3];
  std::unique_ptr<TH1D> h_utau[3];
 
  std::unique_ptr<TH1D> h_ue_real[3];
  std::unique_ptr<TH1D> h_umu_real[3];
  std::unique_ptr<TH1D> h_utau_real[3];
 
  std::unique_ptr<TH1D> h_ue_imag[3];
  std::unique_ptr<TH1D> h_umu_imag[3];
  std::unique_ptr<TH1D> h_utau_imag[3];
 
  std::unique_ptr<TH2D> h_ue_s2th12[3];
  std::unique_ptr<TH2D> h_umu_s2th12[3];
  std::unique_ptr<TH2D> h_utau_s2th12[3];
 
  std::unique_ptr<TH2D> h_ue_s2th13[3];
  std::unique_ptr<TH2D> h_umu_s2th13[3];
  std::unique_ptr<TH2D> h_utau_s2th13[3];
 
  std::unique_ptr<TH2D> h_ue_s2th23[3];
  std::unique_ptr<TH2D> h_umu_s2th23[3];
  std::unique_ptr<TH2D> h_utau_s2th23[3];
 
  std::unique_ptr<TH2D> h_ue_dcp[3];
  std::unique_ptr<TH2D> h_umu_dcp[3];
  std::unique_ptr<TH2D> h_utau_dcp[3];
 
  for (int iU = 0; iU < 3; iU++) {
    h_ue[iU] = std::make_unique<TH1D>(Form("h_ue%d", iU+1), Form(";|U_{e%d}|", iU+1), n_bins, 0., 1.);
    h_ue[iU]->SetDirectory(nullptr);
    h_umu[iU] = std::make_unique<TH1D>(Form("h_umu%d", iU+1), Form(";|U_{#mu%d}|", iU+1), n_bins, 0., 1.);
    h_umu[iU]->SetDirectory(nullptr);
    h_utau[iU] = std::make_unique<TH1D>(Form("h_utau%d", iU+1), Form(";|U_{#tau%d}|", iU+1), n_bins, 0., 1.);
    h_utau[iU]->SetDirectory(nullptr);
 
    h_ue_real[iU] = std::make_unique<TH1D>(Form("h_ue%d_real", iU+1), Form(";Re(U_{e%d})", iU+1), n_bins, -1., 1.);
    h_ue_real[iU]->SetDirectory(nullptr);
    h_umu_real[iU] = std::make_unique<TH1D>(Form("h_umu%d_real", iU+1), Form(";Re(U_{#mu%d})", iU+1), n_bins, -1., 1.);
    h_umu_real[iU]->SetDirectory(nullptr);
    h_utau_real[iU] = std::make_unique<TH1D>(Form("h_utau%d_real", iU+1), Form(";Re(U_{#tau%d})", iU+1), n_bins, -1., 1.);
    h_utau_real[iU]->SetDirectory(nullptr);
 
    h_ue_imag[iU] = std::make_unique<TH1D>(Form("h_ue%d_imag", iU+1), Form(";Im(U_{e%d})", iU+1), n_bins, -1., 1.);
    h_ue_imag[iU]->SetDirectory(nullptr);
    h_umu_imag[iU] = std::make_unique<TH1D>(Form("h_umu%d_imag", iU+1), Form(";Im(U_{#mu%d})", iU+1), n_bins, -1., 1.);
    h_umu_imag[iU]->SetDirectory(nullptr);
    h_utau_imag[iU] = std::make_unique<TH1D>(Form("h_utau%d_imag", iU+1), Form(";Im(U_{#tau%d})", iU+1), n_bins, -1., 1.);
    h_utau_imag[iU]->SetDirectory(nullptr);
 
    h_ue_s2th12[iU] = std::make_unique<TH2D>(Form("h_ue%d_s2th12", iU+1), Form(";|U_{e%d}|;sin^{2}(#theta_{12})", iU+1),
                                             n_bins, 0., 1., n_bins, 0.15, 0.5);
    h_ue_s2th12[iU]->SetDirectory(nullptr);
    h_umu_s2th12[iU] = std::make_unique<TH2D>(Form("h_umu%d_s2th12", iU+1), Form(";|U_{#mu%d}|;sin^{2}(#theta_{12})", iU+1),
                                              n_bins, 0., 1., n_bins, 0.15, 0.5);
    h_umu_s2th12[iU]->SetDirectory(nullptr);
    h_utau_s2th12[iU] = std::make_unique<TH2D>(Form("h_utau%d_s2th12", iU+1), Form(";|U_{#tau%d}|;sin^{2}(#theta_{12})", iU+1),
                                               n_bins, 0., 1., n_bins, 0.15, 0.5);
    h_utau_s2th12[iU]->SetDirectory(nullptr);
 
 
    h_ue_s2th13[iU] = std::make_unique<TH2D>(Form("h_ue%d_s2th13", iU+1), Form(";|U_{e%d}|;sin^{2}(#theta_{13})", iU+1),
                                             n_bins, 0., 1., n_bins, 0., 0.1);
    h_ue_s2th13[iU]->SetDirectory(nullptr);
    h_umu_s2th13[iU] = std::make_unique<TH2D>(Form("h_umu%d_s2th13", iU+1), Form(";|U_{#mu%d}|;sin^{2}(#theta_{13})", iU+1),
                                              n_bins, 0., 1., n_bins, 0., 0.1);
    h_umu_s2th13[iU]->SetDirectory(nullptr);
    h_utau_s2th13[iU] = std::make_unique<TH2D>(Form("h_utau%d_s2th13", iU+1), Form(";|U_{#tau%d}|;sin^{2}(#theta_{13})", iU+1),
                                               n_bins, 0., 1., n_bins, 0., 0.1);
    h_utau_s2th13[iU]->SetDirectory(nullptr);
 
 
    h_ue_s2th23[iU] = std::make_unique<TH2D>(Form("h_ue%d_s2th23", iU+1), Form(";|U_{e%d}|;sin^{2}(#theta_{23})", iU+1),
                                             n_bins, 0., 1., n_bins, 0.3, 0.8);
    h_ue_s2th23[iU]->SetDirectory(nullptr);
    h_umu_s2th23[iU] = std::make_unique<TH2D>(Form("h_umu%d_s2th23", iU+1), Form(";|U_{#mu%d}|;sin^{2}(#theta_{23})", iU+1),
                                              n_bins, 0., 1., n_bins, 0.3, 0.8);
    h_umu_s2th23[iU]->SetDirectory(nullptr);
    h_utau_s2th23[iU] = std::make_unique<TH2D>(Form("h_utau%d_s2th23", iU+1), Form(";|U_{#tau%d}|;sin^{2}(#theta_{23})", iU+1),
                                               n_bins, 0., 1., n_bins, 0.3, 0.8);
    h_utau_s2th23[iU]->SetDirectory(nullptr);
 
 
    h_ue_dcp[iU] = std::make_unique<TH2D>(Form("h_ue%d_dcp", iU+1), Form(";|U_{e%d}|;#delta_{CP}", iU+1),
                                          n_bins, 0., 1., n_bins, -TMath::Pi(), TMath::Pi());
    h_ue_dcp[iU]->SetDirectory(nullptr);
    h_umu_dcp[iU] = std::make_unique<TH2D>(Form("h_umu%d_dcp", iU+1), Form(";|U_{#mu%d}|;#delta_{CP}", iU+1),
                                           n_bins, 0., 1., n_bins, -TMath::Pi(), TMath::Pi());
    h_umu_dcp[iU]->SetDirectory(nullptr);
    h_utau_dcp[iU] = std::make_unique<TH2D>(Form("h_utau%d_dcp", iU+1), Form(";|U_{#tau%d}|;#delta_{CP}", iU+1),
                                            n_bins, 0., 1., n_bins, -TMath::Pi(), TMath::Pi());
    h_utau_dcp[iU]->SetDirectory(nullptr);
  }
 
  // MP: 2D histograms for all unique pairs of |U_{alpha i}| vs |U_{beta j}| (no self-pairs)
  std::vector<std::unique_ptr<TH2D>> h_UU;
  std::string U_names[9] = {"ue1", "ue2", "ue3", "umu1", "umu2", "umu3", "utau1", "utau2", "utau3"};
  std::string U_tex[9]  = {"|U_{e1}|", "|U_{e2}|", "|U_{e3}|", "|U_{#mu1}|", "|U_{#mu2}|", "|U_{#mu3}|", "|U_{#tau1}|", "|U_{#tau2}|", "|U_{#tau3}|"};
  for(int i=0; i < 9; ++i){
    for(int j = i+1; j < 9; ++j){
      std::string name = "h_" + U_names[i] + "_" + U_names[j];
      std::string title = ";" + U_tex[i] + ";" + U_tex[j];
      h_UU.push_back(std::make_unique<TH2D>(name.c_str(), title.c_str(), n_bins, 0., 1., n_bins, 0., 1.));
      h_UU.back()->SetDirectory(nullptr);
    }
  }
  const Long64_t countwidth = nEntries/5;
  for(int i = 0; i < nEntries; i++) {
    if (i % countwidth == 0) {
      M3::Utils::PrintProgressBar(i, nEntries);
      M3::Utils::EstimateDataTransferRate(Chain, i);
    } else {
      Chain->GetEntry(i);
    }
 
    if(step < BurnInCut) continue; // burn-in cut
 
    const std::array<std::array<TComplex, 3>, 3> U = CalculatePMNSElements(s2th13, s2th23, s2th12, dcp);
    // KS: Calculate total weight
    weight = 1.0;
    for (const auto& w : reweights) {
      weight *= w;
    }
    for(int iU = 0; iU < 3; iU++){
      h_ue[iU]->Fill(TComplex::Abs(U[0][iU]), weight);
      h_umu[iU]->Fill(TComplex::Abs(U[1][iU]), weight);
      h_utau[iU]->Fill(TComplex::Abs(U[2][iU]), weight);
 
      h_ue_real[iU]->Fill(U[0][iU].Re(), weight);
      h_umu_real[iU]->Fill(U[1][iU].Re(), weight);
      h_utau_real[iU]->Fill(U[2][iU].Re(), weight);
 
      h_ue_imag[iU]->Fill(U[0][iU].Im(), weight);
      h_umu_imag[iU]->Fill(U[1][iU].Im(), weight);
      h_utau_imag[iU]->Fill(U[2][iU].Im(), weight);
 
      h_ue_s2th12[iU]->Fill(TComplex::Abs(U[0][iU]), s2th12, weight);
      h_umu_s2th12[iU]->Fill(TComplex::Abs(U[1][iU]), s2th12, weight);
      h_utau_s2th12[iU]->Fill(TComplex::Abs(U[2][iU]), s2th12, weight);
 
      h_ue_s2th13[iU]->Fill(TComplex::Abs(U[0][iU]), s2th13, weight);
      h_umu_s2th13[iU]->Fill(TComplex::Abs(U[1][iU]), s2th13, weight);
      h_utau_s2th13[iU]->Fill(TComplex::Abs(U[2][iU]), s2th13, weight);
 
      h_ue_s2th23[iU]->Fill(TComplex::Abs(U[0][iU]), s2th23, weight);
      h_umu_s2th23[iU]->Fill(TComplex::Abs(U[1][iU]), s2th23, weight);
      h_utau_s2th23[iU]->Fill(TComplex::Abs(U[2][iU]), s2th23, weight);
 
      h_ue_dcp[iU]->Fill(TComplex::Abs(U[0][iU]), dcp, weight);
      h_umu_dcp[iU]->Fill(TComplex::Abs(U[1][iU]), dcp, weight);
      h_utau_dcp[iU]->Fill(TComplex::Abs(U[2][iU]), dcp, weight);
    }
 
    // MP: Store in an array for easy access
    double U_mod[9] = {TComplex::Abs(U[0][0]), TComplex::Abs(U[0][1]), TComplex::Abs(U[0][2]), TComplex::Abs(U[1][0]), TComplex::Abs(U[1][1]), TComplex::Abs(U[1][2]), TComplex::Abs(U[2][0]), TComplex::Abs(U[2][1]), TComplex::Abs(U[2][2])};
    int idx = 0;
    for(int ix = 0; ix < 9; ++ix){
      for(int iy = ix+1; iy<9; ++iy){
        h_UU[idx]->Fill(U_mod[ix], U_mod[iy],  weight);
        ++idx;
      }
    }
  } // end loop over steps
 
  // Now we save
  PMNSElementsDir->cd();
 
  for (int iU = 0; iU < 3; iU++) {
    h_ue[iU]->Write(Form("h_ue%d", iU+1));
    h_umu[iU]->Write(Form("h_umu%d", iU+1));
    h_utau[iU]->Write(Form("h_utau%d", iU+1));
 
    h_ue_real[iU]->Write(Form("h_ue%d_real", iU+1));
    h_umu_real[iU]->Write(Form("h_umu%d_real", iU+1));
    h_utau_real[iU]->Write(Form("h_utau%d_real", iU+1));
 
    h_ue_imag[iU]->Write(Form("h_ue%d_imag", iU+1));
    h_umu_imag[iU]->Write(Form("h_umu%d_imag", iU+1));
    h_utau_imag[iU]->Write(Form("h_utau%d_imag", iU+1));
 
    h_ue_s2th12[iU]->Write(Form("h_ue%d_s2th12", iU+1));
    h_umu_s2th12[iU]->Write(Form("h_umu%d_s2th12", iU+1));
    h_utau_s2th12[iU]->Write(Form("h_utau%d_s2th12", iU+1));
 
    h_ue_s2th13[iU]->Write(Form("h_ue%d_s2th13", iU+1));
    h_umu_s2th13[iU]->Write(Form("h_umu%d_s2th13", iU+1));
    h_utau_s2th13[iU]->Write(Form("h_utau%d_s2th13", iU+1));
 
    h_ue_s2th23[iU]->Write(Form("h_ue%d_s2th23", iU+1));
    h_umu_s2th23[iU]->Write(Form("h_umu%d_s2th23", iU+1));
    h_utau_s2th23[iU]->Write(Form("h_utau%d_s2th23", iU+1));
 
    h_ue_dcp[iU]->Write(Form("h_ue%d_dcp", iU+1));
    h_umu_dcp[iU]->Write(Form("h_umu%d_dcp", iU+1));
    h_utau_dcp[iU]->Write(Form("h_utau%d_dcp", iU+1));
  }
 
  // MP: Write all |U_{alpha i}| vs |U_{beta j}| 2D histograms (no self-pairs)
  for(size_t iPlot = 0; iPlot < h_UU.size(); ++iPlot){
    h_UU[iPlot]->Write();
  }
 
  PMNSElementsDir->Close();
  delete PMNSElementsDir;
 
  Chain->SetBranchStatus("*", true);
  OutputFile->cd();
}
 
// ***************
// MP: Produce unitarity triangles from PMNS matrix elements
void OscProcessor::ProduceUnitarityTriangles() {
// ***************
  if(!OscEnabled ||
    Sin2Theta13Index == M3::_BAD_INT_ ||
    Sin2Theta12Index == M3::_BAD_INT_ ||
    Sin2Theta23Index == M3::_BAD_INT_ ||
    DeltaCPIndex == M3::_BAD_INT_)
  {
    MACH3LOG_WARN("Will not {}, as oscillation parameters are missing", __func__);
    return;
  }
  MACH3LOG_INFO("Starting {}", __func__);
 
  double s2th13, s2th23, s2th12, dcp = M3::_BAD_DOUBLE_;
 
  TComplex tr_emu_num, tr_etau_num, tr_mutau_num, tr_12_num, tr_13_num, tr_23_num;
  TComplex tr_emu_denom, tr_etau_denom, tr_mutau_denom, tr_12_denom, tr_13_denom, tr_23_denom;
  TComplex tr_emu, tr_etau, tr_mutau, tr_12, tr_13, tr_23;
 
  TDirectory *UnitarityTrianglesDir = OutputFile->mkdir("UnitarityTriangles");
  UnitarityTrianglesDir->cd();
 
  unsigned int step = 0;
  Chain->SetBranchStatus("*", false);
 
  Chain->SetBranchStatus(Sin2Theta13Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta13Name.c_str(), &s2th13);
 
  Chain->SetBranchStatus(Sin2Theta23Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta23Name.c_str(), &s2th23);
 
  Chain->SetBranchStatus(Sin2Theta12Name.c_str(), true);
  Chain->SetBranchAddress(Sin2Theta12Name.c_str(), &s2th12);
 
  Chain->SetBranchStatus(DeltaCPName.c_str(), true);
  Chain->SetBranchAddress(DeltaCPName.c_str(), &dcp);
 
  Chain->SetBranchStatus("step", true);
  Chain->SetBranchAddress("step", &step);
 
  double weight = 1.0;
  std::vector<double> reweights(ReweightNames.size(), 1.0);
  bool anyMissing = false;
  for (size_t i = 0; i < ReweightNames.size(); ++i) {
    const auto& name = ReweightNames[i];
    if (Chain->GetBranch(name.c_str())) {
      Chain->SetBranchStatus(name.c_str(), true);
      Chain->SetBranchAddress(name.c_str(), &reweights[i]);
    } else {
      anyMissing = true;
    }
  }
  if (anyMissing) {
    MACH3LOG_WARN("Some reweight branches were missing — using partial weight");
    reweights.clear();
  }
  constexpr int n_bins = 1000;
 
  std::unique_ptr<TH2D> h_tr_emu;
  std::unique_ptr<TH2D> h_tr_etau;
  std::unique_ptr<TH2D> h_tr_mutau;
 
  std::unique_ptr<TH2D> h_tr_12;
  std::unique_ptr<TH2D> h_tr_13;
  std::unique_ptr<TH2D> h_tr_23;
 
  h_tr_emu = std::make_unique<TH2D>("h_tr_emu",";#rho_{e#mu};#eta_{e#mu}",n_bins,-5,5,n_bins,-5,5);
  h_tr_emu->SetDirectory(nullptr);
  h_tr_etau = std::make_unique<TH2D>("h_tr_etau",";#rho_{e#tau};#eta_{e#tau}",n_bins,0,2,n_bins,-1,1);
  h_tr_etau->SetDirectory(nullptr);
  h_tr_mutau = std::make_unique<TH2D>("h_tr_mutau",";#rho_{#mu#tau};#eta_{#mu#tau}",n_bins,0,2,n_bins,-1,1);
  h_tr_mutau->SetDirectory(nullptr);
 
  h_tr_12 = std::make_unique<TH2D>("h_tr_12",";#rho_{12};#eta_{12}",n_bins,0,3.2,n_bins,-1.5,1.5);
  h_tr_12->SetDirectory(nullptr);
  h_tr_13 = std::make_unique<TH2D>("h_tr_13",";#rho_{13};#eta_{13}",n_bins,0,2,n_bins,-1,1);
  h_tr_13->SetDirectory(nullptr);
  h_tr_23 = std::make_unique<TH2D>("h_tr_23",";#rho_{23};#eta_{23}",n_bins,-8,8,n_bins,-8,8);
  h_tr_23->SetDirectory(nullptr);
 
  const Long64_t countwidth = nEntries/5;
  for(int i = 0; i < nEntries; i++) {
    if (i % countwidth == 0) {
      M3::Utils::PrintProgressBar(i, nEntries);
      M3::Utils::EstimateDataTransferRate(Chain, i);
    } else {
      Chain->GetEntry(i);
    }
 
    if(step < BurnInCut) continue; // burn-in cut
 
    // KS: Calculate total weight
    weight = 1.0;
    for (const auto& w : reweights) {
      weight *= w;
    }
 
    const std::array<std::array<TComplex, 3>, 3> U = CalculatePMNSElements(s2th13, s2th23, s2th12, dcp);
 
    // Calculate the sides of the unitarity triangles
    tr_emu_num   = U[0][0].operator*(TComplex::Conjugate(U[1][0]));
    tr_emu_denom = U[0][2].operator*(TComplex::Conjugate(U[1][2]));
    tr_emu       = - tr_emu_num.operator/(tr_emu_denom);
 
    tr_etau_num   = U[0][1].operator*(TComplex::Conjugate(U[2][1]));
    tr_etau_denom = U[0][0].operator*(TComplex::Conjugate(U[2][0]));
    tr_etau       = - tr_etau_num.operator/(tr_etau_denom);
 
    tr_mutau_num   = U[1][2].operator*(TComplex::Conjugate(U[2][2]));
    tr_mutau_denom = U[1][1].operator*(TComplex::Conjugate(U[2][1]));
    tr_mutau       = - tr_mutau_num.operator/(tr_mutau_denom);
 
    tr_12_num   = U[0][0].operator*(TComplex::Conjugate(U[0][1]));
    tr_12_denom = U[1][0].operator*(TComplex::Conjugate(U[1][1]));
    tr_12       = - tr_12_num.operator/(tr_12_denom);
 
    tr_13_num   = U[1][0].operator*(TComplex::Conjugate(U[1][2]));
    tr_13_denom = U[2][0].operator*(TComplex::Conjugate(U[2][2]));
    tr_13       = - tr_13_num.operator/(tr_13_denom);
 
    tr_23_num   = U[2][1].operator*(TComplex::Conjugate(U[2][2]));
    tr_23_denom = U[0][1].operator*(TComplex::Conjugate(U[0][2]));
    tr_23       = - tr_23_num.operator/(tr_23_denom);
 
    h_tr_emu->Fill(tr_emu.Re(), tr_emu.Im(), weight);
    h_tr_etau->Fill(tr_etau.Re(), tr_etau.Im(), weight);
    h_tr_mutau->Fill(tr_mutau.Re(), tr_mutau.Im(), weight);
 
    h_tr_12->Fill(tr_12.Re(), tr_12.Im(), weight);
    h_tr_13->Fill(tr_13.Re(), tr_13.Im(), weight);
    h_tr_23->Fill(tr_23.Re(), tr_23.Im(), weight);
  } // end loop over steps
 
  // Now we save
  UnitarityTrianglesDir->cd();
 
  h_tr_emu->Write("h_tr_emu");
  h_tr_etau->Write("h_tr_etau");
  h_tr_mutau->Write("h_tr_mutau");
  
  h_tr_12->Write("h_tr_12");
  h_tr_13->Write("h_tr_13");
  h_tr_23->Write("h_tr_23");
 
  UnitarityTrianglesDir->Close();
  delete UnitarityTrianglesDir;
 
  Chain->SetBranchStatus("*", true);
  OutputFile->cd();
}