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

Question

我有一个 data.table 包含单位立方体中 10,000 个点（对于本例）的 x、y、z 值，每个点都有一个对应的属性（称为 P）。我使用 RANN 包中的 nn2 来查找距离原始 data.frame （返回作为矩阵）。

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 从 DATA 中获取相应的 P 值并计算平均值 P 然后分配给 [=20= 中的新列] 称为 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 特定的 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 以下。

Answer 2

这是另一个使用 melt() 以长格式重塑索引矩阵、连接和聚合的解决方案：

long <- melt(as.data.table(nn.idx)[, pt := .I], measure.vars = patterns("V"))
tmp <- long[DATA[, pt := .I], on = .(value = pt)][, mean(P), by = .(pt)][order(pt), V1]
DATA[, mean.P := tmp][, pt := NULL][]

说明

索引矩阵nn.idx转换为data.table并获得列pt，这是点的行id。然后矩阵从宽格式重塑为长格式。

tmp是相邻点均值的向量。这些是通过右连接 DATA 和 long 来匹配最近的相邻点（在列 value 中）的索引与预先附加到 DATA 的点索引找到的。

最后一步是将结果作为新列追加到 DATA。

变体 2

或者，可以使用第二个连接附加中间结果：

long <- melt(as.data.table(nn.idx)[, pt := .I], measure.vars = patterns("V"))
    long[DATA[, pt := .I], on = .(value = pt)][, mean(P), by = .(pt)][DATA, on = "pt"]

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

Mean of data.table column values as specified using a matrix

r

matrix

nearest-neighbor

data.table

说明

变体 2