r - 使用 Rcpp 和 openMP 从截断正态分布快速采样

Question

更新:

我尝试实现德克的建议。评论？我现在在 JSM 很忙，但我想在为画廊编织 Rmd 之前得到一些反馈。我从 Armadillo 切换回普通 Rcpp，因为它没有增加任何值(value)。带 R::的标量版本非常好。如果将平均值/标准差作为标量输入，而不是作为所需输出长度的向量，我也许应该输入一个参数 n 来表示绘制次数。

<小时/>

有很多 MCMC 应用程序需要从截断正态分布中抽取样本。我在 TN 的现有实现的基础上构建并添加了并行计算。

问题:

1. 有人认为速度还有进一步提升的潜力吗？ 在基准测试的最后一个案例中，rtruncnorm 有时更快。 Rcpp 实现总是比现有包更快，但它可以进一步改进吗？
2. 我在一个无法共享的复杂模型中运行它，并且我的 R session 崩溃了。但是，我无法系统地重现它，因此它可能是代码的另一部分。如果有人正在使用 TN，请对其进行测试并告诉我。更新:更新后的代码没有遇到任何问题，请告诉我。

我如何将事情组合在一起:据我所知，最快的实现不在CRAN上，但源代码可以下载OSU stat 。在我的基准测试中，msm 和 trunco​​rm 中的竞争实现速度较慢。诀窍是有效地调整提议分布，其中指数对于截断正态分布的尾部效果很好。因此，我采用了 Chris 的代码，对其进行了“Rcpp”编辑，并向其中添加了一些 openMP 香料。动态计划在这里是最佳的，因为采样可能需要更多或更少的时间，具体取决于边界。我发现一件令人讨厌的事情:当我想使用 double 时，许多统计分布都是基于 NumericVector 类型。我只是用我的方式解决了这个问题。

这是 Rcpp 代码:

#include <Rcpp.h>
#include <omp.h>


// norm_rs(a, b)
// generates a sample from a N(0,1) RV restricted to be in the interval
// (a,b) via rejection sampling.
// ======================================================================

// [[Rcpp::export]]

double norm_rs(double a, double b)
{
   double  x;
   x = Rf_rnorm(0.0, 1.0);
   while( (x < a) || (x > b) ) x = norm_rand();
   return x;
}

// half_norm_rs(a, b)
// generates a sample from a N(0,1) RV restricted to the interval
// (a,b) (with a > 0) using half normal rejection sampling.
// ======================================================================

// [[Rcpp::export]]

double half_norm_rs(double a, double b)
{
   double   x;
   x = fabs(norm_rand());
   while( (x<a) || (x>b) ) x = fabs(norm_rand());
   return x;
}

// unif_rs(a, b)
// generates a sample from a N(0,1) RV restricted to the interval
// (a,b) using uniform rejection sampling. 
// ======================================================================

// [[Rcpp::export]]

double unif_rs(double a, double b)
{
   double xstar, logphixstar, x, logu;

   // Find the argmax (b is always >= 0)
   // This works because we want to sample from N(0,1)
   if(a <= 0.0) xstar = 0.0;
   else xstar = a;
   logphixstar = R::dnorm(xstar, 0.0, 1.0, 1.0);

   x = R::runif(a, b);
   logu = log(R::runif(0.0, 1.0));
   while( logu > (R::dnorm(x, 0.0, 1.0,1.0) - logphixstar))
   {
      x = R::runif(a, b);
      logu = log(R::runif(0.0, 1.0));
   }
   return x;
}

// exp_rs(a, b)
// generates a sample from a N(0,1) RV restricted to the interval
// (a,b) using exponential rejection sampling.
// ======================================================================

// [[Rcpp::export]]

double exp_rs(double a, double b)
{
  double  z, u, rate;

//  Rprintf("in exp_rs");
  rate = 1/a;
//1/a

   // Generate a proposal on (0, b-a)
   z = R::rexp(rate);
   while(z > (b-a)) z = R::rexp(rate);
   u = R::runif(0.0, 1.0);

   while( log(u) > (-0.5*z*z))
   {
      z = R::rexp(rate);
      while(z > (b-a)) z = R::rexp(rate);
      u = R::runif(0.0,1.0);
   }
   return(z+a);
}




// rnorm_trunc( mu, sigma, lower, upper)
//
// generates one random normal RVs with mean 'mu' and standard
// deviation 'sigma', truncated to the interval (lower,upper), where
// lower can be -Inf and upper can be Inf.
//======================================================================

// [[Rcpp::export]]
double rnorm_trunc (double mu, double sigma, double lower, double upper)
{
int change;
 double a, b;
 double logt1 = log(0.150), logt2 = log(2.18), t3 = 0.725;
 double z, tmp, lograt;

 change = 0;
 a = (lower - mu)/sigma;
 b = (upper - mu)/sigma;

 // First scenario
 if( (a == R_NegInf) || (b == R_PosInf))
   {
     if(a == R_NegInf)
       {
     change = 1;
     a = -b;
     b = R_PosInf;
       }

     // The two possibilities for this scenario
     if(a <= 0.45) z = norm_rs(a, b);
     else z = exp_rs(a, b);
     if(change) z = -z;
   }
 // Second scenario
 else if((a * b) <= 0.0)
   {
     // The two possibilities for this scenario
     if((R::dnorm(a, 0.0, 1.0,1.0) <= logt1) || (R::dnorm(b, 0.0, 1.0, 1.0) <= logt1))
       {
     z = norm_rs(a, b);
       }
     else z = unif_rs(a,b);
   }
 // Third scenario
 else
   {
     if(b < 0)
       {
     tmp = b; b = -a; a = -tmp; change = 1;
       }

     lograt = R::dnorm(a, 0.0, 1.0, 1.0) - R::dnorm(b, 0.0, 1.0, 1.0);
     if(lograt <= logt2) z = unif_rs(a,b);
     else if((lograt > logt1) && (a < t3)) z = half_norm_rs(a,b);
     else z = exp_rs(a,b);
     if(change) z = -z;
   }
   double output;
   output = sigma*z + mu;
 return (output);
}


// rtnm( mu, sigma, lower, upper, cores)
//
// generates one random normal RVs with mean 'mu' and standard
// deviation 'sigma', truncated to the interval (lower,upper), where
// lower can be -Inf and upper can be Inf.
// mu, sigma, lower, upper are vectors, and vectorized calls of this function
// speed up computation
// cores is an intege, representing the number of cores to be used in parallel
//======================================================================


// [[Rcpp::export]]

Rcpp::NumericVector rtnm(Rcpp::NumericVector mus, Rcpp::NumericVector sigmas, Rcpp::NumericVector lower, Rcpp::NumericVector upper, int cores){
  omp_set_num_threads(cores);
  int nobs = mus.size();
  Rcpp::NumericVector out(nobs);
  double logt1 = log(0.150), logt2 = log(2.18), t3 = 0.725;
    double a,b, z, tmp, lograt;

     int  change;

  #pragma omp parallel for schedule(dynamic)   
  for(int i=0;i<nobs;i++) {  

     a = (lower(i) - mus(i))/sigmas(i);
     b = (upper(i) - mus(i))/sigmas(i);
     change=0;
     // First scenario
     if( (a == R_NegInf) || (b == R_PosInf))
       {
         if(a == R_NegInf)
           {
              change = 1;
              a = -b;
              b = R_PosInf;
           }

         // The two possibilities for this scenario
         if(a <= 0.45) z = norm_rs(a, b);
         else z = exp_rs(a, b);
         if(change) z = -z;
       }
     // Second scenario
     else if((a * b) <= 0.0)
       {
         // The two possibilities for this scenario
         if((R::dnorm(a, 0.0, 1.0,1.0) <= logt1) || (R::dnorm(b, 0.0, 1.0, 1.0) <= logt1))
           {
                z = norm_rs(a, b);
           }
         else z = unif_rs(a,b);
       }

     // Third scenario
     else
       {
         if(b < 0)
           {
                tmp = b; b = -a; a = -tmp; change = 1;
           }

         lograt = R::dnorm(a, 0.0, 1.0, 1.0) - R::dnorm(b, 0.0, 1.0, 1.0);
         if(lograt <= logt2) z = unif_rs(a,b);
         else if((lograt > logt1) && (a < t3)) z = half_norm_rs(a,b);
         else z = exp_rs(a,b);
         if(change) z = -z;
       }
    out(i)=sigmas(i)*z + mus(i);          
  }

return(out);
}

这是基准:

libs=c("truncnorm","msm","inline","Rcpp","RcppArmadillo","rbenchmark")
if( sum(!(libs %in% .packages(all.available = TRUE)))>0){ install.packages(libs[!(libs %in% .packages(all.available = TRUE))])}
for(i in 1:length(libs)) {library(libs[i],character.only = TRUE,quietly=TRUE)}


#needed for openMP parallel
Sys.setenv("PKG_CXXFLAGS"="-fopenmp")
Sys.setenv("PKG_LIBS"="-fopenmp")

#no of cores for openMP version
cores = 4

#surce code from same dir
Rcpp::sourceCpp('truncnorm.cpp')


#sample size
nn=1000000


bb= 100
aa=-100
benchmark( rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),cores), rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),1),rtnorm(nn,rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn)),rtruncnorm(nn, a=aa, b=100, mean = 0, sd = 1) , order="relative", replications=3    )[,1:4]

aa=0 
benchmark( rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),cores), rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),1),rtnorm(nn,rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn)),rtruncnorm(nn, a=aa, b=100, mean = 0, sd = 1) , order="relative", replications=3    )[,1:4]

aa=2
benchmark( rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),cores), rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),1),rtnorm(nn,rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn)),rtruncnorm(nn, a=aa, b=100, mean = 0, sd = 1) , order="relative", replications=3    )[,1:4]

aa=50
benchmark( rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),cores), rtnm(rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn),1),rtnorm(nn,rep(0,nn),rep(1,nn),rep(aa,nn),rep(100,nn)),rtruncnorm(nn, a=aa, b=100, mean = 0, sd = 1) , order="relative", replications=3    )[,1:4]

由于速度取决于上/下边界，因此需要进行多次基准测试。对于不同的情况，算法的不同部分会起作用。

Answer 1

非常快速的评论:

1. 如果您包含 RcppArmadillo.h您不需要包括 Rcpp.h -- 事实上，您不应该这样做，我们甚至对此进行了测试

2. rep(oneDraw, n)调用 n 次电话。我会编写一个调用一次的函数，返回 n 次绘制——它会更快，因为你可以节省 n-1 次函数调用开销

3. 您对许多统计分布的评论均基于 NumericVector类型，当我想使用 double 时可能会揭示一些误解:NumericVector是我们方便的内部 R 类型的代理类:没有副本。您可以随意使用std::vector<double>或您喜欢的任何形式。

4. 我对截断法线知之甚少，因此无法评论您的算法的具体细节。

5. 一旦你解决了这个问题，请考虑发布一个关于 Rcpp Gallery 的帖子。 .