[Numpy-discussion] ndrange, like range but multidimensiontal

Mark Harfouche mark.harfouche at gmail.com
Sun Oct 7 10:32:11 EDT 2018

Hi All,

I've been using numpy array objects to store collections of 2D (and soon
ND) variables. When iterating through these collections, I often found it
useful to use `ndindex`, which for `for loops` behaves much like `range`
with only a `stop` parameter.

That said, it lacks a few features that are now present in `range` are
missing from `ndindex`, most notably the ability to iterate over a subset
of the ndindex.

I found myself often writing `itertools.product(range(1, data.shapep[0]),
range(3, data.shape[2]))` for custom iterations. While it does flatten out
the for loop, it is arguable less readable than having 1 or 2 levels of
nested for loops.

It is quite possible that `nditer` would solve my problems, but
unfortunately I am still not able to make sense of then numerous options it

I propose an `ndrange` class that can be used to iterate over
nd-collections mimicking the API of `range` as much as possible and
adapting it to the ND case (i.e. returning tuples instead of singletons).

Since this is an enhancement proposal, I am bringing the idea to the
mailing list for reactions.

The implementation in this PR https://github.com/numpy/numpy/pull/12094 is
based on keeping track of a tuple of python `range` range objects. The
`__iter__` method returns the result of `itertools.product(*self._ranges)`

By leveraging python's `range` implementation, operations like
`containement` `index`, `reversed`, `equality` and most importantly slicing
of the ndrange object are possible to offer to the general numpy audiance.

For example, iterating through a 2D collection but avoiding indexing the
first and last column used to look like this:

c = np.empty((4, 4), dtype=object)
# ... compute on c
for j in range(c.shape[0]):
     for i in range(1, c.shape[1]-1):
         c[j, i] # = compute on c[j, i] that depends on the index i, j

With `np.ndrange` it can look something like this:

c = np.empty((4, 4), dtype=object)
# ... compute on c
for i in np.ndrange(c.shape)[:, 1:-1]:
    c[i] # = some operation on c[i] that depends on the index i

very pythonic, very familiar to numpy users

Thank you for the feedback,


An issue requesting expansion to the ndindex API on github:
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20181007/3611d2f3/attachment.html>

More information about the NumPy-Discussion mailing list