使用SymPy计算不同宽度均匀分布卷积过早陷入计算困境的问题

问题背景 #

我有多个独立的均匀分布，每个分布服从Uniform(0, width)（width取值各不相同），想要推导它们之和的概率分布。当所有宽度相同时，结果是Irwin-Hall分布，但宽度不同时没有现成的解析解。于是尝试用SymPy迭代计算卷积积分来求解，前2-3次卷积计算正常，但第三次之后直接陷入停滞，代码如下：

from time import time
from sympy.functions.elementary.piecewise import Piecewise
from sympy import Symbol, oo, simplify, integrate

t = Symbol('t')
widths = [7, 5, 14, 32, 3]
uniforms = [Piecewise((1 / width, (t <= width) & (t >= 0)), (0, True)) for width in widths]


def convolution(f, g, t, lower_limit=-oo, upper_limit=oo):
    tau = Symbol('tau')
    return integrate(f.subs(t, tau) * g.subs(t, t - tau),
                     (tau, lower_limit, upper_limit))


pdf = uniforms[0]
for uniform in uniforms[1:]:
    time_start = time()
    pdf = simplify(convolution(pdf, uniform, t))
    print(f"Convolution took {round(time() - time_start, 3)}s.")

为什么第三次卷积会卡住？ #

1. 分段函数爆炸式增长：每一次卷积后，分段函数的段数会大幅增加——n个均匀分布卷积后的分段数是O(n²)级别。第三次卷积后，分段数已经足够多，SymPy处理这些分段的符号积分和化简会变得异常缓慢。
2. 全局化简拖慢速度：每次卷积后调用的simplify()是全局化简函数，会对整个分段函数做全面的表达式优化，当分段数变多后，这个操作的成本会指数上升。
3. 无效积分区间浪费算力：用-oo和oo作为积分上下限，SymPy需要额外判断大量无效的零值区间，完全是做无用功。

优化方案 #

1. 限定积分的有效区间 #

均匀分布的非零区间是有限的，卷积时可以手动限定积分变量的有效范围，不用全局的正负无穷。比如当前PDF的非零区间是[a, b]，新增的均匀分布区间是[0, w]，那么积分变量tau的有效范围是max(0, t - w)到min(b, t)，直接缩小积分范围能大幅减少计算量。

2. 改用分段专用化简函数 #

把全局的simplify()换成piecewise_simplify()，这个函数专门针对分段函数做化简，效率比全局化简高得多；甚至可以暂时不化简，等所有卷积完成后再统一处理。

3. 手动管理分段区间 #

不用完全依赖SymPy的Piecewise对象，自己用列表存储每个分段的区间和表达式，手动处理区间合并，避免SymPy内部的额外开销。

优化后的代码示例 #

from time import time
from sympy.functions.elementary.piecewise import Piecewise
from sympy import Symbol, simplify, integrate, piecewise_simplify

t = Symbol('t')
widths = [7, 5, 14, 32, 3]

# 初始化单个均匀分布的分段表示：[(区间下限, 区间上限, 表达式)]
def init_uniform(w):
    return [(0, w, 1/w)]

# 自定义卷积逻辑，基于分段区间计算
def convolve_segments(segments_f, w_g):
    new_segments = []
    # 遍历原PDF的每个分段
    for a_f, b_f, expr_f in segments_f:
        # 计算卷积后可能的分界点，减少无效区间判断
        breakpoints = sorted({a_f, b_f, a_f + w_g, b_f + w_g})
        # 遍历每个分界点划分的t区间
        for i in range(len(breakpoints)-1):
            t_low, t_high = breakpoints[i], breakpoints[i+1]
            # 确定当前t区间内积分变量tau的有效范围
            tau_low = max(a_f, t - w_g)
            tau_high = min(b_f, t)
            # 计算积分
            tau = Symbol('tau')
            integrand = expr_f.subs(t, tau) * (1/w_g)
            integral = integrate(integrand, (tau, tau_low, tau_high))
            new_segments.append((t_low, t_high, integral))
    
    # 合并相邻且表达式相同的区间
    new_segments.sort(key=lambda x: x[0])
    merged_segments = []
    for seg in new_segments:
        if not merged_segments:
            merged_segments.append(seg)
        else:
            last_a, last_b, last_expr = merged_segments[-1]
            curr_a, curr_b, curr_expr = seg
            # 检查相邻区间是否可以合并
            if curr_a == last_b and simplify(last_expr - curr_expr) == 0:
                merged_segments[-1] = (last_a, curr_b, last_expr)
            else:
                merged_segments.append(seg)
    
    # 转换为SymPy的Piecewise对象
    piecewise_parts = []
    for a, b, expr in merged_segments:
        piecewise_parts.append((expr, (t >= a) & (t < b)))
    piecewise_parts.append((0, True))
    return Piecewise(*piecewise_parts), merged_segments

# 执行迭代卷积
start = time()
pdf, segments = convolve_segments(init_uniform(widths[0]), widths[1])
print(f"第一次卷积耗时：{round(time()-start, 3)}s")

start = time()
pdf, segments = convolve_segments(segments, widths[2])
pdf = piecewise_simplify(pdf)
print(f"第二次卷积耗时：{round(time()-start, 3)}s")

start = time()
pdf, segments = convolve_segments(segments, widths[3])
pdf = piecewise_simplify(pdf)
print(f"第三次卷积耗时：{round(time()-start, 3)}s")

start = time()
pdf, segments = convolve_segments(segments, widths[4])
pdf = piecewise_simplify(pdf)
print(f"第四次卷积耗时：{round(time()-start, 3)}s")

额外建议 #

如果不需要严格的符号解析解，只是要数值结果，直接用SciPy的convolve函数做数值卷积会高效得多。当分布数量超过4-5个时，符号解的复杂度会急剧上升，数值方法是更务实的选择。

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