diff options
| -rw-r--r-- | numpy/core/numeric.py | 19 | ||||
| -rw-r--r-- | numpy/core/src/multiarraymodule.c | 2 | ||||
| -rw-r--r-- | numpy/lib/index_tricks.py | 11 |
3 files changed, 23 insertions, 9 deletions
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index 557c8e7f4..5dbd1c232 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -1,8 +1,8 @@ __all__ = ['newaxis', 'ndarray', 'flatiter', 'ufunc', 'arange', 'array', 'zeros', 'empty', 'broadcast', 'dtype', 'fromstring', 'fromfile', 'frombuffer','newbuffer', - 'getbuffer', - 'where', 'concatenate', 'fastCopyAndTranspose', 'lexsort', + 'getbuffer', 'where', 'argwhere', + 'concatenate', 'fastCopyAndTranspose', 'lexsort', 'register_dtype', 'set_numeric_ops', 'can_cast', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'isfortran', 'empty_like', 'zeros_like', @@ -129,6 +129,21 @@ def asfortranarray(a, dtype=None): def isfortran(a): return a.flags.fnc +def argwhere(a): + """Return a 2-d array of shape N x a.ndim where each row + is a sequence of indices into a. This sequence must be + converted to a tuple in order to be used to index into a. + """ + if a.ndim == 0: + return array([],dtype=intp) + else: + b = a.nonzero() + retarr = empty((b[0].shape[0],a.ndim),dtype=intp) + for k in xrange(a.ndim): + retarr[:,k] = b[k] + return retarr + + _mode_from_name_dict = {'v': 0, 's' : 1, 'f' : 2} diff --git a/numpy/core/src/multiarraymodule.c b/numpy/core/src/multiarraymodule.c index 99e98cd35..aeb9860ee 100644 --- a/numpy/core/src/multiarraymodule.c +++ b/numpy/core/src/multiarraymodule.c @@ -5801,7 +5801,7 @@ PyArray_Where(PyObject *condition, PyObject *x, PyObject *y) static char doc_where[] = "where(condition, | x, y) is shaped like condition"\ " and has elements of x and y where condition is respectively true or"\ " false. If x or y are not given, then it is equivalent to"\ - " nonzero(condition)."; + " condition.nonzero()."; static PyObject * array_where(PyObject *ignored, PyObject *args) diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py index 444c67e03..6f9d14e3e 100644 --- a/numpy/lib/index_tricks.py +++ b/numpy/lib/index_tricks.py @@ -1,6 +1,6 @@ ## Automatically adapted for numpy Sep 19, 2005 by convertcode.py -__all__ = ['fromflat', +__all__ = ['unravel_index', 'mgrid', 'ogrid', 'r_', 'c_', 's_', @@ -16,16 +16,15 @@ import function_base import numpy.core.defmatrix as matrix makemat = matrix.matrix - -# fromflat contributed by Stefan van der Walt -def fromflat(x,dims): +# contributed by Stefan van der Walt +def unravel_index(x,dims): """Convert a flat index into an index tuple for a matrix of given shape. - e.g. for a 2x2 matrix, fromflat(2,(2,2)) translates to (1,0). + e.g. for a 2x2 matrix, unravel_index(2,(2,2)) returns (1,0). Example usage: p = x.argmax() - idx = fromflat(p) + idx = unravel_index(p) x[idx] == x.max() Note: x.flat[p] == x.max() |