numpy の array で、特定の値を持つ要素を探す

PyFITS を使ってある 3 次元の FITS ファイルを開いたところ、観測データの存在しないピクセルは NaN が詰まっていました。この配列に対して処理をしようとすると、演算が不正になってしまうので、以下の方法で NaN のピクセルだけ値を 0 に置き換えました。

import pyfits

hdulist = pyfits.open("foo.fits")
data = hdulist[0].data

index = numpy.where(data != data) # check if NaN

for i in range(len(index[0])):
    data[index[0][i]][index[1][i]][index[2][i]] = 0.

numpy.where という関数を使うことで、条件に一致する配列要素の (i, j, k...) の組を tuple で返してくれます。

入門 Python 3

入門 Python 3

以下の行では、NaN と NaN を != 演算子で比べた場合の戻り値が True であるという特殊な振る舞いを利用しています。

index = numpy.where(data != data) # check if NaN

もしも値が 7.5 のピクセルだけ処理したい場合は、次のようにすれば OK です。

index = numpy.where(data == 7.5)

次の書き方では期待通りに動作しないので注意が必要です。

index = numpy.where(data == float("nan"))