【matplotlib.pyplot.stem】stem plotで残差を表現したい：回帰直線からの残差【python】

前回の記事ではstemを使って残差を表現しようとしたが、yが定数の時しか対応できなかった。

chemstat.hatenablog.com

今回はこんな感じで回帰直線から残差を引くことを目指す。

スマートな答えは見つからず結局それぞれの点でstemプロットをするのがよさそうということになった。これだけでもわざわざ関数を定義してやるよりは大分楽なので良しとしよう。

import matplotlib.pyplot as plt
import numpy as np

#点のデータ
x = np.array([1,2,3,4,5])
y = np.array([2,6,6,9,6])
#回帰直線のデータ
y_reg = 1.1 * x + 2.5

fig, ax = plt.subplots()

#それぞれの点でstemプロットを引く
for n in np.arange(len(x)):

    ax.stem(
        [x[n]],
        [y[n]],
        bottom=y_reg[n],
        linefmt=('-','r')
    )

#回帰直線は別でプロット
ax.plot(x,y_reg)

plt.show()