potisanのプログラミングメモ

プログラミング素人です。昔の自分を育ててくれたネット情報に少しでも貢献できるよう、情報を貯めていこうと思っています。Windows環境のC++やC#がメインです。

Python&SciPy&matplotlib 2つの正規分布の交点を求めて垂線を描画する。

SciPyとmatplotlibの学習で書いたコードです。実行すると2つの正規分布の解(根)を求め、グラフに交点から横軸への垂線を描画します。根の存在を前提としてコードを簡略化しています。ここではscipy.optimize.root_scalarを使用しますが、sympyを使う方法もあります。

import numpy as np
from scipy.stats import norm
from scipy.optimize import root_scalar
import matplotlib.pyplot as plt

#正規分布のPDF
f1 = lambda x: norm.pdf(x, -5, 5)
f2 = lambda x: norm.pdf(x, 8, 3)

#交点の計算
#0を始点としてSecant法(割線法)で解(根)を検索します。
f = lambda x: f1(x) - f2(x)
root_x = root_scalar(f, x0=0, x1=1, method="secant").root
root_y = f2(root_x)
print([root_x, root_y]) #[2.549188246445655, 0.025522845926062094]

#グラフのプロット
plt.rcParams["font.family"] = "MS Gothic" #日本語表示対応
fig, ax = plt.subplots()
x = np.linspace(-20, 20, 1000)
#2つの正規分布のPDFを描画する。
ax.plot(x, f1(x), color="black")
ax.plot(x, f2(x), color="black")
#交点の垂線を描画する。
#ax.axvlineは座標系が異なるのでax.vlinesを使います。
ax.vlines(root_x, 0, root_y, color="black")

ax.set_ylim(bottom=0) #y軸を0開始正方向にする。
ax.set_title("2つの正規分布のPDFと交点の垂線")
ax.set_xlabel("確率変数")
ax.set_ylabel("確率")
plt.show()