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

【问题标题】：Is there a way to improve the performance of this fractal calculation algorithm?有没有办法提高这种分形计算算法的性能？
【发布时间】：2021-12-06 03:37:28
【问题描述】：

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

【问题讨论】：

3B1B 在 github 上有 his library。您也可以尝试使用它，但不确定它是否有助于提高性能
您是否考虑过通过 PyOpenCL 进行 GPU 加速？您需要了解 C++，但 PyOpenCL 文档提供了一个很好的入门示例：documen.tician.de/pyopencl/#codecell0

标签： python numpy performance optimization fractals

【解决方案1】：

通过使用单一复数 (np.complex64) 精度，我得到了改进（快了约 40%）。

(...)
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)
(...)

3Blue1Brown 在描述中添加了此链接：https://codepen.io/mherreshoff/pen/RwZPazd 你可以看看那里是怎么做的（旁注：这支笔的作者也使用了单精度）

【讨论】：

【解决方案2】：

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

如果你有足够的内存，你可以尝试将这个循环向量化，并用矩阵计算存储每个迭代步骤。

【讨论】：

很抱歉，但我不确定我是否理解您的意思。由于接下来的每一步都依赖于前一步，那么该操作如何向量化？
你不应该对平面上的每个点都做这个循环吗？然后你可以对这些进行矢量化。

【解决方案3】：

尝试 pypy3、numba 或 cython。

Pypy3 是一个快速的 cpython 替代品。对于纯 python 代码，这可以大大加快。

Numba 的 nopython 模式可以显着加快 cpython 中的数学运算速度。不过，它只做数学。

Cython 是一种 python 方言，可以转换为 c，并允许您混合 cpython 和 c 类型。您使用的 c 类型越多越好（通常）。看看你的 cython -a opti9n。

【讨论】：