Jupyter Notebook と Cython で検証を高速に回したい

# 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` を添えるだけです。

```python
%%cython
def function_to_be_compiled_by_cython():
    ...
```

# 4. 検証①：調和平均

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

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

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

```python
from statistics import harmonic_mean
```

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

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

```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` 添えるだけです。

```python
%%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` で回した方が速くなりました。

```python
%%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回計算したときの処理時間の中央値を取得します。

```python
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` を使用します。

```python
import Levenshtein

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

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

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

```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 でコンパイル

```python
%%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 用に最適化

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

```python
%%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 用に最適化

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

```python
%%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回計算したときの処理時間の中央値を取得します。

```python
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 ライフを！