有没有办法提高这种分形计算算法的性能？

Question

昨天，我看到了有关牛顿分形的新 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.polynomials.Polynomial 类来评估函数，但它比这个版本慢。

Answer 1

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 您可以看看它是如何完成的（旁注：这支笔的作者也使用单精度）