np.whereで効率的に値を出し分ける

今回もnumpyのテクニックの紹介です。np.whereというメソッドを解説します。

参考: numpy.where — NumPy v2.0 Manual

これは何かというと、第1引数にTrue/Falseで評価できるデータの配列を渡すとその評価に応じてTrueなら第2引数、Falseなら第3引数の値を返す、というものです。

第2, 第3引数に渡すのは第1引数に渡した配列と同じ長さ(多次元なら全て同じ)でも良いし、定数であったり、ブロードキャストすれば同じ形にできるものなら何でも良いです。

一番シンプルな例としては、条件を満たすかどうかでそれぞれ異なる定数を返すようなものでしょうか。

import numpy as np


scores = np.array([45, 85, 72, 50, 90])
results = np.where(scores >= 60, '合格', '不合格')
print(results)
# ['不合格' '合格' '合格' '不合格' '合格']

説明いらないと思いますが、60点以上なら合格、と判定するメソッドですね。

上記の例のように、事前にTrue/False の配列を作っておくのではなく、何かしらの条件式を代1引数に渡すような使い方になると思います。条件に応じて何かしらの演算を行いたい場合は、第2, 第3引数に計算式を入れて結果を渡すような形になります。例えば、偶数なら1/2, 奇数なら 3倍して1を足す、みたいな処理をするならこうです。

np.where(scores%2 == 0, scores/2, 3*scores+1)
# array([136., 256.,  36.,  25.,  45.])

1次元配列の場合は、内包表記でもほぼ同じことができるのでありがたみが薄いですが、np.whereは多次元配列で便利なことがあります。(単純に、内包表記の方が不便になるだけという見方もできますが。)

自分が最近使った例としては、欠損値がある行列Aと別の行列Bがあった時に、欠損値以外は元の行列Aの値、欠損してる部分はBの値、で埋めたいというものでした。

これが次のようにして簡単に行えます。

A = np.array(
        [[1, 2, 3,], [np.nan, 5, 6], [7, np.nan, 9]]
    )
B = np.array(
        [[11, 12, 13,], [14, 15, 16], [17, 18, 19]]
    )

print(np.where(np.isnan(A), B, A))
"""
[[ 1.  2.  3.]
 [14.  5.  6.]
 [ 7. 18.  9.]]
"""

コードがすっきり書けること以外にもベクトル処理が行えることによるパフォーマンス面のメリットなど、利点があるので機会があれば使ってみてください。

Nanを含むnumpy配列のデータを専用メソッドで手軽に集計する

numpyのちょっとしたテクニックの話です。僕は最近まで知らなかったのですが、numpyには np.nansum など nan + 集計関数名 という命名規則のメソッド群が用意されています。これの紹介をします。

前提として、 numpy配列の値を合計したり平均を取ったりする時、データ中にnanがあると結果もnanになります。pandasのSeriesの場合と挙動が違うのですね。例えば以下のような感じです。(Seriesと挙動違うよという話は以前どこかの記事で書いた覚えがあります)

import numpy as np
import pandas as pd


# nanを含むデータを作る
ary = np.array([1, 1, 2, 3, np.nan, 8])
print(ary)
# [ 1.  1.  2.  3. nan  8.]

# 合計するとnanになる
print(ary.sum())
# nan

# 平均も同様
print(ary.mean())
# nan

# Series はnanを無視してくれる
print(pd.Series(ary).sum())
# 15.0
print(pd.Series(ary).mean())
# 3.0

欠損値の存在に気づくきっかけになったりしてありがたいこともありますし、仕様としてどうあるべきかを考えたらnullの伝播が実装されているこの作りが正しいと思えるのですが、この挙動が不便なことが多いのも事実です。

僕はこういう時大体Seriesに変換してしまって集計していました。

ただ、実は numpyにもNanに対応したメソッドがちゃんとあり、それが冒頭に書いたnansumです。maxにはnanmax, stdにはnanstd のように多くのメソッドに対して実装されています。

dir()で探すと一覧額作成できます。

for m in dir(np):
    if m.startswith("nan"):  # メソッド名がnanで始まるか
        if m.replace("nan", "") in dir(np):  # nanの部分を除外した場合に同じ名前のメソッドがあるか
            print(m)
"""
nanargmax
nanargmin
nancumprod
nancumsum
nanmax
nanmean
nanmedian
nanmin
nanpercentile
nanprod
nanquantile
nanstd
nansum
nanvar
"""

これらを使うと、エラーが起きずにnanを無視して無視して残りの要素について集計してくれます。

print(np.nansum(ary))
# 15.0
print(np.nanmean(ary))
# 3

1次元配列の場合は内包表記での対応とか色々やり方もあるのですが多次元になってくると面倒だし集計のために補完するのも面倒なのでありがたいですね。使い方がnp.nansum(ary)であって、ary.nansum() では無いので注意してください。

もう一点、 np.nan ではなく、Noneを含めてるとこれは数値の欠損値では無いので相わらずエラーになります。ここも注意です。

ary2 = np.array([1, 1, 2, 3, None, 8])

try:
    np.nansum(ary2)

except Exception as e:
    print(e)
# unsupported operand type(s) for +: 'int' and 'NoneType'