如何用Numpy向量化实现二元泊松分布PMF的求和部分？

我需要编写函数计算二元泊松分布的PMF，要求支持向量参数且无循环：

• theta0 为标量
• theta1、theta2 长度为 l
• x、y 长度均为 n
  输出形状为 (l, n, n)，其中切片 [j, :, :] 对应 theta1[j]、theta2[j] 的PMF矩阵。

已完成求和前的常数部分代码：

import numpy as np
from scipy.special import factorial, comb

def constant(theta1, theta2, theta0, x, y):
    exponential_part = np.exp(-(theta1 + theta2 + theta0)).reshape(-1, 1, 1)
    # 生成x的(n,n)矩阵：行重复x，列对应每个x_i
    x_mat = np.tile(x, (len(x), 1)).T
    # 生成y的(n,n)矩阵：列重复y，行对应每个y_j
    y_mat = np.tile(y, (len(y), 1))
    
    double_factorial = (np.power(np.array(theta1).reshape(-1, 1, 1), x_mat)/factorial(x_mat)) * \
                       (np.power(np.array(theta2).reshape(-1, 1, 1), y_mat)/factorial(y_mat))
    
    return exponential_part * double_factorial, x_mat, y_mat

向量化处理求和项 #

二元泊松PMF的完整公式为：
$$P(X=x,Y=y) = e^{-(\theta_1+\theta_2+\theta_0)} \frac{\theta_1^x}{x!} \frac{\theta_2^y}{y!} \sum_{k=0}^{\min(x,y)} \binom{x}{k}\binom{y}{k}k! \left( \frac{\theta_0}{\theta_1\theta_2} \right)^k$$

要向量化实现求和，步骤如下：

1. 计算每个(x,y)对的求和上限 min_xy = np.minimum(x_mat, y_mat)
2. 生成所有可能的k值（从0到最大的min(x,y)）
3. 利用广播生成所有k对应的项，再通过掩码筛选有效项后求和

完整实现代码：

def bivariate_poisson_pmf(theta1, theta2, theta0, x, y):
    # 获取常数部分和x、y的矩阵形式
    const_part, x_mat, y_mat = constant(theta1, theta2, theta0, x, y)
    l = len(theta1)
    n = len(x)
    
    # 计算每个(x,y)对的求和上限
    min_xy = np.minimum(x_mat, y_mat)
    max_k = np.max(min_xy)
    k = np.arange(0, max_k + 1)  # shape (K,)
    
    # 扩展各参数的维度以支持广播
    k_reshaped = k.reshape(1, 1, 1, -1)  # (1,1,1,K)
    x_reshaped = x_mat.reshape(1, n, n, 1)  # (1,n,n,1)
    y_reshaped = y_mat.reshape(1, n, n, 1)  # (1,n,n,1)
    
    # 计算ratio = theta0/(theta1*theta2)，扩展维度
    ratio = theta0 / (theta1 * theta2)
    ratio_reshaped = ratio.reshape(l, 1, 1, 1)  # (l,1,1,1)
    
    # 计算求和项的每一部分
    # 组合数C(x,k)，当k>x时为0
    comb_x_k = comb(x_reshaped, k_reshaped, exact=True, default=0)
    comb_y_k = comb(y_reshaped, k_reshaped, exact=True, default=0)
    k_fact = factorial(k_reshaped)
    ratio_pow_k = np.power(ratio_reshaped, k_reshaped)
    
    # 计算每个k对应的项
    sum_terms = comb_x_k * comb_y_k * k_fact * ratio_pow_k  # shape (l,n,n,K)
    
    # 生成掩码：仅保留k <= min_xy的项
    mask = k_reshaped <= min_xy.reshape(1, n, n, 1)
    sum_terms = sum_terms * mask.astype(float)
    
    # 对k维度求和
    sum_part = np.sum(sum_terms, axis=-1)  # shape (l,n,n)
    
    # 最终PMF
    pmf = const_part * sum_part
    return pmf

代码说明 #

• 广播机制：通过调整各数组的维度，让l、n、K三个维度自动扩展，完全避免循环
• 掩码筛选：确保每个(x,y)对只求和到min(x,y)，超出上限的项被置为0，不影响求和结果
• 高效计算：利用scipy的comb函数直接计算组合数，并通过default=0处理k>x或k>y的边界情况

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