numpy · mwiebe · Feb 8, 2011 · Feb 10, 2011
diff --git a/doc/source/reference/routines.indexing.rst b/doc/source/reference/routines.indexing.rst
@@ -20,6 +20,7 @@ Generating index arrays
    indices
    ix_
    ogrid
+   ravel_coords
    unravel_index
    diag_indices
    diag_indices_from

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
@@ -4337,6 +4337,104 @@
 
     """)
 
+add_newdoc('numpy.lib._compiled_base', 'ravel_coords',
+    """
+    ravel_coords(coords, dims, mode='raise', order='C')
+
+    Converts a tuple of coordinate arrays into an array of flat
+    indices, applying boundary modes to the coordinates.
+
+    Parameters
+    ----------
+    coords : tuple of array_like
+        A tuple of integer arrays, one array for each dimension.
+    dims : tuple of ints
+        The shape of array into which the indices from ``coords`` apply.
+    mode : {'raise', 'wrap', 'clip'}, optional
+        Specifies how out-of-bounds indices are handled.  Can specify
+        either one mode or a tuple of modes, one mode per index.
+
+        * 'raise' -- raise an error (default)
+        * 'wrap' -- wrap around
+        * 'clip' -- clip to the range
+
+        In 'clip' mode, a negative index which would normally
+        wrap will clip to 0 instead.
+    order : {'C', 'F'}, optional
+        Determines whether the coords should be viewed as indexing in
+        C (row-major) order or FORTRAN (column-major) order.
+
+    Returns
+    -------
+    raveled_indices : ndarray
+        An array of indices into the flattened version of an array
+        of dimensions ``dims``.
+
+    See Also
+    --------
+    unravel_index
+
+    Notes
+    -----
+    .. versionadded:: 1.6.0
+
+    Examples
+    --------
+    >>> arr = np.array([[3,6,6],[4,5,1]])
+    >>> np.ravel_coords(arr, (7,6))
+    array([22, 41, 37])
+    >>> np.ravel_coords(arr, (7,6), order='F')
+    array([31, 41, 13])
+    >>> np.ravel_coords(arr, (4,6), mode='clip')
+    array([22, 23, 19])
+    >>> np.ravel_coords(arr, (4,4), mode=('clip','wrap'))
+    array([12, 13, 13])
+
+    >>> np.ravel_coords((3,1,4,1), (6,7,8,9))
+    1621
+    """)
+
+add_newdoc('numpy.lib._compiled_base', 'unravel_index',
+    """
+    unravel_index(indices, dims, order='C')
+
+    Converts a flat index or array of flat indices into a tuple
+    of coordinate arrays.
+
+    Parameters
+    ----------
+    indices : array_like
+        An integer array whose elements are indices into the flattened
+        version of an array of dimensions ``dims``. Before version 1.6.0,
+        this function accepted just one index value.
+    dims : tuple of ints
+        The shape of the array to use for unravelling ``indices``.
+    order : {'C', 'F'}, optional
+        .. versionadded:: 1.6.0
+        Determines whether the indices should be viewed as indexing in
+        C (row-major) order or FORTRAN (column-major) order.
+
+    Returns
+    -------
+    unraveled_coords : tuple of ndarray
+        Each array in the tuple has the same shape as the ``indices``
+        array.
+
+    See Also
+    --------
+    ravel_coords
+
+    Examples
+    --------
+    >>> np.unravel_index([22, 41, 37], (7,6))
+    (array([3, 6, 6]), array([4, 5, 1]))
+    >>> np.unravel_index([31, 41, 13], (7,6), order='F')
+    (array([3, 6, 6]), array([4, 5, 1]))
+
+    >>> np.unravel_index(1621, (6,7,8,9))
+    (3, 1, 4, 1)
+    """)
+
 add_newdoc('numpy.lib._compiled_base', 'add_docstring',
     """
     docstring(obj, docstring)

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
@@ -300,18 +300,20 @@
     'NpyIter_GetAxisStrideArray':           264,
     'NpyIter_RequiresBuffering':            265,
     'NpyIter_GetInitialDataPtrArray':       266,
+    'NpyIter_CreateCompatibleStrides':      267,
     #
-    'PyArray_CastingConverter':             267,
-    'PyArray_CountNonzero':                 268,
-    'PyArray_PromoteTypes':                 269,
-    'PyArray_MinScalarType':                270,
-    'PyArray_ResultType':                   271,
-    'PyArray_CanCastArrayTo':               272,
-    'PyArray_CanCastTypeTo':                273,
-    'PyArray_EinsteinSum':                  274,
-    'PyArray_FillWithZero':                 275,
-    'PyArray_NewLikeArray':                 276,
-    'PyArray_GetArrayParamsFromObject':     277,
+    'PyArray_CastingConverter':             268,
+    'PyArray_CountNonzero':                 269,
+    'PyArray_PromoteTypes':                 270,
+    'PyArray_MinScalarType':                271,
+    'PyArray_ResultType':                   272,
+    'PyArray_CanCastArrayTo':               273,
+    'PyArray_CanCastTypeTo':                274,
+    'PyArray_EinsteinSum':                  275,
+    'PyArray_FillWithZero':                 276,
+    'PyArray_NewLikeArray':                 277,
+    'PyArray_GetArrayParamsFromObject':     278,
+    'PyArray_ConvertClipmodeSequence':      279,
 }
 
 ufunc_types_api = {

diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
@@ -1311,9 +1311,53 @@ PyArray_ClipmodeConverter(PyObject *object, NPY_CLIPMODE *val)
     return PY_FAIL;
 }
 
+/*NUMPY_API
+ * Convert an object to an array of n NPY_CLIPMODE values.
+ * This is intended to be used in functions where a different mode
+ * could be applied to each axis, like in ravel_coords.
+ */
+NPY_NO_EXPORT int
+PyArray_ConvertClipmodeSequence(PyObject *object, NPY_CLIPMODE *modes, int n)
+{
+    int i;
+    /* Get the clip mode(s) */
+    if (object && (PyTuple_Check(object) || PyList_Check(object))) {
+        if (PySequence_Size(object) != n) {
+            PyErr_Format(PyExc_ValueError,
+                    "list of clipmodes has wrong length (%d instead of %d)",
+                    (int)PySequence_Size(object), n);
+            return PY_FAIL;
+        }
+
+        for (i = 0; i < n; ++i) {
+            PyObject *item = PySequence_GetItem(object, i);
+            if(item == NULL) {
+                return PY_FAIL;
+            }
+
+            if(PyArray_ClipmodeConverter(item, &modes[i]) != PY_SUCCEED) {
+                Py_DECREF(item);
+                return PY_FAIL;
+            }
+
+            Py_DECREF(item);
+        }
+    }
+    else if (PyArray_ClipmodeConverter(object, &modes[0]) == PY_SUCCEED) {
+        for (i = 1; i < n; ++i) {
+            modes[i] = modes[0];
+        }
+    }
+    else {
+        return PY_FAIL;
+    }
+    return PY_SUCCEED;
+}
+
 /*
- * compare the field dictionary for two types
- * return 1 if the same or 0 if not
+ * Compare the field dictionaries for two types.
+ *
+ * Return 1 if the contents are the same, 0 if not.
  */
 static int
 _equivalent_fields(PyObject *field1, PyObject *field2) {

diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
@@ -2316,9 +2316,18 @@ NpyIter_GetIterIndexRange(NpyIter *iter,
 }
 
 /*NUMPY_API
- * Gets the broadcast shape if coords are enabled, otherwise
- * gets the shape of the iteration as Fortran-order (fastest-changing
- * coordinate first)
+ * Gets the broadcast shape if coords are being tracked by the iterator,
+ * otherwise gets the shape of the iteration as Fortran-order
+ * (fastest-changing coordinate first).
+ *
+ * The reason Fortran-order is returned when coords
+ * are not enabled is that this is providing a direct view into how
+ * the iterator traverses the n-dimensional space. The iterator organizes
+ * its memory from fastest coordinate to slowest coordinate, and when
+ * coords are enabled, it uses a permutation to recover the original
+ * order.
+ *
+ * Returns NPY_SUCCEED or NPY_FAIL.
  */
 NPY_NO_EXPORT int
 NpyIter_GetShape(NpyIter *iter, npy_intp *outshape)
@@ -2358,6 +2367,74 @@ NpyIter_GetShape(NpyIter *iter, npy_intp *outshape)
     return NPY_SUCCEED;
 }
 
+/*NUMPY_API
+ * Builds a set of strides which are the same as the strides of an
+ * output array created using the NPY_ITER_ALLOCATE flag, where NULL
+ * was passed for op_axes.  This is for data packed contiguously,
+ * but not necessarily in C or Fortran order. This should be used
+ * together with NpyIter_GetShape and NpyIter_GetNDim.
+ *
+ * A use case for this function is to match the shape and layout of
+ * the iterator and tack on one or more dimensions.  For example,
+ * in order to generate a vector per input value for a numerical gradient,
+ * you pass in ndim*itemsize for itemsize, then add another dimension to
+ * the end with size ndim and stride itemsize.  To do the Hessian matrix,
+ * you do the same thing but add two dimensions, or take advantage of
+ * the symmetry and pack it into 1 dimension with a particular encoding.
+ *
+ * This function may only be called if the iterator is tracking coordinates
+ * and if NPY_ITER_DONT_NEGATE_STRIDES was used to prevent an axis from
+ * being iterated in reverse order.
+ *
+ * If an array is created with this method, simply adding 'itemsize'
+ * for each iteration will traverse the new array matching the
+ * iterator.
+ *
+ * Returns NPY_SUCCEED or NPY_FAIL.
+ */
+NPY_NO_EXPORT int
+NpyIter_CreateCompatibleStrides(NpyIter *iter,
+                            npy_intp itemsize, npy_intp *outstrides)
+{
+    npy_uint32 itflags = NIT_ITFLAGS(iter);
+    npy_intp idim, ndim = NIT_NDIM(iter);
+    npy_intp niter = NIT_NITER(iter);
+
+    npy_intp sizeof_axisdata;
+    NpyIter_AxisData *axisdata;
+    char *perm;
+
+    if (!(itflags&NPY_ITFLAG_HASCOORDS)) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "Iterator CreateCompatibleStrides may only be called "
+                "if coordinates are being tracked");
+        return NPY_FAIL;
+    }
+
+    axisdata = NIT_AXISDATA(iter);
+    sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, niter);
+
+    perm = NIT_PERM(iter);
+    for(idim = 0; idim < ndim; ++idim) {
+        char p = perm[idim];
+        if (p < 0) {
+            PyErr_SetString(PyExc_RuntimeError,
+                    "Iterator CreateCompatibleStrides may only be called "
+                    "if DONT_NEGATE_STRIDES was used to prevent reverse "
+                    "iteration of an axis");
+            return NPY_FAIL;
+        }
+        else {
+            outstrides[ndim-p-1] = itemsize;
+        }
+
+        itemsize *= NAD_SHAPE(axisdata);
+        NIT_ADVANCE_AXISDATA(axisdata, 1);
+    }
+
+    return NPY_SUCCEED;
+}
+
 /*NUMPY_API
  * Get the array of data pointers (1 per object being iterated)
  *

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
@@ -1,4 +1,5 @@
-__all__ = ['unravel_index',
+__all__ = ['ravel_coords',
+           'unravel_index',
            'mgrid',
            'ogrid',
            'r_', 'c_', 's_',
@@ -16,70 +17,9 @@
 import function_base
 import numpy.matrixlib as matrix
 from function_base import diff
+from _compiled_base import ravel_coords, unravel_index
 makemat = matrix.matrix
 
-# contributed by Stefan van der Walt
-def unravel_index(x,dims):
-    """
-    Convert a flat index to an index tuple for an array of given shape.
-
-    Parameters
-    ----------
-    x : int
-        Flattened index.
-    dims : tuple of ints
-        Input shape, the shape of an array into which indexing is
-        required.
-
-    Returns
-    -------
-    idx : tuple of ints
-        Tuple of the same shape as `dims`, containing the unraveled index.
-
-    Notes
-    -----
-    In the Examples section, since ``arr.flat[x] == arr.max()`` it may be
-    easier to use flattened indexing than to re-map the index to a tuple.
-
-    Examples
-    --------
-    >>> arr = np.arange(20).reshape(5, 4)
-    >>> arr
-    array([[ 0,  1,  2,  3],
-           [ 4,  5,  6,  7],
-           [ 8,  9, 10, 11],
-           [12, 13, 14, 15],
-           [16, 17, 18, 19]])
-    >>> x = arr.argmax()
-    >>> x
-    19
-    >>> dims = arr.shape
-    >>> idx = np.unravel_index(x, dims)
-    >>> idx
-    (4, 3)
-    >>> arr[idx] == arr.max()
-    True
-
-    """
-    if x > _nx.prod(dims)-1 or x < 0:
-        raise ValueError("Invalid index, must be 0 <= x <= number of elements.")
-
-    idx = _nx.empty_like(dims)
-
-    # Take dimensions
-    # [a,b,c,d]
-    # Reverse and drop first element
-    # [d,c,b]
-    # Prepend [1]
-    # [1,d,c,b]
-    # Calculate cumulative product
-    # [1,d,dc,dcb]
-    # Reverse
-    # [dcb,dc,d,1]
-    dim_prod = _nx.cumprod([1] + list(dims)[:0:-1])[::-1]
-    # Indices become [x/dcb % a, x/dc % b, x/d % c, x/1 % d]
-    return tuple(x//dim_prod % dims)
-
 def ix_(*args):
     """
     Construct an open mesh from multiple sequences.