Hi Robert 2009/2/27 Robert Kern <[email protected]>: >> a[ix_([2,3,6],range(a.shape[1]),[3,2])] >> >> If anyone knows a better way? > > One could probably make ix_() take slice objects, too, to generate the > correct arange() in the appropriate place.
I was wondering how one would implement this, since the ix_ function has no knowledge of the dimensions of "a". The best I could do was to allow a[ix_[[2,3,6], :3, [3, 2]] to work (see attached patch). Cheers Stéfan
From b2076197b94e3539e804685f332d2e6a769ee1ed Mon Sep 17 00:00:00 2001 From: Stefan van der Walt <[email protected]> Date: Tue, 3 Mar 2009 11:08:52 +0200 Subject: [PATCH] Allow fully specified ranges in ix_. --- numpy/lib/index_tricks.py | 77 +++++++++++++++++++++++++++------------------ 1 files changed, 46 insertions(+), 31 deletions(-) diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py index fcd3909..c7dfd29 100644 --- a/numpy/lib/index_tricks.py +++ b/numpy/lib/index_tricks.py @@ -71,37 +71,51 @@ def unravel_index(x,dims): # Indices become [x/dcb % a, x/dc % b, x/d % c, x/1 % d] return tuple(x/dim_prod % dims) -def ix_(*args): - """ Construct an open mesh from multiple sequences. - - This function takes n 1-d sequences and returns n outputs with n - dimensions each such that the shape is 1 in all but one dimension and - the dimension with the non-unit shape value cycles through all n - dimensions. - - Using ix_() one can quickly construct index arrays that will index - the cross product. - - a[ix_([1,3,7],[2,5,8])] returns the array - - a[1,2] a[1,5] a[1,8] - a[3,2] a[3,5] a[3,8] - a[7,2] a[7,5] a[7,8] - """ - out = [] - nd = len(args) - baseshape = [1]*nd - for k in range(nd): - new = _nx.asarray(args[k]) - if (new.ndim != 1): - raise ValueError, "Cross index must be 1 dimensional" - if issubclass(new.dtype.type, _nx.bool_): - new = new.nonzero()[0] - baseshape[k] = len(new) - new = new.reshape(tuple(baseshape)) - out.append(new) - baseshape[k] = 1 - return tuple(out) +class IndexMesh: + def __call__(self, *args): + return self.__getitem__(args) + + def __getitem__(self, *args): + """ Construct an open mesh from multiple sequences. + + This function takes n 1-d sequences and returns n outputs with n + dimensions each such that the shape is 1 in all but one dimension and + the dimension with the non-unit shape value cycles through all n + dimensions. + + Using ix_() one can quickly construct index arrays that will index + the cross product. + + a[ix_([1,3,7],[2,5,8])] returns the array + + a[1,2] a[1,5] a[1,8] + a[3,2] a[3,5] a[3,8] + a[7,2] a[7,5] a[7,8] + """ + args, = args # unpack args tuple + + out = [] + nd = len(args) + baseshape = [1]*nd + for k in range(nd): + if isinstance(args[k], slice): + s = args[k] + if s.stop is None: + raise ValueError("ix_ only accepts full range " + "specifications") + out.append(_nx.arange(s.start or 0, s.stop, s.step)) + continue + + new = _nx.asarray(args[k]) + if (new.ndim != 1): + raise ValueError, "Cross index must be 1 dimensional" + if issubclass(new.dtype.type, _nx.bool_): + new = new.nonzero()[0] + baseshape[k] = len(new) + new = new.reshape(tuple(baseshape)) + out.append(new) + baseshape[k] = 1 + return tuple(out) class nd_grid(object): """ @@ -210,6 +224,7 @@ class nd_grid(object): def __len__(self): return 0 +ix_ = IndexMesh() mgrid = nd_grid(sparse=False) ogrid = nd_grid(sparse=True) mgrid.__doc__ = None # set in numpy.add_newdocs -- 1.5.6.3
_______________________________________________ Numpy-discussion mailing list [email protected] http://projects.scipy.org/mailman/listinfo/numpy-discussion
