Create an array where each element stores its indices

Question

I want to create a 2d numpy array where every element is a tuple of its indices.

Example (4x5):

array([[[0, 0],
        [0, 1],
        [0, 2],
        [0, 3],
        [0, 4]],

       [[1, 0],
        [1, 1],
        [1, 2],
        [1, 3],
        [1, 4]],

       [[2, 0],
        [2, 1],
        [2, 2],
        [2, 3],
        [2, 4]],

       [[3, 0],
        [3, 1],
        [3, 2],
        [3, 3],
        [3, 4]]])

I would create an python list with the following list comprehension:

[[(y,x) for x in range(width)] for y in range(height)]

Is there a faster way to achieve the same, maybe with numpy methods?

Technically you are showing a 3d array (4,5,2). – hpaulj Feb 17 '17 at 17:26 — hpaulj, Feb 17 '17 at 17:26

Paul Panzer · Answer 1 · 2017-02-17T19:33:22.087

9

Do you do this because you need it or just for sport? In the former case:

np.moveaxis(np.indices((4,5)), 0, -1)

np.indices does precisely what its name suggests. Only it arranges axes differently to you. So we move them with moveaxis

As @Eric points out one attractive feature of this method is that it works unmodified at arbitrary number of dimensions:

dims = tuple(np.multiply.reduceat(np.zeros(16,int)+2, np.r_[0, np.sort(np.random.choice(16, np.random.randint(10)))]))
# len(dims) == ?
np.moveaxis(np.indices(dims), 0, -1) # works

edited Feb 17 '17 at 19:33

answered Feb 17 '17 at 16:12

Paul Panzer

51,835
3
54
99

This is the only answer here that works for `ndim != 2` – Eric Feb 17 '17 at 16:35
Or with the base transpose method: `np.indices((4,5)).transpose(1,2,0)` – hpaulj Feb 17 '17 at 17:24
@hpaulj as Eric points out `moveaxis(..., 0, -1)` works for any number of dimensions. With the `transpose` you might find yourself reduced to `roll`ing an `arange` just to set up the shuffle axes ;-) – Paul Panzer Feb 17 '17 at 17:30

Divakar · Accepted Answer · 2017-02-17T17:46:10.300

6

Here's an initialization based method -

def create_grid(m,n):
    out = np.empty((m,n,2),dtype=int) #Improvement suggested by @AndrasDeak
    out[...,0] = np.arange(m)[:,None]
    out[...,1] = np.arange(n)
    return out

Sample run -

In [47]: create_grid(4,5)
Out[47]: 
array([[[0, 0],
        [0, 1],
        [0, 2],
        [0, 3],
        [0, 4]],

       [[1, 0],
        [1, 1],
        [1, 2],
        [1, 3],
        [1, 4]],

       [[2, 0],
        [2, 1],
        [2, 2],
        [2, 3],
        [2, 4]],

       [[3, 0],
        [3, 1],
        [3, 2],
        [3, 3],
        [3, 4]]])

Runtime test for all approaches posted thus far on (4,5) grided and bigger sizes -

In [111]: %timeit np.moveaxis(np.indices((4,5)), 0, -1)
     ...: %timeit np.mgrid[:4, :5].swapaxes(2, 0).swapaxes(0, 1)
     ...: %timeit np.mgrid[:4,:5].transpose(1,2,0)
     ...: %timeit create_grid(4,5)
     ...: 
100000 loops, best of 3: 11.1 µs per loop
100000 loops, best of 3: 17.1 µs per loop
100000 loops, best of 3: 17 µs per loop
100000 loops, best of 3: 2.51 µs per loop

In [113]: %timeit np.moveaxis(np.indices((400,500)), 0, -1)
     ...: %timeit np.mgrid[:400, :500].swapaxes(2, 0).swapaxes(0, 1)
     ...: %timeit np.mgrid[:400,:500].transpose(1,2,0)
     ...: %timeit create_grid(400,500)
     ...: 
1000 loops, best of 3: 351 µs per loop
1000 loops, best of 3: 1.01 ms per loop
1000 loops, best of 3: 1.03 ms per loop
10000 loops, best of 3: 190 µs per loop

edited Feb 17 '17 at 17:46

answered Feb 17 '17 at 16:35

Divakar

218,885
19
262
358

`.swapaxes(-1, 0)` gives a different output, so comparing the timing is not very meaningfull – Eric Feb 17 '17 at 16:42
That second `:,:` could be replaced with a `...` to match the first, right? – Eric Feb 17 '17 at 16:45
@Eric Yup, had to fix on that runtime test. – Divakar Feb 17 '17 at 16:47
The implementation of `np.indices` looks pretty weird to me, which explains why your solution is so much faster – Eric Feb 17 '17 at 16:49
Ok, I've [submitted a patch](https://github.com/numpy/numpy/pull/8629) that should bring these timings more in line – Eric Feb 17 '17 at 17:13
@Eric Awesome ! – Divakar Feb 17 '17 at 17:14
Would `np.empty` vs `np.zeros` make any difference? – Andras Deak -- Слава Україні Feb 17 '17 at 17:40
Update: it indeed seems to, `465 mus` -> `317 mus` on my system if I change `np.zeros` to `np.empty`. – Andras Deak -- Слава Україні Feb 17 '17 at 17:45
1

@AndrasDeak Yup! Some noticeable improvement, thanks a lot! :) – Divakar Feb 17 '17 at 17:46
1

yay, my answer isn't the slowest :-D Probably not worth it to create a numba-function here but it's tempting if it could beat yours. :) – MSeifert Feb 17 '17 at 18:13

score 5 · Answer 3 · answered Feb 17 '17 at 16:11

You can abuse numpy.mgrid or meshgrid for this purpose:

>>> import numpy as np
>>> np.mgrid[:4,:5].transpose(1,2,0)
array([[[0, 0],
        [0, 1],
        [0, 2],
        [0, 3],
        [0, 4]],

       [[1, 0],
        [1, 1],
        [1, 2],
        [1, 3],
        [1, 4]],

       [[2, 0],
        [2, 1],
        [2, 2],
        [2, 3],
        [2, 4]],

       [[3, 0],
        [3, 1],
        [3, 2],
        [3, 3],
        [3, 4]]])

MSeifert · Answer 4 · 2017-02-17T18:30:55.580

You can use numpy.mgrid and swap it's axes:

>>> # assuming a 3x3 array
>>> np.mgrid[:3, :3].swapaxes(-1, 0)
array([[[0, 0],
        [1, 0],
        [2, 0]],

       [[0, 1],
        [1, 1],
        [2, 1]],

       [[0, 2],
        [1, 2],
        [2, 2]]])

That still differs a bit from your desired array so you can roll your axes:

>>> np.mgrid[:3, :3].swapaxes(2, 0).swapaxes(0, 1)
array([[[0, 0],
        [0, 1],
        [0, 2]],

       [[1, 0],
        [1, 1],
        [1, 2]],

       [[2, 0],
        [2, 1],
        [2, 2]]])

Given that someone timed the results I also want to present a manual numba based version that "beats 'em all":

import numba as nb
import numpy as np

@nb.njit
def _indexarr(a, b, out):
    for i in range(a):
        for j in range(b):
            out[i, j, 0] = i
            out[i, j, 1] = j
    return out

def indexarr(a, b):
    arr = np.empty([a, b, 2], dtype=int)
    return _indexarr(a, b, arr)

Timed:

a, b = 400, 500
indexarr(a, b)  # numba needs a warmup run
%timeit indexarr(a, b)                                  # 1000 loops, best of 3: 1.5 ms per loop
%timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1)  # 100 loops, best of 3: 7.17 ms per loop
%timeit np.mgrid[:a, :b].transpose(1,2,0)               # 100 loops, best of 3: 7.47 ms per loop
%timeit create_grid(a, b)                               # 100 loops, best of 3: 2.26 ms per loop

and on a smaller array:

a, b = 4, 5
indexarr(a, b)
%timeit indexarr(a, b)                                 # 100000 loops, best of 3: 13 µs per loop
%timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1) # 10000 loops, best of 3: 181 µs per loop
%timeit np.mgrid[:a, :b].transpose(1,2,0)              # 10000 loops, best of 3: 182 µs per loop
%timeit create_grid(a, b)                              # 10000 loops, best of 3: 32.3 µs per loop

As promised it "beats 'em all" in terms of performance :-)

Interesting that `swapaxes` is a method but `moveaxis` is not — Eric, Feb 17 '17 at 16:37
Well, never fails! I don't think it's fair though comparing against NumPy ones ;) — Divakar, Feb 17 '17 at 18:46
Nice one! Now you only have to generalise to arbitrary ndim ;-) — Paul Panzer, Feb 17 '17 at 18:52
@PaulPanzer I think almost all answers have some problem with arbitary dimensions (especially for `ndim >= 3` it's not really clear in what order the indices should appear). — MSeifert, Feb 17 '17 at 19:13
What's wrong with row-major? Or should that be hyper-row-major? — Paul Panzer, Feb 17 '17 at 19:32
@PaulPanzer Well currently `[[(y,x) for x in range(width)] for y in range(height)]` could mean that the indices should be reversed _or_ that first and the last axis should be moved to first. In the second case your answer is correct but that still leaves the first case. — MSeifert, Feb 17 '17 at 20:22

Create an array where each element stores its indices

4 Answers4

Linked