Hi all,

Attached is a newer version of my patch that adds new mgrid / ogrid
functionality for working with arrays in addition to slices.  In fact, I
have attached two versions of the patch: index_tricks.patch, that is
just the last version of the patch I sent, and index_tricks.new.patch,
that has been modified so that it is backward compatible.  In the last
version, mgrid calls where all arguments are slices will return an
array, otherwise it returns a list as ogrid does.  This is the only
reasonable way to have the new functionality and maintain backwards
compatibility.  

My 2 cents - I personally think the version that always returns a list
will ultimately be more transparent and cause fewer problems than the
newer version.  In either case, the plan should be to eventually have it
always return a list as that is the only fully consistent option, the
question is just when that switch should be made and by who.  If it is
done at the next major release, someone else will have to remember to ax
the additional code and correct the documentation....

Other changes that would be nice: add a __call__ method, create an
instance called ndgrid for matlab compatibility, and have meshgrid be
reimplimented using an nd_grid instance.

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 5834)
+++ 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 5834)
+++ 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,14 +118,28 @@
     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
 
+    ***IMPORTANT NOTE*** Indexing an nd_grid instance with
+    sparse=False will currently return an array N+1 axis array if all
+    arguments are slices (i.e., something like -4:5:20j or :20:0.5)
+    and there are N arguments.  However, if any of the arguments is
+    not a slice (e.g., is an array or list), then the return is a list
+    of arrays.  This is to maintain backwards compatibility.  However,
+    this functionality will disappear during the next major release
+    (after today's date: 2008-09-19) so that returns will always be
+    lists in the future.
+
     Examples
     --------
     >>> mgrid = np.lib.index_tricks.nd_grid()
-    >>> mgrid[0:5,0:5]
+    >>> mgrid[0:5,0:5] # NOTE currently returns array, but will become a list
     array([[[0, 0, 0, 0, 0],
             [1, 1, 1, 1, 1],
             [2, 2, 2, 2, 2],
@@ -139,6 +153,27 @@
             [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 +185,81 @@
     """
     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
+
+            # IMPORTANT NOTE: Delete stuff below before next major
+            # release.  This stuff makes it so that return argument
+            # will be an array if all arguments are slices.  This is
+            # for backwards compatibility, but this functionality will
+            # be removed at next major release.  Note that
+            # documentation must also be fixed when this is done.
+            allslice = True
+            for k in range(len(key)):
+                allslice = allslice and isinstance(key[k],slice)
+            if allslice: nn = asarray(nn)
+
+        return nn
+
     def __getslice__(self,i,j):
         return _nx.arange(i,j)
 
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
[email protected]
http://projects.scipy.org/mailman/listinfo/numpy-discussion

Reply via email to