[Numpy-discussion] fast iteration (I think I've got it)

Neal Becker ndbecker2@gmail....
Tue Jan 1 13:07:33 CST 2008


This is a c-api question.

I'm trying to get iterators that are both fast and reasonably general.  I
did confirm that iterating using just the general PyArrayIterObject
protocol is not as fast as using c-style pointers for contiguous arrays.

Please confirm if my understanding is correct.  There are 2 cases to
consider.  There are plain old dense arrays, which will be contiguous, and
there are array 'views' or 'slices'.  The views will not be contiguous. 
However, in all cases, the strides between elements in one dimension are
constant.  This last point is key.  As long as the stride in the last
dimension is a constant, a fast iterator is possible.

I put together a small test, which seems to work for the cases I tried.
The idea is to use the general iterator (via PyArray_IterAllButAxis), but
iteration over the last dimension is done with a c-style pointer.

I have tested it with these cases, and the results appear correct:

a = array ((xrange(4),xrange(4,8),xrange(8,12)), int)
b = a[:,::2]
c = a[:,1::2]
d = a[:,-1::-1]
sum3(a)
0 1 2 3 
4 5 6 7 
8 9 10 11 
sum3(b)
0 2 
4 6 
8 10 
sum3(c) 
1 3 
5 7 
9 11 
sum3 (d)
3 2 1 0 
7 6 5 4 
11 10 9 8 

template<typename T>
inline T sum3 (object const& o) {
  PyArrayObject* arr = (PyArrayObject*)(o.ptr());
  int last_dim = arr->nd - 1;
  PyArrayIterObject* it = (PyArrayIterObject*)(PyArray_IterAllButAxis
(o.ptr(), &last_dim));

  for (int i = 0 ; i < arr->nd; ++i)
    std::cout << arr->dimensions[i] << ' ';
  std::cout << '\n';
  while (it->index < it->size) {
    T * dptr = (T*)(it->dataptr);
    const int stride = arr->strides[last_dim];
    for (int i = 0; i < arr->dimensions[last_dim]; ++i) {
      std::cout << *dptr << ' ';
      dptr += stride/sizeof(T);
    }
    std::cout << '\n';
      
    PyArray_ITER_NEXT(it);
  }
  return 0;
}




More information about the Numpy-discussion mailing list