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

【问题标题】：Mean of data.table column values as specified using a matrix使用矩阵指定的 data.table 列值的平均值
【发布时间】：2017-09-24 01:03:04
【问题描述】：

我有一个 data.table，其中包含单位立方体中 10,000 个点（对于此示例）的 x、y、z 值，每个点都有一个相应的属性（称为P）。我已经使用RANN 包中的nn2 来查找距原始data.frame 0.075 单位半径内的每个点的k-neighbors（最多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 从DATA 中获取对应的P 值并计算平均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 中是否有一种方法可以使其规模化和更快。提前致谢！

【问题讨论】：

如果可以使用更大的数据在内存方面进行处理，您可以一次提取所有值 -x = DATA[c(nn.idx), P]- 并通过 by = row(nn.idx)[as.logical(nn.idx)] 找到平均值：meanP = c(rowsum(x, by)) / tabulate(by)
到目前为止，这两种解决方案似乎都是可行的，因此我需要进行更多测试。在我的机器（2016 Dell w/Xeon E5-2620 2.10GHz）上对它们进行基准测试，Uwe 的解决方案以牺牲第二个 data.table（确实变得巨大）为代价最快，并且 bdemarest 足够快更便宜的载体。所以 1Mil pts 和 k=100：# bdemarest solution; elapsed = 16.22 和 # Uwe solution; elapsed = 4.94 对应大小的对象：# Int_vec = 8,000,040 long = 1,200,007,392

标签： r matrix data.table nearest-neighbor

【解决方案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 以下。

【讨论】：

【解决方案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"]

【讨论】：