[FIXED] randomly sample from a high dimensional array along with a specific dimension

Issue

There has a 3-dimensional array x of shape (2000,60,5). If we think it represents a video, the 2000 can represent 2000 frames. I would like to randomly sample it along with the first dimension, i.e., get a set of frame samples. For instance, how to get an array of (500,60,5) which is randomly sampled from x along with the first dimension?

Solution

You can pass x as the first argument of the choice method. If you don't want repeated frames in your sample, use replace=False.

For example,

In [10]: x = np.arange(72).reshape(9, 2, 4)  # Small array for the demo.

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],
        [20, 21, 22, 23]],

       [[24, 25, 26, 27],
        [28, 29, 30, 31]],

       [[32, 33, 34, 35],
        [36, 37, 38, 39]],

       [[40, 41, 42, 43],
        [44, 45, 46, 47]],

       [[48, 49, 50, 51],
        [52, 53, 54, 55]],

       [[56, 57, 58, 59],
        [60, 61, 62, 63]],

       [[64, 65, 66, 67],
        [68, 69, 70, 71]]])

Sample "frames" from x with the choice method of NumPy random generator instance.

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

In [13]: rng.choice(x, size=3)
Out[13]: 
array([[[40, 41, 42, 43],
        [44, 45, 46, 47]],

       [[40, 41, 42, 43],
        [44, 45, 46, 47]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])


In [14]: rng.choice(x, size=3, replace=False)
Out[14]: 
array([[[ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[32, 33, 34, 35],
        [36, 37, 38, 39]],

       [[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]]])

Note that the frames will be in random order; if you want to preserve the order, you could use choice to generate an array of indices, then use the sorted indices to pull the frames out of x.

Answered By - Warren Weckesser