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

Reply via email to