1. はじめに

分析をするときも検証をするときも、とりあえず Jupyter Notebook で始めることが多いんじゃないかと思います。
でも Python は複雑な処理を書くと遅くて、それでぐるぐる検証を回すと待ち時間が増えてしまう...。
そんなとき Cython が Jupyter Notebook で簡単に使えることを知ったので、どのくらい高速化できるのか試してみました。

2. 環境

検証に使用した環境は以下の通りです。

Python: 3.12.8
Cython: 3.0.11

   OS: macOS Sonoma
チップ: Apple M1
メモリ: 16GB

3. Jupyter Notebook での Cython の使い方

使い方は簡単で、まず Jupyter Notebook で以下のマジックコマンドを実行します。

%load_ext Cython

そして以下のように Cython でコンパイルしたい関数に %%cython を添えるだけです。

%%cython
def function_to_be_compiled_by_cython():
    ...

4. 検証①：調和平均

まず調和平均の計算にかかる時間を計測してみます。

4-1. 実装１：標準ライブラリ

標準ライブラリに調和平均を計算する関数があります。
ただ引数に重みも与えられる実装になっているので、その分処理時間は長くなります。

from statistics import harmonic_mean

4-2. 実装２：素直に Python で実装

組み込み関数を用いて素直に実装します。

def harmonic_mean(data: list[float]) -> float:
    try:
        return len(data) / sum([1 / item for item in data])
    except ZeroDivisionError:
        return 0.0

4-3. 実装３：２をそのまま Cython でコンパイル

上の関数に %%cython 添えるだけです。

%%cython
def harmonic_mean(data: list[float]) -> float:
    try:
        return len(data) / sum([1 / item for item in data])
    except ZeroDivisionError:
        return 0.0

4-4. 実装４：２を Cython 用に最適化

まずデータ型を指定します。
また sum を使うよりも for で回した方が速くなりました。

%%cython
def harmonic_mean(data: list[float]) -> float:
    cdef int length = len(data)
    cdef double item, denominator = 0.0

    try:
        for item in data:
            denominator += 1 / item
    except ZeroDivisionError:
        return 0.0

    return length / denominator

4-5. 検証結果

1から99までの値が100個、10,000個、1,000,000個あるリストに対して、調和平均を1,000回計算したときの処理時間の中央値を取得します。

import random
import statistics
import time
from typing import Callable

size2data = {
    'small': random.choices(range(1, 100), k=100),
    'medium': random.choices(range(1, 100), k=10000),
    'large': random.choices(range(1, 100), k=1000000)
}

def timeit_harmonic_mean(func: Callable[list[float], float]) -> None:
    for size, data in size2data.items():
        elapsed_times = []
        for _ in range(1000):
            start_time = time.perf_counter()
            func(data)
            elapsed_time = time.perf_counter() - start_time
            elapsed_times.append(elapsed_time)

        median = statistics.median(elapsed_times)
        print(f'{size: >6}: {median}')

結果は以下のようになりました。表の単位は秒です。
標準ライブラリは思った以上に遅く、素直に Python で実装するだけでかなりの高速化ができました。
また、そのまま %%cython を付けるでも少し速くなりますが、Cython 用に書き換えるとそこからさらに2倍くらい速くなりました。

実装	100個	10,000個	1,000,000個
実装１：標準ライブラリ	`0.0001010`	`0.00690`	`0.698`
実装２：素直に Python で実装	`0.0000060`	`0.00062`	`0.067`
実装３：２をそのまま Cython でコンパイル	`0.0000050`	`0.00039`	`0.043`
実装４：２を Cython 用に最適化	`0.0000023`	`0.00020`	`0.020`

5. 検証②：編集距離

ループをたくさんするコードで試してみたかったので、編集距離を実装してみます。
もちろん python-Levenshtein には足元にも及びませんが、参考までに載せています。

5-1. 実装１：ライブラリ

参考用に python-Levenshtein を使用します。

import Levenshtein

def edit_distance(str1: str, str2: str) -> int:
    return Levenshtein.distance(str1, str2)

5-2. 実装２：素直に Python で実装

編集距離の定義に沿って、素直に実装します。

def edit_distance(str1: str, str2: str) -> int:
    len1 = len(str1)
    len2 = len(str2)
    dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]

    for i in range(len1 + 1):
        for j in range(len2 + 1):
            if i == 0:
                dp[i][j] = j
            elif j == 0:
                dp[i][j] = i
            elif str1[i-1] == str2[j-1]:
                dp[i][j] = dp[i-1][j-1]
            else:
                dp[i][j] = 1 + min(
                    dp[i-1][j],
                    dp[i][j-1],
                    dp[i-1][j-1]
                )

    return dp[len1][len2]

5-3. 実装３：２をそのまま Cython でコンパイル

%%cython
def edit_distance(str1: str, str2: str) -> int:
    len1 = len(str1)
    len2 = len(str2)
    dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]

    for i in range(len1 + 1):
        for j in range(len2 + 1):
            if i == 0:
                dp[i][j] = j
            elif j == 0:
                dp[i][j] = i
            elif str1[i-1] == str2[j-1]:
                dp[i][j] = dp[i-1][j-1]
            else:
                dp[i][j] = 1 + min(
                    dp[i-1][j],
                    dp[i][j-1],
                    dp[i-1][j-1]
                )

    return dp[len1][len2]

5-4. 実装４：２を Cython 用に最適化

まずデータ型を指定します。

%%cython
def edit_distance(str1: str, str2: str) -> int:
    cdef int i, j, len1 = len(str1), len2 = len(str2)
    cdef list dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]

    for i in range(len1 + 1):
        for j in range(len2 + 1):
            if i == 0:
                dp[i][j] = j
            elif j == 0:
                dp[i][j] = i
            elif str1[i-1] == str2[j-1]:
                dp[i][j] = dp[i-1][j-1]
            else:
                dp[i][j] = 1 + min(
                    dp[i-1][j],
                    dp[i][j-1],
                    dp[i-1][j-1]
                )

    return dp[len1][len2]

5-5. 実装５：２をもっと Cython 用に最適化

データ型の指定に加えて、文字列とリストを配列として処理するようにします。

%%cython
from libc.stdlib cimport free, malloc


def edit_distance(str1: str, str2: str) -> int:
    cdef int i, j, distance, len1 = len(str1), len2 = len(str2)
    cdef Py_UCS4* str1_utf32
    cdef Py_UCS4* str2_utf32
    cdef int** dp

    str1_utf32 = <Py_UCS4*>malloc(len1 * sizeof(Py_UCS4))
    for i in range(len1):
        str1_utf32[i] = str1[i]

    str2_utf32 = <Py_UCS4*>malloc(len2 * sizeof(Py_UCS4))
    for i in range(len2):
        str2_utf32[i] = str2[i]

    dp = <int**>malloc((len1 + 1) * sizeof(int*))
    for i in range(len1 + 1):
        dp[i] = <int*>malloc((len2 + 1) * sizeof(int))

    for i in range(len1 + 1):
        for j in range(len2 + 1):
            if i == 0:
                dp[i][j] = j
            elif j == 0:
                dp[i][j] = i
            elif str1_utf32[i-1] == str2_utf32[j-1]:
                dp[i][j] = dp[i-1][j-1]
            else:
                dp[i][j] = 1 + min(
                    dp[i-1][j],
                    dp[i][j-1],
                    dp[i-1][j-1]
                )

    distance = dp[len1][len2]

    for i in range(len1 + 1):
        free(dp[i])
    free(dp)

    return distance

5-6. 検証結果

長さが10、100、1,000の英語のアルファベットをランダムに並べた文字列に対して、編集距離を1,000回計算したときの処理時間の中央値を取得します。

import random
import statistics
import string
import time
from typing import Callable

size2data = {
    'small': [''.join(random.choices(string.ascii_letters, k=10)) for _ in range(2)],
    'medium': [''.join(random.choices(string.ascii_letters, k=100)) for _ in range(2)],
    'large': [''.join(random.choices(string.ascii_letters, k=1000)) for _ in range(2)]
}

def timeit_edit_distance(func: Callable[[str, str], int]) -> None:
    for size, data in size2data.items():
        elapsed_times = []
        for _ in range(1000):
            start_time = time.perf_counter()
            func(data[0], data[1])
            elapsed_time = time.perf_counter() - start_time
            elapsed_times.append(elapsed_time)

        median = statistics.median(elapsed_times)
        print(f'{size: >6}: {median}')

結果は以下のようになりました。表の単位は秒です。
実装２から実装４で、つまり %%cython をつけてデータ型を指定するだけで、およそ85%も処理速度が削減されました。
そこから文字列とリストを配列にして処理を行うと、追加で90%前後、全体で99%以上の処理速度が削減できました。
そして python-Levenshtein は文字列が長くなっても速い。すごい。

実装	長さ10	長さ100	長さ1,000
実装１：ライブラリ	`0.0000015`	`0.0000027`	`0.000042`
実装２：素直に Python で実装	`0.0000661`	`0.0047116`	`0.670170`
実装３：２をそのまま Cython でコンパイル	`0.0000265`	`0.0016129`	`0.348899`
実装４：２を Cython 用に最適化	`0.0000113`	`0.0004852`	`0.079378`
実装５：２をもっと Cython 用に最適化	`0.0000014`	`0.0000426`	`0.002885`

6. 終わりに

Cython で快適な Jupyter Notebook ライフを！