Rcpp调用sample()实现矩阵行概率加权采样报错排查

有一个存储概率值的矩阵，共4列，每列依次对应取值为0到4的整数分值。现需要以每行存储的概率值作为采样权重，为每一行采样得到1个对应分值：若某行存在NA值（即部分列无有效概率），则仅在非NA列对应的分值范围内采样，例如某行概率为0.45,0.55,NA,NA时，仅在0、1两个分值中采样。
编写Rcpp代码实现该逻辑时接连出现报错。

• 最初直接调用Rcpp的sample()方法时报错：

error: no matching function for call to 'as<Rcpp::IntegerVector>(Rcpp::Matrix<14>::Sub&)'
     score[i] = sample(scrs,1,true,as<IntegerVector>(probs));

• 查阅资料提示可使用RcppArmadillo实现，在cppFunction前加载RcppArmadillo包，将sample调用替换为score[i] = Rcpp::RcppArmadillo::sample(scrs,1,true,probs);后报错：

error: 'Rcpp::RcppArmadillo' has not been declared
     score[i] = Rcpp::RcppArmadillo::sample(scrs,1,true,probs);

• 进一步在代码顶部添加#include <RcppArmadilloExtensions/sample.h>头文件引入声明后报错：

fatal error: RcppArmadilloExtensions/sample.h: No such file or directory
   #include <RcppArmadilloExtensions/sample.h>

p.vals <- matrix(c(0.44892077,0.55107923,NA,NA,
                 0.37111195,0.62888805,NA,NA,
                 0.04461714,0.47764478,0.303590351,1.741477e-01,
                 0.91741642,0.07968127,0.002826406,7.589714e-05,
                 0.69330800,0.24355559,0.058340934,4.795468e-03,
                 0.43516823,0.43483784,0.120895859,9.098067e-03,
                 0.73680809,0.22595438,0.037237525,NA,
                 0.89569365,0.10142719,0.002879163,NA),nrow=8,ncol=4,byrow=TRUE)

step.vals <- c(1,1,3,3,3,3,2,2)

require(Rcpp)
cppFunction('IntegerVector scores_cpp(NumericMatrix p, IntegerVector steps){

  int prows = p.nrow();

  IntegerVector score(prows);
  
  for(int i=0;i<prows;i++){
    int step = steps[i];
    
    IntegerVector scrs = seq(0,step);
    
    int start = 0;
    int end = step;
    
    NumericMatrix::Sub probs = p(Range(i,i),Range(start,end));

    score[i] = sample(scrs,1,true,probs);
  }
  
  return score;
  
}')

test <- scores_cpp(p.vals,step.vals)
test

step.vals中每行的取值始终等于对应行包含有效概率的列数减1，因此向函数传入step.vals属于冗余参数。

报错核心原因有两点：

1. Rcpp原生sample要求概率参数是向量类型，原代码里取的p(Range(i,i),Range(start,end))是矩阵切片（Matrix::Sub类型），本质是1行多列的子矩阵，不能直接转换为向量传给sample。
2. 调用RcppArmadillo接口时不需要手动写头文件路径，只要在cppFunction里加depends = "RcppArmadillo"参数，编译时会自动定位对应头文件，手动写include路径很容易因为路径配置错误报找不到文件的问题。

以下是不依赖额外拓展包、自动识别行内有效概率、可直接运行的实现：

library(Rcpp)
cppFunction('
IntegerVector scores_cpp(NumericMatrix p){
  int n_rows = p.nrow();
  int n_cols = p.ncol();
  IntegerVector res(n_rows);
  
  for(int i = 0; i < n_rows; i++){
    NumericVector cur_row = p(i, _);
    std::vector<double> valid_p;
    std::vector<int> valid_score;
    double p_sum = 0.0;
    
    // 提取当前行所有非NA的概率和对应分值
    for(int j = 0; j < n_cols; j++){
      if(!Rcpp::traits::is_na<REALSXP>(cur_row[j])){
        valid_p.push_back(cur_row[j]);
        valid_score.push_back(j);
        p_sum += cur_row[j];
      }
    }
    
    // 概率归一化，避免浮点误差导致和不为1
    for(int k = 0; k < valid_p.size(); k++){
      valid_p[k] = valid_p[k] / p_sum;
    }
    
    // 转换为Rcpp原生类型后采样
    NumericVector p_vec = Rcpp::wrap(valid_p);
    IntegerVector s_vec = Rcpp::wrap(valid_score);
    res[i] = Rcpp::sample(s_vec, 1, true, p_vec)[0];
  }
  return res;
}
')

# 测试运行
test <- scores_cpp(p.vals)
print(test)

内容的提问来源于stack exchange，提问作者MJC