IO factorization for Interpolator
[u/mrichter/AliRoot.git] / STAT / Macros / testKNN.C
1 Float_t p[]={1.4, -.6}; //}{1.7, -.4};
2 //______________________________________________________________
3 void testKNN(const Float_t *p, const int kNN=20)
4 {
5 // Draw "kNN" nearest neighbours of point "p". The distance (in the L1
6 // metric) is encoded in the color code.
7 // To build the data refere to function build().
8
9         TFile::Open("2D_Gauss.root");
10         TKDInterpolator in(db, "x0:x1", "x0>-1.5&&x0<2.&&x1>-2.&&x1<2.", 300);
11         in.DrawNode(in.FindNode(p)-in.GetNNodes());
12         
13         TMarker *mp = new TMarker(p[0], p[1], 3);
14         mp->SetMarkerColor(4);
15         mp->Draw();
16                 
17         Int_t *index, color;
18         Float_t d, d0, pknn[2];
19         in.FindNearestNeighbors(p, kNN, index, d0);
20         TMarker *ma = new TMarker[kNN];
21         for(int ip=0; ip<kNN; ip++){
22                 in.GetDataPoint(index[ip], pknn);
23                 d = TMath::Abs(p[0]-pknn[0]) + TMath::Abs(p[1]-pknn[1]);
24                 ma[ip].SetMarkerStyle(4);
25                 color = 101 - Int_t(50. * d/d0);
26                 ma[ip].SetMarkerColor(color);
27                 ma[ip].DrawMarker(pknn[0], pknn[1]);
28         }
29 }
30
31 void performanceKNN(const int np = 10000)
32 {
33         const int ntimes = 1000; // times to repeate kNN search
34         const int ndim = 5;      // no of dimensions for data
35
36         if(!gSystem->FindFile(".", Form("%dD_Gauss.root", ndim))) build(ndim);
37         TFile::Open(Form("%dD_Gauss.root", ndim));
38         TString var="x0";
39         for(int i=1; i<ndim; i++) var+=Form(":x%d", i);
40         TKDInterpolator in(db, var.Data(), "", 30, np);
41         
42         Int_t *index, fNN = 20;
43         Float_t *d, p[ndim];
44         for(int idim=0; idim<ndim; idim++) p[idim] = gRandom->Gaus();
45         TStopwatch time;
46         time.Start();
47         for(int itime=0; itime<ntimes; itime++) in.FindNearestNeighbors(p, fNN, index, d);
48         time.Stop();
49         printf("cpu = %f [mus] real = %f [mus]\n", 1.E6*time.CpuTime()/ntimes, 1.E6*time.RealTime()/ntimes);
50         for(int i=0; i<fNN; i++){
51                 printf("%5d ", index[i]);
52                 if(i%5==4) printf("\n");
53         }
54         
55         // draw kNN
56         TH2 *h2 = new TH2I("h2", "", 100, -2., 2., 100, -2., 2.);
57         h2->Draw();
58         TGraph *gKD = new TGraph(fNN);
59         gKD->SetMarkerStyle(24);
60         gKD->SetMarkerColor(2);
61         gKD->SetMarkerSize(.5);
62         Float_t pknn[ndim];
63         for(int ip=0; ip<fNN; ip++){
64                 in.GetDataPoint(index[ip], pknn);
65                 gKD->SetPoint(ip, pknn[0], pknn[1]);
66         }
67         gKD->Draw("p");
68         TMarker *mp = new TMarker(p[0], p[1], 3);
69         mp->SetMarkerColor(4);
70         mp->Draw();
71
72         
73         // STAND ALONE
74         const Int_t kNN = fNN;
75         Float_t ftmp, gtmp, dist;
76         Int_t itmp, jtmp;
77         Int_t   fkNN[kNN];
78         Float_t fkNNdist[kNN];
79         for(int i=0; i<kNN; i++) fkNNdist[i] = 9999.;
80
81         // read data in memory
82         Int_t npoints = db->Draw("x0", "", "goff", np);
83         Float_t **x = new Float_t*[ndim];
84         Double_t *v;
85         for(int idim=0; idim<ndim; idim++){
86                 x[idim] = new Float_t[npoints];
87                 db->Draw(Form("x%d", idim), "", "goff", np); v = db->GetV1();
88                 for(int i=0; i<npoints; i++) x[idim][i] = v[i];
89         }
90         // calculate
91         printf("stand alone calculation for %d neighbors in %d points\n", kNN, npoints);
92         time.Start(kTRUE);
93         for(int idx=0; idx<npoints; idx++){
94                 // calculate distance in the L1 metric
95                 dist = 0.;
96                 for(int idim=0; idim<ndim; idim++) dist += TMath::Abs(p[idim] - x[idim][idx]);
97                 if(dist >= fkNNdist[kNN-1]) continue;
98
99                 //insert neighbor
100                 int iNN=0;
101                 while(dist >= fkNNdist[iNN]) iNN++;
102                 itmp = fkNN[iNN]; ftmp = fkNNdist[iNN];
103                 fkNN[iNN]     = idx;
104                 fkNNdist[iNN] = dist;
105                 for(int ireplace=iNN+1; ireplace<kNN; ireplace++){
106                         jtmp = fkNN[ireplace]; gtmp = fkNNdist[ireplace];
107                         fkNN[ireplace] = itmp; fkNNdist[ireplace] = ftmp;
108                         itmp = jtmp; ftmp = gtmp;
109                         if(ftmp == 9999.) break;
110                 }
111         }
112         time.Stop();
113         printf("cpu = %f [s] real = %f [s]\n", time.CpuTime(), time.RealTime());
114         
115         for(int i=0; i<kNN; i++){
116                 printf("%5d ", fkNN[i]);
117                 if(i%5==4) printf("\n");
118         }
119         
120         TGraph *gSA = new TGraph(kNN);
121         gSA->SetMarkerStyle(24);
122         for(int ip=0; ip<kNN; ip++) gSA->SetPoint(ip, x[0][fkNN[ip]], x[1][fkNN[ip]]);
123
124         gSA->Draw("p");
125 }
126
127 //______________________________________________________________
128 void build(const Int_t ndim = 2, const Int_t nstat = 1000000)
129 {
130 // Build "nstat" data points in "ndim" dimensions taken from an
131 // uncorrelated 2D Gauss distribution.
132         
133
134         printf("build data ... \n");
135         Double_t pntTrain[ndim];
136         f = TFile::Open(Form("%dD_Gauss.root", ndim), "RECREATE");
137         t = new TTree("db", "gauss database");
138         for(int idim=0; idim<ndim; idim++) t->Branch(Form("x%d", idim), &pntTrain[idim], Form("x%d/D", idim));
139         
140         for(int istat=0; istat<nstat; istat++){
141                 for(int idim=0; idim<ndim; idim++) pntTrain[idim] = gRandom->Gaus();
142                 t->Fill();
143         }
144         f->cd();
145         t->Write();
146         f->Close();
147         delete f;
148 }