Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse files

ENH: Allow newaxis in combination with fancy indexing.

This commit allows the use of None when also using fancy indexing.
To achieve this mit->consec is used to note the denote the
transposition axis (which means that if its consecutive but at the
start, mit->consec==0, and no transposition is necessary).
  • Loading branch information...
commit 18ee645bf53185ff4f00ad91861356f2eb1222a0 1 parent 29c8ae8
@seberg authored
View
3  numpy/core/include/numpy/ndarraytypes.h
@@ -1275,7 +1275,8 @@ typedef struct {
npy_intp bscoord[NPY_MAXDIMS];
PyObject *indexobj; /* creating obj */
- int consec;
+ int consec; /* axes to swap if index
+ * arrays are consecutive */
char *dataptr;
} PyArrayMapIterObject;
View
49 numpy/core/src/multiarray/mapping.c
@@ -226,7 +226,8 @@ PyArray_MapIterSwapAxes(PyArrayMapIterObject *mit, PyArrayObject **ret, int getm
* (n2,...,n1+n2-1,0,...,n2-1,n1+n2,...n3-1)
*/
n1 = mit->iters[0]->nd_m1 + 1;
- n2 = mit->iteraxes[0];
+ /* We have to normalize for newaxes */
+ n2 = mit->consec;
n3 = mit->nd;
/* use n1 as the boundary if getting but n2 if setting */
@@ -303,9 +304,7 @@ PyArray_GetMap(PyArrayMapIterObject *mit)
/* check for consecutive axes */
if ((mit->subspace != NULL) && (mit->consec)) {
- if (mit->iteraxes[0] > 0) { /* then we need to swap */
- PyArray_MapIterSwapAxes(mit, &ret, 1);
- }
+ PyArray_MapIterSwapAxes(mit, &ret, 1);
}
return (PyObject *)ret;
}
@@ -338,11 +337,9 @@ PyArray_SetMap(PyArrayMapIterObject *mit, PyObject *op)
return -1;
}
if ((mit->subspace != NULL) && (mit->consec)) {
- if (mit->iteraxes[0] > 0) { /* then we need to swap */
- PyArray_MapIterSwapAxes(mit, &arr, 0);
- if (arr == NULL) {
- return -1;
- }
+ PyArray_MapIterSwapAxes(mit, &arr, 0);
+ if (arr == NULL) {
+ return -1;
}
}
@@ -1666,7 +1663,7 @@ _nonzero_indices(PyObject *myBool, PyArrayIterObject **iters)
}
/* convert an indexing object to an INTP indexing array iterator
- if possible -- otherwise, it is a Slice or Ellipsis object
+ if possible -- otherwise, it is a Slice, Ellipsis or None object
and has to be interpreted on bind to a particular
array so leave it NULL for now.
*/
@@ -1676,7 +1673,7 @@ _convert_obj(PyObject *obj, PyArrayIterObject **iter)
PyArray_Descr *indtype;
PyObject *arr;
- if (PySlice_Check(obj) || (obj == Py_Ellipsis)) {
+ if (PySlice_Check(obj) || (obj == Py_Ellipsis) || (obj == Py_None)) {
return 0;
}
else if (PyArray_Check(obj) && PyArray_ISBOOL((PyArrayObject *)obj)) {
@@ -1812,7 +1809,7 @@ PyArray_MapIterBind(PyArrayMapIterObject *mit, PyArrayObject *arr)
{
int subnd;
PyObject *sub, *obj = NULL;
- int i, j, n, curraxis, ellipexp, noellip;
+ int i, j, n, curraxis, ellipexp, noellip, subdim, newaxes;
PyArrayIterObject *it;
npy_intp dimsize;
npy_intp *indptr;
@@ -1841,7 +1838,8 @@ PyArray_MapIterBind(PyArrayMapIterObject *mit, PyArrayObject *arr)
* all indexing arrays have been converted to 0
* therefore we can extract the subspace with a simple
* getitem call which will use view semantics, but
- * without index checking.
+ * without index checking since all original normal
+ * indexes are checked later as fancy ones.
*
* But, be sure to do it with a true array.
*/
@@ -1867,19 +1865,21 @@ PyArray_MapIterBind(PyArrayMapIterObject *mit, PyArrayObject *arr)
goto fail;
}
/* Expand dimensions of result */
- n = PyArray_NDIM(mit->subspace->ao);
- for (i = 0; i < n; i++) {
+ subdim = PyArray_NDIM(mit->subspace->ao);
+ for (i = 0; i < subdim; i++) {
mit->dimensions[mit->nd+i] = PyArray_DIMS(mit->subspace->ao)[i];
}
- mit->nd += n;
+ mit->nd += subdim;
/*
- * Now, we still need to interpret the ellipsis and slice objects
- * to determine which axes the indexing arrays are referring to
+ * Now, we still need to interpret the ellipsis, slice and None
+ * objects to determine which axes the indexing arrays are
+ * referring to
*/
n = PyTuple_GET_SIZE(mit->indexobj);
/* The number of dimensions an ellipsis takes up */
- ellipexp = PyArray_NDIM(arr) - n + 1;
+ newaxes = subdim - (PyArray_NDIM(arr) - mit->numiter);
+ ellipexp = PyArray_NDIM(arr) + newaxes - n + 1;
/*
* Now fill in iteraxes -- remember indexing arrays have been
* converted to 0's in mit->indexobj
@@ -1888,7 +1888,7 @@ PyArray_MapIterBind(PyArrayMapIterObject *mit, PyArrayObject *arr)
j = 0;
/* Only expand the first ellipsis */
noellip = 1;
- memset(mit->bscoord, 0, sizeof(npy_intp)*PyArray_NDIM(arr));
+ memset(mit->bscoord, 0, sizeof(npy_intp)*(PyArray_NDIM(arr)));
for (i = 0; i < n; i++) {
/*
* We need to fill in the starting coordinates for
@@ -1902,6 +1902,11 @@ PyArray_MapIterBind(PyArrayMapIterObject *mit, PyArrayObject *arr)
curraxis += ellipexp;
noellip = 0;
}
+ else if (obj == Py_None) {
+ if (j == 0 && mit->consec) {
+ mit->consec += 1; /* count newaxes before iter axes */
+ }
+ }
else {
npy_intp start = 0;
npy_intp stop, step;
@@ -1926,6 +1931,10 @@ PyArray_MapIterBind(PyArrayMapIterObject *mit, PyArrayObject *arr)
curraxis += 1;
}
}
+ if (mit->consec) {
+ /* if set, it was 1 + newaxes before first iter axis */
+ mit->consec += mit->iteraxes[0] - 1;
+ }
finish:
/* Here check the indexes (now that we have iteraxes) */
Please sign in to comment.
Something went wrong with that request. Please try again.