Updates in order to enable the '2D' PID for the TRD developed by Daniel Lohner.
[u/mrichter/AliRoot.git] / STEER / STEERBase / AliTRDTKDInterpolator.cxx
1 #include "AliTRDTKDInterpolator.h"
2
3 #include "TClonesArray.h"
4 #include "TLinearFitter.h"
5 #include "TMath.h"
6 #include "TRandom.h"
7
8 #include "iostream"
9 using namespace std;
10
11 ClassImp(AliTRDTKDInterpolator)
12
13 //_________________________________________________________________
14 AliTRDTKDInterpolator::AliTRDTKDInterpolator() :
15 TKDTreeIF(),
16 fNDataNodes(0),
17 fNodes(NULL),
18 fLambda(0),
19 fNPointsI(0),
20 fUseHelperNodes(kFALSE),
21 fUseWeights(kFALSE),
22 fPDFMode(kInterpolation)
23 {
24   // default constructor
25 }
26
27 //_________________________________________________________________
28 AliTRDTKDInterpolator::AliTRDTKDInterpolator(Int_t npoints, Int_t ndim, UInt_t bsize, Float_t **data) :
29 TKDTreeIF(npoints, ndim, bsize, data),
30 fNDataNodes(0),
31 fNodes(NULL),
32 fLambda(1 + ndim + (ndim*(ndim+1)>>1)),
33 fNPointsI(100),
34 fUseHelperNodes(kFALSE),
35 fUseWeights(kFALSE),
36 fPDFMode(kInterpolation)
37 {
38 }
39
40 //_________________________________________________________________
41 AliTRDTKDInterpolator::~AliTRDTKDInterpolator()
42 {
43     if(fNodes){
44         fNodes->Delete();
45         delete fNodes;
46         fNodes=NULL;
47     }
48 }
49 //_________________________________________________________________
50 AliTRDTKDInterpolator::AliTRDTKDInterpolator(const AliTRDTKDInterpolator &ref):
51 TKDTreeIF(),
52 fNDataNodes(ref.fNDataNodes),
53 fNodes(ref.fNodes),
54 fLambda(ref.fLambda),
55 fNPointsI(ref.fNPointsI),
56 fUseHelperNodes(ref.fUseHelperNodes),
57 fUseWeights(ref.fUseWeights),
58 fPDFMode(ref.fPDFMode)
59 {
60     // Copy constructor
61 }
62
63 //____________________________________________________________
64
65 AliTRDTKDInterpolator &AliTRDTKDInterpolator::operator=(const AliTRDTKDInterpolator &ref){
66     //
67     // Assignment operator
68     //
69     if(this == &ref) return *this;
70
71     // Make copy
72     TObject::operator=(ref);
73
74     return *this;
75 }
76
77 //_________________________________________________________________
78 Bool_t AliTRDTKDInterpolator::Build()
79 {
80     TKDTreeIF::Build();
81     if(!fBoundaries) MakeBoundaries();
82
83     // allocate interpolation nodes
84     fNDataNodes = fNPoints/fBucketSize + ((fNPoints%fBucketSize)?1:0);/*TKDTreeIF::GetNTNodes();*/
85
86     if(fNodes){
87         Warning("AliTRDTKDInterpolator::Build()", "Data already allocated.");
88         fNodes->Delete();
89     } else {
90         fNodes = new TClonesArray("AliTRDTKDInterpolator::AliTRDTKDNodeInfo", fNDataNodes);
91         fNodes->SetOwner();
92     }
93     for(int in=0; in<fNDataNodes; in++) new ((*fNodes)[in]) AliTRDTKDNodeInfo(fNDim);
94
95     // Set Interpolator nodes
96
97     for(int inode=0, tnode = fNNodes; inode<fNDataNodes-1; inode++, tnode++){
98         AliTRDTKDNodeInfo *node =GetNodeInfo(inode);
99         memcpy(node->fBounds,GetBoundary(tnode),2*fNDim*sizeof(Float_t));
100         node->fVal[0] =  Float_t(fBucketSize)/fNPoints;
101         for(int idim=0; idim<fNDim; idim++) node->fVal[0] /= (node->fBounds[2*idim+1] - node->fBounds[2*idim]);
102         node->fVal[1] =  node->fVal[0]/TMath::Sqrt(float(fBucketSize));
103
104         Int_t *indexPoints = GetPointsIndexes(tnode);
105         // loop points in this terminal node
106         for(int idim=0; idim<fNDim; idim++){
107             node->fData[idim] = 0.;
108             for(int ip = 0; ip<fBucketSize; ip++) node->fData[idim] += fData[idim][indexPoints[ip]];
109             node->fData[idim] /= fBucketSize;
110         }
111     }
112
113     // Analyze last (incomplete) terminal node
114
115     Int_t counts = fNPoints%fBucketSize;
116     counts = counts ? counts : fBucketSize;
117     Int_t inode = fNDataNodes - 1, tnode = inode + fNNodes;
118     AliTRDTKDNodeInfo *ftnode = GetNodeInfo(inode);
119     ftnode->fVal[0] =  Float_t(counts)/fNPoints;
120     memcpy(ftnode->fBounds,GetBoundary(tnode),2*fNDim*sizeof(Float_t));
121  
122     for(int idim=0; idim<fNDim; idim++){
123         Float_t dx = ftnode->fBounds[2*idim+1]-ftnode->fBounds[2*idim];
124         if(dx < 1.e-30){
125             Warning("AliTRDTKDInterpolator::Build()", "Terminal bucket index[%d] too narrow on the %d dimension.", inode, idim);
126             continue;
127         }
128         ftnode->fVal[0] /= (ftnode->fBounds[2*idim+1] - ftnode->fBounds[2*idim]);
129     }
130     ftnode->fVal[1] =  ftnode->fVal[0]/TMath::Sqrt(float(counts));
131
132     // loop points in this terminal node
133     Int_t *indexPoints = GetPointsIndexes(tnode);
134     for(int idim=0; idim<fNDim; idim++){
135         ftnode->fData[idim] = 0.;
136         for(int ip = 0; ip<counts; ip++) ftnode->fData[idim] += fData[idim][indexPoints[ip]];
137         ftnode->fData[idim] /= counts;
138     }
139
140     delete [] fBoundaries;
141     fBoundaries = NULL;
142     // Add Helper Nodes
143     if(fUseHelperNodes){BuildBoundaryNodes();}
144
145     if(fNPointsI>GetNTNodes()){fNPointsI=GetNTNodes();}
146
147     BuildInterpolation();
148
149     return kTRUE;
150 }
151
152
153 //_________________________________________________________________
154 Bool_t AliTRDTKDInterpolator::Eval(const Double_t *point, Double_t &result, Double_t &error)
155 {
156     Float_t pointF[fNDim]; // local Float_t conversion for "point"
157     for(int idim=0; idim<fNDim; idim++) pointF[idim] = (Float_t)point[idim];
158     Int_t nodeIndex = GetNodeIndex(pointF);
159     if(nodeIndex<0){
160         Error("AliTRDTKDInterpolator::Eval()", "Can not retrive node for data point.");
161         result = 0.;
162         error = 1.E10;
163         return kFALSE;
164     }
165     AliTRDTKDNodeInfo *node =GetNodeInfo(nodeIndex);
166
167     switch(fPDFMode){
168
169     case kInterpolation:
170         return node->CookPDF(point, result, error);
171     case kMinError:
172         node->CookPDF(point, result, error);
173         if(error<node->fVal[1]){
174             return kTRUE;
175         }
176         error=node->fVal[1];
177         result=node->fVal[0];
178         return kTRUE;
179     case kNodeVal:
180         error=node->fVal[1];
181         result=node->fVal[0];
182         return kTRUE;
183     default:
184         return kFALSE;
185     }
186 }
187
188 //__________________________________________________________________
189 void AliTRDTKDInterpolator::Print(const Option_t */*opt*/) const
190 {
191     for(Int_t i=GetNTNodes(); i--;){
192         printf("Node %d of %d: \n",i,GetNTNodes());
193         GetNodeInfo(i)->Print();
194     }
195
196 }
197
198 //__________________________________________________________________
199 Int_t AliTRDTKDInterpolator::GetNodeIndex(const Float_t *p)
200 {
201     Int_t inode=FindNode(p)-fNDataNodes+1;
202     if(GetNodeInfo(inode)->Has(p))return inode;
203
204     // Search extra nodes
205
206     for(inode=fNDataNodes;inode<GetNTNodes();inode++){
207         if(GetNodeInfo(inode)->Has(p)){return inode;}
208     }
209
210     // Search for nearest neighbor
211     Float_t dist;
212     Float_t closestdist=10000;
213     inode=-1;
214     for(Int_t ii=0;ii<GetNTNodes();ii++){
215         AliTRDTKDNodeInfo *node=GetNodeInfo(ii);
216         dist=0;
217         for(Int_t idim=0;idim<fNDim;idim++){
218             dist+=TMath::Power((node->fData[idim]-p[idim]),2);
219         }
220         dist=TMath::Sqrt(dist);
221         if(dist<closestdist){closestdist=dist;inode=ii;}
222     }
223     return inode;
224 }
225
226 //_________________________________________________________________
227 AliTRDTKDInterpolator::AliTRDTKDNodeInfo* AliTRDTKDInterpolator::GetNodeInfo(Int_t inode) const
228 {
229   if(!fNodes || inode >= GetNTNodes()) return NULL;
230   return (AliTRDTKDNodeInfo*)(*fNodes)[inode];
231 }
232
233 //_________________________________________________________________
234 Int_t AliTRDTKDInterpolator::GetNTNodes() const 
235 {
236   return fNodes?fNodes->GetEntriesFast():0;
237 }
238
239 //_________________________________________________________________
240 Bool_t AliTRDTKDInterpolator::GetRange(Int_t idim,Float_t range[2]) const
241 {
242     if(!fNodes) return kFALSE;
243     if(idim<0 || idim>=fNDim){
244     range[0]=0.; range[1]=0.;
245     return kFALSE;
246     }
247     range[0]=1.e10; range[1]=-1.e10;
248     for(Int_t in=GetNTNodes(); in--; ){
249         AliTRDTKDNodeInfo *node = GetNodeInfo(in);
250
251         if(node->fBounds[2*idim]<range[0]) range[0] = node->fBounds[2*idim];
252         if(node->fBounds[2*idim+1]>range[1]) range[1] = node->fBounds[2*idim+1];
253     }
254
255     return kTRUE;
256 }
257
258 //_________________________________________________________________
259 TH2Poly *AliTRDTKDInterpolator::Projection(Int_t xdim,Int_t ydim)
260 {
261     Float_t rangex[2],rangey[2];
262     GetRange(xdim,rangex);
263     GetRange(ydim,rangey);
264
265     TH2Poly* h2 = new TH2Poly("hTKDnodes","hTKDnodes", rangex[0],rangex[1],rangey[0],rangey[1]);
266     h2->GetXaxis()->SetTitle(Form("Q_{%d}", xdim));
267     h2->GetYaxis()->SetTitle(Form("Q_{%d}", ydim));
268
269     for(Int_t inode=0;inode<GetNTNodes();inode++){
270
271         AliTRDTKDNodeInfo* node=GetNodeInfo(inode);
272         h2->AddBin(node->fBounds[2*xdim],node->fBounds[2*ydim],node->fBounds[2*xdim+1],node->fBounds[2*ydim+1]);
273         h2->SetBinContent(inode+1, node->fVal[0]);
274     }
275     return h2;
276 }
277
278 //_________________________________________________________________
279 void AliTRDTKDInterpolator::BuildInterpolation()
280 {
281
282     // Calculate Interpolation
283
284     Double_t buffer[fLambda];
285
286     Float_t **refPoints = new Float_t*[fNDim];
287     for(int id=0; id<fNDim; id++){
288         refPoints[id] = new Float_t[GetNTNodes()];
289         for(int in=0; in<GetNTNodes(); in++) refPoints[id][in] = GetNodeInfo(in)->fData[id];
290     }
291     TKDTreeIF *KDhelper = new TKDTreeIF(GetNTNodes(), fNDim, 30, refPoints);
292     KDhelper->Build();
293     KDhelper->MakeBoundariesExact();
294
295     Float_t dist[fNPointsI];
296     Int_t ind[fNPointsI];
297
298     TLinearFitter fitter(fLambda, Form("hyp%d", fLambda-1));
299
300     Int_t nodeIndex(0); Float_t param[6], *pp(NULL);
301     nodeIndex=GetNTNodes(); pp=&param[0];
302     while(nodeIndex--){
303
304         fitter.ClearPoints();
305
306         AliTRDTKDNodeInfo *node = GetNodeInfo(nodeIndex);
307         // find nearest neighbors
308         KDhelper->FindNearestNeighbors(node->fData,fNPointsI, &ind[0], &dist[0]);
309
310         for(int in=0;in<fNPointsI;in++){
311             AliTRDTKDNodeInfo *nnode = GetNodeInfo(ind[in]);
312
313             Float_t w=1; //weight
314             // calculate tri-cubic weighting function
315             if(fUseWeights){
316                 Float_t d = dist[in]/dist[fNPointsI-1];
317                 Float_t w0 = (1. - d*d*d);
318                 w = w0*w0*w0;
319                 if(w<1.e-30) continue;
320             }
321             Int_t ipar=0;
322             for(int idim=0; idim<fNDim; idim++){
323                 buffer[ipar++] = nnode->fData[idim];
324                 for(int jdim=idim; jdim<fNDim; jdim++) buffer[ipar++] = nnode->fData[idim]*nnode->fData[jdim];
325             }
326             fitter.AddPoint(buffer,nnode->fVal[0], nnode->fVal[1]/w);
327
328             // Ensure Boundary Condition
329             for(Int_t kdim=0;kdim<fNDim;kdim++){
330                 if(node->fBounds[2*kdim]==0){
331                     Float_t zdata[fNDim];
332                     memcpy(&zdata[0],node->fData,fNDim*sizeof(Float_t));
333                     zdata[kdim]=0;
334                     ipar=0;
335                     for(int idim=0; idim<fNDim; idim++){
336                         buffer[ipar++] = zdata[idim];
337                         for(int jdim=idim; jdim<fNDim; jdim++) buffer[ipar++] = zdata[idim]*zdata[jdim];
338                     }
339                     fitter.AddPoint(buffer,0,1);
340                 }
341             }
342         }
343
344         fitter.Eval();
345
346         // retrive fitter results
347         TMatrixD cov(fLambda, fLambda);
348         TVectorD par(fLambda);
349         fitter.GetCovarianceMatrix(cov);
350         fitter.GetParameters(par);
351
352         // store results
353         node->Store(&par,&cov);
354     }
355
356     delete KDhelper;
357     for(int id=0; id<fNDim; id++){
358         delete refPoints[id];
359     }
360     delete[] refPoints;
361 }
362
363 //_________________________________________________________________
364 void AliTRDTKDInterpolator::BuildBoundaryNodes(){
365
366     Int_t nnew=0;
367
368     Float_t treebounds[2*fNDim];
369     for(Int_t idim=0;idim<fNDim;idim++){
370         GetRange(idim,&treebounds[2*idim]);
371     }
372
373     for(int inode=0; inode<GetNTNodes(); inode++){
374
375         AliTRDTKDNodeInfo *node=GetNodeInfo(inode);
376
377         for(Int_t vdim=0;vdim<fNDim;vdim++){
378
379             // Try expansion to lower and higher values
380             for(Int_t iter=0;iter<2;iter++){
381                 if(node->fBounds[2*vdim+iter]==treebounds[2*vdim+iter]){
382
383                     // Add new Node
384                     new ((*fNodes)[GetNTNodes()]) AliTRDTKDNodeInfo(fNDim);
385
386                     AliTRDTKDNodeInfo *newnode = GetNodeInfo(GetNTNodes()-1);
387                     if(iter==0)newnode->fBounds[2*vdim+iter]=0;
388                     if(iter==1)newnode->fBounds[2*vdim+iter]=2*treebounds[2*vdim+iter];
389                     newnode->fBounds[2*vdim+!iter]=node->fBounds[2*vdim+iter];
390                     for(Int_t idim=0;idim<fNDim;idim++){
391                         if(idim==vdim)continue;
392                         newnode->fBounds[2*idim]=node->fBounds[2*idim];
393                         newnode->fBounds[2*idim+1]=node->fBounds[2*idim+1];
394                     }
395                     newnode->fVal[0]=0;
396                     newnode->fVal[1]=Float_t(1)/fNPoints;
397                     for(int idim=0; idim<fNDim; idim++){
398                         newnode->fVal[1] /= (newnode->fBounds[2*idim+1] - newnode->fBounds[2*idim]);
399                         newnode->fData[idim]=0.5*(newnode->fBounds[2*idim+1] + newnode->fBounds[2*idim]);
400                     }
401                     nnew++;
402                 }
403             }
404         }
405     }
406     printf("%d Boundary Nodes Added \n",nnew);
407 }
408
409 //_________________________________________________________________
410 AliTRDTKDInterpolator::AliTRDTKDNodeInfo::AliTRDTKDNodeInfo(Int_t ndim):
411   TObject()
412   ,fNDim(ndim)
413   ,fNBounds(2*ndim)
414   ,fNPar(1 + ndim + (ndim*(ndim+1)>>1))
415   ,fNCov(fNPar*fNPar)
416   ,fData(NULL)
417   ,fBounds(NULL)
418   ,fPar(NULL)
419   ,fCov(NULL)
420 {
421   // Default constructor
422     fVal[0] = 0.; fVal[1] = 0.;
423     fData=new Float_t[fNDim];
424     fBounds=new Float_t[fNBounds];
425 }
426
427 //_________________________________________________________________
428 AliTRDTKDInterpolator::AliTRDTKDNodeInfo::AliTRDTKDNodeInfo(const AliTRDTKDInterpolator::AliTRDTKDNodeInfo &ref):
429 TObject(ref)
430    ,fNDim(ref.fNDim)
431    ,fNBounds(ref.fNBounds)
432    ,fNPar(ref.fNPar)
433    ,fNCov(ref.fNCov)
434    ,fData(NULL)
435    ,fBounds(NULL)
436    ,fPar(NULL)
437    ,fCov(NULL)
438 {
439   // Copy constructor
440
441     if(ref.fData){
442         fData = new Float_t[fNDim];
443         memcpy(fData, ref.fData, fNDim*sizeof(Float_t));
444     }
445     if(ref.fBounds){
446         fBounds = new Float_t[2*fNDim];
447         memcpy(fBounds, ref.fBounds, 2*fNDim*sizeof(Float_t));
448     }
449
450     fVal[0] = ref.fVal[0];
451     fVal[1] = ref.fVal[1];
452
453     if(ref.fPar){
454         fPar=new Double_t[fNPar];
455         memcpy(fPar, ref.fPar, fNPar*sizeof(Double_t));
456     }
457     if(ref.fCov){
458         fCov=new Double_t[fNCov];
459         memcpy(fCov, ref.fCov, fNCov*sizeof(Double_t));
460     }
461 }
462
463 //_________________________________________________________________
464 AliTRDTKDInterpolator::AliTRDTKDNodeInfo::~AliTRDTKDNodeInfo()
465 {
466   // Destructor
467   if(fData) delete [] fData;
468   if(fBounds) delete [] fBounds;
469   if(fPar) delete [] fPar;
470   if(fCov) delete [] fCov;
471 }
472
473 //_________________________________________________________________
474 AliTRDTKDInterpolator::AliTRDTKDNodeInfo &AliTRDTKDInterpolator::AliTRDTKDNodeInfo::operator=(const AliTRDTKDInterpolator::AliTRDTKDNodeInfo &ref)
475 {
476     //
477     // Assignment operator
478     //
479
480     if(&ref==this) return *this;
481     memcpy(fData, ref.fData, fNDim*sizeof(Float_t));
482     memcpy(fBounds, ref.fBounds, 2*fNDim*sizeof(Float_t));
483
484     fVal[0] = ref.fVal[0];
485     fVal[1] = ref.fVal[1];
486
487     if(ref.fPar){
488         fPar=new Double_t[fNPar];
489         memcpy(fPar, ref.fPar, fNPar*sizeof(Double_t));
490     }
491     if(ref.fCov){
492         fCov=new Double_t[fNCov];
493         memcpy(fCov, ref.fCov, fNCov*sizeof(Double_t));
494     }
495     return *this;
496 }
497
498
499 //_________________________________________________________________
500 void AliTRDTKDInterpolator::AliTRDTKDNodeInfo::Print(const Option_t *opt) const
501 {
502   // Print the content of the node
503   printf("x [");
504   for(int idim=0; idim<fNDim; idim++) printf("%f ", fData?fData[idim]:0.);
505   printf("] f = [%e +- %e]\n", fVal[0], fVal[1]);
506
507   if(fBounds){
508       printf("range [");
509       for(int idim=0; idim<fNDim; idim++) printf("(%f %f) ", fBounds[2*idim], fBounds[2*idim+1]);
510     printf("]\n");
511   }
512   if(strcmp(opt, "a")!=0) return;
513
514   if(fPar){
515     printf("Fit parameters : \n");
516     for(int ip=0; ip<fNPar; ip++) printf("p%d[%f] ", ip, fPar[ip]);
517     printf("\n");
518   }
519   if(!fCov) return;
520   for(int ip(0), n(0); ip<fNPar; ip++){
521     for(int jp(ip); jp<fNPar; jp++) printf("c(%d %d)[%f] ", ip, jp, fCov[n++]);
522     printf("\n");
523   }
524 }
525
526 //_________________________________________________________________
527 void AliTRDTKDInterpolator::AliTRDTKDNodeInfo::Store(TVectorD const *par, TMatrixD const *cov)
528 {
529 // Store the parameters and the covariance in the node
530
531     if(!fPar){fPar = new Double_t[fNPar];}
532     for(int ip=0; ip<fNPar; ip++) fPar[ip] = (*par)[ip];
533
534     if(!cov) return;
535     if(!fCov){fCov = new Double_t[fNCov];}
536     for(int ip(0), np(0); ip<fNPar; ip++)
537         for(int jp=ip; jp<fNPar; jp++) fCov[np++] = (*cov)(ip,jp);
538 }
539
540 //_________________________________________________________________
541 Bool_t AliTRDTKDInterpolator::AliTRDTKDNodeInfo::CookPDF(const Double_t *point, Double_t &result, Double_t &error) const
542 {
543     // Recalculate the PDF for one node from the results of interpolation (parameters and covariance matrix)
544
545     result =0.; error = 1.;
546     if(!fPar) return kFALSE;
547
548     Double_t fdfdp[fNDim];
549     Int_t ipar = 0;
550     fdfdp[ipar++] = 1.;
551     for(int idim=0; idim<fNDim; idim++){
552         fdfdp[ipar++] = point[idim];
553         for(int jdim=idim; jdim<fNDim; jdim++) fdfdp[ipar++] = point[idim]*point[jdim];
554     }
555
556     // calculate estimation
557     for(int i=0; i<fNPar; i++) result += fdfdp[i]*fPar[i];
558
559     if(!fCov)return kTRUE;
560     // calculate error
561     error=0;
562
563     for(int i(0), n(0); i<fNPar; i++){
564         error += fdfdp[i]*fdfdp[i]*fCov[n++];
565         for(int j(i+1); j<fNPar; j++) error += 2.*fdfdp[i]*fdfdp[j]*fCov[n++];
566     }
567     if(error>0)error = TMath::Sqrt(error);
568     else{error=100;}
569
570     // Boundary condition
571     if(result<0){
572         result=fVal[0];
573         error=fVal[1];
574     }
575
576     return kTRUE;
577 }
578
579 //_____________________________________________________________________
580 Bool_t AliTRDTKDInterpolator::AliTRDTKDNodeInfo::Has(const Float_t *p) const
581 {
582   Int_t n(0);
583   for(int id=0; id<fNDim; id++)
584       if(p[id]>=fBounds[2*id] && p[id]<fBounds[2*id+1]) n++;
585   if(n==fNDim) return kTRUE;
586   return kFALSE;
587 }
588