Оптимизация кода Python с Cython и Numba

# Пример простейшего .pyx файла (example.pyx)
def say_hello(name):
    print(f"Hello, {name}!")

# Python-версия
def calculate_sum(n):
    total = 0
    for i in range(n):
        total += i
    return total

from setuptools import setup
from Cython.Build import cythonize
 
setup(
    ext_modules = cythonize("sum.pyx")
)

import sum
result = sum.calculate_sum_typed(1000000)
print(result)

cpdef int calculate_sum_c_loop(int n):
    cdef int total = 0
    cdef int i = 0
    while i < n:  # C-style loop
        total += i
        i += 1
    return total

# profiling.pyx
cimport cython
 
@cython.profile(True)  # Включаем профилирование для этой функции
def intensive_calculation(int n):
    cdef int i, j
    cdef double result = 0
    for i in range(n):
        for j in range(n):
            result += i / (j + 1)
    return result

myproject/
├── myproject/
│   ├── __init__.py
│   ├── core.pyx
│   ├── utils.pyx
│   └── _dependencies.pxd
├── setup.py
└── tests/

from setuptools import setup, find_packages, Extension
from Cython.Build import cythonize
import numpy
 
# Определяем расширения
extensions = [
    Extension(
        "myproject.core",
        ["myproject/core.pyx"],
        include_dirs=[numpy.get_include()],  # Для работы с NumPy
        extra_compile_args=["-O3"],  # Включаем оптимизации компилятора
        language="c++"  # Используем C++ вместо C
    ),
    Extension("myproject.utils", ["myproject/utils.pyx"])
]
 
setup(
    name="myproject",
    version="0.1.0",
    packages=find_packages(),
    ext_modules=cythonize(
        extensions,
        compiler_directives={
            "language_level": 3,
            "boundscheck": False,
            "wraparound": False,
            "initializedcheck": False
        }
    )
)

pip install numba

from numba import jit
 
@jit
def calculate_sum(n):
    total = 0
    for i in range(n):
        total += i
    return total
 
# Первый вызов скомпилирует функцию
result = calculate_sum(1000000)

from numba import njit
 
@njit  # Эквивалентно @jit(nopython=True)
def fast_function(x, y):
    result = 0.0
    for i in range(len(x)):
        result += x[i] * y[i]
    return result

import numpy as np
from numba import jit, njit
import time
 
def python_function(arr):
    result = np.zeros_like(arr)
    for i in range(len(arr)):
        result[i] = np.sqrt(arr[i])
    return result
 
@jit
def jit_function(arr):
    result = np.zeros_like(arr)
    for i in range(len(arr)):
        result[i] = np.sqrt(arr[i])
    return result
 
@njit
def njit_function(arr):
    result = np.zeros_like(arr)
    for i in range(len(arr)):
        result[i] = np.sqrt(arr[i])
    return result
 
# Тестирование
arr = np.random.random(1000000)
 
start = time.time()
python_result = python_function(arr)
print(f"Python: {time.time() - start:.4f} сек")
 
start = time.time()
jit_result = jit_function(arr)
print(f"JIT: {time.time() - start:.4f} сек")
 
start = time.time()
njit_result = njit_function(arr)
print(f"NJIT: {time.time() - start:.4f} сек")

from numba import njit, prange
 
@njit(parallel=True)
def parallel_sum(arr):
    result = 0.0
    # prange вместо range для параллельного выполнения
    for i in prange(len(arr)):
        result += arr[i]
    return result

from numba import cuda
import numpy as np
 
@cuda.jit
def increment_by_one(arr):
    i = cuda.grid(1)  # Получаем индекс текущего потока
    if i < arr.size:  # Проверяем границы массива
        arr[i] += 1
        
# Создаем массив на CPU
data = np.ones(1000000, dtype=np.float32)
 
# Копируем данные на GPU
d_data = cuda.to_device(data)
 
# Запускаем ядро с 256 потоками в блоке и достаточным количеством блоков
threads_per_block = 256
blocks_per_grid = (data.size + threads_per_block - 1) // threads_per_block
increment_by_one[blocks_per_grid, threads_per_block](d_data)
 
# Копируем результат обратно на CPU
result = d_data.copy_to_host()

import cProfile
 
def slow_function(n):
    result = 0
    for i in range(n):
        result += i * i
    return result
 
cProfile.run('slow_function(1000000)')

import timeit
import numpy as np
from numba import njit
 
def python_version(x):
    return np.sin(x) * np.sqrt(x)
 
@njit
def numba_version(x):
    return np.sin(x) * np.sqrt(x)
 
x = np.random.random(1000000)
 
# Прогрев для компиляции
numba_version(x)
 
# Измерение производительности
py_time = timeit.timeit('python_version(x)', globals=globals(), number=100)
nb_time = timeit.timeit('numba_version(x)', globals=globals(), number=100)
 
print(f"Python: {py_time:.4f} сек")
print(f"Numba: {nb_time:.4f} сек")
print(f"Ускорение: {py_time / nb_time:.1f}x")

@njit(cache=True)  # Включить кэширование (по умолчанию)
def cached_function(x):
    return x * x

from numba import int32, float64
from numba.experimental import jitclass
 
# Определяем спецификацию полей класса
spec = [
    ('x', float64),
    ('y', float64),
    ('count', int32)
]
 
@jitclass(spec)
class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.count = 0
        
    def move(self, dx, dy):
        self.x += dx
        self.y += dy
        self.count += 1
        
    def distance(self, other):
        return ((self.x - other.x)[B]2 + (self.y - other.y)[/B]2)**0.5
 
# Использование
@njit
def compute_distances(points):
    n = len(points)
    result = np.zeros(n * (n-1) // 2)
    idx = 0
    for i in range(n):
        for j in range(i+1, n):
            result[idx] = points[i].distance(points[j])
            idx += 1
    return result

from numba import njit
 
# Эта функция компилируется в машинный код
@njit
def fast_calculation(data):
    result = np.zeros_like(data)
    for i in range(len(data)):
        result[i] = data[i] ** 2 + np.sqrt(data[i])
    return result
 
# А здесь используется обычный Python
def process_dataset(filename):
    # Сложная логика чтения файла и предобработки
    data = read_and_preprocess(filename)
    
    # Критичная к производительности часть делегируется Numba
    result = fast_calculation(data)
    
    # Постобработка и сохранение результатов
    return postprocess_and_save(result)

from numba import njit
import numpy as np
 
# Python-версия использует встроенные функции
def python_stats(data):
    return np.mean(data), np.median(data), np.std(data)
 
# Numba-версия реализует их вручную
@njit
def numba_stats(data):
    n = len(data)
    # Среднее
    mean = 0.0
    for i in range(n):
        mean += data[i]
    mean /= n
    
    # Стандартное отклонение
    var = 0.0
    for i in range(n):
        var += (data[i] - mean) ** 2
    std = np.sqrt(var / n)
    
    # Медиана (требует сортировки)
    sorted_data = np.sort(data)  # Numba поддерживает np.sort
    if n % 2 == 0:
        median = (sorted_data[n//2-1] + sorted_data[n//2]) / 2
    else:
        median = sorted_data[n//2]
    
    return mean, median, std

import dask.array as da
from numba import njit
 
@njit
def complex_function(x):
    # Сложные вычисления с каждым элементом
    return x**3 - 2*x**2 + 3*x - 4
 
# Создаем большой распределенный массив
x = da.random.random((10000, 10000), chunks=(1000, 1000))
 
# Применяем скомпилированную функцию ко всему массиву
result = x.map_blocks(lambda block: complex_function(block))
 
# Запускаем вычисления
result.compute()

from numba import njit, vectorize, float64
 
# Стандартная функция с циклом
@njit
def scalar_function(a, b):
    result = np.empty_like(a)
    for i in range(len(a)):
        result[i] = a[i] * b[i] + np.sin(a[i])
    return result
 
# Векторизованная версия
@vectorize([float64(float64, float64)])
def vector_function(a, b):
    return a * b + np.sin(a)
 
# Тестирование
a = np.random.random(1000000)
b = np.random.random(1000000)
 
# Сравниваем производительность

@njit(debug=True)
def problematic_function(x):
    # Ваш код
    return result

from numba import njit, int64, float64
 
@njit
def smart_sum(x):
    total = 0  # Тип определяется на этапе компиляции
    for i in range(len(x)):
        total += x[i]
    return total
 
# Автоматически создаются две версии функции
ints = np.array([1, 2, 3], dtype=np.int64)
floats = np.array([1.1, 2.2, 3.3], dtype=np.float64)
 
smart_sum(ints)   # Используется целочисленная версия
smart_sum(floats) # Используется версия с плавающей точкой

from numba import njit, types
from numba.core.registry import CPUDispatcher
 
# Создаем функцию
def add(a, b):
    return a + b
 
# Получаем прокси-объект диспетчера
dispatcher = njit(add)
 
# Компилируем для конкретных типов
compiled_int = dispatcher.compile((types.int64, types.int64))
compiled_float = dispatcher.compile((types.float64, types.float64))
 
# Используем скомпилированные функции
result1 = compiled_int(1, 2)
result2 = compiled_float(1.5, 2.5)

import numpy as np
import time
 
# Исходная Python/NumPy функция
def matrix_multiply_numpy(A, B):
    return A @ B
 
# Numba-версия
from numba import njit
 
@njit
def matrix_multiply_numba(A, B):
    return A @ B
 
# Cython-версия (предполагается, что уже скомпилирована)
# Код в файле matrix_cy.pyx:
"""
import numpy as np
cimport numpy as cnp
 
def matrix_multiply_cython(cnp.ndarray[double, ndim=2] A, cnp.ndarray[double, ndim=2] B):
    return A @ B
"""
 
# Тестирование
A = np.random.rand(1000, 1000)
B = np.random.rand(1000, 1000)
 
# Прогрев для Numba
_ = matrix_multiply_numba(A, B)
 
start = time.time()
C_numpy = matrix_multiply_numpy(A, B)
numpy_time = time.time() - start
 
start = time.time()
C_numba = matrix_multiply_numba(A, B)
numba_time = time.time() - start
 
# Для Cython предполагаем импорт из скомпилированного модуля
# from matrix_cy import matrix_multiply_cython
# start = time.time()
# C_cython = matrix_multiply_cython(A, B)
# cython_time = time.time() - start

# Вычисление расстояния Левенштейна между строками
def levenshtein_python(s1, s2):
    if len(s1) < len(s2):
        return levenshtein_python(s2, s1)
    
    if len(s2) == 0:
        return len(s1)
    
    previous_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row
    
    return previous_row[-1]
 
# Cython-версия будет аналогична, но с добавлением типизации
# Numba-версия
@njit
def levenshtein_numba(s1, s2):
    # Аналогичный код, но с типами, которые Numba может скомпилировать
    # (в данном случае мы должны конвертировать строки в массивы символов)
    a = np.array(list(s1))
    b = np.array(list(s2))
    # ...остальной код...

import pandas as pd
import numpy as np
from numba import njit
 
# Создадим синтетический набор данных
n = 10_000_000
df = pd.DataFrame({
    'id': np.random.randint(0, 1000, n),
    'value': np.random.random(n)
})
 
# Python/Pandas версия
def aggregate_pandas(df):
    return df.groupby('id')['value'].agg(['mean', 'std', 'count'])
 
# Numba-версия требует работы с массивами
@njit
def aggregate_numba(ids, values):
    # Предварительное вычисление количества уникальных id
    # и создание массивов для результатов
    # ...код здесь...
    return means, stds, counts
 
# Для Cython потребуется более сложный код с явной типизацией и,
# возможно, использованием C-структур данных

import numpy as np
from numba import njit
 
@njit(parallel=True)
def knn_predict_numba(train_data, train_labels, test_data, k):
    n_test = len(test_data)
    n_train = len(train_data)
    predictions = np.zeros(n_test, dtype=np.int64)
    
    for i in range(n_test):
        # Вычисляем расстояния от текущей тестовой точки до всех тренировочных
        distances = np.zeros(n_train)
        for j in range(n_train):
            dist = 0.0
            for d in range(train_data.shape[1]):
                diff = train_data[j, d] - test_data[i, d]
                dist += diff * diff
            distances[j] = np.sqrt(dist)
        
        # Находим индексы k ближайших точек
        nearest_indices = np.argsort(distances)[:k]
        
        # Голосование: определяем наиболее частый класс
        nearest_labels = train_labels[nearest_indices]
        unique_labels, counts = np.unique(nearest_labels, return_counts=True)
        predictions[i] = unique_labels[np.argmax(counts)]
    
    return predictions

import time
 
def benchmark(func, *args, repeats=5):
    times = []
    for _ in range(repeats):
        start = time.time()
        result = func(*args)
        end = time.time()
        times.append(end - start)
    
    return {
        'mean': np.mean(times),
        'min': np.min(times),
        'max': np.max(times),
        'std': np.std(times)
    }
 
# Использование
results_python = benchmark(python_function, data, repeats=10)
results_numba = benchmark(numba_function, data, repeats=10)
results_cython = benchmark(cython_function, data, repeats=10)

from flask import Flask, request, jsonify
import numpy as np
from numba import njit
 
app = Flask(__name__)
 
@njit(cache=True)  # Кэширование критично для веб-приложений!
def process_data(data):
    # Тяжелые вычисления
    result = np.zeros_like(data)
    for i in range(len(data)):
        result[i] = np.sqrt(data[i]) * np.log(1 + data[i])
    return result
 
@app.route('/analyze', methods=['POST'])
def analyze():
    data = np.array(request.json['data'], dtype=np.float64)
    result = process_data(data)
    return jsonify({'result': result.tolist()})

# При инициализации приложения
dummy_data = np.random.random(100)
_ = process_data(dummy_data)  # Прогреваем JIT-компилятор

@njit
def factorial_numba(n):
    if n <= 1:
        return 1
    return n * factorial_numba(n - 1)

from memory_profiler import profile
 
@profile
def memory_hungry_function(size):
    # Создаем большие массивы
    array1 = np.random.random((size, size))
    array2 = np.random.random((size, size))
    # Выполняем операции
    result = np.dot(array1, array2)
    return result

@cuda.jit
def cuda_kernel(array_in, array_out):
    # CUDA-оптимизированный код
    
def process_with_gpu(data):
    # Время копирования данных на GPU
    d_data = cuda.to_device(data)
    # Выполнение вычислений
    d_result = cuda.device_array_like(d_data)
    cuda_kernel[blocks, threads](d_data, d_result)
    # Время копирования результата обратно
    result = d_result.copy_to_host()