Python 将带索引的切片混合到NumPy中的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-D元组，提供了使用

np.mgrid

的切片+索引的组合。

使用of@hpaulj解决的任务使用

使用@hpaulj解决的任务

---

探索

meshgrid

，例如

np.meshgrid（np.arange（1,3），[1,2,3]，index='ij'）

@hpaulj谢谢！探索

meshgrid

，例如

np.meshgrid（np.arange（1,3），[1,2,3]，index='ij'）

@hpaulj谢谢！

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)