2022年12月29日木曜日

Pythonでラベル付きの数直線を描く

意外とサンプルが存在しないmatplotlibを使った数直線の描き方.

import numpy as np
import matplotlib.pyplot as plt

必要なライブラリはこれだけ.

b = np.array(["$B_1$", "$B_2$", "$B_3$", "$B_4$", "$B_5$", ])
x = np.array([-0.86, 0.90, -0.88, -0.27, 1.10])
y = np.zeros_like(x)
x_ano = x + 15
x_ano[0] += 10 # B1
x_ano[2] -= 10 # B3

B_1,B_2はデータ点につけるラベル.x_anoはラベルの水平位置で$B_1$と$B_3$の位置(b[0]とb[2])だけ重なってしまうので手動で左右にずらしている. こういうのを自動でやってくれるとありがたいけどそういう機能は見当たらないのであきらめた.

#数直線
figsize_px = np.array([1024, 102])
dpi = 200
figsize_inch = figsize_px / dpi
fig,ax=plt.subplots(figsize=figsize_inch,dpi=dpi)

ここは普通にplt.subplotsを読んで問題ない.画素数をしていしたかったのでこんな感じに.

# メモリと枠線を消す
ax.tick_params(labelbottom=True, bottom=False)
ax.tick_params(labelleft=False, left=False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['bottom'].set_visible(False)
plt.gca().spines['left'].set_visible(False)

数直線を描くためにmatplotlibデフォルトの枠線とメモリを全部消去する.

# 自前でメモリと水平軸を描画
xmin, xmax= -1.2,1.2
ys = 0.005
yl = 0.01
xs = np.linspace(int(xmin),int(xmax), int(xmax)-int(xmin)+1)
xl = np.linspace(int(xmin),int(xmax), (int(xmax)-int(xmin))*10+1)
plt.hlines(y=0, xmin=xmin, xmax=xmax, zorder=1) #横軸
plt.vlines(x=xs, ymin=-yl,ymax=yl, zorder=1) #目盛り大
plt.vlines(x=xl, ymin=-ys, ymax=ys, zorder=1) #目盛り小

ここが数直線のほぼすべて.水平線と垂直線の組み合わせ手で自前で軸を描いている.

# データ点と注釈を描画
plt.scatter(x,y,c='r') #散布図
for i in range(x.shape[0]):
    ax.annotate('%s' % b[i],
                 size=12,
                 xy=(x[i],y[i]),
                 xytext=(x_ano[i], y[i]+25.0),
                 textcoords='offset points',
                 arrowprops=dict(arrowstyle="->",connectionstyle="arc3,rad=.3")
                )
plt.savefig('plot_out.png', bbox_inches="tight", transparent=True)

データ点はscatterで注釈はannotateで描くだけ.