Hi, A while back, I sent some changes to index_tricks.py that would allow mgrid and ogrid to mesh things other than slices. For example:
>>> mgrid[['a','b'],[float,int],:3] [array([[['a', 'a', 'a'], ['a', 'a', 'a']], [['b', 'b', 'b'], ['b', 'b', 'b']]], dtype='|S1'), array([[[<type 'float'>, <type 'float'>, <type 'float'>], [<type 'int'>, <type 'int'>, <type 'int'>]], [[<type 'float'>, <type 'float'>, <type 'float'>], [<type 'int'>, <type 'int'>, <type 'int'>]]], dtype=object), array([[[0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2]]])] At the time, there wasn't much follow-up, but I am hoping that there is still interest in this functionality, as I have gone ahead and finished the patch including documentation changes and updates to test_index_tricks.py. Attached is a patch set to the latest subversion of the numpy trunk. I don't think I am allowed to commit the changes myself - correct me if I am wrong. This functionality seems like a nice addition to me as it allows one to mesh things that are not uniformly spaced and potentially not even numbers. The changes don't affect functionality that existed previously except for one minor change - instead of returning a numpy array of arrays, mgrid/ogrid now return a list of arrays. However, this is unlikely to be a problem as the majority of users generally unpack the results of mgrid/ogrid so that each matrix can be used individually. Comments welcome. Cheers, David -- ********************************** David M. Kaplan Charge de Recherche 1 Institut de Recherche pour le Developpement Centre de Recherche Halieutique Mediterraneenne et Tropicale av. Jean Monnet B.P. 171 34203 Sete cedex France Phone: +33 (0)4 99 57 32 27 Fax: +33 (0)4 99 57 32 95 http://www.ur097.ird.fr/team/dkaplan/index.html **********************************
Index: numpy/lib/tests/test_index_tricks.py =================================================================== --- numpy/lib/tests/test_index_tricks.py (revision 5654) +++ numpy/lib/tests/test_index_tricks.py (working copy) @@ -24,15 +24,21 @@ def test_nd(self): c = mgrid[-1:1:10j,-2:2:10j] d = mgrid[-1:1:0.1,-2:2:0.2] - assert(c.shape == (2,10,10)) - assert(d.shape == (2,20,20)) + assert(array(c).shape == (2,10,10)) + assert(array(d).shape == (2,20,20)) assert_array_equal(c[0][0,:],-ones(10,'d')) assert_array_equal(c[1][:,0],-2*ones(10,'d')) assert_array_almost_equal(c[0][-1,:],ones(10,'d'),11) assert_array_almost_equal(c[1][:,-1],2*ones(10,'d'),11) - assert_array_almost_equal(d[0,1,:]-d[0,0,:], 0.1*ones(20,'d'),11) - assert_array_almost_equal(d[1,:,1]-d[1,:,0], 0.2*ones(20,'d'),11) + assert_array_almost_equal(d[0][1,:]-d[0][0,:], 0.1*ones(20,'d'),11) + assert_array_almost_equal(d[1][:,1]-d[1][:,0], 0.2*ones(20,'d'),11) + def test_listargs(self): + e = mgrid[ :2, ['a', 'b', 'c'], [1,5,50,500] ] + assert( array(e).shape == (3,2,3,4) ) + assert_array_equal( e[0][:,1,1].ravel(), r_[:2] ) + assert_array_equal( e[1][1,:,1].ravel(), array(['a','b','c']) ) + assert_array_equal( e[2][1,1,:].ravel(), array([1,5,50,500]) ) class TestConcatenator(TestCase): def test_1d(self): Index: numpy/lib/index_tricks.py =================================================================== --- numpy/lib/index_tricks.py (revision 5654) +++ numpy/lib/index_tricks.py (working copy) @@ -11,7 +11,7 @@ from numpy.core.numerictypes import find_common_type import math -import function_base +import function_base, shape_base import numpy.core.defmatrix as matrix makemat = matrix.matrix @@ -118,6 +118,10 @@ number of points to create between the start and stop values, where the stop value **is inclusive**. + One can also use lists or arrays as indexing arguments, in which case + these will be meshed out themselves instead of generating matrices from + the slice arguments. See examples below. + If instantiated with an argument of sparse=True, the mesh-grid is open (or not fleshed out) so that only one-dimension of each returned argument is greater than 1 @@ -126,19 +130,38 @@ -------- >>> mgrid = np.lib.index_tricks.nd_grid() >>> mgrid[0:5,0:5] - array([[[0, 0, 0, 0, 0], - [1, 1, 1, 1, 1], - [2, 2, 2, 2, 2], - [3, 3, 3, 3, 3], - [4, 4, 4, 4, 4]], - <BLANKLINE> - [[0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4]]]) + [array([[0, 0, 0, 0, 0], + [1, 1, 1, 1, 1], + [2, 2, 2, 2, 2], + [3, 3, 3, 3, 3], + [4, 4, 4, 4, 4]]), array([[0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4]])] >>> mgrid[-1:1:5j] array([-1. , -0.5, 0. , 0.5, 1. ]) + >>> mgrid[:2,[1,5,50],['a','b']] # Example using lists as indexing args + [array([[[0, 0], + [0, 0], + [0, 0]], + <BLANKLINE> + [[1, 1], + [1, 1], + [1, 1]]]), array([[[ 1, 1], + [ 5, 5], + [50, 50]], + <BLANKLINE> + [[ 1, 1], + [ 5, 5], + [50, 50]]]), array([[['a', 'b'], + ['a', 'b'], + ['a', 'b']], + <BLANKLINE> + [['a', 'b'], + ['a', 'b'], + ['a', 'b']]], + dtype='|S1')] >>> ogrid = np.lib.index_tricks.nd_grid(sparse=True) >>> ogrid[0:5,0:5] [array([[0], @@ -150,61 +173,70 @@ """ def __init__(self, sparse=False): self.sparse = sparse + def __getitem__(self,key): - try: - size = [] - typ = int - for k in range(len(key)): + typ = int + + # If not given a tuple, recall with a tuple + if not isinstance(key,tuple): + return self.__getitem__((key,))[0] + + # Now we know we have a tuple + + # Loop and get whether is float or not + for k in range(len(key)): + if isinstance(key[k],slice): # Only check type for slices + if isinstance(key[k].step, complex) or \ + isinstance(key[k].step, float) or \ + isinstance(key[k].start, float) or \ + isinstance(key[k].stop, float): + typ = float + + # Create basic arrays of appropriate sizes + nn = [[1]]*len(key) + for k in range(len(key)): + if isinstance(key[k],slice): step = key[k].step start = key[k].start + stop = key[k].stop if start is None: start=0 if step is None: step=1 if isinstance(step, complex): - size.append(int(abs(step))) - typ = float + # Will auto default to float + nn[k] = function_base.linspace( + start,stop,int(abs(step))) else: - size.append(math.ceil((key[k].stop - start)/(step*1.0))) - if isinstance(step, float) or \ - isinstance(start, float) or \ - isinstance(key[k].stop, float): - typ = float - if self.sparse: - nn = map(lambda x,t: _nx.arange(x, dtype=t), size, \ - (typ,)*len(size)) + nn[k] = _nx.arange(start,stop,step,dtype=typ) else: - nn = _nx.indices(size, typ) - for k in range(len(size)): - step = key[k].step - start = key[k].start - if start is None: start=0 - if step is None: step=1 - if isinstance(step, complex): - step = int(abs(step)) - if step != 1: - step = (key[k].stop - start)/float(step-1) - nn[k] = (nn[k]*step+start) - if self.sparse: - slobj = [_nx.newaxis]*len(size) - for k in range(len(size)): - slobj[k] = slice(None,None) - nn[k] = nn[k][slobj] - slobj[k] = _nx.newaxis - return nn - except (IndexError, TypeError): - step = key.step - stop = key.stop - start = key.start - if start is None: start = 0 - if isinstance(step, complex): - step = abs(step) - length = int(step) - if step != 1: - step = (key.stop-start)/float(step-1) - stop = key.stop+step - return _nx.arange(0, length,1, float)*step + start - else: - return _nx.arange(start, stop, step) + if function_base.isscalar(key[k]): + nn[k] = array([key[k]]) + else: + nn[k] = asarray(key[k]).ravel() + # Return quick if just one input + if len(key)==1: return nn + + # Get sizes for repeating if not sparse + if not self.sparse: + size = [ len(n) for n in nn ] + + # Reshape + slobj = [_nx.newaxis]*len(key) + for k in range(len(key)): + slobj[k] = slice(None,None) + nn[k] = nn[k][slobj] + slobj[k] = _nx.newaxis + + # If not sparse, tile + if not self.sparse: + for k in range(len(key)): + s = size[k] + size[k] = 1 + nn[k] = shape_base.tile( nn[k], size ) + size[k] = s + + return nn + def __getslice__(self,i,j): return _nx.arange(i,j)
_______________________________________________ Numpy-discussion mailing list Numpy-discussion@scipy.org http://projects.scipy.org/mailman/listinfo/numpy-discussion