ENH: Alternative to `random.shuffle`, with an `axis` argument.

It would be nice to have an alternative to numpy.random.shuffle that accepts an axis argument, and that independently shuffles the one-dimensional slices. Here’s an implementation that I’ll call disarrange. It works, but it would be nice to have a more efficient C implementation.

def disarrange(a, axis=-1):
    """
    Shuffle `a` in-place along the given axis.

    Apply numpy.random.shuffle to the given axis of `a`.
    Each one-dimensional slice is shuffled independently.
    """
    b = a.swapaxes(axis, -1)
    # Shuffle `b` in-place along the last axis.  `b` is a view of `a`,
    # so `a` is shuffled in place, too.
    shp = b.shape[:-1]
    for ndx in np.ndindex(shp):
        np.random.shuffle(b[ndx])
    return

Example:

In [156]: a = np.arange(20).reshape(4,5)

In [157]: a
Out[157]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19]])

In [158]: disarrange(a, axis=-1)

In [159]: a
Out[159]: 
array([[ 2,  0,  4,  3,  1],
       [ 8,  6,  7,  9,  5],
       [11, 14, 13, 10, 12],
       [19, 18, 16, 17, 15]])

In [160]: a = np.arange(20).reshape(4,5)

In [161]: disarrange(a, axis=0)

In [162]: a
Out[162]: 
array([[ 5, 11,  7, 13, 14],
       [ 0,  6,  2,  3,  4],
       [10,  1, 17, 18, 19],
       [15, 16, 12,  8,  9]])

This request was motivated by this question on stackoverflow: http://stackoverflow.com/questions/26310346/quickly-calculate-randomized-3d-numpy-array-from-2d-numpy-array/

[jklaise]

Any news on this? I was surprised this functionality doesn’t exist. For now I’m using np.apply_along_axis with np.random.permutation as a workaround.

[WarrenWeckesser]

(Note that while working on the examples here, I found a bug in the new shuffle code. In what follows, I am using the fix proposed in https://github.com/numpy/numpy/pull/14662, which has been merged.)

@wkschwartz, the change in #13829 is useful, but it is not the enhancement requested here. The axis added in #13829 still treats the array as a 1-d sequence to be shuffled. The new axis argument allows the user to specify which axis is viewed as the 1-d axis, but it does not do an independent shuffle within the axis.

For example,

In [1]: import numpy as np                                                      

In [2]: rng = np.random.default_rng()                                           

In [3]: x = np.arange(20).reshape(2, 10)                                        

In [4]: x                                                                       
Out[4]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

In [5]: rng.shuffle(x, axis=1)                                                  

In [6]: x                                                                       
Out[6]: 
array([[ 5,  9,  6,  4,  7,  0,  3,  2,  1,  8],
       [15, 19, 16, 14, 17, 10, 13, 12, 11, 18]])

You can see that the rows have not been independently shuffled. The columns have been rearranged, but the values within each column are the same.

The behavior requested in this issue is to shuffle independently, as in the disarrange code I gave above:

In [10]: x = np.arange(20).reshape(2, 10)                                       

In [11]: x                                                                      
Out[11]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

In [12]: disarrange(x, axis=1)                                                  

In [13]: x                                                                      
Out[13]: 
array([[ 4,  3,  7,  8,  0,  6,  5,  2,  9,  1],
       [12, 15, 19, 17, 18, 14, 10, 13, 11, 16]])