概要
Python で多重配列 numpy.ndarray
もしくは torch.Tensor
からブロック多重配列をスライスするときは関数 numpy.ix_
を使う。
詳細
Julia で
# Julia
# 多重配列を作る
(i, j, k) = (2, 3, 4)
x = reshape(1:i * j * k, i, j, k)
# 多重配列を多重添字でスライスする
(is, js, ks) = ([2, 1], [1, 3], 2:3)
x[is, js, ks] |> println
みたいにして多重配列を多重添字でスライスするのを Python でどうやるのか調べてもなかなか情報にたどり着かなかったのでメモ。numpy.ix_
が所望のものである。 NumPy や PyTorch はデフォルトだと Julia の Array とは多重配列の次元の並べ方が逆であることに注意。
# Python + NumPy
import numpy
# 多重配列を作る
i, j, k = 2, 3, 4
x = (numpy.arange(k * j * i) + 1).reshape(k, j, i)
# 多重配列を多重添字でスライスする
# (Python では is は予約語かなんかのようなので is_ とする)
is_, js, ks = numpy.array([2, 1]) - 1, numpy.array([1, 3]) - 1, numpy.arange(2 - 1, 3)
print(x[numpy.ix_(ks, js, is_)])
# Python + PyTorch
import torch
import numpy
# 多重配列を作る
i, j, k = 2, 3, 4
x = (torch.arange(k * j * i) + 1).reshape(k, j, i)
# 多重配列を多重添字でスライスする
# (Python では is は予約語かなんかのようなので is_ とする)
is_, js, ks = torch.tensor([2, 1]) - 1, torch.tensor([1, 3]) - 1, torch.arange(2 - 1, 3)
print(x[numpy.ix_(ks, js, is_)])