如何在自定义Scipy COM-Poisson分布中缓存归一化常数？

要解决归一化常数重复计算的问题，最直接的方式是用缓存机制复用相同mu和nu组合下的归一化常数结果。以下是具体实现方案：

核心思路 #

COM-Poisson的归一化常数$Z(\mu, \nu) = \sum_{j=0}^{\infty} \frac{\mu^j}{(j!)\nu}$只和mu、nu有关，和k无关。我们可以用Python标准库的functools.lru_cache缓存这个值，避免每次计算PMF时重复求和。

修改后的完整代码 #

import numpy as np
from mpmath import nsum, fac, inf
from scipy import stats
from functools import lru_cache

# 缓存归一化常数，针对不同mu、nu组合只计算一次
@lru_cache(maxsize=None)
def compute_normalization_constant(mu, nu):
    # 处理mu=0的特殊情况：此时只有k=0的概率为1，其余为0
    if mu == 0:
        return 1.0
    return float(nsum(lambda j: (mu ** j) / (fac(j) ** nu), [0, inf]))

def pmf(k, mu, nu):
    z = compute_normalization_constant(mu, nu)
    # 处理k为非负整数的情况，非整数k概率为0
    if not isinstance(k, (int, np.integer)) or k < 0:
        return 0.0
    return float((mu ** k) / ((fac(k) ** nu) * z))

# 向量化PMF函数，适配scipy的数组输入
vectorized_pmf = np.vectorize(pmf)

class com_poisson_gen(stats.rv_discrete):
    def _argcheck(self, mu, nu):
        # 修正原代码的位运算问题，用逻辑与判断参数合法性
        return np.logical_and(mu >= 0, nu >= 0)

    def _pmf(self, k, mu, nu):
        return vectorized_pmf(k, mu, nu)

com_poisson = com_poisson_gen(name="com-poisson", longname='Conway-Maxwell-Poisson')
# 测试生成样本
samples = com_poisson.rvs(2, 4, size=100000)

关键改动说明 #

1. 缓存归一化常数：用@lru_cache装饰compute_normalization_constant函数，自动缓存不同mu、nu组合的计算结果，后续相同参数直接读取缓存。
2. 特殊情况处理：增加mu=0的分支，此时归一化常数为1，避免不必要的无穷求和。
3. 修正参数检查逻辑：原代码用&（位运算）判断参数合法性，改为np.logical_and适配数组输入的场景。
4. PMF逻辑优化：增加对非负整数k的判断，符合离散分布的定义。

性能提升效果 #

当生成大量样本（如10万条）时，缓存机制会避免重复执行上百次的无穷求和，计算速度会有数量级的提升。

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