update DistSampler to support discrete distributions

git-svn-id: http://root.cern.ch/svn/root/trunk@34077 27541ba8-7e3a-0410-8455-c3a389f83636

math/mathcore/inc/Math/DistSampler.h

@@ -129,6 +129,14 @@ public:
    /// set range using DataRange class 
    void SetRange(const ROOT::Fit::DataRange & range); 
 
+   /// set the mode of the distribution (could be useful to some methods)
+   /// implemented by derived classes if needed
+   virtual void SetMode(double  ) {}
+
+   /// set the normalization area of distribution 
+   /// implemented by derived classes if needed
+   virtual void SetArea(double) {}
+
    /// get the parent distribution function (must be called after setting the function) 
    const ROOT::Math::IMultiGenFunction & ParentPdf() const { 
       return *fFunc;

math/mathcore/test/testDistSampler.cxx

@@ -8,13 +8,16 @@
 #include "TStopwatch.h"
 #include "TRandom.h"
 #include "TError.h"
+#include "TCanvas.h"
 // double Pdf(double x) { 
 // }
 
 using namespace ROOT::Math; 
 
-int testDistSampler(int n = 10000) { 
 
+int testCont1D(int n)  { 
+
+   // test gaussian
 
    DistSampler * sampler = Factory::CreateDistSampler("Unuran"); 
    if (!sampler) return -1;
@@ -37,6 +40,8 @@ int testDistSampler(int n = 10000) {
    w.Stop(); 
    double c = 1.E9/double(n);
    std::cout << "Unuran sampling - (ns)/call = " << c*w.RealTime() << "   " << c*w.CpuTime() << std::endl;
+   new TCanvas("Continous test");
+   h1->SetLineColor(kBlue);
    h1->Draw();
 
    // generate ref histogram 
@@ -62,6 +67,70 @@ int testDistSampler(int n = 10000) {
    return 0;
 }
 
+
+int testDisc1D(int n)  { 
+
+   // test a discrete distribution
+
+
+   DistSampler * sampler = Factory::CreateDistSampler("Unuran"); 
+   if (!sampler) return -1;
+   TF1 * f = new TF1("pdf","TMath::Poisson(x,[0])");
+   double mu = 10;
+   f->SetParameter(0,mu);
+
+   TH1D * h1 = new TH1D("h2","h1",50,0,50);
+   TH1D * hr = new TH1D("hd","h2",50,0,50);
+
+   sampler->SetFunction(*f,1);
+   sampler->SetMode(mu);
+   sampler->SetArea(1);
+   bool ret = sampler->Init("DARI"); 
+   if (!ret)      return -1;
+
+
+   TStopwatch w; 
+   w.Start();
+   for (int i = 0; i < n; ++i) { 
+      h1->Fill(sampler->Sample1D() );
+   }
+   w.Stop(); 
+   double c = 1.E9/double(n);
+   std::cout << "Unuran sampling - (ns)/call = " << c*w.RealTime() << "   " << c*w.CpuTime() << std::endl;
+   new TCanvas("Discrete test");
+   h1->SetLineColor(kBlue);
+   h1->Draw();
+
+   // generate ref histogram 
+   w.Start();
+   for (int i = 0; i < n; ++i) { 
+      hr->Fill(gRandom->Poisson(mu) );
+   }
+   w.Stop(); 
+   std::cout << "TRandom::Poisson sampling - (ns)/call = " << c*w.RealTime() << "   " << c*w.CpuTime() << std::endl;
+
+   hr->Draw("SAME");
+
+   // do a Chi2 test 
+   // switch off printing of  info messages from chi2 test
+   gErrorIgnoreLevel = 1001; 
+   double prob = h1->Chi2Test(hr,"UU");
+   if (prob < 1.E-6) { 
+      std::cerr << "Chi2 test of generated histogram failed" << std::endl;
+      return -2;
+   }
+   gErrorIgnoreLevel = 0;
+
+   return 0;
+}
+
+
+int testDistSampler(int n = 10000) { 
+   int iret = 0;
+   iret |= testCont1D(n); 
+   iret |= testDisc1D(n);
+   return iret; 
+}
 int main() { 
    int iret = testDistSampler(); 
    if (iret)  std::cerr << "\ntestDistSampler: ....  FAILED!" << std::endl;

math/unuran/inc/TUnuranDiscrDist.h

@@ -57,7 +57,7 @@ public:
    /** 
       Constructor from a generic function object specifying the pdf
    */ 
-   TUnuranDiscrDist (const ROOT::Math::IGenFunction & func );
+   TUnuranDiscrDist (const ROOT::Math::IGenFunction & func, bool copyFunc = false );
 
    /** 
       Constructor from a TF1 objects specifying the pdf

math/unuran/inc/TUnuranSampler.h

@@ -89,6 +89,22 @@ public:
    */ 
    void SetSeed(unsigned int seed);
 
+   /* 
+      set the mode
+    */
+   void SetMode(double mode) { 
+      fMode = mode; 
+      fHasMode = true; 
+   }
+
+   /*
+     set the area 
+    */
+   void SetArea(double area) { 
+      fArea = area;
+      fHasArea = true; 
+   }
+
    /**
       Get the random engine used by the sampler
     */
@@ -127,6 +143,8 @@ protected:
 
    //initialization for 1D distributions
    bool DoInit1D(const char * algo); 
+   //initialization for 1D discrete distributions
+   bool DoInitDiscrete1D(const char * algo); 
    //initialization for multi-dim distributions
    bool DoInitND(const char * algo); 
    
@@ -135,6 +153,11 @@ private:
 
    // private member
    bool                              fOneDim;      // flag to indicate if the function is 1 dimension
+   bool                              fDiscrete;    // flag to indicate if the function is discrete
+   bool                              fHasMode;     // flag to indicate if a mode is set
+   bool                              fHasArea;     // flag to indicate if a area is set
+   double                            fMode;        // mode of dist
+   double                            fArea;        // area of dist
    const ROOT::Math::IGenFunction *  fFunc1D;      // 1D function pointer
    TUnuran *                         fUnuran;      // unuran engine class
 

math/unuran/src/TUnuranDiscrDist.cxx

@@ -20,7 +20,7 @@
 #include <cassert>
 
 
-TUnuranDiscrDist::TUnuranDiscrDist (const ROOT::Math::IGenFunction & func) : 
+TUnuranDiscrDist::TUnuranDiscrDist (const ROOT::Math::IGenFunction & func, bool copyFunc) : 
    fPmf(&func), 
    fCdf(0), 
    fXmin(1), 
@@ -30,9 +30,13 @@ TUnuranDiscrDist::TUnuranDiscrDist (const ROOT::Math::IGenFunction & func) :
    fHasDomain(0),
    fHasMode(0),
    fHasSum(0), 
-   fOwnFunc(false)
+   fOwnFunc(copyFunc)
 {
    //Constructor from a generic function object
+   if (fOwnFunc) { 
+      fPmf = fPmf->Clone(); 
+      //if (fCdf) fCdf->Clone();
+   }
 } 
 
 

math/unuran/src/TUnuranSampler.cxx

@@ -12,6 +12,7 @@
 #include "TUnuranSampler.h"
 
 #include "TUnuranContDist.h"
+#include "TUnuranDiscrDist.h"
 #include "TUnuranMultiContDist.h"
 #include "TUnuran.h"
 #include "Math/OneDimFunctionAdapter.h"
@@ -29,6 +30,9 @@ ClassImp(TUnuranSampler)
 
 TUnuranSampler::TUnuranSampler() : ROOT::Math::DistSampler(), 
    fOneDim(false), 
+   fDiscrete(false),
+   fHasMode(false), fHasArea(false),
+   fMode(0), fArea(0),
    fFunc1D(0),
    fUnuran(new TUnuran()  )
 {}
@@ -45,7 +49,15 @@ bool TUnuranSampler::Init(const char * algo) {
       Error("TUnuranSampler::Init","Distribution function has not been set ! Need to call SetFunction first.");
       return false;
    }
-   if (NDim() == 1) return DoInit1D(algo); 
+   TString method(algo); 
+   method.ToUpper();
+
+   if (NDim() == 1) { 
+      // check if distribution is discrete by 
+      // using first string in the method name is "D"
+      if (method.First("D") == 0) return DoInitDiscrete1D(algo);
+      return DoInit1D(algo); 
+   }
    else return DoInitND(algo); 
 }
 
@@ -54,25 +66,66 @@ bool TUnuranSampler::DoInit1D(const char * method) {
    // need to create 1D interface from Multidim one 
    // (to do: use directly 1D functions ??)
    fOneDim = true; 
-   TUnuranContDist dist;
+   TUnuranContDist * dist = 0;
    if (fFunc1D == 0) { 
       ROOT::Math::OneDimMultiFunctionAdapter<> function(ParentPdf() ); 
-      dist = TUnuranContDist(function,0,false,true); 
+      dist = new TUnuranContDist(function,0,false,true); 
    }
    else { 
-      dist = TUnuranContDist(*fFunc1D); // no need to copy the function
+      dist = new TUnuranContDist(*fFunc1D); // no need to copy the function
    }
    // set range in distribution (support only one range)
    const ROOT::Fit::DataRange & range = PdfRange(); 
    if (range.Size(0) > 0) { 
       double xmin, xmax; 
       range.GetRange(0,xmin,xmax); 
-      dist.SetDomain(xmin,xmax); 
+      dist->SetDomain(xmin,xmax); 
    }
-   if (method) return fUnuran->Init(dist, method);       
-   return fUnuran->Init(dist);
+   if (fHasMode) dist->SetMode(fMode);
+   if (fHasArea) dist->SetPdfArea(fArea);
+
+   bool ret = false; 
+   if (method) ret =  fUnuran->Init(*dist, method);       
+   else ret =  fUnuran->Init(*dist);
+   delete dist; 
+   return ret; 
 }
 
+bool TUnuranSampler::DoInitDiscrete1D(const char * method) { 
+   // initilize for 1D sampling of discrete distributions
+   fOneDim = true; 
+   fDiscrete = true;
+   TUnuranDiscrDist * dist = 0;
+   if (fFunc1D == 0) { 
+      // need to copy the passed function pointer in this case
+      ROOT::Math::OneDimMultiFunctionAdapter<> function(ParentPdf() ); 
+      dist = new TUnuranDiscrDist(function,true); 
+   }
+   else { 
+      // no need to copy the function since fFunc1D is managed outside
+      dist = new TUnuranDiscrDist(*fFunc1D, false); 
+   }
+   // set range in distribution (support only one range)
+   // otherwise 0, inf is assumed
+   const ROOT::Fit::DataRange & range = PdfRange(); 
+   if (range.Size(0) > 0) { 
+      double xmin, xmax; 
+      range.GetRange(0,xmin,xmax);
+      if (xmin < 0) { 
+         Warning("DoInitDiscrete1D","range starts from negative values - set minimum to zero"); 
+         xmin = 0; 
+      }
+      dist->SetDomain(int(xmin+0.1),int(xmax+0.1)); 
+   }
+   if (fHasMode) dist->SetMode(fMode);
+   if (fHasArea) dist->SetProbSum(fArea);
+
+   bool ret =  fUnuran->Init(*dist, method);       
+   delete dist;
+   return ret;
+}
+
+
 bool TUnuranSampler::DoInitND(const char * method) { 
    // initilize for 1D sampling
    TUnuranMultiContDist dist(ParentPdf()); 
@@ -111,7 +164,7 @@ TRandom * TUnuranSampler::GetRandom() {
 
 double TUnuranSampler::Sample1D() { 
    // sample 1D distributions
-   return fUnuran->Sample(); 
+   return (fDiscrete) ? (double) fUnuran->SampleDiscr() : fUnuran->Sample(); 
 }
 
 bool TUnuranSampler::Sample(double * x) {