]> git.uio.no Git - u/mrichter/AliRoot.git/blobdiff - CORRFW/AliCFUnfolding.cxx
Enabling creation and reading of the CDB snapshot also in simulation
[u/mrichter/AliRoot.git] / CORRFW / AliCFUnfolding.cxx
index 88cec41b363af2ab0022bf0980b5aefc6bdfd7d2..8a998e8acfbf55832781689aa413cf9502d710ad 100644 (file)
  * provided "as is" without express or implied warranty.                  *
  **************************************************************************/
 
-//--------------------------------------------------------------------//
-//                                                                    //
-// AliCFUnfolding Class                                               //
-// Class to handle general unfolding procedure                        // 
-// For the moment only bayesian unfolding is supported                //
-// The next steps are to add chi2 minimisation and weighting methods  //
-//                                                                    //
-// Author : renaud.vernet@cern.ch                                     //
-//--------------------------------------------------------------------//
+//---------------------------------------------------------------------//
+//                                                                     //
+// AliCFUnfolding Class                                                //
+// Class to handle general unfolding procedure                         // 
+// For the moment only bayesian unfolding is supported                 //
+// The next steps are to add chi2 minimisation and weighting methods   //
+//                                                                     //
+//                                                                     //
+//                                                                     //
+// Use :                                                               //
+// -------                                                             //
+// The Bayesian unfolding consists of several iterations.              //
+// At each iteration, an inverse response matrix is calculated, given  //
+// the measured spectrum, the a priori (guessed) spectrum,             //
+// the efficiency spectrum and the response matrix.                    //
+//                                                                     //
+// Then at each iteration, the unfolded spectrum is calculated using   //
+// the inverse response : the goal is to get an unfolded spectrum      //
+// similar (according to some criterion) to the a priori one.          //
+// If the difference is too big, another iteration is performed :      //
+// the a priori spectrum is updated to the unfolded one from the       //
+// previous iteration, and so on so forth, until the maximum number    //
+// of iterations or the similarity criterion is reached.               //
+//                                                                     //
+// Currently the similarity criterion is the Chi2 between the a priori //
+// and the unfolded spectrum (OBSOLETE).                               //
+//                                                                     //
+// Chi2 calculation became absolute with the correlated error          //
+// calculation.                                                        //
+// Errors on the unfolded distribution are not known until the end     //
+// Use the convergence criterion instead                               //
+//                                                                     //
+// Currently the user has to define the max. number of iterations      //
+// (::SetMaxNumberOfIterations)                                        //
+// and                                                                 //
+//     - the chi2 below which the procedure will stop                  //
+// (::SetMaxChi2 or ::SetMaxChi2PerDOF)   (OBSOLETE)                   //
+//     - the convergence criterion below which the procedure will stop //
+// SetMaxConvergencePerDOF(Double_t val);                              //
+//                                                                     //
+// Correlated error calculation can be activated by using:             //
+// SetUseCorrelatedErrors(Bool_t b) in combination with convergence    //
+// criterion                                                           //
+// Documentation about correlated error calculation method can be      //
+// found in AliCFUnfolding::CalculateCorrelatedErrors()                //
+// Author: marta.verweij@cern.ch                                       //
+//                                                                     //
+// An optional possibility is to smooth the unfolded spectrum at the   //
+// end of each iteration, either using a fit function                  //
+// (only if #dimensions <=3)                                           //
+// or a simple averaging using the neighbouring bins values.           //
+// This is possible calling the function ::UseSmoothing                //
+// If no argument is passed to this function, then the second option   //
+// is used.                                                            //
+//                                                                     //
+// IMPORTANT:                                                          //
+//-----------                                                          //
+// With this approach, the efficiency map must be calculated           //
+// with *simulated* values only, otherwise the method won't work.      //
+//                                                                     //
+// ex: efficiency(bin_pt) = number_rec(bin_pt) / number_sim(bin_pt)    //
+//                                                                     //
+// the pt bin "bin_pt" must always be the same in both the efficiency  //
+// numerator and denominator.                                          //
+// This is why the efficiency map has to be created by a method        //
+// from which both reconstructed and simulated values are accessible   //
+// simultaneously.                                                     //
+//                                                                     //
+//                                                                     //
+//---------------------------------------------------------------------//
+// Author : renaud.vernet@cern.ch                                      //
+//---------------------------------------------------------------------//
 
 
 #include "AliCFUnfolding.h"
 #include "TMath.h"
 #include "TAxis.h"
 #include "AliLog.h"
+#include "TF1.h"
+#include "TH1D.h"
+#include "TH2D.h"
+#include "TH3D.h"
+#include "TProfile.h"
+#include "TRandom3.h"
 
 ClassImp(AliCFUnfolding)
 
@@ -37,13 +106,18 @@ AliCFUnfolding::AliCFUnfolding() :
   TNamed(),
   fResponse(0x0),
   fPrior(0x0),
-  fOriginalPrior(0x0),
   fEfficiency(0x0),
   fMeasured(0x0),
-  fMaxNumIterations(0),
+  fMeasuredOrig(0x0),
+  fMaxNumIterations(20),
   fNVariables(0),
-  fMaxChi2(0),
   fUseSmoothing(kFALSE),
+  fSmoothFunction(0x0),
+  fSmoothOption(""),
+  fMaxConvergence(0),
+  fUseCorrelatedErrors(kTRUE),
+  fNRandomIterations(20),
+  fOriginalPrior(0x0),
   fInverseResponse(0x0),
   fMeasuredEstimate(0x0),
   fConditional(0x0),
@@ -51,7 +125,13 @@ AliCFUnfolding::AliCFUnfolding() :
   fUnfolded(0x0),
   fCoordinates2N(0x0),
   fCoordinatesN_M(0x0),
-  fCoordinatesN_T(0x0)
+  fCoordinatesN_T(0x0),
+  fRandomizedDist(0x0),
+  fRandom3(0x0),
+  fRandomUnfolded(0x0),
+  fDeltaUnfoldedP(0x0),
+  fNCalcCorrErrors(0),
+  fRandomSeed(0)
 {
   //
   // default constructor
@@ -65,13 +145,18 @@ AliCFUnfolding::AliCFUnfolding(const Char_t* name, const Char_t* title, const In
   TNamed(name,title),
   fResponse((THnSparse*)response->Clone()),
   fPrior(0x0),
-  fOriginalPrior(0x0),
   fEfficiency((THnSparse*)efficiency->Clone()),
   fMeasured((THnSparse*)measured->Clone()),
+  fMeasuredOrig((THnSparse*)measured->Clone()),
   fMaxNumIterations(0),
   fNVariables(nVar),
-  fMaxChi2(0),
   fUseSmoothing(kFALSE),
+  fSmoothFunction(0x0),
+  fSmoothOption(""),
+  fMaxConvergence(0),
+  fUseCorrelatedErrors(kTRUE),
+  fNRandomIterations(20),
+  fOriginalPrior(0x0),
   fInverseResponse(0x0),
   fMeasuredEstimate(0x0),
   fConditional(0x0),
@@ -79,7 +164,13 @@ AliCFUnfolding::AliCFUnfolding(const Char_t* name, const Char_t* title, const In
   fUnfolded(0x0),
   fCoordinates2N(0x0),
   fCoordinatesN_M(0x0),
-  fCoordinatesN_T(0x0)
+  fCoordinatesN_T(0x0),
+  fRandomizedDist(0x0),
+  fRandom3(0x0),
+  fRandomUnfolded(0x0),
+  fDeltaUnfoldedP(0x0),
+  fNCalcCorrErrors(0),
+  fRandomSeed(0)
 {
   //
   // named constructor
@@ -91,13 +182,24 @@ AliCFUnfolding::AliCFUnfolding(const Char_t* name, const Char_t* title, const In
   else {
     fPrior = (THnSparse*) prior->Clone();
     fOriginalPrior = (THnSparse*)fPrior->Clone();
+    if (fPrior->GetNdimensions() != fNVariables) 
+      AliFatal(Form("The prior matrix should have %d dimensions, and it has actually %d",fNVariables,fPrior->GetNdimensions()));
   }
+
+  if (fEfficiency->GetNdimensions() != fNVariables) 
+    AliFatal(Form("The efficiency matrix should have %d dimensions, and it has actually %d",fNVariables,fEfficiency->GetNdimensions()));
+  if (fMeasured->GetNdimensions() != fNVariables) 
+    AliFatal(Form("The measured matrix should have %d dimensions, and it has actually %d",fNVariables,fMeasured->GetNdimensions()));
+  if (fResponse->GetNdimensions() != 2*fNVariables) 
+    AliFatal(Form("The response matrix should have %d dimensions, and it has actually %d",2*fNVariables,fResponse->GetNdimensions()));
   
+
   for (Int_t iVar=0; iVar<fNVariables; iVar++) {
     AliInfo(Form("prior      matrix has %d bins in dimension %d",fPrior     ->GetAxis(iVar)->GetNbins(),iVar));
     AliInfo(Form("efficiency matrix has %d bins in dimension %d",fEfficiency->GetAxis(iVar)->GetNbins(),iVar));
     AliInfo(Form("measured   matrix has %d bins in dimension %d",fMeasured  ->GetAxis(iVar)->GetNbins(),iVar));
   }
+
   Init();
 }
 
@@ -108,13 +210,18 @@ AliCFUnfolding::AliCFUnfolding(const AliCFUnfolding& c) :
   TNamed(c),
   fResponse((THnSparse*)c.fResponse->Clone()),
   fPrior((THnSparse*)c.fPrior->Clone()),
-  fOriginalPrior((THnSparse*)c.fOriginalPrior->Clone()),
   fEfficiency((THnSparse*)c.fEfficiency->Clone()),
   fMeasured((THnSparse*)c.fMeasured->Clone()),
+  fMeasuredOrig((THnSparse*)c.fMeasuredOrig->Clone()),
   fMaxNumIterations(c.fMaxNumIterations),
   fNVariables(c.fNVariables),
-  fMaxChi2(c.fMaxChi2),
   fUseSmoothing(c.fUseSmoothing),
+  fSmoothFunction((TF1*)c.fSmoothFunction->Clone()),
+  fSmoothOption(c.fSmoothOption),
+  fMaxConvergence(c.fMaxConvergence),
+  fUseCorrelatedErrors(c.fUseCorrelatedErrors),
+  fNRandomIterations(c.fNRandomIterations),
+  fOriginalPrior((THnSparse*)c.fOriginalPrior->Clone()),
   fInverseResponse((THnSparse*)c.fInverseResponse->Clone()),
   fMeasuredEstimate((THnSparse*)fMeasuredEstimate->Clone()),
   fConditional((THnSparse*)c.fConditional->Clone()),
@@ -122,7 +229,13 @@ AliCFUnfolding::AliCFUnfolding(const AliCFUnfolding& c) :
   fUnfolded((THnSparse*)c.fUnfolded->Clone()),
   fCoordinates2N(new Int_t(*c.fCoordinates2N)),
   fCoordinatesN_M(new Int_t(*c.fCoordinatesN_M)),
-  fCoordinatesN_T(new Int_t(*c.fCoordinatesN_T))
+  fCoordinatesN_T(new Int_t(*c.fCoordinatesN_T)),
+  fRandomizedDist((THnSparse*)c.fRandomizedDist->Clone()),
+  fRandom3((TRandom3*)c.fRandom3->Clone()),
+  fRandomUnfolded((THnSparse*)c.fRandomUnfolded->Clone()),
+  fDeltaUnfoldedP((TProfile*)c.fDeltaUnfoldedP),
+  fNCalcCorrErrors(c.fNCalcCorrErrors),
+  fRandomSeed(c.fRandomSeed)
 {
   //
   // copy constructor
@@ -140,13 +253,18 @@ AliCFUnfolding& AliCFUnfolding::operator=(const AliCFUnfolding& c) {
     TNamed::operator=(c);
     fResponse = (THnSparse*)c.fResponse->Clone() ;
     fPrior = (THnSparse*)c.fPrior->Clone() ;
-    fOriginalPrior = (THnSparse*)c.fOriginalPrior->Clone() ;
     fEfficiency = (THnSparse*)c.fEfficiency->Clone() ;
     fMeasured = (THnSparse*)c.fMeasured->Clone() ;
+    fMeasuredOrig = ((THnSparse*)c.fMeasuredOrig->Clone()),
     fMaxNumIterations = c.fMaxNumIterations ;
     fNVariables = c.fNVariables ;
-    fMaxChi2 = c.fMaxChi2 ;
+    fMaxConvergence = c.fMaxConvergence ;
     fUseSmoothing = c.fUseSmoothing ;
+    fSmoothFunction = (TF1*)c.fSmoothFunction->Clone();
+    fSmoothOption = c.fSmoothOption ;
+    fUseCorrelatedErrors = c.fUseCorrelatedErrors;
+    fNRandomIterations = c.fNRandomIterations;
+    fOriginalPrior = (THnSparse*)c.fOriginalPrior->Clone() ;
     fInverseResponse = (THnSparse*)c.fInverseResponse->Clone() ;
     fMeasuredEstimate = (THnSparse*)fMeasuredEstimate->Clone() ;
     fConditional = (THnSparse*)c.fConditional->Clone() ;
@@ -155,6 +273,12 @@ AliCFUnfolding& AliCFUnfolding::operator=(const AliCFUnfolding& c) {
     fCoordinates2N  = new Int_t(*c.fCoordinates2N)  ;
     fCoordinatesN_M = new Int_t(*c.fCoordinatesN_M) ;
     fCoordinatesN_T = new Int_t(*c.fCoordinatesN_T) ;
+    fRandomizedDist = (THnSparse*)c.fRandomizedDist->Clone();
+    fRandom3 = (TRandom3*)c.fRandom3->Clone();
+    fRandomUnfolded = (THnSparse*)c.fRandomUnfolded->Clone();
+    fDeltaUnfoldedP = (TProfile*)c.fDeltaUnfoldedP;
+    fNCalcCorrErrors = c.fNCalcCorrErrors;
+    fRandomSeed = c.fRandomSeed ;
   }
   return *this;
 }
@@ -167,9 +291,11 @@ AliCFUnfolding::~AliCFUnfolding() {
   //
   if (fResponse)           delete fResponse;
   if (fPrior)              delete fPrior;
-  if (fOriginalPrior)      delete fOriginalPrior;
   if (fEfficiency)         delete fEfficiency;
   if (fMeasured)           delete fMeasured;
+  if (fMeasuredOrig)       delete fMeasuredOrig;
+  if (fSmoothFunction)     delete fSmoothFunction;
+  if (fOriginalPrior)      delete fOriginalPrior;
   if (fInverseResponse)    delete fInverseResponse;
   if (fMeasuredEstimate)   delete fMeasuredEstimate;
   if (fConditional)        delete fConditional;
@@ -177,6 +303,10 @@ AliCFUnfolding::~AliCFUnfolding() {
   if (fCoordinates2N)      delete [] fCoordinates2N; 
   if (fCoordinatesN_M)     delete [] fCoordinatesN_M; 
   if (fCoordinatesN_T)     delete [] fCoordinatesN_T; 
+  if (fRandomizedDist)     delete fRandomizedDist;
+  if (fRandom3)            delete fRandom3;
+  if (fRandomUnfolded)     delete fRandomUnfolded;
+  if (fDeltaUnfoldedP)     delete fDeltaUnfoldedP;
 }
 
 //______________________________________________________________
@@ -186,6 +316,8 @@ void AliCFUnfolding::Init() {
   // initialisation function : creates internal settings
   //
 
+  fRandom3 = new TRandom3(fRandomSeed);
+
   fCoordinates2N  = new Int_t[2*fNVariables];
   fCoordinatesN_M = new Int_t[fNVariables];
   fCoordinatesN_T = new Int_t[fNVariables];
@@ -197,8 +329,15 @@ void AliCFUnfolding::Init() {
   fInverseResponse  = (THnSparse*) fResponse->Clone();
   // create the frame of the unfolded spectrum
   fUnfolded = (THnSparse*) fPrior->Clone();
+  // create the frame of the random unfolded spectrum
+  fRandomUnfolded = (THnSparse*) fPrior->Clone();
   // create the frame of the measurement estimate spectrum
   fMeasuredEstimate = (THnSparse*) fMeasured->Clone();
+  // create the frame of the original measurement spectrum
+  fMeasuredOrig = (THnSparse*) fMeasured->Clone();
+
+  InitDeltaUnfoldedProfile();
+
 }
 
 //______________________________________________________________
@@ -206,29 +345,34 @@ void AliCFUnfolding::Init() {
 void AliCFUnfolding::CreateEstMeasured() {
   //
   // This function creates a estimate (M) of the reconstructed spectrum 
-  // given the a priori distribution (T) and the conditional matrix (COND)
+  // given the a priori distribution (T), the efficiency (E) and the conditional matrix (COND)
   //
   // --> P(M) = SUM   { P(M|T)    * P(T) }
-  // --> M(i) = SUM_k { COND(i,k) * T(k) }
+  // --> M(i) = SUM_k { COND(i,k) * T(k) * E (k)}
   //
   // This is needed to calculate the inverse response matrix
   //
 
 
   // clean the measured estimate spectrum
-  for (Long64_t i=0; i<fMeasuredEstimate->GetNbins(); i++) {
-    fMeasuredEstimate->GetBinContent(i,fCoordinatesN_M);
-    fMeasuredEstimate->SetBinContent(fCoordinatesN_M,0.);
-  }
-  
+  fMeasuredEstimate->Reset();
+
+  THnSparse* priorTimesEff = (THnSparse*) fPrior->Clone();
+  priorTimesEff->Multiply(fEfficiency);
+
   // fill it
-  for (Int_t iBin=0; iBin<fConditional->GetNbins(); iBin++) {
+  for (Long_t iBin=0; iBin<fConditional->GetNbins(); iBin++) {
     Double_t conditionalValue = fConditional->GetBinContent(iBin,fCoordinates2N);
     GetCoordinates();
-    Double_t priorValue = fPrior->GetBinContent(fCoordinatesN_T);
-    Double_t fill = fMeasuredEstimate->GetBinContent(fCoordinatesN_M) + conditionalValue * priorValue * fEfficiency->GetBinContent(fCoordinatesN_T);
-    if (fill>0.) fMeasuredEstimate->SetBinContent(fCoordinatesN_M,fill);
+    Double_t priorTimesEffValue = priorTimesEff->GetBinContent(fCoordinatesN_T);
+    Double_t fill = conditionalValue * priorTimesEffValue ;
+    
+    if (fill>0.) {
+      fMeasuredEstimate->AddBinContent(fCoordinatesN_M,fill);
+      fMeasuredEstimate->SetBinError(fCoordinatesN_M,0.);
+    }
   }
+  delete priorTimesEff ;
 }
 
 //______________________________________________________________
@@ -242,14 +386,21 @@ void AliCFUnfolding::CreateInvResponse() {
   // --> INV(i,j) = COND(i,j) * T(j) * E(j)   / SUM_k { COND(i,k) * T(k) }
   //
 
-  for (Int_t iBin=0; iBin<fConditional->GetNbins(); iBin++) {
+  THnSparse* priorTimesEff = (THnSparse*) fPrior->Clone();
+  priorTimesEff->Multiply(fEfficiency);
+
+  for (Long_t iBin=0; iBin<fConditional->GetNbins(); iBin++) {
     Double_t conditionalValue = fConditional->GetBinContent(iBin,fCoordinates2N);
     GetCoordinates();
-    Double_t priorValue = fPrior->GetBinContent(fCoordinatesN_T);
-    Double_t estimatedMeasured = fMeasuredEstimate->GetBinContent(fCoordinatesN_M);
-    Double_t fill = (estimatedMeasured>0. ? conditionalValue * priorValue * fEfficiency->GetBinContent(fCoordinatesN_T) / estimatedMeasured : 0. ) ;
-    if (fill>0. || fInverseResponse->GetBinContent(fCoordinates2N)>0.) fInverseResponse->SetBinContent(fCoordinates2N,fill);
-  }
+    Double_t estMeasuredValue   = fMeasuredEstimate->GetBinContent(fCoordinatesN_M);
+    Double_t priorTimesEffValue = priorTimesEff    ->GetBinContent(fCoordinatesN_T);
+    Double_t fill = (estMeasuredValue>0. ? conditionalValue * priorTimesEffValue / estMeasuredValue : 0. ) ;
+    if (fill>0. || fInverseResponse->GetBinContent(fCoordinates2N)>0.) {
+      fInverseResponse->SetBinContent(fCoordinates2N,fill);
+      fInverseResponse->SetBinError  (fCoordinates2N,0.);
+    }
+  } 
+  delete priorTimesEff ;
 }
 
 //______________________________________________________________
@@ -257,27 +408,66 @@ void AliCFUnfolding::CreateInvResponse() {
 void AliCFUnfolding::Unfold() {
   //
   // Main routine called by the user : 
-  // it calculates the unfolded spectrum from the response matrix and the measured spectrum
-  // several iterations are performed until a reasonable chi2 is reached
+  // it calculates the unfolded spectrum from the response matrix, measured spectrum and efficiency
+  // several iterations are performed until a reasonable chi2 or convergence criterion is reached
   //
 
-  Int_t iIterBayes=0 ;
-  Double_t chi2=0 ;
+  Int_t iIterBayes     = 0 ;
+  Double_t convergence = 0.;
 
   for (iIterBayes=0; iIterBayes<fMaxNumIterations; iIterBayes++) { // bayes iterations
     CreateEstMeasured();
     CreateInvResponse();
     CreateUnfolded();
-    chi2 = GetChi2();
-    //printf("chi2 = %e\n",chi2);
-    if (fMaxChi2>0. && chi2<fMaxChi2) {
+    if (fUseCorrelatedErrors) {
+      convergence = GetConvergence();
+      AliDebug(0,Form("convergence at iteration %d is %e",iIterBayes,convergence));
+    } 
+    else AliWarning("No errors will be calculated. Activate SetUseCorrelatedErrors(kTRUE)\n");
+
+    if (fMaxConvergence>0. && convergence<fMaxConvergence && fNCalcCorrErrors<1) {
+      fNRandomIterations=iIterBayes;
       break;
     }
+
     // update the prior distribution
-    if (fUseSmoothing) Smooth();
-    fPrior = (THnSparse*)fUnfolded->Clone() ; // this should be changed (memory)
+    if (fUseSmoothing) {
+      if (Smooth()) {
+       AliError("Couldn't smooth the unfolded spectrum!!");
+       if (fUseCorrelatedErrors) {
+         if (fNCalcCorrErrors>0) {
+           AliInfo(Form("=======================\nUnfolding of randomized distribution finished at iteration %d with convergence %e \n",iIterBayes,convergence));
+         }
+         else {
+           AliInfo(Form("\n\n=======================\nFinished at iteration %d : convergence is %e and you required it to be < %e\n=======================\n\n",iIterBayes,convergence,fMaxConvergence));
+         }
+       }
+       return;
+      }
+    }
+    if(fNCalcCorrErrors>0) {
+      if (fPrior) delete fPrior ;
+      fPrior = (THnSparse*)fRandomUnfolded->Clone() ;
+    } 
+    else {
+      if (fPrior) delete fPrior ;
+      fPrior = (THnSparse*)fUnfolded->Clone() ;
+    }
+  }
+
+  if (fUseCorrelatedErrors && fNCalcCorrErrors==0) {
+    fNCalcCorrErrors=1;
+    CalculateCorrelatedErrors();
+  }
+
+  if (fUseCorrelatedErrors) {
+    if (fNCalcCorrErrors>1) {
+      AliInfo(Form("\n\n=======================\nFinished at iteration %d : convergence is %e and you required it to be < %e\n=======================\n\n",iIterBayes,convergence,fMaxConvergence));
+    }
+    else if(fNCalcCorrErrors>0) {
+      AliInfo(Form("=======================\nUnfolding of randomized distribution finished at iteration %d with convergence %e \n",iIterBayes,convergence));
+    }
   }
-  AliInfo(Form("Finished at iteration %d : Chi2 is %e and you required it to be < %e",iIterBayes,chi2,fMaxChi2));
 }
 
 //______________________________________________________________
@@ -291,20 +481,125 @@ void AliCFUnfolding::CreateUnfolded() {
 
 
   // clear the unfolded spectrum
-  for (Long64_t i=0; i<fUnfolded->GetNbins(); i++) {
-    fUnfolded->GetBinContent(i,fCoordinatesN_T);
-    fUnfolded->SetBinContent(fCoordinatesN_T,0.);
+  if(fNCalcCorrErrors>0) {
+    //unfold randomized dist
+    fRandomUnfolded->Reset();
+  }
+  else {  
+    //unfold measured dist
+    fUnfolded->Reset();
   }
   
-  for (Int_t iBin=0; iBin<fInverseResponse->GetNbins(); iBin++) {
+  for (Long_t iBin=0; iBin<fInverseResponse->GetNbins(); iBin++) {
     Double_t invResponseValue = fInverseResponse->GetBinContent(iBin,fCoordinates2N);
     GetCoordinates();
-    Double_t effValue = fEfficiency->GetBinContent(fCoordinatesN_T);
-    Double_t fill = fUnfolded->GetBinContent(fCoordinatesN_T) + (effValue>0. ? invResponseValue*fMeasured->GetBinContent(fCoordinatesN_M)/effValue : 0.) ;
-    if (fill>0.) fUnfolded->SetBinContent(fCoordinatesN_T,fill);
+    Double_t effValue      = fEfficiency->GetBinContent(fCoordinatesN_T);
+    Double_t measuredValue = fMeasured  ->GetBinContent(fCoordinatesN_M);
+    Double_t fill = (effValue>0. ? invResponseValue * measuredValue / effValue : 0.) ;
+
+    if (fill>0.) {
+        
+      Double_t err = 0.;
+      
+      if(fNCalcCorrErrors>0) {
+       fRandomUnfolded->SetBinError(fCoordinatesN_T,err);
+       fRandomUnfolded->AddBinContent(fCoordinatesN_T,fill);
+      }
+      else {
+       fUnfolded->SetBinError(fCoordinatesN_T,err);
+       fUnfolded->AddBinContent(fCoordinatesN_T,fill);
+      }
+    }
   }
 }
-    
+
+//______________________________________________________________
+
+void AliCFUnfolding::CalculateCorrelatedErrors() {
+
+  fRandomizedDist     = (THnSparse*) fMeasuredOrig->Clone(); 
+  fPrior              = (THnSparse*) fOriginalPrior->Clone();
+
+  // Step 1: Create randomized distribution (fRandomizedDist) of each bin of the measured spectrum to calculate correlated errors. Poisson statistics: mean = measured value of bin
+  // Step 2: Unfold randomized distribution
+  // Step 3: Store difference of unfolded spectrum from measured distribution and unfolded distribution from randomized distribution -> fDeltaUnfoldedP (TProfile with option "S")
+  // Step 4: Repeat Step 1-3 several times (fNRandomIterations)
+  // Step 5: The spread of fDeltaUnfoldedP for each bin is the error on the unfolded spectrum of that specific bin
+
+    //Do fNRandomIterations = bayes iterations performed
+    for(int i=0; i<fNRandomIterations; i++) {
+      if (fPrior)          delete fPrior ;
+      if (fRandomizedDist) delete fRandomizedDist ;
+      fPrior          = (THnSparse*) fOriginalPrior->Clone();
+      fRandomizedDist = (THnSparse*) fMeasuredOrig->Clone();
+      CreateRandomizedDist();
+      if (fMeasured) delete fMeasured ;
+      fMeasured = (THnSparse*) fRandomizedDist->Clone();
+      //Unfold fRandomizedDist
+      Unfold();
+      FillDeltaUnfoldedProfile();
+    }
+
+    // Get statistical errors for final unfolded spectrum
+    // ie. spread of each pt bin in fDeltaUnfoldedP
+    Double_t sigma       = 0.;
+    Double_t dummy       = 0.;
+    for (Long_t iBin=0; iBin<fRandomUnfolded->GetNbins(); iBin++) {
+      dummy = fUnfolded->GetBinContent(iBin,fCoordinatesN_M);
+      sigma = fDeltaUnfoldedP->GetBinError(fCoordinatesN_M[0]);
+      fUnfolded->SetBinError(fCoordinatesN_M,sigma);
+      fNCalcCorrErrors = 2;
+    }
+
+}
+
+//______________________________________________________________
+void AliCFUnfolding::InitDeltaUnfoldedProfile() {
+  //
+  //Initialize the fDeltaUnfoldedP profile
+  //Errors will be filled with spread between unfolded measured and unfolded randomized spectra
+  //
+
+  Int_t nbinsx = fResponse->GetAxis(0)->GetNbins();
+  Double_t xbins[nbinsx];
+  for(int ix=0; ix<nbinsx; ix++) {
+    xbins[ix] = fResponse->GetAxis(0)->GetBinLowEdge(ix+1);
+  }
+  xbins[nbinsx] = fResponse->GetAxis(0)->GetBinUpEdge(nbinsx);
+  fDeltaUnfoldedP = new TProfile("fDeltaUnfoldedP","Profile of pTUnfolded with spread in error",nbinsx,xbins,"S");
+}
+//______________________________________________________________
+void AliCFUnfolding::CreateRandomizedDist() {
+  //
+  // Create randomized dist from measured distribution
+  //
+
+  Double_t random = 0.;
+  Double_t measuredValue = 0.;
+  Double_t measuredError = 0.;
+  for (Long_t iBin=0; iBin<fRandomizedDist->GetNbins(); iBin++) {
+    measuredValue = fMeasuredOrig->GetBinContent(iBin,fCoordinatesN_M); //used as mean
+    measuredError = fMeasuredOrig->GetBinError(fCoordinatesN_M); //used as sigma
+    // random        = fRandom3->PoissonD(measuredValue); //doesn't work for normalized spectra, use Gaus (assuming raw counts in bin is large >10)
+    random        = fRandom3->Gaus(measuredValue,measuredError);
+    fRandomizedDist->SetBinContent(iBin,random);
+  }
+}
+
+//______________________________________________________________
+void AliCFUnfolding::FillDeltaUnfoldedProfile() {
+  //
+  // Store difference of unfolded spectrum from measured distribution and unfolded distribution from randomized distribution
+  //
+
+  for (Long_t iBin2=0; iBin2<fRandomUnfolded->GetNbins(); iBin2++) {
+    Double_t delta = fUnfolded->GetBinContent(iBin2,fCoordinatesN_M)-fRandomUnfolded->GetBinContent(iBin2,fCoordinatesN_M);
+    fDeltaUnfoldedP->Fill(fDeltaUnfoldedP->GetBinCenter(fCoordinatesN_M[0]),delta);
+  }
+}
+
+//______________________________________________________________
+
 void AliCFUnfolding::GetCoordinates() {
   //
   // assign coordinates in Measured and True spaces (dim=N) from coordinates in global space (dim=2N)
@@ -327,81 +622,143 @@ void AliCFUnfolding::CreateConditional() {
   fConditional     = (THnSparse*) fResponse->Clone(); // output of this function
   fProjResponseInT = (THnSparse*) fPrior->Clone();    // output denominator : 
                                                       // projection of the response matrix on the TRUE axis
-  
-  // set in fProjResponseInT zero everywhere
-  for (Int_t iBin=0; iBin<fProjResponseInT->GetNbins(); iBin++) {
-    fProjResponseInT->GetBinContent(iBin,fCoordinatesN_T);
-    fProjResponseInT->SetBinContent(fCoordinatesN_T,0.);
-  }
+  Int_t* dim = new Int_t [fNVariables];
+  for (Int_t iDim=0; iDim<fNVariables; iDim++) dim[iDim] = fNVariables+iDim ; //dimensions corresponding to TRUE values (i.e. from N to 2N-1)
+  fProjResponseInT = fConditional->Projection(fNVariables,dim,"E"); //project
+  delete [] dim; 
 
-  // calculate the response projection on T axis
-  for (Int_t iBin=0; iBin<fResponse->GetNbins(); iBin++) {
-    Double_t responseValue = fResponse->GetBinContent(iBin,fCoordinates2N);
-    GetCoordinates();
-    Double_t fill = fProjResponseInT->GetBinContent(fCoordinatesN_T) + responseValue ;
-    if (fill>0.) fProjResponseInT->SetBinContent(fCoordinatesN_T,fill);
-  }
-  
   // fill the conditional probability matrix
-  for (Int_t iBin=0; iBin<fResponse->GetNbins(); iBin++) {
+  for (Long_t iBin=0; iBin<fResponse->GetNbins(); iBin++) {
     Double_t responseValue = fResponse->GetBinContent(iBin,fCoordinates2N);
     GetCoordinates();
-    Double_t fill = responseValue / fProjResponseInT->GetBinContent(fCoordinatesN_T) ;
-    if (fill>0. || fConditional->GetBinContent(fCoordinates2N)) fConditional->SetBinContent(fCoordinates2N,fill);
+    Double_t projValue = fProjResponseInT->GetBinContent(fCoordinatesN_T);
+   
+    Double_t fill = responseValue / projValue ;
+    if (fill>0. || fConditional->GetBinContent(fCoordinates2N)>0.) {
+      fConditional->SetBinContent(fCoordinates2N,fill);
+      Double_t err = 0.;
+      fConditional->SetBinError  (fCoordinates2N,err);
+    }
   }
 }
+//______________________________________________________________
+
+Int_t AliCFUnfolding::GetDOF() {
+  //
+  // number of dof = number of bins
+  //
+
+  Int_t nDOF = 1 ;
+  for (Int_t iDim=0; iDim<fNVariables; iDim++) {
+    nDOF *= fPrior->GetAxis(iDim)->GetNbins();
+  }
+  AliDebug(0,Form("Number of degrees of freedom = %d",nDOF));
+  return nDOF;
+}
 
 //______________________________________________________________
 
 Double_t AliCFUnfolding::GetChi2() {
   //
   // Returns the chi2 between unfolded and a priori spectrum
+  // This function became absolute with the correlated error calculation.
+  // Errors on the unfolded distribution are not known until the end
+  // Use the convergence criterion instead
   //
 
-  Double_t chi2 = 0. ;
-  for (Int_t iBin=0; iBin<fPrior->GetNbins(); iBin++) {
-    Double_t priorValue = fPrior->GetBinContent(iBin);
-    chi2 += (priorValue>0. ? TMath::Power(fUnfolded->GetBinContent(iBin) - priorValue,2) / priorValue : 0.) ;
+  Double_t chi2      = 0. ;
+  Double_t error_unf = 0.;
+  for (Long_t iBin=0; iBin<fPrior->GetNbins(); iBin++) {
+    Double_t priorValue = fPrior->GetBinContent(iBin,fCoordinatesN_T);
+    error_unf  = fUnfolded->GetBinError(fCoordinatesN_T);
+    chi2 += (error_unf > 0. ? TMath::Power((fUnfolded->GetBinContent(fCoordinatesN_T) - priorValue)/error_unf,2) / priorValue : 0.) ;
   }
   return chi2;
 }
 
 //______________________________________________________________
 
-void AliCFUnfolding::SetMaxChi2PerDOF(Double_t val) {
+Double_t AliCFUnfolding::GetConvergence() {
+  //
+  // Returns convergence criterion = \sum_t ((U_t^{n-1}-U_t^n)/U_t^{n-1})^2
+  // U is unfolded spectrum, t is the bin, n = current, n-1 = previous
+  //
+  Double_t convergence = 0.;
+  Double_t priorValue  = 0.;
+  Double_t currentValue = 0.;
+  for (Long_t iBin=0; iBin < fPrior->GetNbins(); iBin++) {
+    priorValue = fPrior->GetBinContent(iBin,fCoordinatesN_T);
+    if (fNCalcCorrErrors > 0) 
+      currentValue = fRandomUnfolded->GetBinContent(fCoordinatesN_T);
+    else 
+      currentValue = fUnfolded->GetBinContent(fCoordinatesN_T);
+
+    if (priorValue > 0.)
+      convergence += ((priorValue-currentValue)/priorValue)*((priorValue-currentValue)/priorValue);
+    else {
+      AliWarning(Form("priorValue = %f. Adding 0 to convergence criterion.",priorValue)); 
+      convergence += 0.;
+    }
+  }
+  return convergence;
+}
+
+//______________________________________________________________
+
+void AliCFUnfolding::SetMaxConvergencePerDOF(Double_t val) {
   //
-  // Max. chi2 per degree of freedom : user setting
+  // Max. convergence criterion per degree of freedom : user setting
+  // convergence criterion = DOF*val; DOF = number of bins
+  // In Jan-Fiete's multiplicity note: Convergence criterion = DOF*0.001^2
   //
 
-  Int_t nDOF = 1 ;
-  for (Int_t iDim=0; iDim<fNVariables; iDim++) {
-    nDOF *= fPrior->GetAxis(iDim)->GetNbins();
-  }
-  AliInfo(Form("Number of degrees of freedom = %d",nDOF));
-  fMaxChi2 = val * nDOF ;
+  fUseCorrelatedErrors = kTRUE;
+  Int_t nDOF = GetDOF() ;
+  fMaxConvergence = val * nDOF ;
+  AliInfo(Form("MaxConvergence = %e. Number of degrees of freedom = %d",fMaxConvergence,nDOF));
 }
 
 //______________________________________________________________
 
-void AliCFUnfolding::Smooth() {
+Short_t AliCFUnfolding::Smooth() {
   //
   // Smoothes the unfolded spectrum
-  // Each cell content is replaced by the average with the neighbouring bins (but not diagonally-neighbouring bins)
+  //
+  // By default each cell content is replaced by the average with the neighbouring bins (but not diagonally-neighbouring bins)
+  // However, if a specific function fcn has been defined in UseSmoothing(fcn), the unfolded will be fit and updated using fcn 
   //
   
-  Int_t* numBins = new Int_t[fNVariables];
-  for (Int_t iVar=0; iVar<fNVariables; iVar++) numBins[iVar]=fUnfolded->GetAxis(iVar)->GetNbins();
+  if (fSmoothFunction) {
+    AliDebug(2,Form("Smoothing spectrum with fit function %p",fSmoothFunction));
+    return SmoothUsingFunction();
+  }
+  else return SmoothUsingNeighbours(fUnfolded);
+}
+
+//______________________________________________________________
+
+Short_t AliCFUnfolding::SmoothUsingNeighbours(THnSparse* hist) {
+  //
+  // Smoothes the unfolded spectrum using neighouring bins
+  //
+
+  Int_t  const nDimensions = hist->GetNdimensions() ;
+  Int_t* coordinates = new Int_t[nDimensions];
+
+  Int_t* numBins = new Int_t[nDimensions];
+  for (Int_t iVar=0; iVar<nDimensions; iVar++) numBins[iVar] = hist->GetAxis(iVar)->GetNbins();
   
-  //need a copy because fUnfolded will be updated during the loop, and this creates problems
-  THnSparse* copy = (THnSparse*)fUnfolded->Clone();
+  //need a copy because hist will be updated during the loop, and this creates problems
+  THnSparse* copy = (THnSparse*)hist->Clone();
 
-  for (Int_t iBin=0; iBin<copy->GetNbins(); iBin++) { //loop on non-empty bins
-    Double_t content = copy->GetBinContent(iBin,fCoordinatesN_T);
+  for (Long_t iBin=0; iBin<copy->GetNbins(); iBin++) { //loop on non-empty bins
+    Double_t content = copy->GetBinContent(iBin,coordinates);
+    Double_t error2  = TMath::Power(copy->GetBinError(iBin),2);
 
     // skip the under/overflow bins...
     Bool_t isOutside = kFALSE ;
-    for (Int_t iVar=0; iVar<fNVariables; iVar++) {
-      if (fCoordinatesN_T[iVar]<1 || fCoordinatesN_T[iVar]>numBins[iVar]) {
+    for (Int_t iVar=0; iVar<nDimensions; iVar++) {
+      if (coordinates[iVar]<1 || coordinates[iVar]>numBins[iVar]) {
        isOutside=kTRUE;
        break;
       }
@@ -410,29 +767,86 @@ void AliCFUnfolding::Smooth() {
     
     Int_t neighbours = 0; // number of neighbours to average with
 
-    for (Int_t iVar=0; iVar<fNVariables; iVar++) {
-      if (fCoordinatesN_T[iVar] > 1) { // must not be on low edge border
-       fCoordinatesN_T[iVar]-- ; //get lower neighbouring bin 
-       Double_t contentNeighbour = copy->GetBinContent(fCoordinatesN_T);
-       content += contentNeighbour;
+    for (Int_t iVar=0; iVar<nDimensions; iVar++) {
+      if (coordinates[iVar] > 1) { // must not be on low edge border
+       coordinates[iVar]-- ; //get lower neighbouring bin 
+       content += copy->GetBinContent(coordinates);
+       error2  += TMath::Power(copy->GetBinError(coordinates),2);
        neighbours++;
-       fCoordinatesN_T[iVar]++ ; //back to initial coordinate
+       coordinates[iVar]++ ; //back to initial coordinate
       }
-      if (fCoordinatesN_T[iVar] < numBins[iVar]) { // must not be on up edge border
-       fCoordinatesN_T[iVar]++ ; //get upper neighbouring bin
-       Double_t contentNeighbour = copy->GetBinContent(fCoordinatesN_T);
-       content += contentNeighbour ;
+      if (coordinates[iVar] < numBins[iVar]) { // must not be on up edge border
+       coordinates[iVar]++ ; //get upper neighbouring bin
+       content += copy->GetBinContent(coordinates);
+       error2  += TMath::Power(copy->GetBinError(coordinates),2);
        neighbours++;
-       fCoordinatesN_T[iVar]-- ; //back to initial coordinate
+       coordinates[iVar]-- ; //back to initial coordinate
       }
     }
-    content /= (1+neighbours) ; // make an average
-    fUnfolded->SetBinContent(fCoordinatesN_T,content);
+    // make an average
+    hist->SetBinContent(coordinates,content/(1.+neighbours));
+    hist->SetBinError  (coordinates,TMath::Sqrt(error2)/(1.+neighbours));
   }
   delete [] numBins;
+  delete [] coordinates ;
   delete copy;
+  return 0;
 }
 
+//______________________________________________________________
+
+Short_t AliCFUnfolding::SmoothUsingFunction() {
+  //
+  // Fits the unfolded spectrum using the function fSmoothFunction
+  //
+
+  AliDebug(0,Form("Smooth function is a %s with option \"%s\" and has %d parameters : ",fSmoothFunction->ClassName(),fSmoothOption,fSmoothFunction->GetNpar()));
+
+  for (Int_t iPar=0; iPar<fSmoothFunction->GetNpar(); iPar++) AliDebug(0,Form("par[%d]=%e",iPar,fSmoothFunction->GetParameter(iPar)));
+
+  Int_t fitResult = 0;
+
+  switch (fNVariables) {
+  case 1 : fitResult = fUnfolded->Projection(0)    ->Fit(fSmoothFunction,fSmoothOption); break;
+  case 2 : fitResult = fUnfolded->Projection(1,0)  ->Fit(fSmoothFunction,fSmoothOption); break; // (1,0) instead of (0,1) -> TAxis issue
+  case 3 : fitResult = fUnfolded->Projection(0,1,2)->Fit(fSmoothFunction,fSmoothOption); break;
+  default: AliFatal(Form("Cannot handle such fit in %d dimensions",fNVariables)) ; return 1;
+  }
+
+  if (fitResult != 0) {
+    AliWarning(Form("Fit failed with status %d, stopping the loop",fitResult));
+    return 1;
+  }
+
+  Int_t nDim = fNVariables;
+  Int_t* bins = new Int_t[nDim]; // number of bins for each variable
+  Long_t nBins = 1; // used to calculate the total number of bins in the THnSparse
+
+  for (Int_t iVar=0; iVar<nDim; iVar++) {
+    bins[iVar] = fUnfolded->GetAxis(iVar)->GetNbins();
+    nBins *= bins[iVar];
+  }
+
+  Int_t *bin  = new Int_t[nDim];    // bin to fill the THnSparse (holding the bin coordinates)
+  Double_t x[3] = {0,0,0} ;         // value in bin center (max dimension is 3 (TF3))
+
+  // loop on the bins and update of fUnfolded
+  // THnSparse::Multiply(TF1*) doesn't exist, so let's do it bin by bin
+  for (Long_t iBin=0; iBin<nBins; iBin++) {
+    Long_t bin_tmp = iBin ;
+    for (Int_t iVar=0; iVar<nDim; iVar++) {
+      bin[iVar] = 1 + bin_tmp % bins[iVar] ;
+      bin_tmp /= bins[iVar] ;
+      x[iVar] = fUnfolded->GetAxis(iVar)->GetBinCenter(bin[iVar]);
+    }
+    Double_t functionValue = fSmoothFunction->Eval(x[0],x[1],x[2]) ;
+    fUnfolded->SetBinError  (bin,fUnfolded->GetBinError(bin)*functionValue/fUnfolded->GetBinContent(bin));
+    fUnfolded->SetBinContent(bin,functionValue);
+  }
+  delete [] bins;
+  delete [] bin ;
+  return 0;
+}
 
 //______________________________________________________________
 
@@ -446,6 +860,9 @@ void AliCFUnfolding::CreateFlatPrior() {
   // create the frame of the THnSparse given (for example) the one from the efficiency map
   fPrior = (THnSparse*) fEfficiency->Clone();
 
+  if (fNVariables != fPrior->GetNdimensions()) 
+    AliFatal(Form("The prior matrix should have %d dimensions, and it has actually %d",fNVariables,fPrior->GetNdimensions()));
+
   Int_t nDim = fNVariables;
   Int_t* bins = new Int_t[nDim]; // number of bins for each variable
   Long_t nBins = 1; // used to calculate the total number of bins in the THnSparse
@@ -465,6 +882,7 @@ void AliCFUnfolding::CreateFlatPrior() {
       bin_tmp /= bins[iVar] ;
     }
     fPrior->SetBinContent(bin,1.); // put 1 everywhere
+    fPrior->SetBinError  (bin,0.); // put 0 everywhere
   }
   
   fOriginalPrior = (THnSparse*)fPrior->Clone();