I have an n row, m column numpy array, and would like to create a new k x m array by selecting k random elements from each column of the array. I wrote the following python function to do this, but would like to implement something more efficient and faster:

def sample_array_cols(MyMatrix, nelements):
vmat = []
TempMat = MyMatrix.T
for v in TempMat:
    v = np.ndarray.tolist(v)
    subv = random.sample(v, nelements)
    vmat = vmat   [subv]
return(np.array(vmat).T) 

One question is whether there's a way to loop over each column without transposing the array (and then transposing back). More importantly, is there some way to map the random sample onto each column that would be faster than having a for loop over all columns? I don't have that much experience with numpy objects, but I would guess that there should be something analogous to apply/mapply in R that would work?

CodePudding user response：

One alternative is to randomly generate the indices first, and then use take_along_axis to map them to the original array:

arr = np.random.randn(1000, 5000)  # arbitrary
k = 10  # arbitrary
n, m = arr.shape
idx = np.random.randint(0, n, (m, k))
new = np.take_along_axis(arr, idx, axis=0)

Output (shape):

in [215]: new.shape    
out[215]: (10, 500)  # (k x m)

CodePudding user response：

I would

1. Pre-allocate the result array, and fill in columns, and
2. Use numpy index based indexing

def sample_array_cols(matrix, n_result):
    (n,m) = matrix.shape
    vmat = numpy.array([n_result, m], dtype= matrix.dtype)
    for c in range(m):
        random_indices = numpy.random.randint(0, n, n_result)
        vmat[:,c] = matrix[random_indices, c]
    return vmat

Not quite fully vectorized, but better than building up a list, and the code scans just like your description.