动态生成任意峰数高斯混合模型时numpy.vectorize引发性能问题的优化咨询

以下是提问者的问题描述：

我正在用最大似然估计优化高斯混合模型，最初我用的是下面的模型：

def normal(x, mu, sigma):
    """
    Gaussian (normal) probability density function.
    
    Args:
        x (np.ndarray): Data points.
        mu (float): Mean of the distribution.
        sigma (float): Standard deviation of the distribution.
    
    Returns:
        np.ndarray: Probability density values.
    """
    return (1 / (np.sqrt(2 * np.pi) * sigma)) * np.exp(-0.5 * ((x - mu) / sigma) ** 2)

def model(x, a, mu1, s1, mu2, s2):
    return a*normal(x, mu1, s1) + (1-a)*normal(x, mu2, s2)

这个模型表现很好，不到一秒就能找到合适的拟合。现在我想动态生成支持任意峰数的类似函数：

def generate_gaussian_mix(n):
    def gaussian_mix(x, *params):
    
        if len(params) != 3 * n - 1:
            print(params)
            raise ValueError(f"Expected {3 * n - 1} parameters, but got {len(params)}.")

        params = np.asarray(params)
        mu = params[0::3]  # Means
        sigma = params[1::3]  # Standard deviations
        a = params[2::3]  # Weights
        a = np.hstack((a, 1 - np.sum(a)))

        return np.sum((a / (np.sqrt(2 * np.pi) * sigma))*np.exp(-0.5 * ((x - mu) / sigma) ** 2))

    return np.vectorize(gaussian_mix)

这个模型在相同峰数和数据量的情况下，在我的笔记本上要运行三分多钟。我可以采取哪些优化步骤来提升第二个函数的性能？有没有好的方法避免向量化？你有什么办法避免重复切片吗？

补充一下我的优化函数：

def neg_log_event_likelyhood(model, event, theta):
    x = -np.log(model(event, *theta))
    return x

def fit_distribution_anneal(model, events, bounds, data_range = None, **kwargs):
    def total_log_likelyhood(theta, model, events):
        return np.sum(neg_log_event_likelyhood(model, events, theta))

    if data_range is not None:
        events = np.copy(events)
        events = events[np.logical_and(events > data_range[0], events < data_range[1])]
    
    result = dual_annealing(total_log_likelyhood, bounds, args=(model, events), **kwargs)
    params = result.x

    return params

由于问题的非凸性，必须使用模拟退火而不是普通的最小化方法。

看起来你遇到的核心问题是np.vectorize带来的性能瓶颈——这个函数本质上是Python循环的包装，完全没有利用NumPy的向量化运算优势，这也是为什么你的动态生成版本比手写双高斯版本慢这么多的主要原因。下面我给你几个针对性的优化方向，应该能把性能拉回接近手写版本的水平：

1. 彻底抛弃np.vectorize，用全向量化实现高斯混合计算 #

np.vectorize是很多NumPy新手容易踩的坑，它的名字很有迷惑性，但实际是Python级别的循环，对于大量数据点来说效率极低。我们可以直接让函数处理整个x数组，利用NumPy的广播机制完成所有计算：

修改后的generate_gaussian_mix函数如下：

import numpy as np

def generate_gaussian_mix(n):
    def gaussian_mix(x, *params):
        x = np.asarray(x)
        # 验证参数数量是否正确
        if len(params) != 3 * n - 1:
            raise ValueError(f"Expected {3 * n - 1} parameters, but got {len(params)}.")

        params = np.asarray(params)
        # 提取参数：按你的逻辑，每3个参数为一组（mu, sigma, a），最后一组没有a
        mu = params[0::3]  # 所有均值，共n个
        sigma = params[1::3]  # 所有标准差，共n个
        a = params[2::3]  # 前n-1个权重
        # 补全最后一个权重，确保所有权重和为1
        a = np.hstack((a, 1 - np.sum(a)))
        
        # 利用广播机制：将x扩展为(N,1)，mu/sigma/a保持为(1,n)，实现逐点计算所有高斯成分
        x_reshaped = x[:, np.newaxis]
        # 计算每个x对应的所有高斯成分的值
        gaussian_components = (a / (np.sqrt(2 * np.pi) * sigma)) * np.exp(-0.5 * ((x_reshaped - mu) / sigma) ** 2)
        # 对每个x点，求和所有高斯成分
        return np.sum(gaussian_components, axis=1)
    
    return gaussian_mix

这个版本完全基于NumPy的向量化运算，没有任何Python循环，性能会和你手写的双高斯版本几乎一致。

2. 修正原函数的逻辑错误 #

你原来的gaussian_mix函数里，np.sum没有指定轴，会把所有x点和所有高斯成分的结果加总成一个标量，这明显是错误的——我们需要对每个x点单独求和所有高斯成分，优化后的代码用axis=1实现了这一点，确保返回的结果和输入x的形状一致。

3. 简化参数操作，避免冗余计算 #

• 移除参数错误时的print(params)，避免优化过程中产生大量冗余输出，只保留ValueError抛出即可。
• 可以提前将x转换为NumPy数组，确保后续广播运算的兼容性。

4. 优化数值稳定性，提升优化效率 #

在你的对数似然计算函数中，加入对模型输出的极小值截断，避免np.log(0)导致的Inf，提升优化器的收敛稳定性：

def neg_log_event_likelyhood(model, event, theta):
    model_output = model(event, *theta)
    # 截断极小值，避免log(0)引发数值错误
    model_output = np.clip(model_output, 1e-15, np.inf)
    return -np.log(model_output)

5. 可选：调整参数顺序适配习惯 #

如果你希望动态生成的模型参数顺序和手写双高斯模型一致（权重在前，均值、标准差在后），可以修改参数提取逻辑，这样不需要调整优化器的bounds配置：

def generate_gaussian_mix(n):
    def gaussian_mix(x, *params):
        x = np.asarray(x)
        if len(params) != 3 * n - 1:
            raise ValueError(f"Expected {3 * n - 1} parameters, but got {len(params)}.")

        params = np.asarray(params)
        # 前n-1个是权重，后面每2个一组对应一个峰的mu和sigma
        a = params[:n-1]
        a = np.hstack((a, 1 - np.sum(a)))
        mu = params[n-1::2]
        sigma = params[n::2]
        
        x_reshaped = x[:, np.newaxis]
        gaussian_components = (a / (np.sqrt(2 * np.pi) * sigma)) * np.exp(-0.5 * ((x_reshaped - mu) / sigma) ** 2)
        return np.sum(gaussian_components, axis=1)
    
    return gaussian_mix

这个版本的参数顺序和你手写的双高斯模型更接近，可能更符合你的使用习惯。

备注：内容来源于stack exchange，提问作者David K.