【高速化】Numbaの使い方や注意点について世界一わかりやすく解説

from numba import jit
@jit
def my_function(x):
    # ここでxについての処理を記述
    return x

import numpy as np
import time

# 配列の作成: 0から999999までの整数を持つNumpy配列を作成
arr = np.arange(1000000)
# arr中身
array([     0,      1,      2, ..., 999997, 999998, 999999])

# 演算開始
%%timeit
result = arr ** 2  # 配列の各要素を2乗
# result中身
array([           0,            1,            4, ..., 999994000009, 999996000004, 999998000001])

# 処理時間
1.23 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

from numba import jit  # Numbaのjitデコレータをインポート
import numpy as np
import time

# 配列の作成: 0から999999までの整数を持つNumpy配列を作成
arr = np.arange(1000000)
# arr中身
array([     0,      1,      2, ..., 999997, 999998, 999999])

# JITコンパイルを適用した関数を定義
@jit(nopython=True)
def square_array(arr):
    result = np.empty_like(arr)  # 入力配列と同じサイズの空の配列を作成
    for i in range(arr.size):    # 各要素に対して処理を行うループ
        result[i] = arr[i] ** 2  # 各要素を2乗して結果を格納
    return result                # 結果の配列を返す

# 演算開始
%%timeit
result = square_array(arr)  # JITコンパイルされた関数を実行
# result中身
array([           0,            1,            4, ..., 999994000009, 999996000004, 999998000001])

# 処理時間
1.1 ms ± 73.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

import numpy as np
import time
# 1000x1000のランダム行列を作成
A = np.random.rand(1000, 1000)
# Numpyで条件付き演算を実行
%%timeit
B = np.where(A > 0.5, A * 2, A / 2) # 要素が0.5より大きければ2倍、小さければ半分に
# 処理時間
7.11 ms ± 818 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

from numba import jit
import numpy as np
import time

# 1000x1000のランダム行列を作成
A = np.random.rand(1000, 1000)

# JITコンパイルを適用した条件付き演算関数を定義
@jit(nopython=True)
def custom_operation(A):
    n = A.shape[0]
    B = np.empty_like(A) # 結果を格納する配列を初期化
    for i in range(n):
        for j in range(n):
            if A[i, j] > 0.5:
                B[i, j] = A[i, j] * 2 # 条件を満たす場合、要素を2倍
            else:
                B[i, j] = A[i, j] / 2 # 条件を満たさない場合、要素を半分に
    return B

# Numbaで条件付き演算を実行
%%timeit
B = custom_operation(A) # JITコンパイルされた関数を実行

# 処理時間
1.48 ms ± 65.3 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

import numpy as np
import time

# 試行回数
num_samples = 1000000  # 100万回の試行

# モンテカルロ法でπを計算
%%timeit
x = np.random.rand(num_samples)  # ランダムなx座標を生成
y = np.random.rand(num_samples)  # ランダムなy座標を生成
inside_circle = np.sum(x**2 + y**2 <= 1.0)  # 原点からの距離が1以下の点を数える
pi_estimate = 4 * inside_circle / num_samples  # πの近似値を計算

# 計算結果
pi_estimate: 3.140288

# 処理時間
56.4 ms ± 10.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

from numba import jit
import numpy as np
import time

# 試行回数
num_samples = 1000000  # 100万回の試行

# JITコンパイルを適用したモンテカルロ法によるπの計算
@jit(nopython=True)
def monte_carlo_pi(num_samples):
    inside_circle = 0
    for _ in range(num_samples):
        x = np.random.rand()  # ランダムなx座標を生成
        y = np.random.rand()  # ランダムなy座標を生成
        if x**2 + y**2 <= 1.0:
            inside_circle += 1  # 原点からの距離が1以下の点をカウント
    return 4 * inside_circle / num_samples  # πの近似値を返す

# Numbaでモンテカルロ法を実行
%%timeit
pi_estimate = monte_carlo_pi(num_samples)  # JITコンパイルされた関数を実行
pi_estimate

# 計算結果
pi_estimate: 3.142132

# 処理時間
21.8 ms ± 5.69 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

from numba import jit
@jit(nopython=True)
def divide_numbers(x, y):
    try:
        return x / y
    except ZeroDivisionError:
        return "Cannot divide by zero"
# この関数を実行すると、Numbaがサポートしていない例外処理が含まれているためエラーが発生します
result = divide_numbers(10, 0)