Matplotlibで線を描画するにはmatplotlib.lines.Line2Dなどがあるが、一度に複数の線を引くことは(forループを使う以外には)できない。

こういう場合にはmatplotlib.collections.linecollectionを用いるといい、という話。

使用例1

$x \in \mathbb{R}$としたときの2点$(x, 0), (0, x)$を結ぶ線分を10本描く。

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as mc
import matplotlib.cm as cm
N = 10 # 線の数
x = np.linspace(-1, 1, N)
 
# 線のリスト : [(x0, y0), (x1, y1)]が1つの線
lines = [[(0, x[i]), (x[i], 0)] for i in range(N)]
colors = [cm.viridis(i/N) for i in range(N)] # 色のリスト
 
lc = mc.LineCollection(lines, colors=colors, linewidths=2)
 
# 描画
fig = plt.figure(figsize=(4,4))
ax = fig.add_subplot(aspect='1')
ax.add_collection(lc)
ax.autoscale()
plt.show()

使用例2

PCAの際に元データから第一主成分軸に下ろした垂線を描く(実際には元データと射影の点を結ぶ)。

# Generate toy data
np.random.seed(0)
n_samples = 50
x0 = np.random.uniform(-3, 3, size=n_samples)
x1 = x0 + np.random.randn(n_samples)
X = np.vstack((x0, x1)).T
# SVD
u, s, v = np.linalg.svd(X)
 
# 主成分上へのXの射影
expand_v = np.expand_dims(v[0], axis=1)
Z = X @ expand_v @ expand_v.T
 
# 元データから射影への直線のリスト
lines = [[(X[i, 0], X[i, 1]),
          (Z[i, 0], Z[i, 1])] for i in range(n_samples)]
lc = mc.LineCollection(lines, colors="k", linewidths=1)
# 主成分軸の描画用変数
xmax = ymax = 3
xmin = ymin = -3
x = np.arange(xmin, xmax, 1e-2)
 
# 描画
fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(aspect='1')
ax.scatter(X[:, 0], X[:, 1], s=5, marker='o', color='k') # 元データのplot
ax.scatter(Z[:, 0], Z[:, 1], s=5, marker='o', color='k') # 射影のplot
ax.add_collection(lc)
 
# 主成分軸
#ax.quiver(0, 0, v[:, 0], v[:, 1], zorder=11, width=0.01, scale=6, color='orange') # ベクトル表示
for i in range(2):
    ax.plot(x, v[i, 1] / v[i, 0] * x)
 
plt.hlines(0, xmin, xmax, linestyle="dashed") # x軸
plt.vlines(0, ymin, ymax, linestyle="dashed") # y軸
plt.xlim(xmin, xmax); plt.ylim(ymin, ymax)
plt.xlabel('x'); plt.ylabel('y')
plt.show()