如何在Plotly中高效绘制带三角形的2D矢量场？

我已在Plotly的GitHub仓库提交了功能请求，希望官方添加这类矢量场绘制函数。

我编写了一段用Plotly绘制矢量场的代码，效果类似3D锥图，代码如下：

import numpy as np
import plotly.graph_objects as go
import plotly.colors as pc

def plot_vector_field(x,y,u,v):
    
    # Function to get colors for the triangles
    def get_colors(values, colorscale='Viridis'):
                        colorscale = pc.get_colorscale(colorscale)
                        unique_magnitudes = np.unique(values)
                        color_map = {val: pc.sample_colorscale(colorscale, val)[0] for val in unique_magnitudes}
                        return np.vectorize(color_map.get)(values)

    # Function to plot the triangles
    def plot_triangle(fig, position,direction,size,color):
                        x = [position[0]+direction[0]*size/2, 
                                position[0]-direction[0]*size/3 + direction[1]*size/4, 
                                position[0]-direction[0]*size/3 - direction[1]*size/4]
                        y = [position[1]+direction[1]*size/2, 
                                position[1]-direction[1]*size/3 - direction[0]*size/4, 
                                position[1]-direction[1]*size/3 + direction[0]*size/4]

                        fig.add_trace(go.Scatter(
                                                x=x,
                                                y=y,
                                                fill='toself',
                                                mode='lines', 
                                                line=dict(color='rgba(0,0,0,0)'),
                                                fillcolor=color,
                                                showlegend=False,
                                                name='' 
                                            ))
        
    # Calculate the magnitude of the vectors
    magnitude = np.sqrt(u**2 + v**2)
    magnitude_normalized = magnitude/np.max(magnitude)

    # Get the colors for the triangles
    colors = get_colors(magnitude_normalized, colorscale='viridis')

    # Get the direction of the vectors
    angle = np.arctan2(v, u)
    direction = np.array([np.cos(angle), np.sin(angle)]).T

    # Create the figure
    fig = go.Figure()

    # For each point in the grid, plot a triangle
    for i in range(len(x)):
        plot_triangle(fig, 
                    position=[x[i],y[i]], 
                    direction=direction[i], 
                    size=0.8*magnitude_normalized[i], 
                    color=colors[i])

    
    # Add a trace for the colorbar
    colorbar_trace = go.Scatter(
                        x=[None],
                        y=[None],
                        mode='markers',
                        marker=dict(
                            colorscale='Viridis',
                            cmin=0,
                            cmax=np.max(magnitude),
                            colorbar=dict(
                                title='' 
                            )
                        ),
                        showlegend=False
                    )
    fig.add_trace(colorbar_trace)

    return fig


x = np.linspace(-2, 2, 10)
y = np.linspace(-2, 2, 10)
X, Y = np.meshgrid(x, y)
u = -1 - X**2 + Y
v = 1 + X - Y**2

fig = plot_vector_field(X.flatten(), Y.flatten(), u.flatten(), v.flatten())
fig.show()

该代码生成的图如下：

问题说明 #

当增加网格点数时，函数性能会急剧下降，原因是每个三角形都被作为单独的trace绘制。尝试过将所有三角形放在单个trace中，但无法实现理想的填充效果和颜色设置，需要更优的实现方式。

优化方案 #

核心优化点是批量创建并添加trace，避免循环中逐个调用add_trace带来的性能损耗，同时优化颜色计算逻辑，提升向量化程度。

优化后的代码如下：

import numpy as np
import plotly.graph_objects as go
import plotly.colors as pc

def plot_vector_field(x, y, u, v):
    # 计算矢量大小与归一化值
    magnitude = np.sqrt(u**2 + v**2)
    magnitude_normalized = magnitude / np.max(magnitude)
    
    # 向量化生成颜色，替代原字典映射逻辑
    colorscale = pc.get_colorscale('Viridis')
    colors = np.array([pc.sample_colorscale(colorscale, val)[0] for val in magnitude_normalized])
    
    # 计算方向向量的x、y分量
    angle = np.arctan2(v, u)
    dir_x = np.cos(angle)
    dir_y = np.sin(angle)
    
    # 向量化生成所有三角形的三个顶点坐标
    size = 0.8 * magnitude_normalized
    tri_vertex_x = np.column_stack([
        x + dir_x * size / 2,
        x - dir_x * size / 3 + dir_y * size / 4,
        x - dir_x * size / 3 - dir_y * size / 4
    ])
    tri_vertex_y = np.column_stack([
        y + dir_y * size / 2,
        y - dir_y * size / 3 - dir_x * size / 4,
        y - dir_y * size / 3 + dir_x * size / 4
    ])
    
    # 批量创建所有三角形的trace列表
    traces = []
    for i in range(len(x)):
        trace = go.Scatter(
            x=tri_vertex_x[i],
            y=tri_vertex_y[i],
            fill='toself',
            mode='lines',
            line=dict(color='rgba(0,0,0,0)'),
            fillcolor=colors[i],
            showlegend=False,
            name=''
        )
        traces.append(trace)
    
    # 创建图并一次性添加所有trace
    fig = go.Figure()
    fig.add_traces(traces)
    
    # 添加颜色条
    fig.add_trace(go.Scatter(
        x=[None], y=[None],
        mode='markers',
        marker=dict(
            colorscale='Viridis',
            cmin=0,
            cmax=np.max(magnitude),
            colorbar=dict(title='矢量大小')
        ),
        showlegend=False
    ))
    
    return fig

# 测试：增加点数验证性能
x = np.linspace(-2, 2, 20)
y = np.linspace(-2, 2, 20)
X, Y = np.meshgrid(x, y)
u = -1 - X**2 + Y
v = 1 + X - Y**2

fig = plot_vector_field(X.flatten(), Y.flatten(), u.flatten(), v.flatten())
fig.show()

优化效果说明 #

1. 批量添加trace：将循环中逐个调用fig.add_trace改为先创建所有trace的列表，再用fig.add_traces一次性添加，大幅减少了Plotly内部的重复处理开销，点数越多性能提升越明显。
2. 向量化颜色计算：直接通过列表推导生成颜色数组，替代原有的字典映射逻辑，提升了颜色计算的效率。
3. 向量化顶点计算：用numpy向量化方式生成所有三角形的顶点坐标，避免了循环中重复的计算逻辑。

需要注意的是：由于Plotly的2D Scatter trace无法为单个trace内的不同填充区域设置独立颜色，因此每个三角形仍需作为单独的trace，但批量添加的方式已经能显著缓解性能问题。

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