如何提升牛顿分形Python计算算法的运行性能？

Great question—getting Newton fractals to render in real-time with Python is totally achievable, but you’ve hit the typical numpy performance ceiling for iterative complex operations. Let’s walk through actionable optimizations to speed up that bottleneck evolution loop:

1. Inline Function Calculations to Cut Overhead #

Your current get_function_value and get_function_derivative_value functions add unnecessary call overhead, even with numpy vectorization. Since your polynomial is fixed, inline these calculations directly in the loop to eliminate cumulative function call costs:

# Replace the evolution loop with this
for _ in range(n_steps):
    f = state**5 + state**2 - state + 1
    df = 5*state**4 + 2*state - 1
    state -= f / df

2. Use Numba JIT Compilation (Huge CPU Speed Boost) #

Numba’s Just-In-Time compiler can compile your numpy-based evolution loop into optimized machine code, often delivering 2-10x speedups on CPU. It’s a low-effort change with massive returns:

First install numba if you haven’t:

pip install numba

Then wrap your evolution logic in a JIT-decorated function:

from numba import njit, complex128

@njit(complex128[:,:](complex128[:,:], int))
def newton_evolution(state, n_steps):
    for _ in range(n_steps):
        f = state**5 + state**2 - state + 1
        df = 5*state**4 + 2*state - 1
        state -= f / df
    return state

# In your static() function, replace the loop with:
state = newton_evolution(state, n_steps)

For your 800x800 grid, this should cut the evolution time from ~0.8s to well under 0.2s after the first compilation run.

3. Leverage GPU Acceleration (Real-Time Territory) #

If you have access to an NVIDIA GPU, switching to a tensor framework like PyTorch or TensorFlow will push your rendering into real-time. GPUs excel at parallelizing element-wise complex operations:

Here’s a quick PyTorch adaptation of your code:

import torch
from time import time

def static_gpu():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    time_start = time()
    s = 4
    c = [0, 0]
    n = 800
    polynomial = [1, 0, 0, 1, -1, 1]
    roots = np.roots(polynomial)
    roots = torch.tensor(roots, device=device)
    
    # Create grid directly on GPU
    x = torch.linspace(c[0]-s/2, c[0]+s/2, n, device=device)
    y = torch.linspace(c[1]-s/2, c[1]+s/2, n, device=device)
    xx, yy = torch.meshgrid(x, y, indexing='ij')
    state = xx + 1j*yy
    
    n_steps = 15
    time_setup = time()
    for _ in range(n_steps):
        f = state**5 + state**2 - state + 1
        df = 5*state**4 + 2*state - 1
        state -= f / df
    time_evolution = time()
    
    # Check closest roots (vectorized on GPU)
    dist = torch.abs(state[None, :, :] - roots[:, None, None])
    original_state = torch.argmin(dist, dim=0) + 1
    
    time_check = time()
    # Convert back to numpy for plotting
    print_fractal(original_state.cpu().numpy())
    
    # Timing output
    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))

With a mid-range GPU, the evolution step will drop to ~0.01 seconds—easily fast enough for real-time panning/zooming like in the 3Blue1Brown video.

4. Reduce Precision to Float32 #

Numpy defaults to 64-bit floats (complex128), but most fractal calculations don’t need that level of precision. Switching to 32-bit (complex64) cuts memory usage in half and speeds up operations, as CPUs/GPUs handle 32-bit data faster:

# Initialize state with complex64
state = np.transpose((np.linspace(c[0]-s/2, c[0]+s/2, n, dtype=np.float32)[:, None] + 
                      1j*np.linspace(c[1]-s/2, c[1]+s/2, n, dtype=np.float32))).astype(np.complex64)
# Roots also cast to complex64
roots = np.roots(polynomial).astype(np.complex64)

This alone can give you a 20-30% speedup on CPU, and pairs well with other optimizations.

5. Minor Optimization: Simplify Root Checking #

Your check_distance function can be streamlined using numpy broadcasting to avoid the loop over roots:

def check_distance(state, roots):
    # Broadcast roots to match state shape automatically
    dist = np.abs(roots[:, None, None] - state)
    return np.argmin(dist, axis=0) + 1

This removes unnecessary array construction, though since this step isn’t your bottleneck, the speed gain will be small.

Start with Numba if you don’t have a GPU—it’s the easiest way to get big gains without rewriting much code. If you do have a GPU, the PyTorch/TensorFlow route will get you to real-time rendering. Combining these optimizations (Numba + float32, or GPU + float32) will push your performance even further.

内容的提问来源于stack exchange，提问作者Andrea Grimaldi