如何用单二维数组生成Numpy meshgrid？支持多子数组场景

你可以通过拆分输入数组为单个一维数组，再调用原生的np.meshgrid，最后将结果堆叠成一个多维数组来实现这个需求。这个方法可以支持任意数量的子数组（2、3、4个都没问题），完全匹配原生meshgrid的行为。

实现代码 #

import numpy as np

def custom_meshgrid(input_2d_arr):
    # 将二维数组按行拆分为多个一维数组
    individual_arrays = [input_2d_arr[i] for i in range(input_2d_arr.shape[0])]
    # 调用numpy原生meshgrid，默认使用xy索引模式（和你的示例一致）
    grid_tuple = np.meshgrid(*individual_arrays)
    # 将tuple中的多个grid数组堆叠为一个多维数组
    return np.array(grid_tuple)

测试示例 #

1. 2个子数组的情况（你的示例）

ab = np.array([[1,2],[3,4]])
result = custom_meshgrid(ab)
print(result)

输出结果：

array([[[1, 2],
        [1, 2]],

       [[3, 3],
        [4, 4]]])

完全和np.meshgrid(ab[0], ab[1])的结果堆叠后的效果一致。

2. 4个子数组的情况

abcd = np.array([[1,2],[3,4],[5,6],[7,8]])
result = custom_meshgrid(abcd)
print(result.shape)  # 输出 (4, 2, 2, 2, 2)

这里每个子数组长度为2，所以每个grid是4维数组（2×2×2×2），最终堆叠后得到5维数组，第一个维度对应输入的4个子数组。

原理说明 #

• 第一步拆分：把输入的二维数组的每一行单独提取出来，作为np.meshgrid的输入参数（用*解包可变参数）。
• 第二步调用meshgrid：原生meshgrid会根据输入的多个一维数组生成笛卡尔积对应的网格数组，返回一个tuple。
• 第三步堆叠：将tuple中的多个网格数组转换为numpy数组，自动在第一个维度上堆叠，得到你需要的格式。

这个方法完全复用了numpy原生的meshgrid逻辑，所以不管输入多少个子数组，行为都和直接调用np.meshgrid(arr1, arr2, ...)后堆叠的效果一致。

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