Оптимизировать which.max по нескольким измерениям массива

Question

У меня есть код с 4-мерным массивом, и мне нужно применить which.max для нескольких измерений. Это довольно медленно, я бы хотел найти способы его ускорить.

Пример:

library(microbenchmark)

array4d <- array( runif(5*500*50*5 ,-1,0),
                  dim = c(5, 500, 50, 5) )

microbenchmark(
    max_idx <- apply(array4d, c(1,2,3), which.max )
    )

Любые подсказки приветствуются, спасибо!

Изменить: я удалось сделать это немного быстрее (хотя и уродливо), напрямую закодировав его в циклах for - но я надеюсь, что у кого-то есть идеи получше!

method1 <- function(z) {
    apply(z, c(1,2,3), which.max)
}
method2 <- function(z){
    result <- array( , dim = dim(z)[1:3] )
    for(i in 1:dim(z)[1]){
        for(j in 1:dim(z)[2]){
            for(k in 1:dim(z)[3]){
                result[i, j, k] <- which.max(z[i,j,k,])
            }
        }
    }
    return(result)
}

microbenchmark(
result1 <- method1(array4d),
result2 <- method2(array4d))

> microbenchmark(
+     result1 <- method1(array4d),
+     result2 <- method2(array4d)
+ )
Unit: milliseconds
                        expr      min       lq     mean   median       uq      max neval cld
 result1 <- method1(array4d) 111.9061 140.1400 165.2441 155.6773 170.3967 384.6425   100   b
 result2 <- method2(array4d) 113.4572 123.2429 136.8583 130.8505 141.9620 215.0968   100  a 

Answer 1

Добавление дополнительных методов. Один использует R, другой код настройки от @ Allan-Cameron:

method4 <- function(z){
  result <- array(integer(1) , dim = head(dim(z), -1))
  n <- prod(head(dim(z), -1))
  j <- seq_len(tail(dim(z),1)) * n - n
  for(i in seq_len(n)) result[i] <- which.max(z[i+j])
  result}

Rcpp::cppFunction("
NumericVector method5(const NumericVector &input){
  std::vector<int> dims = input.attr(\"dim\");
  int last_dim = dims.back();
  int diff = input.size()/last_dim;
  std::vector<int> result(diff);
  dims.pop_back();
  
  for(int i = 0; i < diff; ++i)
  {
    double max_val = input[i];
    int max_ind = 0;
    for(int j = 0; j < last_dim; ++j)
    {
      if(input[i+j*diff] > max_val) {
        max_val = input[i+j*diff];
        max_ind = j;
      }
    }
    result[i] = max_ind + 1;
  }
  
  NumericVector arr = wrap(result);
  arr.attr(\"dim\") = dims;
  return arr ;
}"
)

Время:

set.seed(42)
array4d <- array(runif(5*500*50*5, -1, 0), dim = c(5, 500, 50, 5))

library(microbenchmark)
microbenchmark(
  check = "equal", control=list(order="block")
, method1(array4d)         #Using code from Question
, method2(array4d)         #Using code from Question
, apply_which_max(array4d) #Using code from Allan Cameron
, method4(array4d)
, method5(array4d)
)
#Unit: microseconds
#                     expr        min         lq        mean      median          uq        max neval  cld
#         method1(array4d) 200857.804 228567.850 266815.6275 254530.3050 294578.1125 423838.879   100    d
#         method2(array4d) 144767.680 149616.981 162367.6556 150688.1860 182290.4980 315650.052   100   c 
# apply_which_max(array4d)   3131.482   3153.712   3346.1025   3175.9445   3206.2220   5922.866   100 a   
#         method4(array4d)  58618.275  60777.584  62334.8258  61198.1815  61702.2170 165254.042   100  b  
#         method5(array4d)    894.823    902.862    972.2953    911.9845    927.0885   2643.957   100 a   

Для случайного выбора вместо первого совпадения:

Rcpp::cppFunction("
NumericVector method6(const NumericVector &input){
  std::srand(std::time(nullptr));
  std::vector<int> dims = input.attr(\"dim\");
  int last_dim = dims.back();
  int diff = input.size()/last_dim;
  std::vector<int> result(diff);
  dims.pop_back();
  
  for(int i = 0; i < diff; ++i)
  {
    double max_val = input[i];
    std::vector<int> max_ind = {0};
    for(int j = 1; j < last_dim; ++j)
    {
      if(input[i+j*diff] > max_val) {
        max_val = input[i+j*diff];
        max_ind.clear();
        max_ind.push_back(j);
      } else if(input[i+j*diff] == max_val) max_ind.push_back(j);
    }
    result[i] = max_ind[std::rand() % max_ind.size()] + 1;
  }
  
  NumericVector arr = wrap(result);
  arr.attr(\"dim\") = dims;
  return arr ;
}"
)

Answer 2

• 1000 to Rcpp:cppFunction:

Rcpp::cppFunction("
NumericVector apply_which_max(NumericVector input){
  std::vector<int> dims = input.attr(\"dim\");
  int last_dim = dims.back();
  int diff = input.size()/last_dim;
  std::vector<int> result(diff);
  dims.pop_back();
  
  for(int i = 0; i < diff; ++i)
  {
    double max_val = input[i];
    int max_ind = i;
    for(int j = i; j < input.size(); j += diff)
    {
      if(input[j] > max_val) {
        max_val = input[j];
        max_ind = j;
      }
    }
    result[i] = max_ind / diff + 1;
  }
  
  NumericVector arr = wrap(result);
  arr.attr(\"dim\") = dims;
  return arr ;
}"
)

Нам лучше показать, что он имеет тот же результат:

max_idx1 <- apply(array4d, c(1 ,2, 3), which.max)
max_idx2 <- apply_which_max(array4d)

all(max_idx1 == max_idx2)
#> [1] TRUE

Но быстрее ли?

Используя тесты (на моем медленная машина), получаем:

microbenchmark(apply(array4d, c(1 ,2, 3), which.max))
#> Unit: milliseconds
#>                                   expr      min       lq     mean   median       uq
#>  apply(array4d, c(1, 2, 3), which.max) 243.2283 276.5381 342.5796 330.0815 403.8358
#>       max neval
#>  543.8325   100

microbenchmark(apply_which_max(array4d))
#> Unit: milliseconds
#>                      expr    min     lq     mean median     uq    max neval
#>  apply_which_max(array4d) 3.1761 3.2725 3.713377  3.342 3.4088 12.962   100

То есть не просто быстрее, а примерно в 100 раз быстрее.