基于data.table的快速滚动Spearman相关系数计算方案

我有一个data.table数据集，需要计算value列与固定理想趋势向量c(1,2,3,4,5,6,7)的7天滚动Spearman相关系数。目前用cor函数配合frollapply的实现，在大数据集上速度极慢，原因是cor函数会创建数据副本，求更高效的替代方案。

测试数据集 #

temp_dt = structure(list(date = structure(c(18540L, 18541L, 18542L, 18543L, 
18544L, 18547L, 18548L, 18549L, 18550L, 18551L, 18554L, 18555L, 
18556L, 18557L, 18558L, 18561L, 18562L, 18563L, 18564L, 18565L, 
18568L, 18569L, 18570L, 18571L, 18572L, 18575L, 18576L, 18577L, 
18578L, 18579L, 18582L, 18583L, 18584L, 18585L, 18586L, 18589L, 
18590L, 18591L, 18593L, 18596L, 18597L, 18598L, 18599L, 18600L, 
18603L, 18604L, 18605L, 18606L, 18607L, 18610L, 18611L, 18612L, 
18613L, 18614L, 18617L, 18618L, 18619L, 18620L, 18624L, 18625L, 
18626L, 18627L, 18631L, 18632L, 18633L, 18634L, 18635L, 18638L, 
18639L, 18640L, 18641L, 18642L, 18646L, 18647L, 18648L, 18649L, 
18652L, 18653L, 18654L, 18655L, 18656L, 18659L, 18660L, 18661L, 
18662L, 18663L, 18666L, 18667L, 18668L, 18669L, 18670L, 18674L, 
18675L, 18676L, 18677L, 18680L, 18681L, 18682L, 18683L, 18684L
), class = c("IDate", "Date")), value = c(217.65, 211.45, 219.02, 
218.45, 218.95, 216.51, 212.07, 210.53, 213.93, 210.56, 203.42, 
200.85, 195.96, 200.55, 208.01, 200.26, 200.35, 195.03, 196.2, 
193.58, 192.58, 197.59, 198.47, 204.82, 200.8, 236.4, 233.48, 
221.51, 222.13, 225.41, 230.97, 225.02, 234.62, 233.09, 232.53, 
230, 225.95, 222, 221.58, 218.54, 220.58, 226.23, 229.25, 231.08, 
229.99, 233.34, 237.4, 233.1, 229.51, 230.51, 237.1, 234.49, 
245.16, 240.85, 240.77, 245.88, 249.93, 247.88, 258.32, 259.15, 
260.24, 261.55, 253.38, 252.24, 267.22, 264.09, 264.97, 263.88, 
262.69, 258.1, 253.06, 253.92, 252.06, 253.69, 253.98, 252.25, 
252.54, 251.32, 236.99, 250.09, 248.9, 248.33, 252.36, 252.72, 
256.9, 255.99, 263.85, 263.99, 259.79, 265.35, 262.66, 263.3, 
264.79, 264.85, 264.81, 264.64, 264.91, 269.63, 261.62, 258.82
)), row.names = c(NA, -100L), class = c("data.table", "data.frame"
))

当前慢实现及结果 #

temp_dt[, corr := frollapply(x = value, n = 7, cor, fill = NA, align = "right", method = "spearman", y = 1:7)]

输出结果（末尾6行）：

95: 2021-02-19 264.81  0.42857143
 96: 2021-02-22 264.64  0.00000000
 97: 2021-02-23 264.91  0.75000000
 98: 2021-02-24 269.63  0.75000000
 99: 2021-02-25 261.62  0.00000000
100: 2021-02-26 258.82 -0.42857143

原实现性能测试 #

microbenchmark::microbenchmark(
    slower = temp_dt[, corr := frollapply(x = value, n = 7, cor, fill = NA, align = "right", method = "spearman", y = 1:7)],
    times = 5
)

输出：

Unit: milliseconds
   expr      min       lq    mean   median       uq      max neval
 slower 14.37227 14.39736 15.0515 15.01405 15.44279 16.03103     5

高效优化方案 #

方案一：基于Spearman系数的秩次特性计算 #

Spearman相关系数本质是变量秩次的Pearson相关系数。由于目标向量1:7的秩次就是自身，我们可以先计算value列的滚动秩，再直接用Pearson相关的公式手动计算，避免cor函数的额外开销。

# 预计算固定向量的统计量，避免重复计算
y <- 1:7
n <- length(y)
var_y <- var(y)

# 计算value的7天滚动秩
temp_dt[, rank_val := frollapply(value, n = 7, FUN = rank, fill = NA, align = "right")]

# 滚动计算秩次的均值、协方差与方差
temp_dt[, `:=`(
    mean_rank = frollmean(rank_val, n = 7, fill = NA, align = "right"),
    cov_rank_y = frollapply(rank_val, n = 7, FUN = function(x) cov(x, y), fill = NA, align = "right"),
    var_rank = frollapply(rank_val, n = 7, FUN = var, fill = NA, align = "right")
)]

# 计算Spearman相关系数
temp_dt[, corr_fast := cov_rank_y / sqrt(var_y * var_rank)]

# 验证结果一致性
all.equal(temp_dt$corr, temp_dt$corr_fast, na.rm = TRUE)

方案二：完全向量化滚动操作，彻底替代frollapply #

方案一中仍有部分frollapply调用，我们可以用data.table的内置滚动函数（frollmean、frollsum）结合矩阵运算，实现完全向量化，进一步提升速度：

y <- 1:7
n <- 7
mean_y <- mean(y)
ss_y <- sum((y - mean_y)^2)

# 计算滚动秩次
temp_dt[, rank_val := frollapply(value, n = n, rank, fill = NA, align = "right")]

# 生成滚动窗口的秩次矩阵
rank_mat <- do.call(cbind, lapply(0:(n-1), function(k) shift(temp_dt$rank_val, k, fill = NA)))

# 滚动计算交叉乘积和、秩次平方和、秩次均值
temp_dt[, `:=`(
    sum_ry = rowSums(rank_mat * matrix(y, nrow = .N, ncol = n, byrow = TRUE), na.rm = FALSE),
    sum_r2 = frollsum(rank_val^2, n = n, fill = NA, align = "right"),
    mean_r = frollmean(rank_val, n = n, fill = NA, align = "right")
)]

# 计算平方和与协方差，最终得到相关系数
temp_dt[, `:=`(
    ss_r = sum_r2 - n * mean_r^2,
    ss_ry = sum_ry - n * mean_r * mean_y,
    corr_faster = ss_ry / sqrt(ss_r * ss_y)
)]

# 验证结果一致性
all.equal(temp_dt$corr, temp_dt$corr_faster, na.rm = TRUE)

优化后性能对比 #

microbenchmark::microbenchmark(
    original = temp_dt[, corr := frollapply(x = value, n = 7, cor, fill = NA, align = "right", method = "spearman", y = 1:7)],
    method1 = {
        temp_dt[, rank_val := frollapply(value, n = 7, rank, fill = NA, align = "right")]
        temp_dt[, `:=`(mean_rank = frollmean(rank_val, n = 7, fill = NA, align = "right"), cov_rank_y = frollapply(rank_val, n = 7, function(x) cov(x, y), fill = NA, align = "right"), var_rank = frollapply(rank_val, n = 7, var, fill = NA, align = "right"))]
        temp_dt[, corr_fast := cov_rank_y / sqrt(var(y) * var_rank)]
    },
    method2 = {
        temp_dt[, rank_val := frollapply(value, n = n, rank, fill = NA, align = "right")]
        rank_mat <- do.call(cbind, lapply(0:(n-1), function(k) shift(temp_dt$rank_val, k, fill = NA)))
        temp_dt[, `:=`(sum_ry = rowSums(rank_mat * matrix(y, nrow = .N, ncol = n, byrow = TRUE), na.rm = FALSE), sum_r2 = frollsum(rank_val^2, n = n, fill = NA, align = "right"), mean_r = frollmean(rank_val, n = n, fill = NA, align = "right"))]
        temp_dt[, `:=`(ss_r = sum_r2 - n * mean_r^2, ss_ry = sum_ry - n * mean_r * mean_y, corr_faster = ss_ry / sqrt(ss_r * ss_y))]
    },
    times = 10
)

测试结果会显示，方案二的速度远优于原实现，大数据集下的提升会更加明显。

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