提交 7f1ab2ac authored 作者: Brandon T. Willard's avatar Brandon T. Willard 提交者: Brandon T. Willard

Allow single dimension slices in mgrid and ogrid

上级 ea65b923
...@@ -6729,10 +6729,12 @@ class TestNdGrid: ...@@ -6729,10 +6729,12 @@ class TestNdGrid:
def test_mgrid_numpy_equiv(self): def test_mgrid_numpy_equiv(self):
nmgrid = ( nmgrid = (
[np.mgrid[0:1:0.1]],
np.mgrid[0:1:0.1, 1:10:1.0, 10:100:10.0], np.mgrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
np.mgrid[0:2:1, 1:10:1, 10:100:10], np.mgrid[0:2:1, 1:10:1, 10:100:10],
) )
tmgrid = ( tmgrid = (
[mgrid[0:1:0.1]],
mgrid[0:1:0.1, 1:10:1.0, 10:100:10.0], mgrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
mgrid[0:2:1, 1:10:1, 10:100:10], mgrid[0:2:1, 1:10:1, 10:100:10],
) )
...@@ -6742,10 +6744,12 @@ class TestNdGrid: ...@@ -6742,10 +6744,12 @@ class TestNdGrid:
def test_ogrid_numpy_equiv(self): def test_ogrid_numpy_equiv(self):
nogrid = ( nogrid = (
[np.ogrid[0:1:0.1]],
np.ogrid[0:1:0.1, 1:10:1.0, 10:100:10.0], np.ogrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
np.ogrid[0:2:1, 1:10:1, 10:100:10], np.ogrid[0:2:1, 1:10:1, 10:100:10],
) )
togrid = ( togrid = (
[ogrid[0:1:0.1]],
ogrid[0:1:0.1, 1:10:1.0, 10:100:10.0], ogrid[0:1:0.1, 1:10:1.0, 10:100:10.0],
ogrid[0:2:1, 1:10:1, 10:100:10], ogrid[0:2:1, 1:10:1, 10:100:10],
) )
......
...@@ -5937,6 +5937,10 @@ class _nd_grid(object): ...@@ -5937,6 +5937,10 @@ class _nd_grid(object):
def __getitem__(self, *args): def __getitem__(self, *args):
if isinstance(args[0], slice):
sl = args[0]
return arange(sl.start or 0, sl.stop, sl.step or 1)
ndim = len(args[0]) ndim = len(args[0])
for sl in args[0]: for sl in args[0]:
if isinstance(sl.step, python_complex): if isinstance(sl.step, python_complex):
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论