bayesian_estimator.h

 
#ifndef __bayesian_estimator_H
#define __bayesian_estimator_H
 
#include "KVBase.h"
#include "impact_parameter_distribution.h"
 
#include <vector>
#include <numeric>
#include <TF2.h>
#include <TH1.h>
#include <TCanvas.h>
#include <TGraph.h>
#include <KVValueRange.h>
#include "TMath.h"
#include "TNamed.h"
 
#include "Math/PdfFuncMathCore.h"
 
namespace KVImpactParameters {
   class gamma_kernel {
   public:
      double operator()(double X, double mean, double reduced_variance)
      {
         return ROOT::Math::gamma_pdf(X, mean / reduced_variance, reduced_variance);
      }
 
      ClassDef(gamma_kernel, 0)
   };
 
   class BD_kernel {
   public:
      double operator()(double X, double mean, double reduced_variance)
      {
         return ROOT::Math::binomial_pdf(TMath::Nint(X), 1. - reduced_variance, mean / (1. - reduced_variance));
      }
 
      ClassDef(BD_kernel, 0)
   };
 
   class NBD_kernel {
   public:
      double operator()(double X, double mean, double reduced_variance)
      {
         return ROOT::Math::negative_binomial_pdf(TMath::Nint(X), 1. / reduced_variance, mean / (reduced_variance - 1.));
      }
 
      ClassDef(NBD_kernel, 0)
   };
 
 
   template
   <class FittingFunction, class FluctuationKernel>
   class bayesian_estimator : public KVBase {
 
      FittingFunction theFitter;
      FluctuationKernel theKernel;
 
      TH1* histo;
      TH1* h_selection;
      std::vector<double> sel_rapp;
      TF1 p_X_cb_integrator;//
      TF1 P_X_fit_function;//
      TF1 mean_X_vs_cb_function;
      TF1 mean_X_vs_b_function;
      TF1 p_X_X_integrator;
      TF1 p_X_X_integrator_with_selection;
      TF1 fitted_P_X;
      TF1 Cb_dist_for_X_select;
      TF1 Cb_dist_for_arb_X_select;
      TF1 B_dist_for_X_select;
      TF1 B_dist_for_arb_X_select;
      TF1 B_dist_for_arb_X_select_from_histo;
 
      Bool_t fIntegerVariable;
 
      impact_parameter_distribution fIPDist;// used to convert centrality to impact parameter
 
      double mean_X_vs_cb(double* x, double* par)
      {
 
         theFitter.fill_params_from_array(par);
         return theFitter.meanX(x[0]);
      }
      double mean_X_vs_b(double* x, double* par)
      {
 
         theFitter.fill_params_from_array(par);
         return theFitter.meanX(fIPDist.GetCentrality().Eval(x[0]));
      }
      double P_X_cb(double X, double cb)
      {
         return fIntegerVariable ? theKernel(TMath::Nint(X), theFitter.meanX(cb), theFitter.redVar(cb))
                : theKernel(X, theFitter.meanX(cb), theFitter.redVar(cb));
      }
      double P_X_cb_for_TF2_obs_vs_b(double* x, double*)
      {
 
         auto X = x[1];
         auto cb = fIPDist.GetCentrality().Eval(x[0]);
         return x[0] * theKernel(X, theFitter.meanX(cb), theFitter.redVar(cb));
      }
      double P_X_cb_for_integral(double* x, double* par)
      {
         return P_X_cb(par[0], x[0]);
      }
      double P_X_cb_for_X_integral(double* x, double* par)
      {
         return P_X_cb(x[0], par[0]);
      }
      double P_X_cb_for_X_integral_with_selection(double* x, double* par)
      {
 
         int bin = histo->FindBin(x[0]);
         if (bin < 1 || bin > histo->GetNbinsX()) return 0;
         return sel_rapp[bin - 1] * P_X_cb(x[0], par[0]);
      }
      double cb_integrated_P_X(double* x, double* p)
      {
 
         theFitter.fill_params_from_array(p);
         p_X_cb_integrator.SetParameter(0, x[0]); // set value of X in integrand
         return p_X_cb_integrator.Integral(0, 1, 1.e-4);
      }
      double P_X_from_fit(double* x, double* par)
      {
         p_X_cb_integrator.SetParameter(0, x[0]); // set value of X in integrand
         return par[0] * p_X_cb_integrator.Integral(0, 1, 1.e-4);
      }
      double cb_dist_for_X_selection(double* x, double* p)
      {
 
         p_X_X_integrator.SetParameter(0, x[0]);
         double num = p_X_X_integrator.Integral(p[0], p[1], 1.e-4);
         double den = fitted_P_X.Integral(p[0], p[1], 1.e-4);
         if (den > 0) return num / den;
         return 0;
      }
      double cb_dist_for_arb_X_selection(double* x, double* p)
      {
 
         p_X_X_integrator_with_selection.SetParameter(0, x[0]);
         double num = p_X_X_integrator_with_selection.Integral(p[0], p[1], 1.e-4);
 
         return num;
      }
      double b_dist_for_X_selection(double* x, double* p)
      {
 
         p_X_X_integrator.SetParameter(0, fIPDist.GetCentrality().Eval(x[0]));
         double num =  p_X_X_integrator.Integral(p[0], p[1], 1.e-4);
         return num * fIPDist.GetDifferentialCrossSection(x[0]);
      }
      double b_dist_for_arb_X_selection(double* x, double* p)
      {
 
         p_X_X_integrator_with_selection.SetParameter(0, fIPDist.GetCentrality().Eval(x[0]));
         double num =  p_X_X_integrator_with_selection.Integral(p[0], p[1], 1.e-4);
         return num * fIPDist.GetDifferentialCrossSection(x[0]);
      }
      double b_dist_for_arb_X_selection_from_histo(double* x, double* p)
      {
         TH1 *hh_centb  = fIPDist.GetCentralityFromHisto();
         Double_t centb = hh_centb->GetBinContent(hh_centb->FindBin(x[0]));
 
         p_X_X_integrator_with_selection.SetParameter(0, centb);
         double num =  p_X_X_integrator_with_selection.Integral(p[0], p[1], 1.e-4);
         return num * fIPDist.GetDifferentialCrossSectionFromHisto(x[0]);
      }
 
   public:
      bayesian_estimator(Bool_t integer_variable = false)
         : KVBase(),
           theFitter(),
           p_X_cb_integrator("p_X_cb_integrator", this, &bayesian_estimator::P_X_cb_for_integral, 0, 1, 1),
           P_X_fit_function("P_X_fit_function", this, &bayesian_estimator::cb_integrated_P_X, 0, 1, theFitter.npar()),
           mean_X_vs_cb_function("mean_X_vs_cb", this, &bayesian_estimator::mean_X_vs_cb, 0, 1, theFitter.npar()),
           mean_X_vs_b_function("mean_X_vs_b", this, &bayesian_estimator::mean_X_vs_b, 0, 20, theFitter.npar()),
           p_X_X_integrator("p_X_X_integrator", this, &bayesian_estimator::P_X_cb_for_X_integral, 0, 1000, 1),
           p_X_X_integrator_with_selection("p_X_X_integrator_with_selection", this, &bayesian_estimator::P_X_cb_for_X_integral_with_selection, 0, 1000, 1),
           fitted_P_X("fitted_P_X", this, &bayesian_estimator::P_X_from_fit, 0, 1000, 1),
           Cb_dist_for_X_select("Cb_dist_for_X_select", this, &bayesian_estimator::cb_dist_for_X_selection, 0, 1, 2),
           Cb_dist_for_arb_X_select("Cb_dist_for_arb_X_select", this, &bayesian_estimator::cb_dist_for_arb_X_selection, 0, 1, 2),
           B_dist_for_X_select("b_dist_for_X_select", this, &bayesian_estimator::b_dist_for_X_selection, 0, 20, 2),
           B_dist_for_arb_X_select("b_dist_for_arb_X_select", this, &bayesian_estimator::b_dist_for_arb_X_selection, 0, 20, 2),
           B_dist_for_arb_X_select_from_histo("b_dist_for_arb_X_select_from_histo", this, &bayesian_estimator::b_dist_for_arb_X_selection_from_histo, 0, 20, 2),
           fIntegerVariable(integer_variable)
      {
         p_X_cb_integrator.SetParNames("X");
         theFitter.set_par_names(P_X_fit_function);
         theFitter.set_par_names(mean_X_vs_cb_function);
         theFitter.set_par_names(mean_X_vs_b_function);
      }
 
      bayesian_estimator(const FittingFunction& previous_fit, Bool_t integer_variable = false)
         : KVBase(),
           theFitter(previous_fit),
           p_X_cb_integrator("p_X_cb_integrator", this, &bayesian_estimator::P_X_cb_for_integral, 0, 1, 1),
           P_X_fit_function("P_X_fit_function", this, &bayesian_estimator::cb_integrated_P_X, 0, 1, theFitter.npar()),
           mean_X_vs_cb_function("mean_X_vs_cb", this, &bayesian_estimator::mean_X_vs_cb, 0, 1, theFitter.npar()),
           mean_X_vs_b_function("mean_X_vs_b", this, &bayesian_estimator::mean_X_vs_b, 0, 20, theFitter.npar()),
           p_X_X_integrator("p_X_X_integrator", this, &bayesian_estimator::P_X_cb_for_X_integral, 0, 1000, 1),
           p_X_X_integrator_with_selection("p_X_X_integrator_with_selection", this, &bayesian_estimator::P_X_cb_for_X_integral_with_selection, 0, 1000, 1),
           fitted_P_X("fitted_P_X", this, &bayesian_estimator::P_X_from_fit, 0, 1000, 1),
           Cb_dist_for_X_select("Cb_dist_for_X_select", this, &bayesian_estimator::cb_dist_for_X_selection, 0, 1, 2),
           Cb_dist_for_arb_X_select("Cb_dist_for_arb_X_select", this, &bayesian_estimator::cb_dist_for_arb_X_selection, 0, 1, 2),
           B_dist_for_X_select("b_dist_for_X_select", this, &bayesian_estimator::b_dist_for_X_selection, 0, 20, 2),
           B_dist_for_arb_X_select("b_dist_for_arb_X_select", this, &bayesian_estimator::b_dist_for_arb_X_selection, 0, 20, 2),
           B_dist_for_arb_X_select_from_histo("b_dist_for_arb_X_select_from_histo", this, &bayesian_estimator::b_dist_for_arb_X_selection_from_histo, 0, 20, 2),
           fIntegerVariable(integer_variable)
      {
         p_X_cb_integrator.SetParNames("X");
         theFitter.set_par_names(P_X_fit_function);
         theFitter.set_par_names(mean_X_vs_cb_function);
         theFitter.set_par_names(mean_X_vs_b_function);
      }
 
      virtual ~bayesian_estimator() {}
 
      void RenormaliseHisto(TH1* h)
      {
 
         histo = h;
         histo->Scale(1. / h->Integral("width"));
      }
 
      void FitHisto(TH1* h = nullptr)
      {
 
         if (h) histo = h;
         P_X_fit_function.SetRange(histo->GetXaxis()->GetBinLowEdge(1), histo->GetXaxis()->GetBinUpEdge(histo->GetNbinsX()));
         theFitter.set_initial_parameters(histo, P_X_fit_function);
         histo->Fit(&P_X_fit_function);
      }
 
      void SetIPDistParams(double sigmaR, double deltab)
      {
 
         fIPDist.SetDeltaB_WithConstantCrossSection(deltab, sigmaR);
      }
 
      void SetIPDistFromHisto(TH1 *ip_histo)
      {
 
          fIPDist.FitIPDist(ip_histo);
      }
 
      impact_parameter_distribution& GetIPDist()
      {
         return fIPDist;
      }
 
      void DrawMeanXvsCb(const TString& title = "", Color_t color = -1, Option_t* opt = "")
      {
 
         GetMeanXvsCb().Draw();
      }
      TF1& GetMeanXvsCb()
      {
 
         Double_t par[theFitter.npar()];
         theFitter.fill_array_from_params(par);
         mean_X_vs_cb_function.SetParameters(par);
         return mean_X_vs_cb_function;
      }
      TGraph* GraphMeanXvsb(int npts = 500)
      {
 
         Double_t par[theFitter.npar()];
         theFitter.fill_array_from_params(par);
         KVValueRange<double> brange(0, GetIPDist().GetB0() + 4 * GetIPDist().GetDeltaB());
         mean_X_vs_b_function.SetParameters(par);
         auto gr = new TGraph;
         for (int i = 0; i <= npts; ++i) gr->SetPoint(i, brange.ValueIofN(i, npts), mean_X_vs_b_function.Eval(brange.ValueIofN(i, npts)));
         return gr;
      }
      void update_fit_params()
      {
 
         if (!histo) {
            Warning("update_fit_params", "no histogram set with FitHisto(TH1*)");
            return;
         }
         TF1* fit = (TF1*)histo->FindObject("P_X_fit_function");
         if (!fit) {
            Warning("update_fit_params", "no fit function found in histogram");
            return;
         }
         theFitter.fill_params_from_array(fit->GetParameters());
      }
      double DrawCbDistForXSelection(double X1, double X2, Option_t* opt = "", Color_t color = kRed, const TString& title = "")
      {
         Cb_dist_for_X_select.SetParameters(X1, X2);
         TF1* f = Cb_dist_for_X_select.DrawCopy(opt);
         f->SetNpx(500);
         f->SetLineColor(color);
         f->SetMarkerColor(color);
         f->SetLineWidth(2);
         f->SetTitle(title);
         return f->GetMaximum();
      }
      void DrawCbDistForSelection(TH1* sel, TH1* incl, double& mean_cb, double& sigma_cb, Option_t* opt = "", Color_t color = kRed, const TString& title = "")
      {
 
         assert(sel->GetNbinsX() == incl->GetNbinsX());
 
         sel_rapp.assign((std::vector<double>::size_type)incl->GetNbinsX(), 0.0);
         int first_bin(0), last_bin(0);
         for (int i = 1; i <= incl->GetNbinsX(); ++i) {
            if (incl->GetBinContent(i) > 0) {
               sel_rapp[i - 1] = sel->GetBinContent(i) / incl->GetBinContent(i);
               if (sel->GetBinContent(i) > 0) {
                  if (!first_bin) first_bin = i;
                  last_bin = i;
               }
            }
         }
         double Xmin = incl->GetBinLowEdge(first_bin);
         double Xmax = incl->GetXaxis()->GetBinUpEdge(last_bin);
 
         histo = incl;
         h_selection = sel;
 
         Cb_dist_for_arb_X_select.SetParameters(Xmin, Xmax);
         Cb_dist_for_arb_X_select.SetRange(0., 1.);
         Cb_dist_for_arb_X_select.GetHistogram();
 
         double cb_mean(0), cb_sqrmean(0), sum_pcb(0);
         TGraph* f = new TGraph;
         for (int i = 0; i < 500; ++i) {//graph with 500 points
            double cb  = i / 499.;
            double p_cb = Cb_dist_for_arb_X_select.Eval(cb);
 
            cb_mean    += p_cb * cb;
            cb_sqrmean += p_cb * cb * cb;
            sum_pcb    += p_cb;
 
            f->SetPoint(i, cb, p_cb);
         }
 
         cb_sqrmean /= sum_pcb;
         cb_mean /= sum_pcb;
         mean_cb = cb_mean;
         sigma_cb = TMath::Sqrt(cb_sqrmean - cb_mean * cb_mean);
 
         f->SetLineColor(color);
         f->SetMarkerColor(color);
         f->SetLineWidth(2);
         f->SetTitle(title);
         if (TString(opt) == "same") f->Draw("l");
         else f->Draw("al");
      }
      double DrawBDistForXSelection(KVValueRange<double> Xrange, Option_t* opt = "", Color_t color = kRed, const TString& title = "")
      {
 
         B_dist_for_X_select.SetParameters(Xrange.Min(), Xrange.Max());
         TGraph* f = new TGraph;
         double maxS = 0;
         for (int i = 0; i < 500; ++i) {
            double b = 2 * i * GetIPDist().GetB0() / 499.;
            double sig = B_dist_for_X_select.Eval(b);
            if (sig > maxS) maxS = sig;
            f->SetPoint(i, b, sig);
         }
         f->SetLineColor(color);
         f->SetMarkerColor(color);
         f->SetLineWidth(2);
         f->SetTitle(title);
         if (TString(opt) == "same") f->Draw("l");
         else f->Draw("al");
         return maxS;
      }
      TF1& GetB_dist_for_X_select()
      {
         return B_dist_for_X_select;
      }
      TGraph* GraphBDistForXSelection(KVValueRange<double> Xrange, int npts = 500)
      {
 
 
         B_dist_for_X_select.SetParameters(Xrange.Min(), Xrange.Max());
         TGraph* f = new TGraph;
         for (int i = 0; i < npts; ++i) {
            double b = 2 * i * GetIPDist().GetB0() / 499.;
            double sig = B_dist_for_X_select.Eval(b);
            f->SetPoint(i, b, sig);
         }
         f->SetLineWidth(2);
         return f;
      }
      void DrawBDistForSelection(TH1* sel, TH1* incl, double& mean, double& sigma, Option_t* opt = "", Color_t color = kRed, const TString& title = "")
      {
 
 
         assert(sel->GetNbinsX() == incl->GetNbinsX());
 
         sel_rapp.assign((std::vector<double>::size_type)incl->GetNbinsX(), 0.0);
         int first_bin(0), last_bin(0);
         for (int i = 1; i <= incl->GetNbinsX(); ++i) {
            if (incl->GetBinContent(i) > 0) {
               sel_rapp[i - 1] = sel->GetBinContent(i) / incl->GetBinContent(i);
               if (sel->GetBinContent(i) > 0) {
                  if (!first_bin) first_bin = i;
                  last_bin = i;
               }
            }
         }
         double Xmin = incl->GetBinLowEdge(first_bin);
         double Xmax = incl->GetXaxis()->GetBinUpEdge(last_bin);
 
         histo = incl;
         h_selection = sel;
 
         B_dist_for_arb_X_select.SetParameters(Xmin, Xmax);
         B_dist_for_arb_X_select.SetRange(0, 2 * GetIPDist().GetB0());
         B_dist_for_arb_X_select.GetHistogram();
 
         double bmean(0), bsqrmean(0), sigtot(0);
         TGraph* f = new TGraph;
         for (int i = 0; i < 500; ++i) {
            double b = 2 * i * GetIPDist().GetB0() / 499.;
            double sig = B_dist_for_arb_X_select.Eval(b);
            bmean += sig * b;
            bsqrmean += sig * b * b;
            sigtot += sig;
            f->SetPoint(i, b, sig);
         }
         f->SetLineColor(color);
         f->SetMarkerColor(color);
         f->SetLineWidth(2);
         f->SetTitle(title);
         mean = bmean / sigtot;
         bsqrmean /= sigtot;
         sigma = TMath::Sqrt(bsqrmean - mean * mean);
         if (TString(opt) == "same") f->Draw("l");
         else f->Draw("al");
      }
 
      void DrawBDistForSelectionFromIPHisto(TH1* sel, TH1* incl, double& mean, double& sigma, Option_t* opt = "", Color_t color = kRed, const TString& title = "")
      {
 
         assert(sel->GetNbinsX() == incl->GetNbinsX());
 
         sel_rapp.assign((std::vector<double>::size_type)incl->GetNbinsX(), 0.0);
         int first_bin(0), last_bin(0);
         for (int i = 1; i <= incl->GetNbinsX(); ++i) {
            if (incl->GetBinContent(i) > 0) {
               sel_rapp[i - 1] = sel->GetBinContent(i) / incl->GetBinContent(i);
               if (sel->GetBinContent(i) > 0) {
                  if (!first_bin) first_bin = i;
                  last_bin = i;
               }
            }
         }
         double Xmin = incl->GetBinLowEdge(first_bin);
         double Xmax = incl->GetXaxis()->GetBinUpEdge(last_bin);
 
         histo = incl;
         h_selection = sel;
 
         B_dist_for_arb_X_select_from_histo.SetParameters(Xmin, Xmax);
         B_dist_for_arb_X_select_from_histo.SetRange(0, 2 * GetIPDist().GetB0()); //if SetIPDistFromHisto() was called, a fit have been applied thus B0 was estimated
         B_dist_for_arb_X_select_from_histo.GetHistogram();
 
         double bmean(0), bsqrmean(0), sigtot(0);
         TGraph* f = new TGraph;
         for (int i = 0; i < 500; ++i) {
            double b = 2 * i * GetIPDist().GetB0() / 499.;
            double sig = B_dist_for_arb_X_select_from_histo.Eval(b);
            bmean += sig * b;
            bsqrmean += sig * b * b;
            sigtot += sig;
            f->SetPoint(i, b, sig);
         }
         f->SetLineColor(color);
         f->SetMarkerColor(color);
         f->SetLineWidth(2);
         f->SetTitle(title);
         mean = bmean / sigtot;
         bsqrmean /= sigtot;
         sigma = TMath::Sqrt(bsqrmean - mean * mean);
         if (TString(opt) == "same") f->Draw("l");
         else f->Draw("al");
      }
 
      void GetMeanAndSigmaBDistForSelection(TH1* sel, TH1* incl, double& mean, double& sigma)
      {
 
 
         assert(sel->GetNbinsX() == incl->GetNbinsX());
 
         sel_rapp.assign((std::vector<double>::size_type)incl->GetNbinsX(), 0.0);
         int first_bin(0), last_bin(0);
         for (int i = 1; i <= incl->GetNbinsX(); ++i) {
            if (incl->GetBinContent(i) > 0) {
               sel_rapp[i - 1] = sel->GetBinContent(i) / incl->GetBinContent(i);
               if (sel->GetBinContent(i) > 0) {
                  if (!first_bin) first_bin = i;
                  last_bin = i;
               }
            }
         }
         double Xmin = incl->GetBinLowEdge(first_bin);
         double Xmax = incl->GetXaxis()->GetBinUpEdge(last_bin);
 
         histo = incl;
         h_selection = sel;
 
         B_dist_for_arb_X_select.SetParameters(Xmin, Xmax);
         mean = B_dist_for_arb_X_select.Mean(0, 20);
         double var = B_dist_for_arb_X_select.Variance(0, 20);
         sigma = TMath::Sqrt(var);
      }
 
      void GetMeanAndSigmaBDistForXSelection(double X1, double X2, double& mean, double& sigma)
      {
 
 
         B_dist_for_X_select.SetParameters(X1, X2);
         mean = B_dist_for_X_select.Mean(0, 20);
         sigma = TMath::Sqrt(B_dist_for_X_select.Variance(0, 20));
      }
      void DrawFittedP_X(double norm = 1.0, Option_t* opt = "", Color_t color = kRed, const TString& title = "")
      {
 
         TF1* f = GetFittedP_X(norm)->DrawCopy(opt);
         f->SetNpx(500);
         f->SetLineColor(color);
         f->SetLineWidth(2);
         f->SetTitle(title);
      }
      TF1* GetFittedP_X(double norm = 1.0)
      {
         fitted_P_X.SetParameter(0, norm);
         return &fitted_P_X;
      }
      TGraph* GraphP_XForGivenB(double b, KVValueRange<double> Xrange, int npts = 500)
      {
         p_X_X_integrator.SetParameter(0, fIPDist.GetCentrality().Eval(b));
         auto graph = new TGraph;
         for (int i = 0; i < npts; ++i) graph->SetPoint(i, Xrange.ValueIofN(i, npts), p_X_X_integrator.Eval(Xrange.ValueIofN(i, npts)));
         return graph;
      }
      void Print(Option_t* = "") const
      {
         theFitter.print_fit_params();
      }
      void DrawNormalisedMeanXvsb(const TString& title, Color_t color, Option_t* opt)
      {
 
 
         theFitter.backup_params();
         theFitter.normalise_shape_function();
         Double_t par[5];
         theFitter.fill_array_from_params(par);
         mean_X_vs_b_function.SetRange(0, GetIPDist().GetB0() + 2 * GetIPDist().GetDeltaB());
         mean_X_vs_b_function.SetParameters(par);
         TF1* copy = mean_X_vs_b_function.DrawCopy(opt);
         if (!title.IsNull()) copy->SetTitle(title);
         if (color >= 0) copy->SetLineColor(color);
         theFitter.restore_params();
      }
 
      TF2* GetJointProbabilityDistribution(KVValueRange<double> b_range, KVValueRange<double> X_range)
      {
 
         auto jpd = new TF2("joint_proba_dist", this, &bayesian_estimator::P_X_cb_for_TF2_obs_vs_b,
                            b_range.Min(), b_range.Max(), X_range.Min(), X_range.Max(), 0);
         return jpd;
      }
 
      ClassDef(bayesian_estimator, 1) //Estimate impact parameter distribution by fits to data
   };
 
}
 
#endif