如何为scipy.stats.rv_continuous编写含可变形状参数的子类？

要让rv_continuous把列表/数组形式的k当成单个分布的参数，需要调整参数校验逻辑，同时确保内部方法正确处理这个参数。以下是修改后的实现：

import numpy as np
from scipy.stats import rv_continuous, chi

class Cee(rv_continuous):
    """A random variable representing the maximum of multiple chi distributions.

    Each chi distribution can have a different ``df``.

    .. note::
        You probably do not need to use this class directly. Instead work with
        the instance :data:`cee`.

    Parameters
    ----------
    k : list of int, tuple of int, or int
        List/tuple of degrees of freedom of the chi-distributed variables to take the
        maximum of. Pass a single int for a single chi distribution.
    """

    def _argcheck(self, k):
        # 校验k是否为非负整数（或整数序列）
        k_arr = np.atleast_1d(k)
        return (np.issubdtype(k_arr.dtype, np.integer) and (k_arr > 0).all())

    def _cdf(self, x, k):
        k_arr = np.atleast_1d(k)
        # 利用广播机制批量计算chi的cdf，再取乘积
        cdfs = chi.cdf(x[..., np.newaxis], df=k_arr)
        return np.prod(cdfs, axis=-1)

    def _pdf(self, x, k):
        k_arr = np.atleast_1d(k)
        # 最大值分布的pdf公式：sum(fi(x) * product_{j≠i} Fj(x))
        cdfs = chi.cdf(x[..., np.newaxis], df=k_arr)
        pdfs = chi.pdf(x[..., np.newaxis], df=k_arr)
        
        prod_cdfs = np.prod(cdfs, axis=-1)
        terms = pdfs * (prod_cdfs[..., np.newaxis] / cdfs)
        # 处理cdfs为0的情况，避免除以0
        terms = np.where(cdfs == 0, 0, terms)
        return np.sum(terms, axis=-1)

# 使用这个实例
cee = Cee(name="cee", a=0)

关键修改点 #

• 重写_argcheck方法：明确告诉SciPy，k是单个形状参数，校验其所有元素为正整数（支持单个int或序列）。
• 统一参数处理逻辑：用np.atleast_1d(k)兼容单个整数和序列输入，结合numpy广播机制替代循环，保证计算结果对应单个分布。
• 手动实现_pdf：相比SciPy自动数值微分，手动实现最大值分布的pdf公式更高效准确。

使用示例 #

调用时传入序列参数，将得到单个结果：

# 传入列表作为k参数
print(cee(k=[1,2,3]).pdf(0.5))  # 输出单个数值
print(cee(k=[1,2,3]).cdf(0.5))  # 输出单个数值

# 支持单个int参数（等价于对应自由度的chi分布）
print(cee(k=2).pdf(0.5))

原代码问题原因 #

原代码中，SciPy的参数解析逻辑会把k=[1,2,3]拆成三个独立的形状参数，相当于创建了三个分别以k=1、k=2、k=3的Cee实例，因此调用pdf(0.5)会返回三个结果。通过重写_argcheck并统一参数处理，让SciPy识别k是单个参数而非多个独立参数。

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