numbaのvectorizeを使いこなす


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

# vectorizeによるベクトル化

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

```python
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'`を渡すといいと思います。

```python
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`を計算する処理が次のように簡潔に書けたりします。

```python
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`](https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html "numpy.searchsorted")という関数をnumbaで自作してみます。この関数は1次元の配列`a`とスカラー`v`を受け取って`a[i-1] < v <= a[i]`となる`i`を返します。この関数をnumbaで実装すると次のようになります（numpyの`searchsorted`は二分探索を行いますが、ここでは簡単のため線形探索で実装しています）。

```python
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に配列を渡した場合、シグネチャーに従って次のルールでブロードキャストがなされます。

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

例として、シグネチャーが`'(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`より少しパワーアップしています。

```python
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)
```

# 参考文献
* [Generalized Universal Function API](https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html "Generalized Universal Function API")
* [Universal functions (ufunc)](https://numpy.org/doc/stable/reference/ufuncs.html "Universal functions (ufunc)")
* [numba Reference Manual/Just-in-Time compilation/guvectorize](https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#numba.guvectorize "Just-in-Time compilation")