Pythonのコードを機械語にコンパイルして処理を高速化するライブラリnumbaにはvectorizeとguvectorizeというデコレーターが用意されており、これを使うことで簡単に関数をベクトル化することができます。このデコレーターの使い方を解説します。

vectorizeによるベクトル化

numbaのvectorizeというデコレーターを関数に付けるだけで、その関数をベクトル化することができます。例えば2つの数$x$と$y$を受け取って$x^2 + y^2$を返す関数dist_sqの場合は次のようになります。

from numba import vectorize, float64

@vectorize([float64(float64, float64)])
def dist_sq(x, y):
    return x**2 + y**2

次のように、これだけでdist_sqは配列を受け取り要素ごとのdist_sqの結果を返せるようになっています。コンパイルもされているので高速に動きます。もっと高速に動かしたい場合にはvectorizeの引数にtarget='parallel'を渡したりtarget='cuda'を渡すといいと思います。

import numpy as np

x = np.arange(3)
y = np.arange(3)
dist_sq(x, y)

# => array([0., 2., 8.])

vectorizeによって修飾された関数はnumpyのufuncというものになっているので、関数の引数にブロードキャストが効きます。このことを利用すると、例えば
配列xと配列yの要素の組み合わせごとにdist_sqを計算する処理が次のように簡潔に書けたりします。

x = np.arange(3)
y = np.arange(4)
z = dist_sq(x[:, np.newaxis], y[np.newaxis, :])

# => array([[ 0.,  1.,  4.,  9.],
#           [ 1.,  2.,  5., 10.],
#           [ 4.,  5.,  8., 13.]])
# 結果の配列の(i, j)番目の要素は x[i]**2 + y[j]**2 と等しい。

numpyのブロードキャストに慣れている人ならjitデコレーターよりもvectorizeの方が使い勝手がいい場面も多いのではないかと思います。

guvectorizeによる拡張版ベクトル化

引数にスカラーだけではなく配列も受け取るような関数をベクトル化するにはguvectorizeというデコレーターを使います。例として、numpyのsearchsortedという関数をnumbaで自作してみます。この関数は1次元の配列aとスカラーvを受け取ってa[i-1] < v <= a[i]となるiを返します。この関数をnumbaで実装すると次のようになります（numpyのsearchsortedは二分探索を行いますが、ここでは簡単のため線形探索で実装しています）。

import numpy as np
from numba import guvectorize, float64, int64

@guvectorize([(float64[:], float64, int64[:])], '(n),()->()')
def my_searchsorted(a, v, res):
    if v >= a[-1]:
        res[0] = a.shape[0]
        return

    for i in range(0, a.shape[0]):
        if v <= a[i]:
            res[0] = i
            return

a = np.array([0, 1, 2])

my_searchsorted(a, 0.5)
# => 1

my_searchsorted(a, np.array([0.5, -1, 2.1]))
# => array([1, 0, 3])

デコレーターguvectorizeに渡している最初の引数[(float64[:], float64, int64[:])]は修飾する関数my_searchsortedの引数と戻り値の型と形状を表しています。戻り値の形状には、スカラーを返す場合でも1次元配列のように[:]を付けます。そして関数本体の中ではres[0]に戻り値を書き込みます。numbaのドキュメントではこの引数をシグネチャーと呼んでいるのですが、numpyの言葉の使い方と異なるので、ここではこの引数のことをシグネチャーとは呼ばないことにします。

デコレーターguvectorizeに渡している2番目の引数'(n),()->()'は、修飾する関数の引数と戻り値の次元と形状を表す文字列です。この場合だと、最初の引数は1次元かつそのサイズは任意、次の引数はスカラー、戻り値もスカラーということを表しています。numpyのドキュメントではこの文字列のことをシグネチャーと呼んでいるので、ここでもそう呼びます。numbaのドキュメントではこの文字列をレイアウト文字列と呼んでいるのですが、ここではnumpyの用語を用います。シグネチャーによって明示された引数や戻り値の次元のことをコア次元と言います。この例だと、最初の引数のコア次元は1、次の引数のコア次元は0ということになります。

デコレーターguvectorizeで修飾された関数はnumpyのufuncというものになります。ufuncに配列を渡した場合、シグネチャーに従って次のルールでブロードキャストがなされます。

引数に渡された配列の次元の後ろの方から次元をコア次元分だけ取り除く。
コア次元が取り除かれた引数の配列の形状がブロードキャストによって揃えられる。揃えられた形状の次元の数のことをループ次元と言います。
戻り値の次元の数を「ループ次元+戻り値のコア次元」と決定する。後ろのほうの次元がコア次元になります。

例として、シグネチャーが'(n),()->()'のufuncに形状が(n)の配列と(a, b)の配列を渡した場合のブロードキャストを考えます。まず、1番目のルールによりコア次元が取り除かれます。取り除かれたあとの形状は()と(a, b)になります。()は(a, b)にブロードキャストできるので、2番目のルールによって形状が(a, b)に揃えられます。ループ次元はこの場合2次元になっています。戻り値のコア次元は0なので、戻り値の次元はループ次元と等しく2、戻り値の形状は(a, b)になります。

単にnumpyのブロードキャストと呼ばれるものは、シグネチャーが(),()->()のufunc（2項演算）に対するブロードキャストだったわけです。

先ほどの例my_searchsortedに戻ると、実は引数aにも次のように多次元配列を代入することができるので、numpyのsearchsortedより少しパワーアップしています。

a = np.array([[0, 1, 2], [3, 4, 5]])
v = np.array([1.5, 1.5])
my_searchsorted(a, v)

# => array([2, 0])
# 結果の配列のi番目の要素は my_searchsorted(a[i, :], v[i]) と等しい。
# ブロードキャストは次のように動いている：
# ・引数の形状は(2, 3)と(2)
# ・コア次元を取り除くと(2)と(2)になる。形状を揃えた結果も(2)なので、戻り値の形状も(2)

numbaのvectorizeを使いこなす

vectorizeによるベクトル化

guvectorizeによる拡張版ベクトル化

参考文献