From 474042ec32b9603696ca04794cb9d02bdc5ecf5d Mon Sep 17 00:00:00 2001 From: Mark Wiebe Date: Tue, 8 Feb 2011 00:50:38 -0800 Subject: [PATCH 1/2] ENH: index_tricks: Implement unravel_index and ravel_coords functions in C --- doc/source/reference/routines.indexing.rst | 1 + numpy/add_newdocs.py | 94 ++++ numpy/core/code_generators/numpy_api.py | 24 +- numpy/core/src/multiarray/multiarraymodule.c | 36 ++ numpy/core/src/multiarray/new_iterator.c.src | 70 +++ numpy/lib/index_tricks.py | 66 +-- numpy/lib/src/_compiled_base.c | 488 +++++++++++++++++++ numpy/lib/tests/test_index_tricks.py | 67 ++- 8 files changed, 767 insertions(+), 79 deletions(-) diff --git a/doc/source/reference/routines.indexing.rst b/doc/source/reference/routines.indexing.rst index 9d8fde882046..98b7817b82e4 100644 --- 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 index a66eb6de3dd2..487ffef00746 100644 --- a/numpy/add_newdocs.py +++ b/numpy/add_newdocs.py @@ -4337,6 +4337,100 @@ """) +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 an array into which indices from `coords` are for. + mode : {'raise', 'wrap', 'clip'}, optional + Specifies how out-of-bounds indices will behave. Can specify + either one mode or a tuple of modes, with length matching that + of ``dims``. + + * 'raise' -- raise an error (default) + * 'wrap' -- wrap around + * 'clip' -- clip to the range + + Note that 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 + This is a new array with the same shape as the broadcast shape + of the arrays in ``coords``. + + See Also + -------- + unravel_index + + 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 type array whose elements are indices for a + flattened array with shape `dims`. + dims : tuple of ints + The shape of an array into which flattened indices from + ``indices`` are for. + order : {'C', 'F'}, optional + 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 is the same shape as the input ``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 index 0beecf8dc6a6..61f56d09cd20 100644 --- 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 index 8931dce26d46..507cd398c803 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1311,6 +1311,42 @@ PyArray_ClipmodeConverter(PyObject *object, NPY_CLIPMODE *val) return PY_FAIL; } +/*NUMPY_API + * Convert an object to an array of n NPY_CLIPMODE values. + */ +NPY_NO_EXPORT int +PyArray_ConvertClipmodeSequence(PyObject *object, NPY_CLIPMODE *vals, 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, &vals[i]) != PY_SUCCEED) { + Py_DECREF(item); + return PY_FAIL; + } + Py_DECREF(item); + } + } else if(PyArray_ClipmodeConverter(object, &vals[0]) == PY_SUCCEED) { + for(i = 1; i < n; ++i) { + vals[i] = vals[0]; + } + } else { + return PY_FAIL; + } + return PY_SUCCEED; +} + /* * compare the field dictionary for two types * return 1 if the same or 0 if not diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src index 1bdc89b4d666..931cb03395d3 100644 --- a/numpy/core/src/multiarray/new_iterator.c.src +++ b/numpy/core/src/multiarray/new_iterator.c.src @@ -2319,6 +2319,8 @@ NpyIter_GetIterIndexRange(NpyIter *iter, * Gets the broadcast shape if coords are enabled, otherwise * gets the shape of the iteration as Fortran-order (fastest-changing * coordinate first) + * + * Returns NPY_SUCCEED or NPY_FAIL. */ NPY_NO_EXPORT int NpyIter_GetShape(NpyIter *iter, npy_intp *outshape) @@ -2358,6 +2360,74 @@ NpyIter_GetShape(NpyIter *iter, npy_intp *outshape) return NPY_SUCCEED; } +/*NUMPY_API + * Builds a set of strides which are compatible with the iteration order + * for data packed contiguously, but not necessarily C or Fortran. + * The strides this produces are equivalent to the strides for an + * automatically allocated output created without an op_axes parameter, + * and 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, + * if you want 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 index 264ebaad0625..360d30c2afe9 100644 --- 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. diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c index 627fc32868d1..a7a736254ab0 100644 --- a/numpy/lib/src/_compiled_base.c +++ b/numpy/lib/src/_compiled_base.c @@ -565,6 +565,489 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict) return NULL; } +static int sequence_to_arrays(PyObject *seq, + PyArrayObject **op, int count, + char *paramname) +{ + int i; + + if (!PySequence_Check(seq) || PySequence_Size(seq) != count) { + PyErr_Format(PyExc_ValueError, + "parameter %s must be a sequence of length %d", + paramname, count); + return -1; + } + + for (i = 0; i < count; ++i) { + PyObject *item = PySequence_GetItem(seq, i); + if (item == NULL) { + while (--i >= 0) { + Py_DECREF(op[i]); + op[i] = NULL; + } + return -1; + } + + op[i] = (PyArrayObject *)PyArray_FromAny(item, NULL, 0, 0, 0, NULL); + if (op[i] == NULL) { + while (--i >= 0) { + Py_DECREF(op[i]); + op[i] = NULL; + } + Py_DECREF(item); + return -1; + } + + Py_DECREF(item); + } + + return 0; +} + +static int +ravel_coords_loop(int ravel_ndim, npy_intp *ravel_dims, + npy_intp *ravel_strides, + npy_intp count, + NPY_CLIPMODE *modes, + char **coords, npy_intp *coords_strides) +{ + int i; + npy_intp j, m; + + while (count--) { + npy_intp raveled = 0; + for (i = 0; i < ravel_ndim; ++i) { + m = ravel_dims[i]; + j = *(npy_intp *)coords[i]; + switch (modes[i]) { + case NPY_RAISE: + if(j < 0 || j>=m) { + PyErr_SetString(PyExc_ValueError, + "invalid entry in coordinates array"); + return NPY_FAIL; + } + break; + case NPY_WRAP: + if(j < 0) { + j += m; + if(j < 0) { + j = j%m; + if(j != 0) { + j += m; + } + } + } + else if(j >= m) { + j -= m; + if(j >= m) { + j = j%m; + } + } + break; + case NPY_CLIP: + if(j < 0) { + j = 0; + } + else if(j >= m) { + j = m-1; + } + break; + + } + raveled += j * ravel_strides[i]; + + coords[i] += coords_strides[i]; + } + *(npy_intp *)coords[ravel_ndim] = raveled; + coords[ravel_ndim] += coords_strides[ravel_ndim]; + } + + return NPY_SUCCEED; +} + +static PyObject * +arr_ravel_coords(PyObject *self, PyObject *args, PyObject *kwds) +{ + int i, s; + PyObject *mode0=NULL, *coords0=NULL, *ret; + PyArray_Dims dimensions={0,0}; + npy_intp ravel_strides[NPY_MAXDIMS]; + PyArray_ORDER order = NPY_CORDER; + NPY_CLIPMODE modes[NPY_MAXDIMS]; + + PyArrayObject *op[NPY_MAXARGS]; + PyArray_Descr *dtype[NPY_MAXARGS]; + npy_uint32 op_flags[NPY_MAXARGS]; + + NpyIter *iter = NULL; + + char *kwlist[] = {"coords", "dims", "mode", "order", NULL}; + + memset(op, 0, sizeof(op)); + dtype[0] = NULL; + + if(!PyArg_ParseTupleAndKeywords(args, kwds, "OO&|OO&:ravel_coords", kwlist, + &coords0, + PyArray_IntpConverter, &dimensions, + &mode0, + PyArray_OrderConverter, &order)) { + goto fail; + } + + if (dimensions.len+1 > NPY_MAXARGS) { + PyErr_SetString(PyExc_ValueError, + "too many dimensions passed to ravel_coords"); + goto fail; + } + + if(!PyArray_ConvertClipmodeSequence(mode0, modes, dimensions.len)) { + goto fail; + } + + switch (order) { + case NPY_CORDER: + s = 1; + for (i = dimensions.len-1; i >= 0; --i) { + ravel_strides[i] = s; + s *= dimensions.ptr[i]; + } + break; + case NPY_FORTRANORDER: + s = 1; + for (i = 0; i < dimensions.len; ++i) { + ravel_strides[i] = s; + s *= dimensions.ptr[i]; + } + break; + default: + PyErr_SetString(PyExc_ValueError, + "only 'C' or 'F' order is permitted"); + goto fail; + } + + /* Get the coords into op */ + if (sequence_to_arrays(coords0, op, dimensions.len, "coords") < 0) { + goto fail; + } + + + for (i = 0; i < dimensions.len; ++i) { + op_flags[i] = NPY_ITER_READONLY| + NPY_ITER_ALIGNED; + } + op_flags[dimensions.len] = NPY_ITER_WRITEONLY| + NPY_ITER_ALIGNED| + NPY_ITER_ALLOCATE; + dtype[0] = PyArray_DescrFromType(NPY_INTP); + for (i = 1; i <= dimensions.len; ++i) { + dtype[i] = dtype[0]; + } + + iter = NpyIter_MultiNew(dimensions.len+1, op, NPY_ITER_BUFFERED| + NPY_ITER_NO_INNER_ITERATION| + NPY_ITER_ZEROSIZE_OK, + NPY_KEEPORDER, + NPY_SAME_KIND_CASTING, + op_flags, dtype, + 0, NULL, 0); + if (iter == NULL) { + goto fail; + } + + if (NpyIter_GetIterSize(iter) != 0) { + NpyIter_IterNext_Fn iternext; + char **dataptr; + npy_intp *strides; + npy_intp *countptr; + + iternext = NpyIter_GetIterNext(iter, NULL); + if (iternext == NULL) { + goto fail; + } + dataptr = NpyIter_GetDataPtrArray(iter); + strides = NpyIter_GetInnerStrideArray(iter); + countptr = NpyIter_GetInnerLoopSizePtr(iter); + + do { + if (ravel_coords_loop(dimensions.len, dimensions.ptr, + ravel_strides, *countptr, modes, + dataptr, strides) != NPY_SUCCEED) { + goto fail; + } + } while(iternext(iter)); + } + + ret = (PyObject *)NpyIter_GetOperandArray(iter)[dimensions.len]; + Py_INCREF(ret); + + Py_DECREF(dtype[0]); + for (i = 0; i < dimensions.len; ++i) { + Py_XDECREF(op[i]); + } + PyDimMem_FREE(dimensions.ptr); + NpyIter_Deallocate(iter); + return PyArray_Return(ret); + +fail: + Py_XDECREF(dtype[0]); + for (i = 0; i < dimensions.len; ++i) { + Py_XDECREF(op[i]); + } + if (dimensions.ptr) { + PyDimMem_FREE(dimensions.ptr); + } + if (iter != NULL) { + NpyIter_Deallocate(iter); + } + return NULL; +} + +static int +unravel_index_loop_corder(int unravel_ndim, npy_intp *unravel_dims, + npy_intp unravel_size, npy_intp count, + char *indices, npy_intp indices_stride, + npy_intp *coords) +{ + int i; + npy_intp val; + + while (count--) { + val = *(npy_intp *)indices; + if (val < 0 || val >= unravel_size) { + PyErr_SetString(PyExc_ValueError, + "invalid entry in index array"); + return NPY_FAIL; + } + for (i = unravel_ndim-1; i >= 0; --i) { + coords[i] = val % unravel_dims[i]; + val /= unravel_dims[i]; + } + coords += unravel_ndim; + indices += indices_stride; + } + + return NPY_SUCCEED; +} + +static int +unravel_index_loop_forder(int unravel_ndim, npy_intp *unravel_dims, + npy_intp unravel_size, npy_intp count, + char *indices, npy_intp indices_stride, + npy_intp *coords) +{ + int i; + npy_intp val; + + while (count--) { + val = *(npy_intp *)indices; + if (val < 0 || val >= unravel_size) { + PyErr_SetString(PyExc_ValueError, + "invalid entry in index array"); + return NPY_FAIL; + } + for (i = 0; i < unravel_ndim; ++i) { + *coords++ = val % unravel_dims[i]; + val /= unravel_dims[i]; + } + indices += indices_stride; + } + + return NPY_SUCCEED; +} + +static PyObject * +arr_unravel_index(PyObject *self, PyObject *args, PyObject *kwds) +{ + PyObject *indices0 = NULL, *ret_tuple = NULL; + PyArrayObject *ret_arr = NULL; + PyArrayObject *indices = NULL; + PyArray_Descr *dtype = NULL; + PyArray_Dims dimensions={0,0}; + PyArray_ORDER order = PyArray_CORDER; + npy_intp unravel_size; + + NpyIter *iter = NULL; + int i, ret_ndim; + npy_intp ret_dims[NPY_MAXDIMS], ret_strides[NPY_MAXDIMS]; + + char *kwlist[] = {"indices", "dims", "order", NULL}; + + if(!PyArg_ParseTupleAndKeywords(args, kwds, "OO&|O&:unravel_index", + kwlist, + &indices0, + PyArray_IntpConverter, &dimensions, + PyArray_OrderConverter, &order)) { + goto fail; + } + + if (dimensions.len == 0) { + PyErr_SetString(PyExc_ValueError, + "dims must have at least one value"); + goto fail; + } + + unravel_size = PyArray_MultiplyList(dimensions.ptr, dimensions.len); + + if(!PyArray_Check(indices0)) { + indices = (PyArrayObject*)PyArray_FromAny(indices0, + NULL, 0, 0, 0, NULL); + if(indices == NULL) { + goto fail; + } + } + else { + indices = (PyArrayObject *)indices0; + Py_INCREF(indices); + } + + dtype = PyArray_DescrFromType(NPY_INTP); + if (dtype == NULL) { + goto fail; + } + + iter = NpyIter_New(indices, NPY_ITER_READONLY| + NPY_ITER_ALIGNED| + NPY_ITER_BUFFERED| + NPY_ITER_ZEROSIZE_OK| + NPY_ITER_DONT_NEGATE_STRIDES| + NPY_ITER_COORDS, + NPY_KEEPORDER, NPY_SAME_KIND_CASTING, + dtype, 0, NULL, 0); + if (iter == NULL) { + goto fail; + } + + /* + * Create the return array with a layout compatible with the indices + * and with a dimension added to the end for the coordinates + */ + ret_ndim = PyArray_NDIM(indices) + 1; + if (NpyIter_GetShape(iter, ret_dims) != NPY_SUCCEED) { + goto fail; + } + ret_dims[ret_ndim-1] = dimensions.len; + if (NpyIter_CreateCompatibleStrides(iter, + dimensions.len*sizeof(npy_intp), ret_strides) != NPY_SUCCEED) { + goto fail; + } + ret_strides[ret_ndim-1] = sizeof(npy_intp); + + /* Remove the coords and inner loop */ + if (NpyIter_RemoveCoords(iter) != NPY_SUCCEED) { + goto fail; + } + if (NpyIter_RemoveInnerLoop(iter) != NPY_SUCCEED) { + goto fail; + } + + ret_arr = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype, + ret_ndim, ret_dims, ret_strides, NULL, 0, NULL); + dtype = NULL; + if (ret_arr == NULL) { + goto fail; + } + + if (order == NPY_CORDER) { + if (NpyIter_GetIterSize(iter) != 0) { + NpyIter_IterNext_Fn iternext; + char **dataptr; + npy_intp *strides; + npy_intp *countptr, count; + npy_intp *coordsptr = (npy_intp *)PyArray_DATA(ret_arr); + + iternext = NpyIter_GetIterNext(iter, NULL); + if (iternext == NULL) { + goto fail; + } + dataptr = NpyIter_GetDataPtrArray(iter); + strides = NpyIter_GetInnerStrideArray(iter); + countptr = NpyIter_GetInnerLoopSizePtr(iter); + + do { + count = *countptr; + if (unravel_index_loop_corder(dimensions.len, dimensions.ptr, + unravel_size, count, *dataptr, *strides, + coordsptr) != NPY_SUCCEED) { + goto fail; + } + coordsptr += count*dimensions.len; + } while(iternext(iter)); + } + } + else if (order == NPY_FORTRANORDER) { + if (NpyIter_GetIterSize(iter) != 0) { + NpyIter_IterNext_Fn iternext; + char **dataptr; + npy_intp *strides; + npy_intp *countptr, count; + npy_intp *coordsptr = (npy_intp *)PyArray_DATA(ret_arr); + + iternext = NpyIter_GetIterNext(iter, NULL); + if (iternext == NULL) { + goto fail; + } + dataptr = NpyIter_GetDataPtrArray(iter); + strides = NpyIter_GetInnerStrideArray(iter); + countptr = NpyIter_GetInnerLoopSizePtr(iter); + + do { + count = *countptr; + if (unravel_index_loop_forder(dimensions.len, dimensions.ptr, + unravel_size, count, *dataptr, *strides, + coordsptr) != NPY_SUCCEED) { + goto fail; + } + coordsptr += count*dimensions.len; + } while(iternext(iter)); + } + } + else { + PyErr_SetString(PyExc_ValueError, + "only 'C' or 'F' order is permitted"); + goto fail; + } + + /* Now make a tuple of views, one per coordinate */ + ret_tuple = PyTuple_New(dimensions.len); + if (ret_tuple == NULL) { + goto fail; + } + for (i = 0; i < dimensions.len; ++i) { + PyArrayObject *view; + + view = (PyArrayObject *)PyArray_New(&PyArray_Type, ret_ndim-1, + ret_dims, NPY_INTP, + ret_strides, + PyArray_BYTES(ret_arr) + i*sizeof(npy_intp), + 0, 0, NULL); + if (view == NULL) { + goto fail; + } + Py_INCREF(ret_arr); + view->base = (PyObject *)ret_arr; + PyTuple_SET_ITEM(ret_tuple, i, PyArray_Return(view)); + } + + Py_DECREF(ret_arr); + Py_XDECREF(indices); + PyDimMem_FREE(dimensions.ptr); + NpyIter_Deallocate(iter); + + return ret_tuple; + +fail: + Py_XDECREF(ret_tuple); + Py_XDECREF(ret_arr); + Py_XDECREF(dtype); + Py_XDECREF(indices); + if (dimensions.ptr) { + PyDimMem_FREE(dimensions.ptr); + } + if (iter != NULL) { + NpyIter_Deallocate(iter); + } + return NULL; +} static PyTypeObject *PyMemberDescr_TypePtr = NULL; @@ -905,6 +1388,7 @@ io_unpack(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds) return pack_or_unpack_bits(obj, axis, 1); } +/* The docstrings for many of these methods are in add_newdocs.py. */ static struct PyMethodDef methods[] = { {"_insert", (PyCFunction)arr_insert, METH_VARARGS | METH_KEYWORDS, arr_insert__doc__}, @@ -914,6 +1398,10 @@ static struct PyMethodDef methods[] = { METH_VARARGS | METH_KEYWORDS, NULL}, {"interp", (PyCFunction)arr_interp, METH_VARARGS | METH_KEYWORDS, NULL}, + {"ravel_coords", (PyCFunction)arr_ravel_coords, + METH_VARARGS | METH_KEYWORDS, NULL}, + {"unravel_index", (PyCFunction)arr_unravel_index, + METH_VARARGS | METH_KEYWORDS, NULL}, {"add_docstring", (PyCFunction)arr_add_docstring, METH_VARARGS, NULL}, {"packbits", (PyCFunction)io_pack, diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index c17ee5d6a2d9..40f75936e9c0 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -4,12 +4,69 @@ ndenumerate, fill_diagonal, diag_indices, diag_indices_from, s_, index_exp ) -class TestUnravelIndex(TestCase): +class TestRavelUnravelIndex(TestCase): def test_basic(self): - assert unravel_index(2,(2,2)) == (1,0) - assert unravel_index(254,(17,94)) == (2, 66) - assert_raises(ValueError, unravel_index, 4,(2,2)) - + assert_equal(np.unravel_index(2,(2,2)), (1,0)) + assert_equal(np.ravel_coords((1,0),(2,2)), 2) + assert_equal(np.unravel_index(254,(17,94)), (2,66)) + assert_equal(np.ravel_coords((2,66),(17,94)), 254) + assert_raises(ValueError, np.unravel_index, -1, (2,2)) + assert_raises(TypeError, np.unravel_index, 0.5, (2,2)) + assert_raises(ValueError, np.unravel_index, 4, (2,2)) + assert_raises(ValueError, np.ravel_coords, (-3,1), (2,2)) + assert_raises(ValueError, np.ravel_coords, (2,1), (2,2)) + assert_raises(ValueError, np.ravel_coords, (0,-3), (2,2)) + assert_raises(ValueError, np.ravel_coords, (0,2), (2,2)) + assert_raises(TypeError, np.ravel_coords, (0.1,0.), (2,2)) + + assert_equal(np.unravel_index((2*3 + 1)*6 + 4, (4,3,6)), [2,1,4]) + assert_equal(np.ravel_coords([2,1,4], (4,3,6)), (2*3 + 1)*6 + 4) + + arr = np.array([[3,6,6],[4,5,1]]) + assert_equal(np.ravel_coords(arr, (7,6)), [22,41,37]) + assert_equal(np.ravel_coords(arr, (7,6), order='F'), [31,41,13]) + assert_equal(np.ravel_coords(arr, (4,6), mode='clip'), [22,23,19]) + assert_equal(np.ravel_coords(arr, (4,4), mode=('clip','wrap')), + [12,13,13]) + assert_equal(np.ravel_coords((3,1,4,1), (6,7,8,9)), 1621) + + assert_equal(np.unravel_index(np.array([22, 41, 37]), (7,6)), + [[3, 6, 6],[4, 5, 1]]) + assert_equal(np.unravel_index(np.array([31, 41, 13]), (7,6), order='F'), + [[3, 6, 6], [4, 5, 1]]) + assert_equal(np.unravel_index(1621, (6,7,8,9)), [3,1,4,1]) + + def test_dtypes(self): + # Test with different data types + for dtype in [np.int16, np.uint16, np.int32, + np.uint32, np.int64, np.uint64]: + coords = np.array([[1,0,1,2,3,4],[1,6,1,3,2,0]], dtype=dtype) + shape = (5,8) + uncoords = 8*coords[0]+coords[1] + assert_equal(np.ravel_coords(coords, shape), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape)) + uncoords = coords[0]+5*coords[1] + assert_equal(np.ravel_coords(coords, shape, order='F'), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) + + coords = np.array([[1,0,1,2,3,4],[1,6,1,3,2,0],[1,3,1,0,9,5]], + dtype=dtype) + shape = (5,8,10) + uncoords = 10*(8*coords[0]+coords[1])+coords[2] + assert_equal(np.ravel_coords(coords, shape), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape)) + uncoords = coords[0]+5*(coords[1]+8*coords[2]) + assert_equal(np.ravel_coords(coords, shape, order='F'), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) + + def test_clipmodes(self): + # Test clipmodes + assert_equal(np.ravel_coords([5,1,-1,2], (4,3,7,12), mode='wrap'), + np.ravel_coords([1,1,6,2], (4,3,7,12))) + assert_equal(np.ravel_coords([5,1,-1,2], (4,3,7,12), + mode=('wrap','raise','clip','raise')), + np.ravel_coords([1,1,0,2], (4,3,7,12))) + assert_raises(ValueError, np.ravel_coords, [5,1,-1,2], (4,3,7,12)) class TestGrid(TestCase): def test_basic(self): From 4c788cd662272452269943399083c36653d2178e Mon Sep 17 00:00:00 2001 From: Mark Wiebe Date: Wed, 9 Feb 2011 21:13:54 -0800 Subject: [PATCH 2/2] STY: index_tricks: Improve comments and documentation strings --- numpy/add_newdocs.py | 28 +++++++++++--------- numpy/core/src/multiarray/multiarraymodule.c | 26 +++++++++++------- numpy/core/src/multiarray/new_iterator.c.src | 25 ++++++++++------- numpy/lib/src/_compiled_base.c | 13 +++++++++ 4 files changed, 62 insertions(+), 30 deletions(-) diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py index 487ffef00746..1a55b6bc77b7 100644 --- a/numpy/add_newdocs.py +++ b/numpy/add_newdocs.py @@ -4349,17 +4349,16 @@ coords : tuple of array_like A tuple of integer arrays, one array for each dimension. dims : tuple of ints - The shape of an array into which indices from `coords` are for. + The shape of array into which the indices from ``coords`` apply. mode : {'raise', 'wrap', 'clip'}, optional - Specifies how out-of-bounds indices will behave. Can specify - either one mode or a tuple of modes, with length matching that - of ``dims``. + 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 - Note that in 'clip' mode, a negative index which would normally + 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 @@ -4368,13 +4367,17 @@ Returns ------- raveled_indices : ndarray - This is a new array with the same shape as the broadcast shape - of the arrays in ``coords``. + 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]]) @@ -4401,19 +4404,20 @@ Parameters ---------- indices : array_like - An integer type array whose elements are indices for a - flattened array with shape `dims`. + 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 an array into which flattened indices from - ``indices`` are for. + 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 is the same shape as the input ``indices`` + Each array in the tuple has the same shape as the ``indices`` array. See Also diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 507cd398c803..b87baeef7cb7 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1313,43 +1313,51 @@ PyArray_ClipmodeConverter(PyObject *object, NPY_CLIPMODE *val) /*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 *vals, int n) +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 (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, &vals[i]) != PY_SUCCEED) { + + if(PyArray_ClipmodeConverter(item, &modes[i]) != PY_SUCCEED) { Py_DECREF(item); return PY_FAIL; } + Py_DECREF(item); } - } else if(PyArray_ClipmodeConverter(object, &vals[0]) == PY_SUCCEED) { - for(i = 1; i < n; ++i) { - vals[i] = vals[0]; + } + else if (PyArray_ClipmodeConverter(object, &modes[0]) == PY_SUCCEED) { + for (i = 1; i < n; ++i) { + modes[i] = modes[0]; } - } else { + } + 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 index 931cb03395d3..4a2bdf14caea 100644 --- a/numpy/core/src/multiarray/new_iterator.c.src +++ b/numpy/core/src/multiarray/new_iterator.c.src @@ -2316,9 +2316,16 @@ 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. */ @@ -2361,15 +2368,15 @@ NpyIter_GetShape(NpyIter *iter, npy_intp *outshape) } /*NUMPY_API - * Builds a set of strides which are compatible with the iteration order - * for data packed contiguously, but not necessarily C or Fortran. - * The strides this produces are equivalent to the strides for an - * automatically allocated output created without an op_axes parameter, - * and should be used together with NpyIter_GetShape and NpyIter_GetNDim. + * 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, - * if you want to generate a vector per input value for a numerical gradient, + * 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 diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c index a7a736254ab0..bca2426e6adf 100644 --- a/numpy/lib/src/_compiled_base.c +++ b/numpy/lib/src/_compiled_base.c @@ -565,6 +565,14 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict) return NULL; } +/* + * Converts a Python sequence into 'count' PyArrayObjects + * + * seq - Input Python object, usually a tuple but any sequence works. + * op - Where the arrays are placed. + * count - How many arrays there should be (errors if it doesn't match). + * paramname - The name of the parameter that produced 'seq'. + */ static int sequence_to_arrays(PyObject *seq, PyArrayObject **op, int count, char *paramname) @@ -604,6 +612,7 @@ static int sequence_to_arrays(PyObject *seq, return 0; } +/* Inner loop for unravel_index */ static int ravel_coords_loop(int ravel_ndim, npy_intp *ravel_dims, npy_intp *ravel_strides, @@ -665,6 +674,7 @@ ravel_coords_loop(int ravel_ndim, npy_intp *ravel_dims, return NPY_SUCCEED; } +/* ravel_coords implementation - see add_newdocs.py */ static PyObject * arr_ravel_coords(PyObject *self, PyObject *args, PyObject *kwds) { @@ -802,6 +812,7 @@ arr_ravel_coords(PyObject *self, PyObject *args, PyObject *kwds) return NULL; } +/* C-order inner loop for unravel_index */ static int unravel_index_loop_corder(int unravel_ndim, npy_intp *unravel_dims, npy_intp unravel_size, npy_intp count, @@ -829,6 +840,7 @@ unravel_index_loop_corder(int unravel_ndim, npy_intp *unravel_dims, return NPY_SUCCEED; } +/* Fortran-order inner loop for unravel_index */ static int unravel_index_loop_forder(int unravel_ndim, npy_intp *unravel_dims, npy_intp unravel_size, npy_intp count, @@ -855,6 +867,7 @@ unravel_index_loop_forder(int unravel_ndim, npy_intp *unravel_dims, return NPY_SUCCEED; } +/* unravel_index implementation - see add_newdocs.py */ static PyObject * arr_unravel_index(PyObject *self, PyObject *args, PyObject *kwds) {