Shuffling NumPy array along a given axis

by

Vectorized solution with rand+argsort trick

We could generate unique indices along the specified axis and index into the the input array with advanced-indexing. To generate the unique indices, we would use random float generation + sort trick, thus giving us a vectorized solution. We would also generalize it to cover generic n-dim arrays and along generic axes with np.take_along_axis. The final implementation would look something like this –

def shuffle_along_axis(a, axis):
    idx = np.random.rand(*a.shape).argsort(axis=axis)
    return np.take_along_axis(a,idx,axis=axis)

Note that this shuffle won’t be in-place and returns a shuffled copy.

Sample run –

In [33]: a
Out[33]: 
array([[18, 95, 45, 33],
       [40, 78, 31, 52],
       [75, 49, 42, 94]])

In [34]: shuffle_along_axis(a, axis=0)
Out[34]: 
array([[75, 78, 42, 94],
       [40, 49, 45, 52],
       [18, 95, 31, 33]])

In [35]: shuffle_along_axis(a, axis=1)
Out[35]: 
array([[45, 18, 33, 95],
       [31, 78, 52, 40],
       [42, 75, 94, 49]])

More Related Contents:

  1. Consistently create same random numpy array
  2. How to specify upper and lower limits when using numpy.random.normal
  3. Binary random array with a specific proportion of ones?
  4. How to get a normal distribution within a range in numpy? [duplicate]
  5. Generate a random sample of points distributed on the surface of a unit sphere
  6. Difference between np.random.seed() and np.random.RandomState()
  7. Should I use `random.seed` or `numpy.random.seed` to control random number generation in `scikit-learn`?
  8. Creating same random number sequence in Python, NumPy and R
  9. Generating Discrete random variables with specified weights using SciPy or NumPy
  10. Better way to shuffle two numpy arrays in unison
  11. How to generate a random normal distribution of integers
  12. Weighted random sample without replacement in python
  13. Vectorizing `numpy.random.choice` for given 2D array of probabilities along an axis
  14. Working with floating point NumPy arrays for comparison and related operations
  15. Generate random locations within a triangular domain
  16. Generate random array of floats between a range
  17. Use a.any() or a.all()
  18. Restart cumsum and get index if cumsum more than value
  19. ValueError: numpy.dtype has the wrong size, try recompiling
  20. subsampling every nth entry in a numpy array
  21. Python’s sum vs. NumPy’s numpy.sum
  22. Pass percentiles to pandas agg function
  23. Weighted choice short and simple [duplicate]
  24. Max in a sliding window in NumPy array
  25. Get first and second highest values in pandas columns
  26. ImportError: No module named ‘pandas.indexes’
  27. Form a big 2d array from multiple smaller 2d arrays
  28. Fit mixture of two gaussian/normal distributions to a histogram from one set of data, python
  29. What’s the best way to generate random strings of a specific length in Python?
  30. NumPy library ImportError: DLL load failed: The specified procedure could not be found

Leave a Comment