如何用Python的rv_discrete实现Polya-Aeppli分布？结果异常排查

实现背景与代码 #

我尝试通过子类化scipy.stats.rv_discrete类、指定PMF（概率质量函数）的方式在Python中实现Polya-Aeppli（几何泊松）分布，代码如下：

from scipy.stats import rv_discrete
import numpy as np
import math

class PolyaAeppli(rv_discrete):

    def _pmf(self, k: np.ndarray, lambda_, theta_) -> np.ndarray:
        """
        Probability mass function for Polya-Aeppli distribution. Extension of Poisson distribution for
        modeling group arrivals.

        :param k: internal parameter of rv_discrete
        :param lambda_: arrival rate param of Poisson dist: [0, inf). The higher the lambda, the more arrivals
        :param theta_: probability param of Geometric dist: [0, 1]. The LOWER the theta, the more arrivals
        :return: probability values for each k
        """

        if isinstance(lambda_, np.ndarray or list):
            lambda_ = lambda_[0]. # 此处存在语法错误，多了一个英文句点
        if isinstance(theta_, np.ndarray or list):
            theta_ = theta_[0]
        k = np.asarray(k, dtype=int)
        res = np.zeros(len(k))
        for ix, k_ in enumerate(k):  # 暂未实现向量化
            if k_ == 0:
                res[ix] = np.exp(-lambda_)
            else:
                res[ix] = np.exp(-lambda_) * np.sum([(np.power(lambda_, i) / math.factorial(i)) *
                                                     ((1 - theta_) ** (k_ - i)) * (theta_ ** i) *
                                                     (comb(k_ - 1, i - 1)) for i in range(1, k_ + 1)])
        return res

问题现象 #

我用以下代码采样并绘制直方图，预期参数lambda_=0.9、theta_=0.9的d组平均到达数最大，但实际结果与预期不符：

PA = PolyaAeppli(name='polya_aeppli')

# 采样
resa = PA.rvs(lambda_=0.5, theta_=0.5, size=1000)
resb = PA.rvs(lambda_=0.5, theta_=0.9, size=1000)
resc = PA.rvs(lambda_=0.9, theta_=0.5, size=1000)
resd = PA.rvs(lambda_=0.9, theta_=0.9, size=1000)

# 绘制直方图
import matplotlib.pyplot as plt
plt.hist(resa, label='a', alpha=0.5)
plt.hist(resb, label='b', alpha=0.5)
plt.hist(resc, label='c', alpha=0.5)
plt.hist(resd, label='d', alpha=0.5)
plt.legend()
plt.show()

问题分析 #

1. 参数理解偏差 #

Polya-Aeppli分布本质是Poisson(λ)个独立几何分布变量的和，其中几何分布定义为「首次成功的次数」（取值为1,2,...），单变量期望为1/θ。因此Polya-Aeppli分布的整体期望公式为：
$$E[K] = \lambda \cdot \frac{1}{\theta}$$
按此公式计算各组期望：

• a组：0.5 / 0.5 = 1
• b组：0.5 / 0.9 ≈ 0.556
• c组：0.9 / 0.5 = 1.8
• d组：0.9 / 0.9 = 1

可见c组期望最大，d组与a组期望相同，这与你的预期完全相反。你代码中对theta的注释是正确的（theta越小，到达数越多），但实际预期时出现了逻辑偏差。

2. 代码中的错误 #

• 语法错误：原代码中lambda_ = lambda_[0].多了一个英文句点，会导致运行报错。
• 类型判断错误：isinstance(lambda_, np.ndarray or list)写法无效，应改为isinstance(lambda_, (np.ndarray, list))，否则只会判断是否为np.ndarray类型。
• 缺失依赖：comb函数未导入，需添加from math import comb（Python3.10+支持）或from scipy.special import comb。

修正后的代码示例 #

from scipy.stats import rv_discrete
import numpy as np
import math
from math import comb

class PolyaAeppli(rv_discrete):

    def _pmf(self, k: np.ndarray, lambda_, theta_) -> np.ndarray:
        """
        Probability mass function for Polya-Aeppli distribution. Extension of Poisson distribution for
        modeling group arrivals.

        :param k: internal parameter of rv_discrete
        :param lambda_: arrival rate param of Poisson dist: [0, inf). The higher the lambda, the more arrivals
        :param theta_: probability param of Geometric dist: [0, 1]. The LOWER the theta, the more arrivals
        :return: probability values for each k
        """

        # 修正类型判断逻辑
        if isinstance(lambda_, (np.ndarray, list)):
            lambda_ = lambda_[0]
        if isinstance(theta_, (np.ndarray, list)):
            theta_ = theta_[0]
        k = np.asarray(k, dtype=int)
        res = np.zeros(len(k))
        
        for ix, k_ in enumerate(k):
            if k_ == 0:
                res[ix] = math.exp(-lambda_)
            else:
                sum_term = 0.0
                for i in range(1, k_ + 1):
                    poisson_term = math.pow(lambda_, i) / math.factorial(i)
                    neg_binom_term = comb(k_ - 1, i - 1) * (theta_ ** i) * ((1 - theta_) ** (k_ - i))
                    sum_term += poisson_term * neg_binom_term
                res[ix] = math.exp(-lambda_) * sum_term
        return res

内容的提问来源于stack exchange，提问作者Nathan Vaartjes