R で二次元配列をなめる処理をするときの速度

R で、検出器上の各ピクセルの分解能を計算しようとした。最初に書いた愚直なコードは以下の通り。

lambda <- 12400 / 7000 # A
distance <- 51.5 * 1000 # um

pixel_size <- 50 # um
buf_size <- 3000 # px

buf_origin <- c(buf_size, buf_size) / 2

reso <- matrix(0, buf_size, buf_size)

system.time({
for (x in 1:buf_size) {
  print(x)
  for (y in 1:buf_size) {
    xy <- c(x, y)
    r <- sqrt(sum((xy - buf_origin) ** 2)) * pixel_size
    theta <- atan2(r, distance) / 2
    reso[y, x] <- lambda / 2 / sin(theta)
  }
}
})

Core i7 3970X 3.50GHz で 37 秒もかかった。

1:buf_size を毎回生成しているのが気になるので、少し変えたが効果なかった。

system.time({
xloop <- 1:buf_size
yloop <- 1:buf_size
for (x in xloop) {
  print(x)
  for (y in yloop) {
    xy <- c(x, y)
    r <- sqrt(sum((xy - buf_origin) ** 2)) * pixel_size
    theta <- atan2(r, distance) / 2
    reso[y, x] <- lambda / 2 / sin(theta)
  }
}
})

二重ループはダサすぎるので、メモリを贅沢に使って、座標のリストを作ってからベクトルで処理してみる。

system.time({
coords <- expand.grid(1:buf_size, 1:buf_size)
coords <- sweep(coords, 2, buf_origin, "-")
r <- sqrt(rowSums(coords ** 2)) * pixel_size
theta <- atan2(r, distance) / 2
reso <- matrix(lambda / 2 / sin(theta), buf_size, buf_size)
rm(coords, r)
})

一気に 2.4 秒まで縮んだ。

しかし、sweep が内部で一列ごとに演算子を eval するのが遅そうである。行列を転置して、sweep を除去してみる。

system.time({
coords <- expand.grid(1:buf_size, 1:buf_size)
coords <- t(coords) - buf_origin
r <- sqrt(colSums(coords ** 2)) * pixel_size
theta <- atan2(r, distance) / 2
reso <- matrix(lambda / 2 / sin(theta), buf_size, buf_size)
rm(coords, r)
})

1 秒未満になった。

言い古されたことだが、R で for ループを回すと遅いことが実感できた。改善版は確かに速いが、本来は不要なメモリを消費すること、本質的でない転置が必要となるのが気に入らない。まあ、R の利便性と引き換えに仕方がないことであろう。