np.vectorizeで関数をベクトル化する

NumPyやScyPyの関数って非常に便利で、NumPy配列(要するにArray)を渡すと空気を読んでその渡したデータの各要素に関数を適用してNumPy配列で結果を返してくれたりします。

自分で定義した関数でもNumPyやSciPyの関数の組み合わせで作った関数であれば結構そのように動いてくれるのですが、文字列操作が入ったりif文による分岐等があると必ずしもそうはならず、スカラー値を受け取ってスカラー値を返すだけの関数になることがあります。

そのような関数を、手軽にベクトルか対応することができる方法があるのでこの記事で紹介します。

それが、記事タイトルのnp.vectorizeです。

ドキュメント: numpy.vectorize — NumPy v1.26 Manual

関数を渡すと戻り値で新しい関数オブジェクトが帰ってきてそれがベクトル対応(配列対応)しています。

基本的な使い方

数学関数だと特にArrayを渡すと元々期待通り動いたりするので、少々無理矢理な例ですが文字列操作の関数を作ってお見せします。これは数値を1個受けとって、その数値に、「回目」っていいう単位をつけて返すだけの関数です。普通に実験、そのまま配列渡してみる、ベクとライズして配列を渡してみる、の3パターンやってみました。

import numpy as np


# 数値に単位をつける関数を実装
def number_format(n):
    return f"{n}回目"


# 数値を渡すと想定通り動く
print(number_format(5))
# 5回目

# 配列を渡すと配列を一個の値とみなして文字列化して単位をつけてしまう。
print(number_format([1, 2, 3]))
# [1, 2, 3]回目

# ベクトル化した関数を作る
number_format_vec = np.vectorize(number_format)

# それに配列を渡すと配列の各要素に元の関数を適用してくれる。
print(number_format_vec([1, 2, 3]))
# ['1回目' '2回目' '3回目']

# Array型もタプルもいける
print(number_format_vec(np.array([1, 2, 3])))
# ['1回目' '2回目' '3回目']
print(number_format_vec((1, 2, 3)))
# ['1回目' '2回目' '3回目']

# もちろん、内包表記で同じことをすることは可能(ただし、この結果はlist)
print([number_format(n) for n in [1, 2, 3]])
# ['1回目', '2回目', '3回目']

ベクトル化した関数を1回しか使わないなら内包表記で済ましちゃっていいんじゃないかな、と思うのですが、何度も利用したい関数であればnp.vectorizeを使うと言う選択肢もあるのかな、と思います。

注意点

NumPyやSciPyで実装されている関数群って並列処理できる部分は並列処理するような賢い実装になっていることがありますが、この np.vectorize はそこまで気が利いたものではありません。どうやら単純にfor文で順次処理するようになるだけらしいので処理の高速化等の効果はありません。ドキュメントにも利便性のためのもので、パフォーマンスのため使うようなものではなく、for loop回してるだけだって書いてありますね。

そのため、本当に頻繁に大規模なベクトルを処理する関数なのであれば別の方法で対応させる必要があるでしょう。

もう一点、細かいですが戻り値がNumPyのarrayであることも注意が必要ですね。と言ってもこれは便利に感じることが多いですが。内包表記であればlistで結果が得られますがvectorizeするとlist渡してもlistではなくarrayで帰ってきます。

引数を複数受け取る関数の場合

この np.vectorize は引数を複数受け取る関数にも対応しています。ドキュメントのサンプルもa, b の2変数受け取っていますしね。一応その例も見ておきましょう。年と月の数値を受け取って何年何月、という文字列返す関数でやってみます。

def month_str(year, month):
    return (f"{year}年{month}月")


month_str_vec = np.vectorize(month_str)

# 元の関数はyear, monthは1個ずつしか値を受け取れない
print(month_str([2020, 2023, 2026], [1, 4, 7]))
# [2020, 2023, 2026]年[1, 4, 7]月

# ベクトル化すると複数ペアをまとめて処理できる。
print(month_str_vec([2020, 2023, 2026], [1, 4, 7]))
# ['2020年1月' '2023年4月' '2026年7月']

# 片方は配列で、片方はスカラーというパターンにも対応する
print(month_str_vec([2020, 2023, 2026], 1))
# ['2020年1月' '2023年1月' '2026年1月']

さいごに

以上が手軽に関数をベクトル化する方法でした。まぁ、内包表記もあればmapを使うやり方もあるのでこれが必須というわけではないのですがいい感じに動く関数を手軽に作る方法として頭の片隅に置いておくと使う場面はあるんじゃないかなと思います。

ちなみに、関数を定義した直後にベクトル化した関数で元の関数名を上書きしておくと、最初っからベクトル化した関数を宣言したのと同じように使えますよ。

def func(x):
    # 何かの処理


func = np.vectorize(func)
# 以降に呼び出されるfuncはベクトル対応した関数。