Fix uses of nonzero and add flatnonzero

7 files changed, 20 insertions, 11 deletions

diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
index 98291f411..5c21edd3b 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -230,8 +230,7 @@ def ravel(m,order='C'):
     return a.ravel(order)
 
 def nonzero(a):
-    """nonzero(a) returns the indices of the elements of a which are not zero,
-    a must be 1d
+    """nonzero(a) returns the indices of the elements of a which are not zero
     """
     try:
         nonzero = a.nonzero
@@ -241,7 +240,6 @@ def nonzero(a):
         res = nonzero()
     return res
     
-
 def shape(a):
     """shape(a) returns the shape of a (as a function call which
        also works on nested sequences).
diff --git a/numpy/core/ma.py b/numpy/core/ma.py
index 5cf4388bf..c1af7ae09 100644
--- a/numpy/core/ma.py
+++ b/numpy/core/ma.py
@@ -445,9 +445,8 @@ fabs = masked_unary_operation(umath.fabs)
 negative = masked_unary_operation(umath.negative)
 
 def nonzero(a):
-    """returns the indices of the elements of a which are not zero and not masked
-
-    a must be 1d
+    """returns the indices of the elements of a which are not zero
+    and not masked
     """
     return asarray(filled(a, 0).nonzero())
 
diff --git a/numpy/core/tests/test_ma.py b/numpy/core/tests/test_ma.py
index 312799323..540e64c8a 100644
--- a/numpy/core/tests/test_ma.py
+++ b/numpy/core/tests/test_ma.py
@@ -306,7 +306,7 @@ class test_ma(NumpyTestCase):
     def check_testMaPut(self):
         (x, y, a10, m1, m2, xm, ym, z, zm, xf, s) = self.d
         m = [1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1]
-        i = numpy.nonzero(m)
+        i = numpy.nonzero(m)[0]
         putmask(xm, m, z)
         assert take(xm, i) == z
         put(ym, i, zm)
diff --git a/numpy/lib/convertnumericA.py b/numpy/lib/convertnumericA.py
index ad1daa329..c3748bae3 100644
--- a/numpy/lib/convertnumericA.py
+++ b/numpy/lib/convertnumericA.py
@@ -13,6 +13,9 @@ Makes the following changes:
  * Converts .flat to .ravel() except for .flat = xxx or .flat[xxx]
  * Replace xxx.spacesaver() with True
  * Convert xx.savespace(?) to pass + ## xx.savespace(?)
+
+ * Converts uses of 'b' to 'B' in the typecode-position of
+   functions and methods
 """
 __all__ = ['fromfile', 'fromstr']
 
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 557383b71..548af953d 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -7,7 +7,8 @@ __all__ = ['logspace', 'linspace',
            'nanargmin', 'nanmin', 'vectorize', 'asarray_chkfinite', 'average',
            'histogram', 'bincount', 'digitize', 'cov', 'corrcoef', 'msort',
            'median', 'sinc', 'hamming', 'hanning', 'bartlett', 'blackman',
-           'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring', 'meshgrid'
+           'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring', 'meshgrid',
+           'flatnonzero'
            ]
 
 import types
@@ -530,7 +531,7 @@ def extract(condition, arr):
 
     Equivalent to compress(ravel(condition), ravel(arr)).
     """
-    return _nx.take(ravel(arr), nonzero(ravel(condition)))
+    return _nx.take(ravel(arr), nonzero(ravel(condition))[0])
 
 def insert(arr, mask, vals):
     """Similar to putmask arr[mask] = vals but the 1D array vals has the
@@ -997,3 +998,11 @@ def meshgrid(x,y):
     y = y.reshape(numRows,1)
     Y = y.repeat(numCols, axis=1)
     return X, Y
+
+def flatnonzero(a):
+    """Return indicies that are not-zero in flattened version of a
+
+    Equivalent to a.ravel().nonzero()[0]
+    """
+    return a.ravel().nonzero()[0]
+    
diff --git a/numpy/lib/polynomial.py b/numpy/lib/polynomial.py
index 206558179..0d66fd6b0 100644
--- a/numpy/lib/polynomial.py
+++ b/numpy/lib/polynomial.py
@@ -90,7 +90,7 @@ def roots(p):
         raise ValueError,"Input must be a rank-1 array."
 
     # find non-zero array entries
-    non_zero = NX.nonzero(NX.ravel(p))
+    non_zero = NX.nonzero(NX.ravel(p))[0]
 
     # find the number of trailing zeros -- this is the number of roots at 0.
     trailing_zeros = len(p) - non_zero[-1] - 1
diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py
index d1de7fea2..d80d36685 100644
--- a/numpy/linalg/linalg.py
+++ b/numpy/linalg/linalg.py
@@ -313,7 +313,7 @@ eigenvalue u[i].  Satisfies the equation dot(a, v[:,i]) = u[i]*v[:,i]
         else:
             w = wr+1j*wi
             v = array(vr, w.dtype)
-            ind = nonzero(wi != 0.0)        # indices of complex e-vals
+            ind = nonzero(wi != 0.0)[0]      # indices of complex e-vals
             for i in range(len(ind)/2):
                 v[ind[2*i]] = vr[ind[2*i]] + 1j*vr[ind[2*i+1]]
                 v[ind[2*i+1]] = vr[ind[2*i]] - 1j*vr[ind[2*i+1]]