ENH: improve PyArray_Nonzero for sparse bool masks

juliantaylor · juliantaylor · commit 61d7a10b7d24 · 2014-03-09T21:32:53.000+01:00
we already count the number of true elements so we can decide when its
sparse and use faster (4-byte) vectorized npy_memchr.
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -2424,10 +2424,7 @@ PyArray_Nonzero(PyArrayObject *self)
     PyObject *ret_tuple;
     npy_intp ret_dims[2];
     PyArray_NonzeroFunc *nonzero = PyArray_DESCR(self)->f->nonzero;
-    char *data;
-    npy_intp stride, count;
     npy_intp nonzero_count;
-    npy_intp *multi_index;
 
     NpyIter *iter;
     NpyIter_IterNextFunc *iternext;
@@ -2454,23 +2451,47 @@ PyArray_Nonzero(PyArrayObject *self)
 
     /* If it's a one-dimensional result, don't use an iterator */
     if (ndim <= 1) {
-        npy_intp j;
+        npy_intp * multi_index = (npy_intp *)PyArray_DATA(ret);
+        char * data = PyArray_BYTES(self);
+        npy_intp stride = (ndim == 0) ? 0 : PyArray_STRIDE(self, 0);
+        npy_intp count = (ndim == 0) ? 1 : PyArray_DIM(self, 0);
 
-        multi_index = (npy_intp *)PyArray_DATA(ret);
-        data = PyArray_BYTES(self);
-        stride = (ndim == 0) ? 0 : PyArray_STRIDE(self, 0);
-        count = (ndim == 0) ? 1 : PyArray_DIM(self, 0);
+        /* nothing to do */
+        if (nonzero_count == 0) {
+            goto finish;
+        }
 
+        /* avoid function call for bool */
         if (PyArray_ISBOOL(self)) {
-            /* avoid function call for bool */
-            for (j = 0; j < count; ++j) {
-                if (*data != 0) {
-                    *multi_index++ = j;
+            /*
+             * use fast memchr variant for sparse data, see gh-4370
+             * the fast bool count is followed by this sparse path is faster
+             * than combining the two loops, even for larger arrays
+             */
+            if (((double)nonzero_count / count) <= 0.1) {
+                npy_intp subsize;
+                npy_intp j = 0;
+                while (1) {
+                    npy_memchr(data + j * stride, 0, stride, count, &subsize, 1);
+                    j += subsize;
+                    if (j >= count) {
+                        break;
+                    }
+                    *multi_index++ = j++;
+                }
+            }
+            else {
+                npy_intp j;
+                for (j = 0; j < count; ++j) {
+                    if (*data != 0) {
+                        *multi_index++ = j;
+                    }
+                    data += stride;
                 }
-                data += stride;
             }
         }
         else {
+            npy_intp j;
             for (j = 0; j < count; ++j) {
                 if (nonzero(data, self)) {
                     *multi_index++ = j;
@@ -2498,6 +2519,7 @@ PyArray_Nonzero(PyArrayObject *self)
     }
 
     if (NpyIter_GetIterSize(iter) != 0) {
+        npy_intp * multi_index;
         /* Get the pointers for inner loop iteration */
         iternext = NpyIter_GetIterNext(iter, NULL);
         if (iternext == NULL) {
@@ -2558,7 +2580,7 @@ PyArray_Nonzero(PyArrayObject *self)
     else {
         for (i = 0; i < ndim; ++i) {
             PyArrayObject *view;
-            stride = ndim*NPY_SIZEOF_INTP;
+            npy_intp stride = ndim * NPY_SIZEOF_INTP;
 
             view = (PyArrayObject *)PyArray_New(Py_TYPE(self), 1,
                                 &nonzero_count,
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
@@ -936,6 +936,20 @@ def test_nonzero_twodim(self):
         assert_equal(np.nonzero(x['a'].T), ([0, 1, 1, 2], [1, 1, 2, 0]))
         assert_equal(np.nonzero(x['b'].T), ([0, 0, 1, 2, 2], [0, 1, 2, 0, 2]))
 
+    def test_sparse(self):
+        # test special sparse condition boolean code path
+        for i in range(20):
+            c = np.zeros(200, dtype=np.bool)
+            c[i::20] = True
+            assert_equal(np.nonzero(c)[0], np.arange(i, 200 + i, 20))
+
+            c = np.zeros(400, dtype=np.bool)
+            c[10 + i:20 + i] = True
+            c[20 + i*2] = True
+            assert_equal(np.nonzero(c)[0],
+                         np.concatenate((np.arange(10 +i, 20 + i), [20 +i*2])))
+
+
 class TestIndex(TestCase):
     def test_boolean(self):
         a = rand(3, 5, 8)
