基于terra::focalCpp的自定义Rcpp中位数函数实现问题

问题场景 #

我要实现一个支持na.rm=TRUE的中位数计算函数，用Rcpp编写后给terra包的focalCpp调用。从GitHub找到的median_dbl代码单独运行没问题，但集成到focalCpp时就报错Error: [focalCpp] test failed。加了NA过滤逻辑后结果还不对，核心问题出在NA处理和focalCpp的函数适配上。

复现代码 #

1. 示例数据与R基准实现 #

library(terra)
# 生成带NA的测试栅格
r <- rast(nrows=5, ncols=5, vals=rnorm(25))
r[sample(1:25, 5)] <- NA

# R原生中位数窗口函数
focal_median_r <- function(x) median(x, na.rm=TRUE)
# 计算基准结果
result_r <- focal(r, matrix(1,3,3), fun=focal_median_r)

2. 原出错的Rcpp代码 #

#include <Rcpp.h>
#include <algorithm>
#include <vector>
using namespace Rcpp;

// [[Rcpp::export]]
double median_dbl(NumericVector x) {
  int n = x.size();
  if (n == 0) return NA_REAL;
  
  std::vector<double> non_na;
  for (int i=0; i<n; ++i) {
    if (!NumericVector::is_na(x[i])) {
      non_na.push_back(x[i]);
    }
  }
  
  int m = non_na.size();
  if (m == 0) return NA_REAL;
  
  std::sort(non_na.begin(), non_na.end());
  
  return (m % 2 == 1) ? non_na[m/2] : (non_na[m/2-1] + non_na[m/2])/2.0;
}

// 错误的包装函数
// [[Rcpp::export]]
NumericVector focal_median_cpp(NumericVector x, IntegerVector dim) {
  int n = x.size();
  NumericVector res(n);
  for (int i=0; i<n; ++i) {
    res[i] = median_dbl(x);
  }
  return res;
}

3. 报错的调用代码 #

# 调用focalCpp时触发错误
result_cpp <- focalCpp(r, matrix(1,3,3), fun=focal_median_cpp)
# 错误信息：Error: [focalCpp] test failed

问题根源 #

1. 函数签名不匹配：focalCpp要求自定义函数必须接受单个窗口的NumericVector，返回一个double结果，而不是返回整个向量。之前写的focal_median_cpp返回NumericVector，完全不符合要求，导致测试失败。
2. NA处理逻辑偏差：原代码只用NumericVector::is_na过滤NA，但R原生median(na.rm=TRUE)会同时忽略NaN，这里没覆盖到，导致结果不一致。

修正方案 #

1. 调整后的Rcpp代码 #

直接实现符合focalCpp要求的单窗口计算函数，同时修正NA/NaN过滤逻辑：

#include <Rcpp.h>
#include <algorithm>
#include <vector>
using namespace Rcpp;

// [[Rcpp::export]]
double median_dbl(NumericVector x) {
  std::vector<double> non_na;
  non_na.reserve(x.size()); // 预分配内存，提升效率
  
  // 过滤NA、NaN和Inf，和R原生median行为一致
  for (double val : x) {
    if (R_FINITE(val)) {
      non_na.push_back(val);
    }
  }
  
  int m = non_na.size();
  if (m == 0) return NA_REAL;
  
  std::sort(non_na.begin(), non_na.end());
  
  return (m % 2 == 1) ? non_na[m / 2] : (non_na[m/2 - 1] + non_na[m/2]) / 2.0;
}

2. 正确调用focalCpp #

不需要额外包装函数，直接传入median_dbl：

result_cpp <- focalCpp(r, matrix(1,3,3), fun=median_dbl)

# 验证结果一致性
all.equal(values(result_r), values(result_cpp), na.rm=TRUE)
# 输出TRUE表示结果一致

关键优化说明 #

• 用R_FINITE替代NumericVector::is_na，确保和R原生median行为完全对齐，同时排除NaN和Inf。
• 调用non_na.reserve(x.size())避免动态扩容，提升大窗口下的计算效率。
• 去掉多余的包装函数，直接使用符合focalCpp签名要求的median_dbl函数。

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