kNN algorithm improved. IO Defined
[u/mrichter/AliRoot.git] / STAT / TKDInterpolator.h
1 #ifndef ROOT_TKDInterpolator
2 #define ROOT_TKDInterpolator
3
4 #ifndef ROOT_TKDTree
5 #include "TKDTree.h"
6 #endif
7 #ifndef ROOT_TVectorD
8 #include "TVectorD.h"
9 #endif
10 #ifndef ROOT_TMatrixD
11 #include "TMatrixD.h"
12 #endif
13
14
15 class TTree;
16 class TLinearFitter;
17 class TKDInterpolator : public TKDTreeIF
18 {
19 public:
20         struct TKDNodeInfo
21         {
22
23                 ClassDef(TKDNodeInfo, 1) // node info for interpolator
24         };
25 public:
26         TKDInterpolator();
27         TKDInterpolator(TTree *t, const Char_t *var, const Char_t *cut = 0, UInt_t bsize = 100, Long64_t nentries = 1000000000, Long64_t firstentry = 0);
28         TKDInterpolator(Int_t npoints, Int_t ndim, UInt_t bsize, Float_t **data);
29         ~TKDInterpolator();
30
31                 Double_t   Eval(const Double_t *point, Double_t &result, Double_t &error);
32         inline Bool_t     GetCOGPoint(Int_t node, Float_t *coord, Float_t &val, Float_t &error) const ;
33         inline Bool_t     GetDataPoint(Int_t n, Float_t *p) const;
34                 TKDTreeIF* GetHelper() {return fKDhelper;}
35         inline Bool_t     GetIntInterpolation() const {return fStatus&1;}
36         inline Bool_t     GetSetStore() const {return fStatus&2;}
37         inline Bool_t     GetUseWeights() const {return fStatus&4;}
38                                         
39                                         void       DrawNodes(UInt_t ax1 = 0, UInt_t ax2 = 1, Int_t depth = -1);
40                 void       DrawNode(Int_t tnode, UInt_t ax1 = 0, UInt_t ax2 = 1);
41                 void       GetStatus();
42                 void       SetIntInterpolation(const Bool_t on = kTRUE);
43                 void       SetSetStore(const Bool_t on = kTRUE);
44                 void       SetUseWeights(const Bool_t on = kTRUE);
45         
46 private:
47         TKDInterpolator(const TKDInterpolator &);
48         TKDInterpolator& operator=(const TKDInterpolator &);    
49                 void       Build();
50                 Double_t   CookPDF(const Double_t *point, const Int_t node, Double_t &result, Double_t &error);
51                                         
52 protected:
53         Int_t     fNTNodes;        //Number of evaluation data points
54         Float_t   **fRefPoints;    //[fNDim]
55         Float_t   *fRefValues;     //[fNTNodes]
56         TMatrixD  *fCov;           //[fNTNodes] cov matrix array for nodes
57         TVectorD  *fPar;           //[fNTNodes] parameters array for nodes
58         Bool_t    *fPDFstatus;     //[fNTNodes] status bit for node's PDF
59
60 private:
61         UChar_t   fStatus;         // status of the interpolator
62         Int_t     fLambda;         // number of parameters in polynom
63         Int_t                   fDepth;          //! depth of the KD Tree structure used
64         Double_t        *fBuffer;        //! working space [2*fLambda]
65         TKDTreeIF *fKDhelper;      //! kNN finder
66         TLinearFitter *fFitter;    //! linear fitter    
67
68         ClassDef(TKDInterpolator, 1)   // data interpolator based on KD tree
69 };
70
71 //__________________________________________________________________
72 Bool_t  TKDInterpolator::GetCOGPoint(Int_t node, Float_t *coord, Float_t &val, Float_t &error) const
73 {
74         if(node < 0 || node > fNTNodes) return kFALSE;
75
76         for(int idim=0; idim<fNDim; idim++) coord[idim] = fRefPoints[idim][node];
77         val   = fRefValues[node];
78         error = fRefValues[node]/TMath::Sqrt(fBucketSize);
79         return kTRUE;
80 }
81
82 //__________________________________________________________________
83 Bool_t  TKDInterpolator::GetDataPoint(Int_t n, Float_t *p) const
84 {
85         if(n < 0 || n >= fNpoints) return kFALSE;
86         
87         for(int i=0; i<fNDim; i++) p[i] = fData[i][n];
88         return kTRUE;
89 }
90
91
92 #endif
93