Python中三维网格参数化函数计算的广播错误与解决方案

import numpy as np

def log_likelihood_function(A, nu_0, alpha, nu, x_i, sigma):
    """
    计算给定参数A, nu_0, alpha和数据nu, x_i的对数似然。
    A, nu_0, alpha 预期为标量。
    nu, x_i 预期为一维数组。
    """
    # 注意：如果nu和x_i已经是NumPy数组，这两行是多余的，但无害。
    # nu = np.array(nu)
    # x_i = np.array(x_i)

    # 模型预测值
    model_prediction = A * (nu / nu_0)**alpha * (1 + nu / nu_0)**(-4 * alpha)

    # 对数似然函数的第一项
    term1 = -len(nu) / 2 * np.log(2 * np.pi * sigma**2)

    # 对数似然函数的第二项（平方误差和）
    term2 = -1 / (2 * sigma**2) * np.sum((x_i - model_prediction)**2)

    return term1 + term2

def model(A, nu_0, alpha, nu):
    """
    根据参数A, nu_0, alpha和输入nu计算模型输出。
    """
    return A * (nu / nu_0)**alpha * (1 + nu / nu_0)**(-4 * alpha)

# 模拟数据
nu = np.linspace(0.05, 1.0, 500)
x_i = model(4.5, 1, 2/3, nu) + np.random.normal(0, 0.05, len(nu))
sigma = 0.05

# 定义参数范围并生成3D网格
A_range = np.arange(0.0, 10.0, 0.1)    # 100个值
nu_0_range = np.arange(0.0, 5.0, 0.1) # 50个值
alpha_range = np.arange(0.0, 5.0, 0.1) # 50个值

# 使用np.meshgrid生成三维参数网格
AA, NU_0_GRID, Alpha_GRID = np.meshgrid(A_range, nu_0_range, alpha_range, indexing="ij")

print(f"AA 形状: {AA.shape}")         # (100, 50, 50)
print(f"NU_0_GRID 形状: {NU_0_GRID.shape}") # (100, 50, 50)
print(f"Alpha_GRID 形状: {Alpha_GRID.shape}") # (100, 50, 50)
print(f"nu 形状: {nu.shape}")         # (500,)
print(f"x_i 形状: {x_i.shape}")       # (500,)

# 尝试直接计算，这将导致ValueError
# L = log_likelihood_function(AA, NU_0_GRID, Alpha_GRID, nu, x_i, sigma)

# 1. 使用 np.vectorize 封装函数
# excluded 参数指定哪些参数不参与向量化，即它们将作为整体传递给每次内部调用。
# 在本例中，nu, x_i, sigma 对于每次似然计算都是固定不变的完整数组或标量。
vectorized_log_likelihood = np.vectorize(log_likelihood_function, 
                                         excluded=['nu', 'x_i', 'sigma'])

# 2. 调用向量化后的函数
# AA, NU_0_GRID, Alpha_GRID 是网格数组，它们将被逐点迭代。
# nu, x_i, sigma 则作为固定参数传递给每次内部调用。
L_vectorized = vectorized_log_likelihood(AA, NU_0_GRID, Alpha_GRID, nu, x_i, sigma)

print("\n--- 使用 np.vectorize ---")
print(f"计算结果 L_vectorized 形状: {L_vectorized.shape}") # 预期: (100, 50, 50)

# 1. 初始化一个数组来存储结果，形状与网格数组相同
L_looped = np.zeros_like(AA)

# 2. 获取网格的维度
dim_A, dim_nu0, dim_alpha = AA.shape

print("\n--- 使用显式循环 ---")
print("开始循环计算...")

# 3. 遍历网格中的每一个点
for i in range(dim_A):
    for j in range(dim_nu0):
        for k in range(dim_alpha):
            # 从网格数组中提取当前点的标量参数值
            current_A = AA[i, j, k]
            current_nu_0 = NU_0_GRID[i, j, k]
            current_alpha = Alpha_GRID[i, j, k]

            # 使用这些标量参数和固定数据计算似然值
            L_looped[i, j, k] = log_likelihood_function(current_A, current_nu_0, current_alpha, nu, x_i, sigma)

print("循环计算完成。")
print(f"计算结果 L_looped 形状: {L_looped.shape}") # 预期: (100, 50, 50)

# 验证两种方法的结果是否一致
# print(f"两种方法结果是否接近: {np.allclose(L_vectorized, L_looped)}")