New class to handle multi-dimensional unfolding + user macros (test/)
[u/mrichter/AliRoot.git] / CORRFW / AliCFUnfolding.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
3  *                                                                        *
4  * Author: The ALICE Off-line Project.                                    *
5  * Contributors are mentioned in the code where appropriate.              *
6  *                                                                        *
7  * Permission to use, copy, modify and distribute this software and its   *
8  * documentation strictly for non-commercial purposes is hereby granted   *
9  * without fee, provided that the above copyright notice appears in all   *
10  * copies and that both the copyright notice and this permission notice   *
11  * appear in the supporting documentation. The authors make no claims     *
12  * about the suitability of this software for any purpose. It is          *
13  * provided "as is" without express or implied warranty.                  *
14  **************************************************************************/
15
16 //--------------------------------------------------------------------//
17 //                                                                    //
18 // AliCFUnfolding Class                                               //
19 // Class to handle general unfolding procedure                        // 
20 // For the moment only bayesian unfolding is supported                //
21 // The next steps are to add chi2 minimisation and weighting methods  //
22 //                                                                    //
23 // Author : renaud.vernet@cern.ch                                     //
24 //--------------------------------------------------------------------//
25
26
27 #include "AliCFUnfolding.h"
28 #include "TMath.h"
29 #include "TAxis.h"
30 #include "AliLog.h"
31
32 ClassImp(AliCFUnfolding)
33
34 //______________________________________________________________
35
36 AliCFUnfolding::AliCFUnfolding() :
37   TNamed(),
38   fResponse(0x0),
39   fPrior(0x0),
40   fOriginalPrior(0x0),
41   fEfficiency(0x0),
42   fMeasured(0x0),
43   fMaxNumIterations(0),
44   fNVariables(0),
45   fMaxChi2(0),
46   fUseSmoothing(kFALSE),
47   fInverseResponse(0x0),
48   fMeasuredEstimate(0x0),
49   fConditional(0x0),
50   fProjResponseInT(0x0),
51   fUnfolded(0x0),
52   fCoordinates2N(0x0),
53   fCoordinatesN_M(0x0),
54   fCoordinatesN_T(0x0)
55 {
56   //
57   // default constructor
58   //
59 }
60
61 //______________________________________________________________
62
63 AliCFUnfolding::AliCFUnfolding(const Char_t* name, const Char_t* title, const Int_t nVar, 
64                                const THnSparse* response, const THnSparse* efficiency, const THnSparse* measured, const THnSparse* prior) :
65   TNamed(name,title),
66   fResponse((THnSparse*)response->Clone()),
67   fPrior(0x0),
68   fOriginalPrior(0x0),
69   fEfficiency((THnSparse*)efficiency->Clone()),
70   fMeasured((THnSparse*)measured->Clone()),
71   fMaxNumIterations(0),
72   fNVariables(nVar),
73   fMaxChi2(0),
74   fUseSmoothing(kFALSE),
75   fInverseResponse(0x0),
76   fMeasuredEstimate(0x0),
77   fConditional(0x0),
78   fProjResponseInT(0x0),
79   fUnfolded(0x0),
80   fCoordinates2N(0x0),
81   fCoordinatesN_M(0x0),
82   fCoordinatesN_T(0x0)
83 {
84   //
85   // named constructor
86   //
87
88   AliInfo(Form("\n\n--------------------------\nCreating an unfolder :\n--------------------------\nresponse matrix has %d dimension(s)",fResponse->GetNdimensions()));
89   
90   if (!prior) CreateFlatPrior(); // if no prior distribution declared, simply use a flat distribution
91   else {
92     fPrior = (THnSparse*) prior->Clone();
93     fOriginalPrior = (THnSparse*)fPrior->Clone();
94   }
95   
96   for (Int_t iVar=0; iVar<fNVariables; iVar++) {
97     AliInfo(Form("prior      matrix has %d bins in dimension %d",fPrior     ->GetAxis(iVar)->GetNbins(),iVar));
98     AliInfo(Form("efficiency matrix has %d bins in dimension %d",fEfficiency->GetAxis(iVar)->GetNbins(),iVar));
99     AliInfo(Form("measured   matrix has %d bins in dimension %d",fMeasured  ->GetAxis(iVar)->GetNbins(),iVar));
100   }
101   Init();
102 }
103
104
105 //______________________________________________________________
106
107 AliCFUnfolding::AliCFUnfolding(const AliCFUnfolding& c) :
108   TNamed(c),
109   fResponse((THnSparse*)c.fResponse->Clone()),
110   fPrior((THnSparse*)c.fPrior->Clone()),
111   fOriginalPrior((THnSparse*)c.fOriginalPrior->Clone()),
112   fEfficiency((THnSparse*)c.fEfficiency->Clone()),
113   fMeasured((THnSparse*)c.fMeasured->Clone()),
114   fMaxNumIterations(c.fMaxNumIterations),
115   fNVariables(c.fNVariables),
116   fMaxChi2(c.fMaxChi2),
117   fUseSmoothing(c.fUseSmoothing),
118   fInverseResponse((THnSparse*)c.fInverseResponse->Clone()),
119   fMeasuredEstimate((THnSparse*)fMeasuredEstimate->Clone()),
120   fConditional((THnSparse*)c.fConditional->Clone()),
121   fProjResponseInT((THnSparse*)c.fProjResponseInT->Clone()),
122   fUnfolded((THnSparse*)c.fUnfolded->Clone()),
123   fCoordinates2N(new Int_t(*c.fCoordinates2N)),
124   fCoordinatesN_M(new Int_t(*c.fCoordinatesN_M)),
125   fCoordinatesN_T(new Int_t(*c.fCoordinatesN_T))
126 {
127   //
128   // copy constructor
129   //
130 }
131
132 //______________________________________________________________
133
134 AliCFUnfolding& AliCFUnfolding::operator=(const AliCFUnfolding& c) {
135   //
136   // assignment operator
137   //
138   
139   if (this!=&c) {
140     TNamed::operator=(c);
141     fResponse = (THnSparse*)c.fResponse->Clone() ;
142     fPrior = (THnSparse*)c.fPrior->Clone() ;
143     fOriginalPrior = (THnSparse*)c.fOriginalPrior->Clone() ;
144     fEfficiency = (THnSparse*)c.fEfficiency->Clone() ;
145     fMeasured = (THnSparse*)c.fMeasured->Clone() ;
146     fMaxNumIterations = c.fMaxNumIterations ;
147     fNVariables = c.fNVariables ;
148     fMaxChi2 = c.fMaxChi2 ;
149     fUseSmoothing = c.fUseSmoothing ;
150     fInverseResponse = (THnSparse*)c.fInverseResponse->Clone() ;
151     fMeasuredEstimate = (THnSparse*)fMeasuredEstimate->Clone() ;
152     fConditional = (THnSparse*)c.fConditional->Clone() ;
153     fProjResponseInT = (THnSparse*)c.fProjResponseInT->Clone() ;
154     fUnfolded = (THnSparse*)c.fUnfolded->Clone() ;
155     fCoordinates2N  = new Int_t(*c.fCoordinates2N)  ;
156     fCoordinatesN_M = new Int_t(*c.fCoordinatesN_M) ;
157     fCoordinatesN_T = new Int_t(*c.fCoordinatesN_T) ;
158   }
159   return *this;
160 }
161
162 //______________________________________________________________
163
164 AliCFUnfolding::~AliCFUnfolding() {
165   //
166   // destructor
167   //
168   if (fResponse)           delete fResponse;
169   if (fPrior)              delete fPrior;
170   if (fOriginalPrior)      delete fOriginalPrior;
171   if (fEfficiency)         delete fEfficiency;
172   if (fMeasured)           delete fMeasured;
173   if (fInverseResponse)    delete fInverseResponse;
174   if (fMeasuredEstimate)   delete fMeasuredEstimate;
175   if (fConditional)        delete fConditional;
176   if (fProjResponseInT)    delete fProjResponseInT;
177   if (fCoordinates2N)      delete [] fCoordinates2N; 
178   if (fCoordinatesN_M)     delete [] fCoordinatesN_M; 
179   if (fCoordinatesN_T)     delete [] fCoordinatesN_T; 
180 }
181
182 //______________________________________________________________
183
184 void AliCFUnfolding::Init() {
185   //
186   // initialisation function : creates internal settings
187   //
188
189   fCoordinates2N  = new Int_t[2*fNVariables];
190   fCoordinatesN_M = new Int_t[fNVariables];
191   fCoordinatesN_T = new Int_t[fNVariables];
192
193   // create the matrix of conditional probabilities P(M|T)
194   CreateConditional();
195   
196   // create the frame of the inverse response matrix
197   fInverseResponse  = (THnSparse*) fResponse->Clone();
198   // create the frame of the unfolded spectrum
199   fUnfolded = (THnSparse*) fPrior->Clone();
200   // create the frame of the measurement estimate spectrum
201   fMeasuredEstimate = (THnSparse*) fMeasured->Clone();
202 }
203
204 //______________________________________________________________
205
206 void AliCFUnfolding::CreateEstMeasured() {
207   //
208   // This function creates a estimate (M) of the reconstructed spectrum 
209   // given the a priori distribution (T) and the conditional matrix (COND)
210   //
211   // --> P(M) = SUM   { P(M|T)    * P(T) }
212   // --> M(i) = SUM_k { COND(i,k) * T(k) }
213   //
214   // This is needed to calculate the inverse response matrix
215   //
216
217
218   // clean the measured estimate spectrum
219   for (Long64_t i=0; i<fMeasuredEstimate->GetNbins(); i++) {
220     fMeasuredEstimate->GetBinContent(i,fCoordinatesN_M);
221     fMeasuredEstimate->SetBinContent(fCoordinatesN_M,0.);
222   }
223   
224   // fill it
225   for (Int_t iBin=0; iBin<fConditional->GetNbins(); iBin++) {
226     Double_t conditionalValue = fConditional->GetBinContent(iBin,fCoordinates2N);
227     GetCoordinates();
228     Double_t priorValue = fPrior->GetBinContent(fCoordinatesN_T);
229     Double_t fill = fMeasuredEstimate->GetBinContent(fCoordinatesN_M) + conditionalValue * priorValue * fEfficiency->GetBinContent(fCoordinatesN_T);
230     if (fill>0.) fMeasuredEstimate->SetBinContent(fCoordinatesN_M,fill);
231   }
232 }
233
234 //______________________________________________________________
235
236 void AliCFUnfolding::CreateInvResponse() {
237   //
238   // Creates the inverse response matrix (INV) with Bayesian method
239   //  : uses the conditional matrix (COND), the prior probabilities (T) and the efficiency map (E)
240   //
241   // --> P(T|M)   = P(M|T)    * P(T) * eff(T) / SUM   { P(M|T)    * P(T) }
242   // --> INV(i,j) = COND(i,j) * T(j) * E(j)   / SUM_k { COND(i,k) * T(k) }
243   //
244
245   for (Int_t iBin=0; iBin<fConditional->GetNbins(); iBin++) {
246     Double_t conditionalValue = fConditional->GetBinContent(iBin,fCoordinates2N);
247     GetCoordinates();
248     Double_t priorValue = fPrior->GetBinContent(fCoordinatesN_T);
249     Double_t estimatedMeasured = fMeasuredEstimate->GetBinContent(fCoordinatesN_M);
250     Double_t fill = (estimatedMeasured>0. ? conditionalValue * priorValue * fEfficiency->GetBinContent(fCoordinatesN_T) / estimatedMeasured : 0. ) ;
251     if (fill>0. || fInverseResponse->GetBinContent(fCoordinates2N)>0.) fInverseResponse->SetBinContent(fCoordinates2N,fill);
252   }
253 }
254
255 //______________________________________________________________
256
257 void AliCFUnfolding::Unfold() {
258   //
259   // Main routine called by the user : 
260   // it calculates the unfolded spectrum from the response matrix and the measured spectrum
261   // several iterations are performed until a reasonable chi2 is reached
262   //
263
264   Int_t iIterBayes=0 ;
265   Double_t chi2=0 ;
266
267   for (iIterBayes=0; iIterBayes<fMaxNumIterations; iIterBayes++) { // bayes iterations
268     CreateEstMeasured();
269     CreateInvResponse();
270     CreateUnfolded();
271     chi2 = GetChi2();
272     //printf("chi2 = %e\n",chi2);
273     if (fMaxChi2>0. && chi2<fMaxChi2) {
274       break;
275     }
276     // update the prior distribution
277     if (fUseSmoothing) Smooth();
278     fPrior = (THnSparse*)fUnfolded->Clone() ; // this should be changed (memory)
279   }
280   AliInfo(Form("Finished at iteration %d : Chi2 is %e and you required it to be < %e",iIterBayes,chi2,fMaxChi2));
281 }
282
283 //______________________________________________________________
284
285 void AliCFUnfolding::CreateUnfolded() {
286   //
287   // Creates the unfolded (T) spectrum from the measured spectrum (M) and the inverse response matrix (INV)
288   // We have P(T) = SUM   { P(T|M)   * P(M) } 
289   //   -->   T(i) = SUM_k { INV(i,k) * M(k) }
290   //
291
292
293   // clear the unfolded spectrum
294   for (Long64_t i=0; i<fUnfolded->GetNbins(); i++) {
295     fUnfolded->GetBinContent(i,fCoordinatesN_T);
296     fUnfolded->SetBinContent(fCoordinatesN_T,0.);
297   }
298   
299   for (Int_t iBin=0; iBin<fInverseResponse->GetNbins(); iBin++) {
300     Double_t invResponseValue = fInverseResponse->GetBinContent(iBin,fCoordinates2N);
301     GetCoordinates();
302     Double_t effValue = fEfficiency->GetBinContent(fCoordinatesN_T);
303     Double_t fill = fUnfolded->GetBinContent(fCoordinatesN_T) + (effValue>0. ? invResponseValue*fMeasured->GetBinContent(fCoordinatesN_M)/effValue : 0.) ;
304     if (fill>0.) fUnfolded->SetBinContent(fCoordinatesN_T,fill);
305   }
306 }
307     
308 void AliCFUnfolding::GetCoordinates() {
309   //
310   // assign coordinates in Measured and True spaces (dim=N) from coordinates in global space (dim=2N)
311   //
312   for (Int_t i = 0; i<fNVariables ; i++) {
313     fCoordinatesN_M[i] = fCoordinates2N[i];
314     fCoordinatesN_T[i] = fCoordinates2N[i+fNVariables];
315   }
316 }
317
318 //______________________________________________________________
319
320 void AliCFUnfolding::CreateConditional() {
321   //
322   // creates the conditional probability matrix (R*) holding the P(M|T), given the reponse matrix R
323   //
324   //  --> R*(i,j) = R(i,j) / SUM_k{ R(k,j) }
325   //
326
327   fConditional     = (THnSparse*) fResponse->Clone(); // output of this function
328   fProjResponseInT = (THnSparse*) fPrior->Clone();    // output denominator : 
329                                                       // projection of the response matrix on the TRUE axis
330   
331   // set in fProjResponseInT zero everywhere
332   for (Int_t iBin=0; iBin<fProjResponseInT->GetNbins(); iBin++) {
333     fProjResponseInT->GetBinContent(iBin,fCoordinatesN_T);
334     fProjResponseInT->SetBinContent(fCoordinatesN_T,0.);
335   }
336
337   // calculate the response projection on T axis
338   for (Int_t iBin=0; iBin<fResponse->GetNbins(); iBin++) {
339     Double_t responseValue = fResponse->GetBinContent(iBin,fCoordinates2N);
340     GetCoordinates();
341     Double_t fill = fProjResponseInT->GetBinContent(fCoordinatesN_T) + responseValue ;
342     if (fill>0.) fProjResponseInT->SetBinContent(fCoordinatesN_T,fill);
343   }
344   
345   // fill the conditional probability matrix
346   for (Int_t iBin=0; iBin<fResponse->GetNbins(); iBin++) {
347     Double_t responseValue = fResponse->GetBinContent(iBin,fCoordinates2N);
348     GetCoordinates();
349     Double_t fill = responseValue / fProjResponseInT->GetBinContent(fCoordinatesN_T) ;
350     if (fill>0. || fConditional->GetBinContent(fCoordinates2N)) fConditional->SetBinContent(fCoordinates2N,fill);
351   }
352 }
353
354 //______________________________________________________________
355
356 Double_t AliCFUnfolding::GetChi2() {
357   //
358   // Returns the chi2 between unfolded and a priori spectrum
359   //
360
361   Double_t chi2 = 0. ;
362   for (Int_t iBin=0; iBin<fPrior->GetNbins(); iBin++) {
363     Double_t priorValue = fPrior->GetBinContent(iBin);
364     chi2 += (priorValue>0. ? TMath::Power(fUnfolded->GetBinContent(iBin) - priorValue,2) / priorValue : 0.) ;
365   }
366   return chi2;
367 }
368
369 //______________________________________________________________
370
371 void AliCFUnfolding::SetMaxChi2PerDOF(Double_t val) {
372   //
373   // Max. chi2 per degree of freedom : user setting
374   //
375
376   Int_t nDOF = 1 ;
377   for (Int_t iDim=0; iDim<fNVariables; iDim++) {
378     nDOF *= fPrior->GetAxis(iDim)->GetNbins();
379   }
380   AliInfo(Form("Number of degrees of freedom = %d",nDOF));
381   fMaxChi2 = val * nDOF ;
382 }
383
384 //______________________________________________________________
385
386 void AliCFUnfolding::Smooth() {
387   //
388   // Smoothes the unfolded spectrum
389   // Each cell content is replaced by the average with the neighbouring bins (but not diagonally-neighbouring bins)
390   //
391   
392   Int_t* numBins = new Int_t[fNVariables];
393   for (Int_t iVar=0; iVar<fNVariables; iVar++) numBins[iVar]=fUnfolded->GetAxis(iVar)->GetNbins();
394   
395   //need a copy because fUnfolded will be updated during the loop, and this creates problems
396   THnSparse* copy = (THnSparse*)fUnfolded->Clone();
397
398   for (Int_t iBin=0; iBin<copy->GetNbins(); iBin++) { //loop on non-empty bins
399     Double_t content = copy->GetBinContent(iBin,fCoordinatesN_T);
400
401     // skip the under/overflow bins...
402     Bool_t isOutside = kFALSE ;
403     for (Int_t iVar=0; iVar<fNVariables; iVar++) {
404       if (fCoordinatesN_T[iVar]<1 || fCoordinatesN_T[iVar]>numBins[iVar]) {
405         isOutside=kTRUE;
406         break;
407       }
408     }
409     if (isOutside) continue;
410     
411     Int_t neighbours = 0; // number of neighbours to average with
412
413     for (Int_t iVar=0; iVar<fNVariables; iVar++) {
414       if (fCoordinatesN_T[iVar] > 1) { // must not be on low edge border
415         fCoordinatesN_T[iVar]-- ; //get lower neighbouring bin 
416         Double_t contentNeighbour = copy->GetBinContent(fCoordinatesN_T);
417         content += contentNeighbour;
418         neighbours++;
419         fCoordinatesN_T[iVar]++ ; //back to initial coordinate
420       }
421       if (fCoordinatesN_T[iVar] < numBins[iVar]) { // must not be on up edge border
422         fCoordinatesN_T[iVar]++ ; //get upper neighbouring bin
423         Double_t contentNeighbour = copy->GetBinContent(fCoordinatesN_T);
424         content += contentNeighbour ;
425         neighbours++;
426         fCoordinatesN_T[iVar]-- ; //back to initial coordinate
427       }
428     }
429     content /= (1+neighbours) ; // make an average
430     fUnfolded->SetBinContent(fCoordinatesN_T,content);
431   }
432   delete [] numBins;
433   delete copy;
434 }
435
436
437 //______________________________________________________________
438
439 void AliCFUnfolding::CreateFlatPrior() {
440   //
441   // Creates a flat prior distribution
442   // 
443
444   AliInfo("Creating a flat a priori distribution");
445   
446   // create the frame of the THnSparse given (for example) the one from the efficiency map
447   fPrior = (THnSparse*) fEfficiency->Clone();
448
449   Int_t nDim = fNVariables;
450   Int_t* bins = new Int_t[nDim]; // number of bins for each variable
451   Long_t nBins = 1; // used to calculate the total number of bins in the THnSparse
452
453   for (Int_t iVar=0; iVar<nDim; iVar++) {
454     bins[iVar] = fPrior->GetAxis(iVar)->GetNbins();
455     nBins *= bins[iVar];
456   }
457
458   Int_t *bin = new Int_t[nDim]; // bin to fill the THnSparse (holding the bin coordinates)
459
460   // loop that sets 1 in each bin
461   for (Long_t iBin=0; iBin<nBins; iBin++) {
462     Long_t bin_tmp = iBin ;
463     for (Int_t iVar=0; iVar<nDim; iVar++) {
464       bin[iVar] = 1 + bin_tmp % bins[iVar] ;
465       bin_tmp /= bins[iVar] ;
466     }
467     fPrior->SetBinContent(bin,1.); // put 1 everywhere
468   }
469   
470   fOriginalPrior = (THnSparse*)fPrior->Clone();
471
472   delete [] bin;
473   delete [] bins;
474 }