Среднее значение столбца data.table, заданное с помощью матрицы

Вот решение, которое дает почти 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() в исходном цикле.

Затем я попытался поместить всю функциональность в функции, специфичные для 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-кратное увеличение скорости по сравнению с исходной петлей.

Кажется, он неплохо масштабируется до 1 миллиона точек данных:

# 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 

Тайминги были выполнены на одном ядре MacBook Pro 2011 года (Sandy Bridge 2.2Ghz). Объем оперативной памяти остался ниже 1,5 ГБ.

Вот еще одно решение, использующее 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 преобразуется в таблицу данных и получает столбец pt, который является идентификатором строки точек. Затем формат матрицы изменяется с широкого на длинный.

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"]