如何优雅地将N个一维数组组合的函数结果存入N维数组？

优化方案：分块处理+减少重复IO #

核心思路是利用my_function中相同a1~aI对应同一模型文件的特点，先按A组参数（a1~aI）分块，每块只加载一次模型，再批量处理该块下所有B组参数（b1~bJ）的组合——既减少了耗时的模型加载操作，又彻底避免了N层嵌套循环。

步骤1：参数分组与结果数组初始化

import numpy as np

# 替换为你的实际数组：A组对应a1~aI，B组对应b1~bJ
A_arrays = [array1, array2, ..., arrayI]
B_arrays = [arrayI_plus_1, ..., arrayN]
data = ...  # 传入my_function的固定data参数

# 构建结果数组的形状（各输入数组的长度依次对应N维的维度）
A_lengths = [len(arr) for arr in A_arrays]
B_lengths = [len(arr) for arr in B_arrays]
result_shape = A_lengths + B_lengths

# 初始化N维结果数组
my_results = np.empty(result_shape, dtype=np.float64)

步骤2：分块处理A组参数，批量计算B组结果

# 遍历A组的所有参数组合（直接生成多维度索引）
for a_indices in np.ndindex(*A_lengths):
    # 获取当前A组的参数值
    a_vals = tuple(arr[idx] for arr, idx in zip(A_arrays, a_indices))
    
    # 仅加载一次模型（关键优化：避免重复IO操作）
    filename = f"file_{'_'.join(map(str, a_vals))}"
    model = my_loadfunction(filename)
    
    # 生成B组参数的全组合网格（indexing='ij'保证与数组索引逻辑一致）
    B_meshes = np.meshgrid(*B_arrays, indexing='ij')
    
    # 方式1：若my_model_manipulator和my_probability_function支持批量/广播输入（最快）
    new_models = my_model_manipulator(model, *B_meshes)
    results_block = my_probability_function(data, new_models)
    my_results[a_indices] = results_block
    
    # 方式2：若上述函数不支持批量，用扁平化循环处理（仍远优于N层嵌套）
    # B_flat = [mesh.ravel() for mesh in B_meshes]
    # results_flat = []
    # for b_vals in zip(*B_flat):
    #     new_model = my_model_manipulator(model, *b_vals)
    #     results_flat.append(my_probability_function(data, new_model))
    # my_results[a_indices] = np.array(results_flat).reshape(B_lengths)

替代方案：全组合循环（无性能优化，但代码极简）

如果不需要优化模型加载次数，可直接生成所有参数组合的索引，用一维循环完成赋值：

import itertools

# 生成所有输入数组的索引范围与参数组合
index_ranges = [range(len(arr)) for arr in [array1, array2, ..., arrayN]]
param_combinations = itertools.product(array1, array2, ..., arrayN)

for indices, params in zip(itertools.product(*index_ranges), param_combinations):
    # 拆分参数为a组和b组，匹配my_function的参数顺序
    a_params = params[:I]
    b_params = params[I:]
    my_results[indices] = my_function(data, *a_params, *b_params)

方案优势 #

• 性能优先：分块处理将模型加载次数从len(array1)*len(array2)*...*len(arrayN)减少到len(array1)*...*len(arrayI)，大幅降低IO耗时；
• 代码简洁：最多2层循环，彻底替代N层嵌套，可读性与可维护性显著提升；
• 灵活性高：适配函数是否支持批量输入的两种场景，无需修改核心业务逻辑。

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