numpyで3次元座標インデックスの配列を作る方法

どういう場合に使うかっていうと、例えば3次元の空間がnp.zeros(100, 100, 100)とかで用意されてるとします。

そしてここに、ある関数によるエフェクトの結果生まれる場を表現したいとします。

えっ何いってんのこいつと思うかもしれませんが、要するに磁力とか、重力とか、座標ごとに値が変化するようなものを計算するのに良いのですね。

そんなの連続座標で座標値があるんだからそれで計算すればええやんと思うかもしれませんが、どうしても離散グリッドが必要な場合というのはたくさんあるのです。(主に組合せ的な最適化問題を探索的に解く場合とかですね…)

ということで、例として10 x 10 x 10の3次元配列に3次元の座標ベクトルが入っている配列を作ってみましょう

import numpy as np
X, Y, Z = np.mgrid[0:10, 0:10, 0:10]
# グリッド中の全セルを走査するindex_listを作る(shape: (1331, 3) )
index_list = np.vstack((X.flatten(), Y.flatten(), Z.flatten())).T
# 3次元の位置ベクトルが格納される3次元配列
index_mesh = np.zeros((10, 10, 10, 3))
# index_listの値でindex_meshをselectし、index_listを流し込む
index_mesh[[index_list [:, 0], index_list [:, 1], index_list [:, 2]]] = index_list

こうすることで、index_mesh[i, j, k]には[i, j k]という値が格納されました。

これならば座標に応じた値を計算するのも簡単ですね。球状に広がるようなエフェクトを作るのも簡単でしょう。

さらに、中心からの方向を示す単位ベクトルに変換してみましょう。まずは中心からのベクトルを計算します。

index_mesh -= np.array([5, 5, 5])

これで中心座標、今回は(5, 5, 5)での値が(0, 0, 0)になりました。

あとは各ベクトルを正規化します

norms = np.linalg.norm(index_mesh , axis=3)
index_mesh = index_mesh / norms[:, :, :, None]
index_mesh [5, 5, 5] = 0

これで、中心から放射状に広がる単位ベクトルの配列となりました。