Блог

Вчера я наткнулся на новое видео 3Blue1Brown о фрактале Ньютона, и я был действительно загипнотизирован его живым представлением фрактала. (Вот ссылка на видео для всех, кто заинтересован, это в 13:40: https://www.youtube.com/watch?v=-RdOwhmqP5s)

Я хотел попробовать сам и попытался закодировать его на python (я думаю, что он тоже использует python).

Я потратил несколько часов, пытаясь улучшить свою наивную реализацию, и дошел до того, что просто не знаю, как я мог бы сделать это быстрее.

Код выглядит так:

 import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from time import time


def print_fractal(state):
    fig = plt.figure(figsize=(8, 8))
    gs = GridSpec(1, 1)
    axs = [fig.add_subplot(gs[0, 0])]
    fig.tight_layout(pad=5)
    axs[0].matshow(state)
    axs[0].set_xticks([])
    axs[0].set_yticks([])
    plt.show()
    plt.close()


def get_function_value(z):
    return z**5   z**2 - z   1


def get_function_derivative_value(z):
    return 5*z**4   2*z - 1


def check_distance(state, roots):
    roots2 = np.zeros((roots.shape[0], state.shape[0], state.shape[1]), dtype=complex)
    for r in range(roots.shape[0]):
        roots2[r] = np.full((state.shape[0], state.shape[1]), roots[r])
    dist_2 = np.abs((roots2 - state))
    original_state = np.argmin(dist_2, axis=0)   1
    return original_state


def static():
    time_start = time()
    s = 4
    c = [0, 0]
    n = 800
    polynomial = [1, 0, 0, 1, -1, 1]
    roots = np.roots(polynomial)
    state = np.transpose((np.linspace(c[0] - s/2, c[0]   s/2, n)[:, None]   1j*np.linspace(c[1] - s/2, c[1]   s/2, n)))
    n_steps = 15
    time_setup = time()
    for _ in range(n_steps):
        state -= (get_function_value(state) / get_function_derivative_value(state))
    time_evolution = time()
    original_state = check_distance(state, roots)
    time_check = time()
    print_fractal(original_state)
    print("{0:<40}".format("Time to setup the initial configuration:"), "{:20.3f}".format(time_setup - time_start))
    print("{0:<40}".format("Time to evolve the state:"), "{:20.3f}".format(time_evolution - time_setup))
    print("{0:<40}".format("Time to check the closest roots:"), "{:20.3f}".format(time_check - time_evolution))
 

Средний результат выглядит следующим образом:

Время настройки начальной конфигурации: 0,004

Время для изменения состояния: 0,796

Время для проверки ближайших корней: 0.094

Ясно, что именно эволюционная часть является узким местом процесса. Это не «медленно», но я думаю, что этого недостаточно, чтобы сделать что-то живое, как в видео. Я уже сделал все, что мог, используя векторы numpy и избегая циклов, но, думаю, этого недостаточно. Какие еще уловки можно здесь применить?

Примечание: Я пытался использовать numpy.многочлены.Полиномиальный класс для оценки функции, но он был медленнее, чем эта версия.

1. Я получил улучшение (~на 40% быстрее), используя точность одного комплекса ( np.complex64 ).

 (...)
state = np.transpose((np.linspace(c[0] - s/2, c[0]   s/2, n)[:, None] 
                        1j*np.linspace(c[1] - s/2, c[1]   s/2, n)))
state = state.astype(np.complex64)
(...)
 

2. 3Blue1Brown добавил эту ссылку в описание: https://codepen.io/mherreshoff/pen/RwZPazd
   Вы можете посмотреть, как это было сделано там (примечание: автор этой ручки также использовал единственную точность).

 for _ in range(n_steps):
    state -= (get_function_value(state) / get_function_derivative_value(state))
 

Если у вас достаточно памяти, вы можете попытаться векторизовать этот цикл и сохранить шаги каждой итерации с помощью матричного вычисления.

Попробуйте pypy3, numba или cython.

Pypy3-это быстрая замена cpython. Для чистого кода на python это может быть значительно быстрее.

Режим nopython Numba может значительно ускорить математику в cpython. Хотя это не намного больше, чем математика.

Cython-это диалект python, который переносится на c и позволяет смешивать типы cpython и c. Чем больше типов c вы используете, тем лучше (обычно). Проверьте yhe cython -a opti9n.