将切片与索引混合到 NumPy 中的 mgrid 输入

Question

np.mgrid 接受切片元组，例如 np.mgrid[1:3, 4:8] 或 np.mgrid[np.s_[1:3, 4:8]]。

但是有没有办法在 mgrid 的元组参数中混合索引切片和数组？例如：

extended_mgrid(np.s_[1:3, 4:8] + (np.array([1,2,3]), np.array([7,8])))

应该给出与

相同的结果

np.mgrid[1:3, 4:8, 1:4, 7:9]

但一般来说，元组内的索引数组可能无法表示为切片。

解决此任务需要能够创建 N 维索引元组，提供切片 + 索引的混合使用 np.mgrid 就像在 this my answer for another question 中一样。

Answer 1

使用np.meshgrid 使用@hpaulj 的help 解决了任务。

Try it online!

import numpy as np

def extended_mgrid(i):
    res = np.meshgrid(*[(
            np.arange(e.start or 0, e.stop, e.step or 1)
            if type(e) is slice else e
        ) for e in {slice: (i,), np.ndarray: (i,), tuple: i}[type(i)]
    ], indexing = 'ij')
    return np.stack(res, 0) if type(i) is tuple else res[0]

# Tests

a = np.mgrid[1:3]
b = extended_mgrid(np.s_[1:3])
assert np.array_equal(a, b), (a, b)

a = np.mgrid[(np.s_[1:3],)]
b = extended_mgrid((np.s_[1:3],))
assert np.array_equal(a, b), (a, b)

a = np.array([[[1,1],[2,2]],[[3,4],[3,4]]])
b = extended_mgrid((np.array([1,2]), np.array([3,4])))
assert np.array_equal(a, b), (a, b)

a = np.mgrid[1:3, 4:8, 1:4, 7:9]
b = extended_mgrid(np.s_[1:3, 4:8] + (np.array([1,2,3]), np.array([7,8])))
assert np.array_equal(a, b), (a, b)