二項分布の確率質量関数をPythonで書き起こしてみる

13399 ワード
Python 統計学 二項分布 Python テキストリンク

数学や統計に興味があるのですが、数式を見ても記号がありすぎてどれがどれなのかわからなくなったり、すぐに忘れてしまったりするので、学んだことをプログラムに書いてみることにしました。

今回は全試行回数n回、成功確率pの二項分布に対して、以下の機能を持つクラスを作成します。

  • 成功回数xの確率質量を取得
  • xの下側累積確率を取得
  • xの上側累積確率を取得
  • グラフを描画

式の内容を覚えることを目的としてるので、パフォーマンスは考慮しません。

今回Pythonで再現する公式

成功確率がpのベルヌーイ試行（二項試行）を全部でn回行ったときの、

成功回数がxになる確率質量を求める関数

f(x,n,p) = {}_n \mathrm{ C }_xp^x(1-p)^{ n-x } \quad (0 \leqq n;\ 0 \leqq x \leqq n;\ 0 \leqq p\leqq 1)\\

成功回数xの下側累積確率を求める関数

P(n,x,p) = \sum_{i=0}^xf(i,n,p)

成功回数xの上側累積確率を求める関数

P(n,x,p) = \sum_{i=x}^nf(i,n,p)

Pythonで書いてみた


import math
import numpy as np
import matplotlib.pyplot as plt

class ToolBox:

  def combination(total, chosen):
      return math.factorial(total) / math.factorial(chosen)/ math.factorial(total-chosen)

class BinomialDistribution:
  """
  　全試行回数nと成功確率pの二項分布オブジェクトを生成する.

    Attributes
    ----------
    total : int
      全試行回数n.
    p_of_sccess : float
      一回の試行における成功確率p.
    p_of_fail : float
      一回の試行における失敗確率1-p.
  """

  def __init__(self, total, p_of_success):
    """
    Parameters
    ----------
    total : int
        全試行回数n.
    p_of_success : float
        一回の試行における成功確率p.
    """
    assert total > 0, "Condition error: 0 <= total"
    assert p_of_success >= 0 and p_of_success <= 1, "Condition error: 0 <= p_of_success <= 1"

    self.total = total
    self.p_of_sccess = p_of_success
    self.p_of_fail = 1 - p_of_success

  def get_probability_mass(self, success):
    """
    x回成功する確率質量を求める

    Parameters
    ----------
    success : int
        成功回数x.

    Returns
    -------
    probability_mass : float
        xの確率質量.
    """
    assert success >= 0 and success <= self.total, "Condition error: 0 <= sccuess <= total"

    fail = self.total - success
    combination = ToolBox.combination(self.total, success)
    probability_mass = combination * (self.p_of_sccess ** success) * (self.p_of_fail) ** fail
    return probability_mass

  def get_lower_culmitive_distribution(self, success):
    """
    xの下側累積確率(0〜x回の確率質量の和)を求める.

    Parameters
    ----------
    success : int
        成功回数x.

    Returns
    -------
    result : float
        xの下側累積確率.
    """
    result = 0
    for i in range (0, success + 1):
      result += binomial_distribution.get_probability_mass(i)
    return result

  def get_upper_culmitive_distribution(self, success):
    """
    xの上側累積確率(x〜全試行回数の確率質量の和、)を求める.

    Parameters
    ----------
    success : int
        成功回数x.

    Returns
    -------
    result : float
        xの上側累積確率.
    """
    result = 0
    for i in range (success, self.total + 1):
      result += binomial_distribution.get_probability_mass(i)
    return result

  def draw_graph(self):
    """
    グラフを描画してpng形式で保存する.
    """
    x = np.arange(0, self.total + 1, 1)
    y = []
    for i in range(0, self.total + 1):
      y.append(self.get_probability_mass(i))
    plt.plot(x, y)
    plt.savefig('graph.png')

実行


# 全試行回数10回、成功確率0.5の二項分布オブジェクトを作成する.
binomial_distribution = BinomialDistribution(10, 0.5)

# 成功回数2回の場合の確率質量を取得する.
print(binomial_distribution.get_probability_mass(2))

# 成功確率2回の場合の下側累積確率を取得する.
print(binomial_distribution.get_lower_culmitive_distribution(2))

# 成功確率2回の場合の上側累積確率を取得する.
print(binomial_distribution.get_upper_culmitive_distribution(2))

# この二項分布のグラフを描画する.
binomial_distribution.draw_graph()

結果

0.0439453125 # 成功回数2回の場合の確率質量
0.0546875 # 成功確率2回の場合の下側累積確率
0.9892578125 # 成功確率2回の場合の上側累積確率

答え合わせはこちらのサイトを使わせていただきました
https://keisan.casio.jp/exec/system/1161228843

グラフは同じディレクトリに保存されます。

2020年4月7日 修正内容

  • 「確率密度関数(Probability Density Function, PDF)」としていた箇所を「確率質量関数(Probability Mass Function, PMF)」に修正。確率変数x(今回の場合success)が離散的な値の場合後者を使うらしい。でも離散的な二項分布に対して「確率密度」を使っているサイトや本もあり、あまり厳密に区別されてないような気がする
  • Matplotlibを用いたグラフ描画機能を追加
  • 成功回数x(successs)をコンストラクタから切り離しメソッド呼び出し時に指定するよう修正
  • 数式の表記ミスを修正

参考にしたサイト

MathJaxの書き方
https://easy-copy-mathjax.nakaken88.com/

二項分布の計算をしてくれるサイト
https://keisan.casio.jp/exec/system/1161228843

確率密度関数と確率質量関数
https://data-science.gr.jp/theory/tbs_pdf_and_pmf.html