Matplotlib绘制多色虚线如何避免启发式计算虚线尺寸

问题描述 #

我想要绘制由不同颜色短划线交替组成的折线，目前的实现代码如下：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm

def plotmulticolor(x, y, cmap, axis, dashes=100):
    norm = BoundaryNorm(np.arange(cmap.N + 1) - 0.5, cmap.N)
    xp = []
    yp = []
    cp = []
    for i in range(len(x) - 1):
        xi = [x[i], x[i + 1]]
        yi = [y[i], y[i + 1]]
        d = np.sqrt(np.diff(ax.transLimits.transform(xi)) ** 2 + np.diff(ax.transLimits.transform(yi)) ** 2)
        n = int(d / np.sqrt(2) * dashes)
        xp += np.linspace(xi[0], xi[1], n).tolist()
        yp += np.linspace(yi[0], yi[1], n).tolist()
        cp += np.tile(np.arange(cmap.N), int(np.ceil(n / cmap.N)))[:n - 1].tolist()
    points = np.array([xp, yp]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    lc = LineCollection(segments, cmap=cmap, norm=norm)
    lc.set_array(np.array(cp))
    return ax.add_collection(lc)

fig, ax = plt.subplots()
ax.set_xlim([0, 10])
ax.set_ylim([0, 10])
plotmulticolor([0, 5, 10], [0, 1, 10], ListedColormap(["r", "b", "g"]), axis=ax)
plt.show()

上述代码可以得到预期效果，但存在缺陷：短划线数量是启发式计算的，需要提前设定坐标轴范围获取刻度比例才能计算虚线尺寸，有没有更巧妙的实现可以避免这种启发式计算逻辑？

解决方案 #

方案1：使用路径效果实现（最优，完全避免手动计算分段） #

直接利用matplotlib.patheffects的描边功能，通过叠加不同颜色、不同偏移的短划线实现交替效果，短划线尺寸直接以显示单位（磅pt）设置，和坐标轴范围、缩放完全无关，无需任何提前计算：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patheffects import Stroke

def plot_multicolor_dash(x, y, colors, dash_len=5, ax=None):
    if ax is None:
        ax = plt.gca()
    # 绘制基础线条
    line, = ax.plot(x, y, color='none')
    n_color = len(colors)
    effects = []
    # 为每个颜色生成对应偏移的短划线描边
    for i, color in enumerate(colors):
        offset = i * dash_len
        # 短划线格式：[显示长度, 空白长度]，总周期为颜色数×短划线长度
        dashes = [dash_len, (n_color - 1) * dash_len]
        effects.append(Stroke(offset=offset, foreground=color, linestyle=(0, dashes)))
    line.set_path_effects(effects)
    return line

# 调用示例
fig, ax = plt.subplots()
# 不需要提前设置轴范围，绘图后会自动适配
plot_multicolor_dash([0,5,10], [0,1,10], colors=["r", "b", "g"], dash_len=8)
plt.show()

这个方案的优势：

• 完全不需要手动拆分线段，无启发式计算逻辑
• 短划线长度固定为显示单位，缩放、平移坐标轴时不会变形
• 不需要提前设置坐标轴范围，会自动适配数据范围

方案2：LineCollection适配动态轴范围 #

如果你需要保留LineCollection的实现逻辑，可以给坐标轴绑定范围变化回调，当轴范围更新时自动重新计算分段，不需要提前固定轴范围：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm

def plotmulticolor(x, y, cmap, ax, dash_density=100):
    norm = BoundaryNorm(np.arange(cmap.N + 1) - 0.5, cmap.N)
    def update_segments(event=None):
        xp, yp, cp = [], [], []
        for i in range(len(x)-1):
            xi = [x[i], x[i+1]]
            yi = [y[i], y[i+1]]
            # 用当前轴的transform计算相对长度
            trans = ax.transLimits.transform
            dx, dy = np.diff(trans(xi))[0], np.diff(trans(yi))[0]
            d = np.sqrt(dx**2 + dy**2)
            n = int(d * dash_density) + 2
            xp_seg = np.linspace(xi[0], xi[1], n)
            yp_seg = np.linspace(yi[0], yi[1], n)
            xp.extend(xp_seg.tolist())
            yp.extend(yp_seg.tolist())
            cp.extend(np.tile(np.arange(cmap.N), int(np.ceil(n/cmap.N)))[:n-1].tolist())
        points = np.array([xp, yp]).T.reshape(-1,1,2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc.set_segments(segments)
        lc.set_array(np.array(cp))
        fig.canvas.draw_idle()
    # 初始化LineCollection
    lc = LineCollection([], cmap=cmap, norm=norm)
    ax.add_collection(lc)
    # 绑定轴范围变化回调
    ax.callbacks.connect('xlim_changed', update_segments)
    ax.callbacks.connect('ylim_changed', update_segments)
    # 首次触发计算
    update_segments()
    return lc

# 调用示例
fig, ax = plt.subplots()
# 不需要提前设置轴范围
plotmulticolor([0,5,10], [0,1,10], ListedColormap(["r","b","g"]), ax=ax)
ax.autoscale() # 自动适配数据范围
plt.show()

内容的提问来源于stack exchange，提问作者Tom de Geus