Efficient permutation of each row (or column) of a numpy array given a permutation matrix

Question

Lets assume, I have two given ndarrays, where the matrix mapping contains information, of how row of the matrix mask should be permuted. We may assume, that the mapping matrix comes from some other algorithm.

import numpy as np
T, K, F = 2, 3, 5
mask = np.random.randint(4, size=(T, K, F))
mapping = np.asarray([
    [0, 1, 2],
    [0, 1, 2],
    [2, 0, 1],
    [0, 1, 2],
    [1, 0, 2]
])

The straight forward way to do this, is by applying a for loop:

out = np.empty_like(mask)
for f in range(F):
    out[:, :, f] = mask[:, mapping[f, :], f]

This seems to be quite efficient, so I looked at Numpy advanced indexing and found this solution:

out = mask[
    np.arange(T)[:, None, None],
    mapping.T[None, :, :],
    np.arange(F)[None, None, :]
]

An answer to a related question suggests the use of ogrid:

ogrid = np.ogrid[:T, :1, :F]
out = mask[
    ogrid[0],
    mapping.T[None, :, :],
    ogrid[2]
]

It seems very uncomfortable to create all the intermediate arrays and broadcast them correctly. So what is the best way, to perform the desired reordering?

Timing information:

To provide meaningful timing information, I used some shapes, closer to my application. The random permutation is just for brevity of the example.

T, K, F = 1000, 3, 257
mask = np.random.randint(4, size=(T, K, F))
mapping = np.stack([list(np.random.permutation(np.arange(3))) for _ in range(F)])

Here are the results:

for loop:                 100 loops, best of 3: 8.4 ms per loop
three times broadcasting: 100 loops, best of 3: 8.37 ms per loop
ogrid:                    100 loops, best of 3: 8.33 ms per loop
swapaxis:                 100 loops, best of 3: 2.43 ms per loop
transpose:                100 loops, best of 3: 2.08 ms per loop

Divakar · Accepted Answer · 2017-06-29 05:50:53Z

1

Defining "best" is debatable, but here's one way with advanced-indexing -

mask[:,mapping, np.arange(F)[:,None]].swapaxes(1,2)

Another way would be to transpose mapping and then use the range array for the last axis without extending to 2D. Each row last axis (axis=-1) of mapping, decides the order of the elements along second last axis (axis=-2) of mask. So, we need that transpose on mapping. In the first approach, we achieved this transposed behaviour through the latter swapping of axes. I would vouch for this one on efficiency.

Thus, we would have the implementation, like so -

mask[:,mapping.T, np.arange(F)]

edited Jun 29, 2017 at 5:50

answered Jun 29, 2017 at 5:39

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Lukas Over a year ago

I do not exactly understand, why this is faster, but I added the timings above, and both of your versions are notably faster.

Collectives™ on Stack Overflow

Efficient permutation of each row (or column) of a numpy array given a permutation matrix

1 Answer 1

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related