Matplotlib LineCollection 中使用线标记

在 Matplotlib 中，LineCollection 是一个非常有用的对象，它可以将多个线段组织成一组，并将它们作为单个物体绘制。然而，如果你想在 LineCollection 中添加一些线标记，就需要一些额外的工作了。

什么是 LineCollection

LineCollection 是一个用于绘制多条线段的容器对象。它可以接受一个 Nx2 或 Nx3 的数组，其中每行表示一个线段或线条，第一列和第二列表示线段的两个端点的 x 和 y 坐标，第三列表示线段的颜色、宽度等属性（可选）。

以下是一个使用 LineCollection 绘制两条线段的示例代码：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# 创建两条线段
lines = np.array([
    [[0, 0], [1, 1]],  # 第一条线段
    [[1, 0], [0, 1]]   # 第二条线段
])

# 创建一个 LineCollection 对象并添加到图表中
col = LineCollection(lines)
fig, ax = plt.subplots()
ax.add_collection(col)

# 设置坐标轴范围
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# 显示图表
plt.show()

添加线标记

如果你想在 LineCollection 中添加一些线标记，例如箭头、圆形或方形等，那么可以通过覆盖 LineCollection 对象的 draw 方法来实现。以下将演示如何向 LineCollection 中添加箭头和圆形线标记。

添加箭头

为了向 LineCollection 中添加箭头线标记，需要继承 LineCollection 类并重写其 draw 方法。下面是一个添加箭头的示例代码：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch

class ArrowLineCollection(LineCollection):
    def __init__(self, lines, arrows=None, *args, **kwargs):
        super().__init__(lines, *args, **kwargs)
        self.arrows = arrows if arrows is not None else []

    def add_arrow(self, arrow):
        self.arrows.append(arrow)

    def draw(self, renderer):
        super().draw(renderer)
        for arrow in self.arrows:
            arrow.draw(renderer)

# 创建两条线段和两个箭头
lines = np.array([
    [[0, 0], [1, 1]],
    [[1, 0], [0, 1]]
])
arrows = [FancyArrowPatch((0, 0), (1, 1), arrowstyle='->', mutation_scale=20),
          FancyArrowPatch((1, 0), (0, 1), arrowstyle='->', mutation_scale=20)]

# 创建一个 ArrowLineCollection 对象
col = ArrowLineCollection(lines, arrows=arrows, linewidth=2)

# 创建一个图形并添加 LineCollection 对象
fig, ax = plt.subplots()
ax.add_collection(col)

# 设置坐标轴范围
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# 显示图表
plt.show()

可以看到，两个箭头已经成功添加到了两条线段的末端。

添加圆形

与添加箭头不同，添加圆形线标记需要对每个线段创建一个单独的圆形，并将其绘制在相应的位置上。下面是一个添加圆形的示例代码：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.patches import Circle

class CircleLineCollection(LineCollection):
    def __init__(self, lines, circles=None, *args, **kwargs):
        super().__init__(lines, *args, **kwargs)
        self.circles = circles if circles is not None else []

    def add_circle(self, circle):
        self.circles.append(circle)

    def draw(self, renderer):
        super().draw(renderer)
        for circle in self.circles:
            circle.draw(renderer)

# 创建两条线段和两个圆形
lines = np.array([
    [[0, 0], [1, 1]],
    [[1, 0], [0, 1]]
])
circles = [Circle((1, 1), radius=0.05),
           Circle((0, 1), radius=0.05)]

# 创建一个 CircleLineCollection 对象
col = CircleLineCollection(lines, circles=circles, linewidth=2)

# 创建一个图形并添加 LineCollection 对象
fig, ax = plt.subplots()
ax.add_collection(col)

# 设置坐标轴范围
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# 显示图表
plt.show()

可以看到，两个圆形已经成功添加到了两条线段的末端。

总结

在 Matplotlib 中，使用 LineCollection 组织多条线段并将其作为单个物体绘制非常方便。但是，如果你想在 LineCollection 中添加一些线标记，例如箭头、圆形或方形等，就需要通过继承 LineCollection 类并重写其 draw 方法来实现。本文介绍了向 LineCollection 中添加箭头和圆形线标记的示例代码，希望能对 Matplotlib 的使用有所帮助。