STAT/Macros/kNNTest.C

   1 #include <malloc.h>
   2
   3 #include "TSystem.h"
   4 #include "TFile.h"
   5 #include "TTree.h"
   6 #include "TH2I.h"
   7 #include "TLegend.h"
   8 #include "TRandom.h"
   9 #include "TString.h"
  10 #include "TGraph.h"
  11 #include "TMarker.h"
  12 #include "TStopwatch.h"
  13 #include "../src/TKDPDF.h"
  14 #include "../src/TKDTree.h"
  15
  16 void kNNTest(const int np = 10000, const int ndim = 2);
  17 void kNNDraw(const Float_t *p, const int kNN=20);
  18 void build(const Int_t ndim = 2, const Int_t nstat = 1000000);
  19 Float_t Mem();
  20
  21
  22 //______________________________________________________________
  23 void kNNTest(const int np, const int ndim)
  24 {
  25 // Test speed and quality of nearest neighbors search.
  26 // The results are compared with a simple loop search through the data.
  27 // The macro should run in compiled mode.
  28
  29         const int ntimes = 1000; // times to repeate kNN search
  30         const Int_t kNN = 1; // number of neighbors
  31
  32         Float_t **x = new Float_t*[ndim];
  33         for(int idim =0 ; idim<ndim; idim++){
  34                 x[idim] = new Float_t[np];
  35                 for(int ip=0; ip<np; ip++) x[idim][ip] = gRandom->Gaus();
  36         }
  37         TKDTreeIF nnFinder(np, ndim, 1, x);
  38
  39         // generate test sample
  40         Float_t **p = new Float_t*[ntimes];
  41         for(int ip =0 ; ip<ntimes; ip++){
  42                 p[ip] = new Float_t[ndim];
  43                 for(int idim=0; idim<ndim; idim++) p[ip][idim] = gRandom->Gaus();
  44         }
  45
  46
  47         Int_t *index;
  48         Float_t *d;
  49         TStopwatch timer;
  50         timer.Start(kTRUE);
  51         for(int itime=0; itime<ntimes; itime++) nnFinder.FindNearestNeighbors(p[itime], kNN, index, d);
  52         timer.Stop();
  53         printf("kDTree NN calculation.\n");
  54         printf("cpu = %f [ms] real = %f [ms]\n", 1.E3*timer.CpuTime()/ntimes, 1.E3*timer.RealTime()/ntimes);
  55         printf("Points indexes in increasing order of their distance to the target point.\n");
  56         for(int i=0; i<kNN; i++){
  57                 printf("%5d [%7.4f] ", index[i], d[i]);
  58                 if(i%5==4) printf("\n");
  59         }
  60         printf("\n");
  61
  62         // draw kNN
  63         TLegend *leg = new TLegend(.7, .7, .9, .9);
  64         leg->SetBorderSize(1);
  65         leg->SetHeader("NN finders");
  66         TH2 *h2 = new TH2I("h2", "", 100, -2., 2., 100, -2., 2.);
  67         h2->Draw(); h2->SetStats(kFALSE);
  68         TMarker *mp = new TMarker(p[ntimes-1][0], p[ntimes-1][1], 3);
  69         mp->SetMarkerColor(4);
  70         mp->Draw(); leg->AddEntry(mp, "Target", "p");
  71         TGraph *gKD = new TGraph(kNN);
  72         gKD->SetMarkerStyle(24);
  73         gKD->SetMarkerColor(2);
  74         gKD->SetMarkerSize(.5);
  75         for(int ip=0; ip<kNN; ip++) gKD->SetPoint(ip, x[0][index[ip]], x[1][index[ip]]);
  76         gKD->Draw("p"); leg->AddEntry(gKD, "kDTree", "p");
  77
  78
  79
  80         // STAND ALONE
  81         Float_t ftmp, gtmp, dist;
  82         Int_t itmp, jtmp;
  83         Int_t   fkNN[kNN];
  84         Float_t fkNNdist[kNN];
  85         for(int i=0; i<kNN; i++) fkNNdist[i] = 9999.;
  86
  87         // calculate
  88         timer.Start(kTRUE);
  89         for(int idx=0; idx<np; idx++){
  90                 // calculate distance in the L1 metric
  91                 dist = 0.;
  92                 for(int idim=0; idim<ndim; idim++) dist += TMath::Abs(p[ntimes-1][idim] - x[idim][idx]);
  93                 if(dist >= fkNNdist[kNN-1]) continue;
  94
  95                 //insert neighbor
  96                 int iNN=0;
  97                 while(dist >= fkNNdist[iNN]) iNN++;
  98                 itmp = fkNN[iNN]; ftmp = fkNNdist[iNN];
  99                 fkNN[iNN]     = idx;
 100                 fkNNdist[iNN] = dist;
 101                 for(int ireplace=iNN+1; ireplace<kNN; ireplace++){
 102                         jtmp = fkNN[ireplace]; gtmp = fkNNdist[ireplace];
 103                         fkNN[ireplace] = itmp; fkNNdist[ireplace] = ftmp;
 104                         itmp = jtmp; ftmp = gtmp;
 105                         if(ftmp == 9999.) break;
 106                 }
 107         }
 108         timer.Stop();
 109         printf("\nStand Alone NN calculation.\n");
 110         printf("cpu = %f [ms] real = %f [ms]\n", 1.E3*timer.CpuTime(), 1.E3*timer.RealTime());
 111         printf("Points indexes in increasing order of their distance to the target point.\n");
 112         for(int i=0; i<kNN; i++){
 113                 printf("%5d [%7.4f] ", fkNN[i], fkNNdist[i]);
 114                 if(i%5==4) printf("\n");
 115         }
 116         printf("\n");
 117
 118         TGraph *gSA = new TGraph(kNN);
 119         gSA->SetMarkerStyle(24);
 120         for(int ip=0; ip<kNN; ip++) gSA->SetPoint(ip, x[0][fkNN[ip]], x[1][fkNN[ip]]);
 121
 122         gSA->Draw("p"); leg->AddEntry(gSA, "Stand Alone", "p");
 123         leg->Draw();
 124
 125         for(int ip=0; ip<ntimes; ip++) delete [] p[ip];
 126         delete [] p;
 127         for(int idim=0; idim<ndim; idim++) delete [] x[idim];
 128         delete [] x;
 129 }
 130
 131
 132 //______________________________________________________________
 133 Float_t p[]={1.4, -.6}; //}{1.7, -.4};
 134 void kNNDraw(const Float_t *p, const int kNN)
 135 {
 136 // Test memory consumption and draw "kNN" nearest neighbours of point "p".
 137 // The distance (in the L1 metric) is encoded in the color code.
 138
 139         const Int_t npoints = 10000;
 140         Float_t *data0 =  new Float_t[npoints*2];
 141   Float_t *data[2];
 142   data[0] = &data0[0];
 143   data[1] = &data0[npoints];
 144   for (Int_t i=0;i<npoints;i++) {
 145     data[1][i]= gRandom->Gaus();
 146     data[0][i]= gRandom->Gaus();
 147   }
 148
 149         TKDPDF pdf(npoints, 2, 100, data);
 150         pdf.DrawNode(pdf.GetNodeIndex(p));
 151
 152         TMarker *mp = new TMarker(p[0], p[1], 3);
 153         mp->SetMarkerColor(4);
 154         mp->Draw();
 155
 156         Int_t *index, color;
 157         Float_t *d, d0, pknn[2];
 158         Float_t start = Mem();
 159         pdf.FindNearestNeighbors(p, kNN, index, d);
 160         Float_t end = Mem();
 161         printf("FindNearestNeighbors memory usage %fKB\n", end-start);
 162         d0 = d[kNN-1];
 163         TMarker *ma = new TMarker[kNN];
 164         for(int ip=0; ip<kNN; ip++){
 165                 pdf.GetDataPoint(index[ip], pknn);
 166                 color = 101 - Int_t(50. * d[ip]/d0);
 167                 ma[ip].SetMarkerStyle(4);
 168                 ma[ip].SetMarkerColor(color);
 169                 ma[ip].DrawMarker(pknn[0], pknn[1]);
 170         }
 171 }
 172
 173
 174 //______________________________________________________________
 175 Float_t Mem()
 176 {
 177   // size in KB
 178   static struct mallinfo memdebug;
 179   memdebug = mallinfo();
 180   return memdebug.uordblks/1024.;
 181 }
 182
 183 //______________________________________________________________
 184 void build(const Int_t ndim, const Int_t nstat)
 185 {
 186 // Build "nstat" data points in "ndim" dimensions taken from an
 187 // uncorrelated 2D Gauss distribution.
 188
 189
 190         printf("build data ... \n");
 191         Double_t pntTrain[ndim];
 192         TFile *f = TFile::Open(Form("%dD_Gauss.root", ndim), "RECREATE");
 193         TTree *t = new TTree("db", "gauss database");
 194         for(int idim=0; idim<ndim; idim++) t->Branch(Form("x%d", idim), &pntTrain[idim], Form("x%d/D", idim));
 195
 196         for(int istat=0; istat<nstat; istat++){
 197                 for(int idim=0; idim<ndim; idim++) pntTrain[idim] = gRandom->Gaus();
 198                 t->Fill();
 199         }
 200         f->cd();
 201         t->Write();
 202         f->Close();
 203         delete f;
 204 }
 205