静かなる名辞

pythonとプログラミングのこと


numbaとnumpyで速いループ処理を書くためのガイド(スレッド並列化のおまけつき)(実はポエム)

はじめに

 この記事は「Pythonおっせーよ」と思っている、そこのあなたのためのものです。

 PythonはLLなので遅いです。その分、楽に書けるし、動的型付けでダイナミックなことができて「楽しぃぃいい」のですが、それでも遅くて困るときがあります。特に数値計算的なことをやらせようとすると*1

 え、numpy使えばいける? numpyは、numpyに組み込まれている関数では処理できなくて、ベクトル演算も効かないようなタイプの処理には基本的に無力です。いけません。

 どうせC/C++で書かれてコンパイルされたバイナリのwrapperと割り切って使うから良い? ライブラリが見つかればそれで良いでしょう。問題はライブラリに頼るまでもない、ちょっとした処理の実装です。コンパイル言語で書いたりすると(Cythonも含みます)、面倒くさくて割りに合わない気がします。なんでこんなことやってるんだという気持ちになれます。

 numbaだと、並のコンパイル言語よりは気楽に書けて、単純な数値計算とかなら並のコンパイル言語くらいの速度が実際に出ます。素晴らしいですね。

 ということで、numpy+numbaで書くときに気をつけることを色々と書いておこうと思います。

 はい。気をつけることが色々あるんです。気をつければ速いんですけどねー。けっこう面倒くさいんですよねー。それでもCythonより楽というのがnumbaの位置づけなので、知っておくとたまに得するかもしれません。どうしても高速に走らないといけないループだけnumbaで書き下すといったプログラミングが可能になります。それくらいのものです。

 目次

コンセプト

 前提として、numpyがどんな仕組みで高速な処理を達成しているのかを知らない人は、この記事を読んでも活用できないと思います。一応、この章で最低限のことは説明しますが、これを読んでわかる人には要らないような説明です。わからない人はCの勉強から。

 先にこの章の結論を書いてしまいます。numpyはぶっちゃけていえば、C配列のwrapperみたいなものです*2

 A = np.zeros(shape=(3, 4), dtype=np.float64)とかすると内部でmalloc(sizeof(double) * 3 * 4)されてメモリ領域が確保され、そしてこれへアクセスするためのPython Objectも作られます。プロキシ・オブジェクトです。PythonでA[i, j] = 1.0すると、内部で確保した配列に対して[i][j] = 1.0が走るわけです。

 ということで、numpy配列は原理的には、その気になればCの配列っぽく扱えます。scikit-learnなどの多くの機械学習ライブラリはnumpy配列で表現されたデータを受け取るように実装されていますが、「すべてがオブジェクト」的な世界観とは無縁のバイナリの世界で処理をするからこそ高速に処理を達成できる訳です。

 逆に、Python側からnumpy配列の中にある数値を取り出したり、numpy配列に数値を書き込んだりする処理は、びっくりするほど遅くなります。それをやる度に、いちいち対応するPython Objectを生成せざるを得ないのです。なにしろPythonの世界は「すべてがオブジェクト」だから、そうならざるを得ない。

 そういう事情から、Pythonでnumpyにバリバリ添字アクセスする処理を書いても、だいたい悲惨な結果になります。

 それでもnumpyは、Pythonで高速な数値計算を実現する方法を、不完全ではありますが提供してくれています。それがベクトル演算で、複数要素にまとめて同じような処理をするのであればPythonの世界とのやり取りはそんなにしなくて済むので、速いという発想です。

 多くの場合ベクトル演算は有用です。が、本質的には予め用意されている処理しか書けないし、どんな処理でもベクトル演算で書ける訳ではないので、限界があります。そういう限界に突き当たった人は添字アクセスの処理を書いて、あまりの遅さに絶望し、なんとかならないかと思ってnumbaにたどり着く訳です*3

 さらに言うと、numpyのベクトル演算にはもう一つ問題っぽいものがあります。素直に書くとベクトル演算した結果が新たな配列として返ることが多いので、けっこう豪快にメモリを使います。ある程度は累積代入文などでカバーできますが、駄目なときもあります。これもなんとかしたい点です。メモリを余計に食うから……ではなくて(GBとかの死ぬほどでかい配列でない限り、メモリ消費自体は普通は問題にはなりません)、メモリを新規に確保するのはそもそも遅い処理だし、キャッシュ効率も悪くなるので、できればインプレースで済むことはインプレースで済ませたいからです。

 ということで、numpy配列の中にあるC配列を直接添字アクセスしてループできると嬉しい、という需要がある訳です。というかそのためのnumpy配列です。もし本格的にやるならC/C++かCythonの出番ですが、もう少し気楽にやりたいという需要もある訳です。それがnumbaの存在する理由です。

 そしてこの記事は、実際にnumbaを使う考え方というか実践的なコツみたいなものを、思いつくままに書いていくというものです。

jitを効かせる

 numbaはJust-in-time compilationという技術でコードの高速化を行います。背後ではLLVMが動きます。私にとってはほとんど謎の技術ですが、ここではnumbaのバックエンドと割り切って気にしないことにします。コンパイラの賢さを信じる。

 重要なのはjitです。というかjitデコレータです。numbaの中核です。極論すればfrom numba import jitだけでnumbaは使えます。

 ただし、日本語Web圏でjitデコレータについて言われていることには注意が必要です。

 検索すると普通に出てくるnumbaの解説記事では「関数に@jitつけるだけでjitコンパイルされて速くなる!」とか書いてあることがあります。それしか内容がなかったら、ブラウザバックしましょう。それかタブを閉じましょう。役に立たない記事です。

 numbaは「オレの書いたイケイケなPythonコードをチョースゴイ技術でJITコンパイルしてくれる」技術ではありません。「色々気を配ってCっぽいPythonを書けばCっぽい速さで動いてくれる」というものです。そこを勘違いした人たちはみんなインターネットの夜空に輝く素晴らしい星々になりました。
(この流れでついでに書いておくと、Cythonは「PythonっぽいC」「Cではないにしても少なくともPythonではないなにか」です。numbaの場合はギリギリ構文まではPythonしています。Pythonパーサで構文解析されてPythonの抽象構文木をJITコンパイルするのです。あれ、それはそれでPythonではないっ?)

 なんで「@jitだけで」が駄目なのかについてもう少し書いておきます。単純な話、そもそもjitコンパイルできていないnumbaの記事が死ぬほどたくさんあります。numba非対応の機能を使うとそうなるのです。その場合は「Python互換モード」みたいなこれまた正体のよくわからないもので動かしてくれますが、こちらは劇的に遅いのが普通です。

 ということで、最低限以下の使い方をしましょう。

@jit(nopython=True)
def f(...):
    ...

 nopython=Trueというオプションが追加されました。こうすると謎の「Python互換モード」とはおさらばでき、コンパイルに失敗したらエラーが出ます。うまく動かないときは動かさない、これが第一歩です。


numpy配列で入出力のやり取りをする

 JITが速いのは実は静的型付けでコンパイルして動かすからです。

 私個人の心情としては静的型付けには窮屈とか面倒くさいといった印象を抱いてしまいがちですが*4、パフォーマンス上の理由なので好き嫌いは置いておいて静的型付けを使いましょう。

 numbaの良いところは型推論してくれることです。a = 42と書いてあったら、aはたぶんint64です、とか。なので型の指定は書かなくて済みます。cdefまみれになるCythonとの違いです。

 ただし、引数と返り値の型は指定した方が何かと高速化の恩恵を受けられます。というか、やらないとたぶんコンパイルエラーでまともに走らないでしょう。

2.2. Just-in-Time compilation — Numba 0.48.0-py3.6-macosx-10.7-x86_64.egg documentation

 指定の仕方も実は3通りもあったりします。1つ目はtupleで引数の型のみ指定する方法、2つ目はnumbaの型オブジェクトを作って渡す方法、3つ目は文字列で書く方法です。文字列で書きたい人が多いと思うので、それを紹介します。

 シグネチャはデコレータの第一引数であり、「使え」という強い意志を感じます。

@jit("i8(f8, f8)", nopython=True)
def f(a, b):
    ...

 "i8(f8, f8)"というのが型の指定になります。カッコの前が関数の返り値、後ろのカッコの中身が引数の型であり、i8は8バイト整数(つまりint64)、f8は8バイト浮動小数点数(float64)です。浮動小数点数を二つ受け取って整数を返す関数です。

 なお、Pythonオブジェクトをこういうnumbaの関数に渡した場合はそれなりに相互変換はしてくれますが、普通に「numpy配列の要素を取り出すとき」の問題が生じるので損な選択になります。

 ということで、numpy配列でやり取りしましょう。配列はスライス表記で表します。

@jit("i8[:](f8[:, :], f8[:, :, :])", nopython=True)
def f(a, b):
    ...

 返り値は整数の一次元配列、第一引数は浮動小数点数の二次元配列、第二引数は浮動小数点数の三次元配列みたいな関数です。

 これが、最初に書いた「numpy配列の中にあるC配列を直接いじりたい」を実現するための魔法です。こうしておけばnumbaが適切にやり取りしてやってくれます。なかなか夢が広がりますね。

 使える型とかの詳細はリファレンスを見てください。

2.1. Types and signatures — Numba 0.48.0-py3.6-macosx-10.7-x86_64.egg documentation

 また、numba.typeofを使うことで任意のオブジェクトに適合するシグネチャを調べることもできます。簡略表記で返してくれないというか、numbaの型オブジェクトを返してくれる(なので型オブジェクトを作って渡す方法向きではあります)のが微妙ですが。このオブジェクトどう書くんだ? というとき使いましょう。

>>> import numba
>>> numba.typeof(1)
int64
>>> numba.typeof(1.0)
float64
>>> @numba.jit("float64(int64)", nopython=True)
... def f(a):
...     return a + 1.5
... 
>>> f(10)
11.5

2.1. Types and signatures — Numba 0.48.0-py3.6-macosx-10.7-x86_64.egg documentation


 なお、当たり前ですが、合致しない型とか形状の配列を渡すと、実行時にエラーになります。JITコンパイルする本丸の関数とは別に、pythonレイヤで形状チェックとかastypeとかしてから本丸を呼び出すようなwrapper関数を書いておくと実用的です。どうせ自分しか使わないしコードのこの箇所では入ってくる型は確定しているしということで割り切るなら、それはそれで実用的です。

 あと、numpy配列を渡せば普通に引数に副作用を及ぼせるので、返り値は必須ではありません。インプレース処理にしてもいいし、返り値用の配列を引数として渡してもいいのです。返り値が要らないときは返り値の型のところにはvoidと書いてください。

jitコンパイルできるコードを書く

 
 ここと、

2.6. Supported Python features — Numba 0.48.0-py3.6-macosx-10.7-x86_64.egg documentation

 ここを、

2.7. Supported NumPy features — Numba 0.48.0-py3.6-macosx-10.7-x86_64.egg documentation


 見て、書いてある機能だけで書いてください。終了。

 では不親切なので、もう少し書きます。

 たとえばnumbaではループを書くのにPythonのrangeが普通に使えます。いえ、使えるようにみえます。

# @jit(nopython=True)しても動く
def f():
    a = 0
    for i in range(10):
        a += i
    print(i)

 Pythonだとrangeオブジェクトを生成して、そこから生成したイテレータから値を一つずつ取り出して……という感じでこのコードは動きます。

>>> range(10)
range(0, 10)

 numbaはまったく違う仕組みで実行します。最終的には機械語のループに変換されます。カウンタ足してってジャンプ命令とかする奴です。

 なので当然ながら制約があります。rangeオブジェクトを引数として受け取ったりはできません。

# @jit(nopython=True)できない
def f(r):  # f(range(10))とか呼び出したいができない
    a = 0
    for i in r:
        a += i
    print(a)

 といった面倒くさい決まりごとが色々あるのがnumbaです。細かく把握して慣れ親しんでいけば、そのうち「こんなのnumbaで書けるんだ!」という驚きが発生するかもしれませんが、普通はだるいのでコンサバティブなコードを書くことになります。

 それでも一昔前のnumbaに比べれば最近のnumbaはかなり立派になっていて、たとえばリスト内包表記なんか昔はできなかったのにいつの間にかできるようになっていて私なんかはかなりびびったのですが、でもPythonオブジェクトのリストを作るんだから遅えんだろうな……という諦観も生じてきて複雑な気持ちになったりして、まあ要するに「Python的な」コード、向いてませんよということです。どうせ相手にするのはnumpy配列、CかFORTRANを書くつもりで書くくらいの気持ちで扱いましょう。そうすると高級なコンパイル言語……? という気がしてきて幸せになれます。

 ということで、手っ取り早いのはrangeループによるインデックスアクセスとnumpyの機能とを組み合わせて使うことなので、numpyへのサポートを見てみます。

 結論を先に言ってしまうと、我々が想像することはだいたいできます。基本的なindexingやslicingはできますし、代表的な関数もだいたい使えます。

 numpyのリッチな機能が使えても、それはつまりnumpyとして走るだけなので、numba的に速くはないのが微妙なのですが。
(たとえば単一のnumpy関数を呼び出して時間がかかっている処理をそのまま関数に入れて@jitしても、高速化されません。つーかそれで速くなったら大慌てだよね。たぶんみんな使うし、それができるならnumpy開発者が最初から組み込むよね)

 コツとしては「Pythonコードで書かれた複雑な処理」をそのまま移すんじゃなくて、ボトルネックになっている書かざるを得ないループをnumbaに移すようにしましょう。短い方がコンパイルを通すのが楽です。あとは走らせるとコンパイルエラーがたくさん出ると思うので、だいたい読んでググればなんとかなります。それと、最初から書いたものを分析やシミュレーションなどの処理の中に組み込で使おうとすると、書き換える度に五分前処理してエラーになるとか辛い事態が発生するので、ダミーデータで走らせるテストも書いておくと良いでしょう(これはnumbaに限らず大切なことなのですが)。

ちょっとだけ実践編

 お待たせしました、実践編です。実際にnumba+numpyでどんなことができるのか、簡単に見ていきます。

 この記事ではfloat64のcumsumを実装します。二次元配列を受け取り、行ごとにcumsumするものを目指します*5

 え、cumsumはnumpyにありますよね。numba要らないって? ブラウザバックしないでください

 私が定義するcumsum(mycumsumと名付けます)は以下のような計算です。

 任意の要素xに対して

  • 0.0 <= x

 普通にcumsm

  • -1.0 <= x < 0

 リセットしてxを0にする

  • -2.0 <= x < -1.0

 前の値を0.5倍したものをxにする(最初なら0)

  • x < -2.0

 前の値に42.0を足してxにする(最初なら42.0)

 はい、cumsumじゃないですよね。名前が思いつきませんでしたが、cumsumに余計な処理が増えたなにかです。

 とりあえずpythonで書いてみましょう。

import random

def mycumsum(lst):
    t = 0.0
    result = []
    for x in lst:
        if 0.0 <= x:
            t += x
        elif -1.0 <= x < 0.0:
            t = 0.0
        elif -2.0 <= x < 1.0:
            t *= 0.5
        else:
            t += 42.0
        result.append(t)
    return result


random.seed(42)
data = [random.uniform(-10.0, 10.0)  # 見づらいので[-10, 10]にする
        for _ in range(10)]

print(data)
print(mycumsum(data))

""" =>
[2.7885359691576745, -9.49978489554666, -4.499413632617615, -5.5357852370235445, 4.729424283280249, 3.533989748458225, 7.843591354096908, -8.261223347411677, -1.561563606294591, -9.404055611238594]
[2.7885359691576745, 44.78853596915768, 86.78853596915768, 128.7885359691577, 133.51796025243794, 137.05195000089617, 144.89554135499307, 186.89554135499307, 93.44777067749654, 135.44777067749652]
"""


 書けたけど……と思いますよね。これは遅いPythonの処理です。上述の理由によりnumpy配列の要素にPythonレイヤからアクセスすると遅いので、numpy配列を渡してもそんなに変わりません(かえって遅い可能性もあります)。そしてこの関数をnumpyのベクトル演算で効率的に実装することはできません(たぶん)。

 まあでも、一応速度を測ってみるか。意外と速いかもしれない。

import time
import random
import numpy as np


def mycumsum_row(lst):
    t = 0.0
    result = []
    for x in lst:
        if 0.0 <= x:
            t += x
        elif -1.0 <= x < 0.0:
            t = 0.0
        elif -2.0 <= x < 1.0:
            t *= 0.5
        else:
            t += 42.0
        result.append(t)
    return result

def mycumsum(A):
    return [mycumsum_row(row) for row in A]

random.seed(42)
# (1000, 1000)配列とします
data = [[random.uniform(-10.0, 10.0) for _ in range(1000)]
        for _ in range(1000)]

# 実行時間計測
times = []
for i in range(10):
    t1 = time.time()
    res = mycumsum(data)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 0.11440501213073731

# numpyのcumsumもやります
data = np.array(data)  # numpy配列にしないと遅いんで・・・

# 実行時間計測
times = []
for i in range(10):
    t1 = time.time()
    res = np.cumsum(data, axis=1)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 0.003912162780761719

# 全体の結果を改めて
""" =>
0.11293444633483887
0.003912162780761719
"""

 知 っ て た 。

 まあ、処理の複雑性を考えると意外と健闘している方かもしれません。約29倍で済んでいます。普通にpythonでループさせるとこれくらいになるという例です。

 じゃあ、次はnumbaで書きましょう。実装の仕方はこんな感じです。

@jit("f8[:, :](f8[:, :])", nopython=True)
def mycumsum_numba(A):
    result = np.empty(A.shape, dtype=np.float64)
    for i in range(A.shape[0]):  # 行のループ
        t = 0.0
        for j in range(A.shape[1]):  # 列のループ
            x = A[i, j]
            if 0.0 <= x:
                t += x
            elif -1.0 <= x < 0.0:
                t = 0.0
            elif -2.0 <= x < 1.0:
                t *= 0.5
            else:
                t += 42.0
            result[i, j] = t
    return result

 特に解説は要らないでしょう。最適化とかはまったく考えていない愚直な実装です。

 また測る。

import time
import random
import numpy as np
from numba import jit

def mycumsum_row(lst):
    t = 0.0
    result = []
    for x in lst:
        if 0.0 <= x:
            t += x
        elif -1.0 <= x < 0.0:
            t = 0.0
        elif -2.0 <= x < 1.0:
            t *= 0.5
        else:
            t += 42.0
        result.append(t)
    return result

def mycumsum(A):
    return [mycumsum_row(row) for row in A]

@jit("f8[:, :](f8[:, :])", nopython=True)
def mycumsum_numba(A):
    result = np.empty(A.shape, dtype=np.float64)
    for i in range(A.shape[0]):  # 行のループ
        t = 0.0
        for j in range(A.shape[1]):  # 列のループ
            x = A[i, j]
            if 0.0 <= x:
                t += x
            elif -1.0 <= x < 0.0:
                t = 0.0
            elif -2.0 <= x < 1.0:
                t *= 0.5
            else:
                t += 42.0
            result[i, j] = t
    return result
            

random.seed(42)
# (1000, 1000)配列とします
data = [[random.uniform(-10.0, 10.0) for _ in range(1000)]
        for _ in range(1000)]

# 実行時間計測
times = []
for i in range(10):
    t1 = time.time()
    res_mycumsum = mycumsum(data)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 0.11440501213073731

# numpyのcumsum
data = np.array(data)

# 実行時間計測
times = []
for i in range(10):
    t1 = time.time()
    res = np.cumsum(data, axis=1)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 0.003912162780761719

# 自作numba版
times = []
for i in range(10):
    t1 = time.time()
    res_mycumsum_numba = mycumsum_numba(data)
    t2 = time.time()
    times.append(t2 - t1)

assert (res_mycumsum_numba == np.array(res_mycumsum)).all()  # 同じ結果か確認

print(np.mean(times)) # => 0.005931735038757324


# 全体の結果を改めて
""" =>
0.11293444633483887
0.003912162780761719
0.005931735038757324
"""

 numpyのcumsumの1.5倍で実行できます。処理の複雑性を考えると大健闘と言えます。

 という感じで使えるのがnumbaで、これはなかなかすごいことだと思います。

GILを切ってマルチスレッドする(おまけ)

 もっと速くしたいときは、numbaではGILを切れます。

 これを使ってマルチスレッド並列化による高速化ができます。Python的には未知の世界です。しかもpython標準のthreadingから叩けば勝手にGILなしで動くので、使いやすいものです。

 ただし先に言っておきますが、スレッド間の複雑な同期処理は難しそうなので、独立に動かせるスレッドがあることを前提にしてください……

 公式の例はこんなのですが、

1.19. Examples — Numba 0.48.0-py3.6-macosx-10.7-x86_64.egg documentation

 私もやってみます。



 まず、並列化する前のコードです。速すぎるので調子に乗って配列を大きくしました。

import time
import numpy as np
from numba import jit


@jit("f8[:, :](f8[:, :])", nopython=True)
def mycumsum_numba(A):
    result = np.empty(A.shape, dtype=np.float64)
    for i in range(A.shape[0]):  # 行のループ
        t = 0.0
        for j in range(A.shape[1]):  # 列のループ
            x = A[i, j]
            if 0.0 <= x:
                t += x
            elif -1.0 <= x < 0.0:
                t = 0.0
            elif -2.0 <= x < 1.0:
                t *= 0.5
            else:
                t += 42.0
            result[i, j] = t
    return result

np.random.seed(42)
data = np.random.uniform(-10.0, 10.0, size=(5000, 50000))

# 自作numba版
times = []
for i in range(10):
    t1 = time.time()
    res_mycumsum_numba = mycumsum_numba(data)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 1.8200997829437255

 そろそろ速度的にきつくなってきました。スレッド並列化でこれを数倍速くします。

 GILを切るにはnogil=Trueを指定します。が、GILを切ることで速くなるわけではありません。切った上でマルチスレッド処理を書くことで速くなるのです。

 実装としてはこんな感じです。

# 並列版の中身
@jit("void(f8[:, :], f8[:, :], i8, i8)", nopython=True, nogil=True)
def mycumsum_numba_pi(A, result, start, stop):
    for i in range(start, stop):  # 行のループ
        t = 0.0
        for j in range(A.shape[1]):  # 列のループ
            x = A[i, j]
            if 0.0 <= x:
                t += x
            elif -1.0 <= x < 0.0:
                t = 0.0
            elif -2.0 <= x < 1.0:
                t *= 0.5
            else:
                t += 42.0
            result[i, j] = t

# 並列版の外側(マルチスレッドで呼び出すためのもの)
def mycumsum_numba_p(A, n_threads=8):  # n_threadsはマシンに合わせて調整してください。当方8C/16Tです。16も試しましたがかえって遅かったので8にしています
    n = A.shape[0]
    result = np.empty(A.shape, dtype=np.float64)
    cl = n // n_threads
    l = [cl * x for x in range(n_threads)] + [n]
    args = [(A, result, start, stop) for start, stop in zip(l, l[1:])]
    threads = [threading.Thread(target=mycumsum_numba_pi, args=a)
               for a in args]
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    return result

 中身にはnogil=Trueしないと何も意味のないことになるので付けましょう。外側はnumbaの関数にしなくてオッケーです。

 結果を受け取るために、返り値を使うのをやめて引数への副作用で結果を取り出すテクを使います。公式の例を見る限り、うまくviewを活かせば、もしかしたらここまで愚直なインデックスアクセスに頼らなくてもできそうな気もしますが、よくわからないのでstartとstopで実装です。

 これによりnumbaで走るスレッドの外側で生成されたresultに対してスレッド並列に書き込みが行われるはずです。

 計測に使ったコードです。

import time
import threading
import numpy as np
from numba import jit

# 愚直な実装
@jit("f8[:, :](f8[:, :])", nopython=True)
def mycumsum_numba(A):
    result = np.empty(A.shape, dtype=np.float64)
    for i in range(A.shape[0]):  # 行のループ
        t = 0.0
        for j in range(A.shape[1]):  # 列のループ
            x = A[i, j]
            if 0.0 <= x:
                t += x
            elif -1.0 <= x < 0.0:
                t = 0.0
            elif -2.0 <= x < 1.0:
                t *= 0.5
            else:
                t += 42.0
            result[i, j] = t
    return result

# 並列版の中身
@jit("void(f8[:, :], f8[:, :], i8, i8)", nopython=True, nogil=True)
def mycumsum_numba_pi(A, result, start, stop):
    for i in range(start, stop):  # 行のループ
        t = 0.0
        for j in range(A.shape[1]):  # 列のループ
            x = A[i, j]
            if 0.0 <= x:
                t += x
            elif -1.0 <= x < 0.0:
                t = 0.0
            elif -2.0 <= x < 1.0:
                t *= 0.5
            else:
                t += 42.0
            result[i, j] = t

# 並列版の外側(マルチスレッドで呼び出すためのもの)
def mycumsum_numba_p(A, n_threads=8):
    n = A.shape[0]
    result = np.empty(A.shape, dtype=np.float64)
    cl = n // n_threads
    l = [cl * x for x in range(n_threads)] + [n]
    args = [(A, result, start, stop) for start, stop in zip(l, l[1:])]
    threads = [threading.Thread(target=mycumsum_numba_pi, args=a)
               for a in args]
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    return result
    
np.random.seed(42)
data = np.random.uniform(-10.0, 10.0, size=(5000, 50000))

# 検証
assert (mycumsum_numba(data) == mycumsum_numba_p(data)).all()

# 自作numba版
times = []
for i in range(10):
    t1 = time.time()
    res = mycumsum_numba(data)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 1.819163990020752

# 自作numba並列版
times = []
for i in range(10):
    t1 = time.time()
    res = mycumsum_numba_p(data)
    t2 = time.time()
    times.append(t2 - t1)

print(np.mean(times)) # => 0.2936941385269165


""" =>
1.8091564416885375
0.2936941385269165
"""

 ということで8スレッドで並列に動かして6倍くらい速くなりました。オーバーヘッドもあれば記憶領域の効率(キャッシュとメモリ帯域)的にも悪いのでこんなものです。それでもmultiprocessingでプロセス並列でやるより、原理的にははるかに効率的な方法です。

まとめ

 こんな感じで使えて速くて便利なnumbaを使わない手はないっ……!

 という訳でもないんですよね。必要になるシチュエーションは限られているので、数値計算寄りのことをしている一般的Pythonユーザでも一年に一回あるかどうかでしょう。私は4年間くらい実用的に使わずに過ごしてきましたが*6、そんなに困りませんでした。

 本格的にアルゴリズムの実装などに使うには何かと機能的に物足りないのは事実で、C/C++やCythonの方がたぶん向いています。また、日常的にインデックスループを書いて走らせる必要がある人はJuliaの方がいいかもしれません。

 じゃあなんで? という話になるのですが、こちらの記事の冒頭にある表を見てください。

qiita.com

 速度的には、CやJuliaと互角に見えます(CとJuliaが互角なのがそもそもびっくりですが、numbaはそれと互角、という意味)。つまりCの速度で動くギリギリPythonです。唯一無二の存在と言っていいでしょう。

 これは強い……

 どんな感じでやるのか知っておいて、頭の引き出しに入れておけば、必要になったときに取り出せます。雑に使うだけなら学習コスト低いのが良いところです*7。それで一年に一回の行き詰まりのときに、諦めるしかないか、突破できるかが決まると思います。

*1:ちなみにこの記事では実践的な数値計算はしません。それでも多少はなにかの手助けになるでしょう。

*2:動的メモリ確保してるから配列じゃないとか面倒くさいこと言わないでください、配列です。

*3:あるいはCythonかもしれないし、Pythonを見限ってJuliaに行ってしまうかもしれませんが

*4:あ、でもCとか古めの静的型付けは好きです。苦手な気がするのはオブジェクト指向以降のクラスと絡みついたモダンな静的型付けです。Javaのインターフェースとか必要なのはわかるけど好きになれる気がしません。新しめのpandasのリファレンスは丁寧にシグネチャにPythonの型ヒントを書いてくれていますが、あれとか正気の沙汰とは思えません。

*5:numpyのaxis=1相当

*6:存在は知ってるけどどう使うのかよくわからない、@jitしても速くならないし……とかそんなレベルでした。

*7:この記事読めばできるでしょという期待を込めて書いています。