r - 使用矩阵指定的 data.table 列值的平均值

Question

我有一个 data.table 包含单位立方体中 10,000 个点(对于本示例)的 x,y,z 值，每个点都有一个相应的属性(称为 P )。我用过 nn2来自 RANN包以从原始 data.frame(以矩阵形式返回)0.075 个单位的半径内查找每个点的 k 邻居(最多 50 个)索引。

library(RANN)
library(data.table)

set.seed(1L) # for reproducible data
DATA <- data.table(runif(10000, 0,1), 
                   runif(10000, 0,1), 
                   runif(10000, 0,1), 
                   runif(10000, 10,30))
colnames(DATA)<-c("x","y","z","P")
nn.idx <- nn2(DATA[,1:3], DATA[,1:3], k=50, 
              treetype = "kd", searchtype = "radius", 
              radius = 0.075)$nn.idx

以下for循环可以完成这项工作，但我想知道是否有任何方法可以通过对其进行矢量化来加快速度，因为这在应用于 > 数百万个点时不会扩展？简单来说，我想用nn.idx得到对应的P来自 DATA 的值并计算平均值 P然后将其分配给 DATA 中的新列叫 mean.P

for(index in 1:nrow(DATA))
  DATA$mean.P[index]<-mean(DATA[nn.idx[index,], P])

出于说明目的，以下代码说明了我要计算的内容 - 对于所有点(灰点)，计算给定点(红点)周围球体中所有点(橙色 + 红点)的平均值并分配它到那一点(红点)。迭代所有点，但要以一种有效的方式进行，以适应大数据集。

library(rgl)
rgl.open()
rgl.points(DATA[1500,1], DATA[1500,2], DATA[1500,3], color ="red")
rgl.points(DATA[nn.idx[1500,],1:3], color ="orange", add=T)
rgl.points(DATA[,1:3], color ="lightgray", alpha=0.1, add=T)

在我的生活中，我从来没有花这么多时间尝试有效地矢量化一个循环!另外，我不反对下注，只是用 c++ 和 Rcpp 来做，但我想我会先在这里问一下，看看 R 中是否有一种方法可以使其扩展和更快。提前致谢!

Answer 1

这是一个将速度提高近 100 倍的解决方案。我不完全理解为什么改进如此之大，但也许真正的 data.table 专家之一可以对此发表评论。

library(RANN)
library(data.table)

set.seed(1L) # for reproducible data
DATA <- data.table(runif(10000, 0,1), 
                   runif(10000, 0,1), 
                   runif(10000, 0,1), 
                   runif(10000, 10,30))
colnames(DATA)<-c("x","y","z","P")
nn.idx <- nn2(DATA[,1:3], DATA[,1:3], k=50, 
              treetype = "kd", searchtype = "radius", 
              radius = 0.075)$nn.idx

# (1)
# Timing for original loop.
system.time(for(index in 1:nrow(DATA)) {
    DATA$mean.P[index] <- mean(DATA[nn.idx[index,], P])
})
#    user  system elapsed 
#   7.830   0.850   8.684 

# (2)
# Use `set()` instead of `$<-` and `[<-`.
system.time({for(index in 1:nrow(DATA)) {
    set(DATA, i=index, j="mean_P_2", value=mean(DATA[nn.idx[index, ], P]))
}})
#    user  system elapsed 
#   3.405   0.008   3.417 

如您所见，仅通过替换 data.table-specific set() 就有 2 倍的改进。原循环中的函数。

接下来，我尝试将所有功能放入特定于 data.table 的函数中(主要在 data.table [] 语法中)。我也放了 P值转换为向量，因为访问向量中的值通常比 data.frames 或 data.tables 上的类似操作快得多。

# (3)
# Add row index.
DATA[, row_idx:=seq(nrow(DATA))]

# Isolate P values in a vector, because vector access is cheaper
# than data.table or data.frame access.
P_vec = DATA$P

system.time({
    # Create a list column where each element is a vector of 50 integer indexes.
    DATA[, nn_idx:=lapply(row_idx, function(i) nn.idx[i, ])]
    # Use `:=` and `by=` to internalize the loop within `[.data.table`.
    DATA[, mean_P_3:=mean(P_vec[nn_idx[[1]]]), by=row_idx]
})
#    user  system elapsed 
#   0.092   0.002   0.095 

# All results are identical.
all.equal(DATA$mean.P, DATA$mean_P_2)
# [1] TRUE
all.equal(DATA$mean.P, DATA$mean_P_3)
# [1] TRUE

与原始循环相比，这使速度提高了近 100 倍。

它似乎可以很好地扩展到 100 万个数据点:

# Try with 1 million data points.
set.seed(1L) # for reproducible data
DATA2 <- data.table(runif(1e6, 0,1), 
                    runif(1e6, 0,1), 
                    runif(1e6, 0,1), 
                    runif(1e6, 10,30))
colnames(DATA2) <- c("x","y","z","P")

system.time({
    nn.idx2 <- nn2(DATA2[,1:3], DATA2[,1:3], k=50, 
                   treetype = "kd", searchtype = "radius", 
                   radius = 0.075)$nn.idx
})
#    user  system elapsed 
# 346.603   1.883 349.708 


DATA2[, row_idx:=seq(nrow(DATA2))]
P_vec = DATA2$P

system.time({
    DATA2[, nn_idx:=lapply(row_idx, function(i) nn.idx2[i, ])]
    DATA2[, mean_P:=mean(P_vec[nn_idx[[1]]]), by=row_idx]
})
#    user  system elapsed 
#  15.685   0.587  16.297 

计时是在 2011 年 macbook pro(Sandy Bridge 2.2Ghz)的单核上完成的。 RAM 使用率保持在 1.5 GB 以下。