对 numpy 数组的许多随机排列进行采样的最快方法

Question

Unlike many other numpy/random functions, numpy.random.Generator.permutation() doesn't provide an obvious way to return multiple results in a single function call. Given a (1d) numpy array x , I want to sample n permutations of x (each of length len(x)), and have the result as a numpy array with shape (n, len(x)) . A naive way of generating many permutations is np.array([rng.permutation(x) for _ in range(n)]) . This is not ideal, mostly because the loop is in Python rather than inside a compiled numpy function.

import numpy as np

rng = np.random.default_rng(1234)
# x is big enough to not want to enumerate all permutations
x = rng.standard_normal(size=20)
n = 10000
perms = np.array([rng.permutation(x) for _ in range(n)])

My use case is for a brute-force search to find permutations that minimise a specific property (constituting a "good enough" search solution). I can calculate the property of interest for each permutation using numpy operations that vectorise/broadcast nicely over the resulting matrix of permutations. It turns out that naively generating the matrix of permutations is the bottleneck in my code. Is there a better way?

Answer 1

You can use rng.permuted instead of rng.permutation and combine it with np.tile so to repeat x multiple times and shuffle each replicates independently. Here is how:

perms = rng.permuted(np.tile(x, n).reshape(n,x.size), axis=1)

This is about 10 times faster on my machine than your initial code.

Answer 2

Just be aware that Jérome's solution provides an array of "n" rows but it may include repetitions. Different rows may have same "x" order (especially if "n" is bigger than "x")

If you need to sample without repetition (as it was my case), you can always do set(list(perm)) and keep unique combinations "x" values