scipy.stats.rvs_continuous为何会修改传入参数的维度？

问题 #

使用scipy.stats.rvs_continuous自定义Phase-Type分布建模合并时间时，遇到参数维度被强制修改的问题：

• Phase-Type分布由矩阵和向量定义，但传入rvs_continuous后，3×3矩阵、3×1向量等参数会被统一转为一维数组
• 目前通过硬编码重塑和切片能让代码运行，但重复操作繁琐，且无法支持通用n×n矩阵场景
• 移除硬编码后报错，此时打印方法内的mat参数形状为(9, 0)，但外部检查是预期的(3,3)

原始代码 #

import numpy as np
from scipy import stats
from scipy import linalg


s1 = np.ones((3,1))


def g_mat(n, p, r):
    c = 1/(2*n*p)
    w = - c - 0.5*r 
    return np.matrix([[w, 0.5*r, 0], [0.5*r,  -1*r, 0.5*r], [0, 0.5*r, w]])
def state_vec(s,g):
    p = (4+3*s)/(4+s)
    p0 = 1/(1+(p**(-1*g/2)))
    return np.array([[p0**2, 2*(1-p0)*p0, (1-p0)**2]])
def ph_cdf(mat, state_vec1, x):
    mat = mat.reshape(3,3)
    return 1 - state_vec1.dot(linalg.expm(x*mat).dot(s1))

def ph_pdf(mat, state_vec1,x):
    mat = mat.reshape(3,3)
    s = -1*mat.dot(s1)
    return state_vec1.dot(linalg.expm(x*mat).dot(s))
def ex_t(mat, state_vec1):
    U = linalg.inv(-1*mat)
    u1= U.dot(s1)
    return state_vec1.dot(u1)[0,0]

class coal_time_dist(stats.rv_continuous):
    def _argcheck(self, mat, vec):
        return (type(mat) == np.ndarray) & (type(vec) == np.ndarray)
    def cdf(self,x, mat, vec):
        x = float(x[0])
        mat=mat.reshape(3,3)
        vec = vec[0:3,]
        return ph_cdf(mat, vec, x)[0]
    def pdf(self,x, mat, vec):
        x = float(x[0])
        mat=mat.reshape(3,3)
        vec = vec[0:3,]
        return ph_pdf(mat, vec, x)[0]

运行代码 #

mat = g_mat(1e4, 0.5, 1e-4)
vec = state_vec(0.5, 10)
dist = coal_time_dist(name="dist", badvalue=6)
dist.cdf(10, mat, vec)

移除硬编码后的报错信息 #

File "<stdin>", line 1, in <module>
File "/home/nzx3cc/.local/lib/python3.11/site-packages/scipy/stats/_distn_infrastructure.py", line 2120, in cdf
place(output, cond, self._cdf(*goodargs))
^^^^^^^^^^^^^^^^^^^^
File "<string>", line 38, in _cdf
File "<string>", line 19, in ph_cdf
File "/local/lib/python3.11/site-packages/scipy/linalg/_matfuncs.py", line 285, in expm
raise LinAlgError('The input array must be at least two-dimensional')
numpy.linalg.LinAlgError: The input array must be at least two-dimensional

解决方案 #

核心原因 #

scipy的rv_continuous会对所有传入的参数做扁平化处理，统一转为一维数组以支持向量化计算，导致矩阵的二维结构丢失。

修改方案 #

1. 新增n参数传递矩阵的阶数，用于动态重塑矩阵和切片向量
2. 优化_argcheck方法，验证矩阵扁平化后的长度是否为n²，向量长度是否为n
3. 在cdf、pdf方法中，利用n动态重塑矩阵、调整向量形状，去掉硬编码

修改后的完整代码 #

import numpy as np
from scipy import stats
from scipy import linalg


def g_mat(n, p, r):
    c = 1/(2*n*p)
    w = - c - 0.5*r 
    # 改用np.array而非np.matrix，避免后续维度问题
    if n == 3:
        return np.array([[w, 0.5*r, 0], [0.5*r,  -1*r, 0.5*r], [0, 0.5*r, w]])
    # 可扩展支持其他n值的矩阵生成逻辑
    else:
        raise ValueError("暂不支持该阶数的矩阵生成")

def state_vec(s, g, n):
    p = (4+3*s)/(4+s)
    p0 = 1/(1+(p**(-1*g/2)))
    if n == 3:
        return np.array([[p0**2, 2*(1-p0)*p0, (1-p0)**2]])
    # 可扩展支持其他n值的状态向量生成逻辑
    else:
        raise ValueError("暂不支持该阶数的状态向量生成")

def ph_cdf(mat, state_vec1, x):
    s1 = np.ones((mat.shape[0], 1))
    return 1 - state_vec1.dot(linalg.expm(x*mat).dot(s1))

def ph_pdf(mat, state_vec1, x):
    s1 = np.ones((mat.shape[0], 1))
    s = -1*mat.dot(s1)
    return state_vec1.dot(linalg.expm(x*mat).dot(s))

def ex_t(mat, state_vec1):
    s1 = np.ones((mat.shape[0], 1))
    U = linalg.inv(-1*mat)
    u1 = U.dot(s1)
    return state_vec1.dot(u1)[0,0]

class coal_time_dist(stats.rv_continuous):
    def _argcheck(self, n, mat_flat, vec_flat):
        # 检查参数合法性
        if not isinstance(n, int) or n <= 0:
            return False
        if len(mat_flat) != n*n or len(vec_flat) != n:
            return False
        return True
    
    def cdf(self, x, n, mat_flat, vec_flat):
        x = float(x[0])
        # 动态重塑矩阵和向量
        mat = mat_flat.reshape(n, n)
        vec = vec_flat.reshape(1, n)
        return ph_cdf(mat, vec, x)[0,0]
    
    def pdf(self, x, n, mat_flat, vec_flat):
        x = float(x[0])
        mat = mat_flat.reshape(n, n)
        vec = vec_flat.reshape(1, n)
        return ph_pdf(mat, vec, x)[0,0]

使用示例 #

n = 3
mat = g_mat(n, 1e4, 0.5)
vec = state_vec(0.5, 10, n)
# 将矩阵和向量扁平化后传入
dist = coal_time_dist(name="dist", badvalue=6)
dist.cdf(10, n, mat.flatten(), vec.flatten())

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