如何为numpy.mgrid实现可变维度网格生成的包装函数？

我希望通过指定变量（例如dim=2）来创建可变维度的meshgrid，而非像示例中那样手动修改表达式来设置二维网格。请问该如何实现对应的包装函数？问题源于我不熟悉mgrid所使用的语法（参考index_tricks）。示例代码如下：

import numpy as np

mgrid = np.mgrid[
                  -5:5:5j,
                  -5:5:5j,
                ]

我查阅了numpy.mgrid的官方文档，但未找到通过变量设置维度的相关说明。

核心思路是先根据维度变量生成重复的切片规则，再用解包操作把这些切片传给np.mgrid。

• 第一步：定义统一的切片规则，比如示例中的-5:5:5j，可以用slice()函数生成等价的切片对象
• 第二步：根据指定的维度dim，生成对应数量的切片列表
• 第三步：用*解包列表，传入np.mgrid的索引中

包装函数实现 #

import numpy as np

def variable_dim_mgrid(dim, start=-5, end=5, num_points=5j):
    # 生成dim个相同的切片规则
    slice_list = [slice(start, end, num_points)] * dim
    # 解包切片并传入mgrid
    return np.mgrid[*slice_list]

调用示例 #

# 创建2维网格
grid_2d = variable_dim_mgrid(dim=2)
print(grid_2d.shape)  # 输出 (2, 5, 5)

# 创建3维网格
grid_3d = variable_dim_mgrid(dim=3)
print(grid_3d.shape)  # 输出 (3, 5, 5, 5)

自定义维度规则 #

如果需要每个维度的范围或点数不同，可以修改函数逻辑，接受一个维度规则列表作为参数：

def custom_dim_mgrid(dim_rules):
    # dim_rules是一个列表，每个元素是(start, end, num_points)
    slice_list = [slice(s, e, n) for s, e, n in dim_rules]
    return np.mgrid[*slice_list]

# 调用示例：3个维度分别设置不同点数
grid_custom = custom_dim_mgrid([(-5,5,5j), (-3,3,3j), (0,10,10j)])
print(grid_custom.shape)  # 输出 (3, 5, 3, 10)

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