numba

Numba - JITによる高性能Python

Numbaを使うと、Pythonコードを高速に実行できます。関数にデコレーターを付けて、Numbaにコンパイルさせます。数値計算の負荷が高いループやNumPy配列の操作を含むコードで特に効果を発揮します。

使用する場面

• NumPyの組み込みベクトル化では、特定のアルゴリズムに十分でない場合。
• 標準的なPythonでは遅くなる、複雑な入れ子のループがある場合。
• 配列の要素ごとに処理するカスタムの「ufunc」（ユニバーサル関数）を作成する必要がある場合。
• 高性能な物理シミュレーション（モンテカルロ、N体、グリッドベースのソルバー）。
• NVIDIA GPU（CUDA）で実行するコードを高速化する場合。
• マルチプロセッシングのオーバーヘッドなしに、すべてのCPUコアを活用する並列コードを作成する場合。

リファレンスドキュメント

公式ドキュメント：https://numba.pydata.org/numba-doc/latest/index.html
ユーザーガイド：https://numba.pydata.org/numba-doc/latest/user/index.html
検索パターン：@njit、@vectorize、prange、cuda.jit、numba.typed

基本原則

nopythonモード（@njit）

これはNumbaの「ゴールドスタンダード」です。このモードでは、NumbaはPython C-APIを使わずにコードをコンパイルするため、最高速度が得られます。コンパイルできない場合（サポートされていないPythonオブジェクトが原因の場合など）は、エラーが発生します。

ジャストインタイム（JIT）コンパイル

コンパイルは、関数を初めて呼び出したときに行われます。その後の呼び出しでは、機械語コードがキャッシュされて使用されます。

配列指向

NumbaはNumPy配列を扱うように設計されています。配列のメモリレイアウトを理解し、配列を処理する高度に最適化されたループを生成できます。

クイックリファレンス

インストール

pip install numba

標準的なインポート

import numpy as np
from numba import njit, prange, vectorize, guvectorize, cuda

基本パターン - ループの高速化

import numpy as np
from numba import njit

# 1. Apply the @njit decorator (alias for @jit(nopython=True))
@njit
def sum_array(arr):
    res = 0.0
    # Standard Python loop that would be slow is now fast as C
    for i in range(arr.shape[0]):
        res += arr[i]
    return res

# 2. Execute
data = np.random.random(1_000_000)
result = sum_array(data) # First call compiles, then runs

重要なルール

✅ 推奨事項

• @njitを優先する - 常にnopython=True（またはその別名である@njit）を使います。これにより、コードが実際に機械語速度で実行されます。
• NumPy配列を使う - NumbaはNumPy向けに最適化されています。JITコンパイルする関数内では、標準のPythonリストを避けます。
• 並列処理を有効にする - 自動マルチスレッド処理には、@njit(parallel=True)を使い、rangeの代わりにprangeを使います。
• コンパイル済みコードをキャッシュする - スクリプトを再起動するたびに再コンパイルしなくて済むよう、@njit(cache=True)を使います。
• ウォームアップする - コンパイルが行われるため、初回の呼び出しは遅くなります。実行時間を測定するベンチマークでは、必ず測定前に関数を一度実行します。
• 型の指定（任意） - シグネチャ（例：(float64[:],)）を指定すると初回の呼び出しを高速化できますが、通常、Numbaは型を適切に推論します。

❌ 禁止事項

• Pythonオブジェクトを使わない - 文字列、辞書、カスタムクラスは、nopythonモードでは遅いか、サポートされていません。必要に応じて、専用コンテナーにはnumba.typedを使います。
• 小さな関数をJITコンパイルしない - PythonからJITコンパイル済み関数を呼び出すオーバーヘッドによって、単純な処理では速度向上の効果が相殺されることがあります。
• サポートされていないライブラリを使わない - @njit関数内では、pandas、matplotlib、requestsを使用できません。
• グローバル状態を変更しない - 安定性を保つため、JITコンパイルする関数は可能な限り「純粋」にします。

アンチパターン（絶対に避ける）

from numba import njit
import pandas as pd

# ❌ BAD: Using Pandas inside @njit (Unsupported)
@njit
def bad_func(df):
    return df['col'].sum() # Will raise a LoweringError

# ✅ GOOD: Pass NumPy arrays instead
@njit
def good_func(arr):
    return arr.sum()

# ❌ BAD: Using @jit without nopython=True
from numba import jit
@jit
def slow_func(x): # This might fall back to "Object Mode" (slow)
    return x + 1

# ✅ GOOD: Always ensure nopython mode
@njit
def fast_func(x):
    return x + 1

# ❌ BAD: Manual loops in Python to call a JIT function
# for i in range(1000):
#     process_element(arr[i]) # Calling JIT overhead 1000 times

# ✅ GOOD: Move the loop INSIDE the @njit function
@njit
def process_all(arr):
    for i in range(arr.shape[0]):
        process_element(arr[i])

並列処理とベクトル化

自動マルチスレッド処理

from numba import njit, prange

@njit(parallel=True)
def parallel_sum(A):
    # Use prange for the loop that should be parallelized
    s = 0
    for i in prange(A.shape[0]):
        s += A[i]
    return s

高速なufuncの作成（@vectorize）

from numba import vectorize

# This creates a NumPy ufunc that supports broadcasting
@vectorize(['float64(float64, float64)'], target='parallel')
def fast_add(x, y):
    return x + y

# Now you can use it on massive arrays
res = fast_add(arr1, arr2)

構造体と型の扱い

配列以外のデータに対するnumba.typed

from numba.typed import List, Dict
from numba import njit

@njit
def use_typed_list():
    l = List()
    l.append(1.0)
    return l

GPUアクセラレーション（numba.cuda）

CUDAカーネルの記述

from numba import cuda

@cuda.jit
def my_kernel(io_array):
    # Calculate thread indices
    pos = cuda.grid(1)
    if pos < io_array.size:
        io_array[pos] *= 2

# Usage
data = np.ones(256)
threadsperblock = 32
blockspergrid = (data.size + (threadsperblock - 1)) // threadsperblock
my_kernel[blockspergrid, threadsperblock](data)

実践的なワークフロー

1. 高速なモンテカルロシミュレーション

import random

@njit(parallel=True)
def monte_carlo_pi(nsamples):
    acc = 0
    for i in prange(nsamples):
        x = random.random()
        y = random.random()
        if (x**2 + y**2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

2. カスタム画像フィルター（ステンシル）

from numba import njit

@njit
def apply_threshold(image, threshold):
    M, N = image.shape
    result = np.zeros_like(image)
    for i in range(M):
        for j in range(N):
            if image[i, j] > threshold:
                result[i, j] = 255
    return result

3. 物理グリッドの解法（ラプラス方程式）

@njit
def solve_laplace(u, niters):
    M, N = u.shape
    for n in range(niters):
        for i in range(1, M-1):
            for j in range(1, N-1):
                u[i, j] = 0.25 * (u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1])
    return u

パフォーマンスの最適化

inspect_types()メソッド

Numbaが高コストなPythonオブジェクトにフォールバックしているか、すべてをネイティブ型に最適化できているかを確認するために使います。

fast_func.inspect_types() # Prints color-coded annotated code

ループ内での配列割り当てを避ける

メモリ管理のオーバーヘッドを避けるため、@njit関数の外で配列を事前に割り当てるか、引数として渡します。

# ✅ GOOD:
@njit
def compute_into(out_arr, in_arr):
    for i in range(in_arr.shape[0]):
        out_arr[i] = in_arr[i] * 2

よくある落とし穴と解決策

「グローバル変数」の問題

Numba はコンパイル時点のグローバル変数の値をキャプチャします。

# ❌ Problem: Changing a global variable won't affect the jitted function
K = 10
@njit
def f(x): return x + K

K = 20
f(1) # Result is still 11!

# ✅ Solution: Pass constants as arguments

オブジェクトモードへのフォールバック

Numba が「オブジェクトモードが有効です」と表示する場合、コードの実行は遅くなります。

# ✅ Solution: Force nopython mode
@njit # If this throws error, fix the code instead of removing @njit

並列処理での乱数シード

parallel=True で np.random を使用する場合、各スレッドで独立した乱数ストリームを確保するために注意が必要です。Numba 内で標準の random.random() または np.random.random() を使用すると、スレッドセーフに動作し、スレッドごとのシード設定も自動的に処理されます。

ベストプラクティス

1. 常に @njit を使用する - nopython=True を指定せずに @jit を使用しないでください
2. 配列を事前に確保する - ホットループ内で配列を作成しないでください
3. 並列処理には prange を使用する - parallel=True と prange を使って自動マルチスレッド処理を有効にします
4. コンパイル済み関数をキャッシュする - cache=True を使って再コンパイルを避けます
5. 関数をウォームアップする - ベンチマークの前に、JIT コンパイルされた関数を一度呼び出します
6. NumPy 配列を渡す - JIT コンパイルされた関数を呼び出す前に、Python のリストを NumPy 配列に変換します
7. Python オブジェクトを避ける - コンテナーが必要な場合は、numba.typed.List と numba.typed.Dict を使用します
8. コンパイルモードを確認する - inspect_types() を使って nopython モードであることを確認します
9. 初回呼び出し時のオーバーヘッドに対処する - 初回の呼び出し時に関数がコンパイルされることを忘れないでください
10. 適切なシグネチャを使用する - 必須ではありませんが、初回コンパイルを高速化できます

Numba は、高性能コンピューティングの分野で Python が C++ や Fortran と競えるようにする橋渡し役です。ループから「Python 税」を取り除き、実行速度を犠牲にすることなく迅速なプロトタイピングを可能にします。