Python Numpy 以不同的方式生成伯努利试验(或为不同的行随机 Select 不同的 p)

我想加速下面的代码，即for循环.有没有办法在numpy做到这一点？

import numpy as np
# define seend and random state
rng = np.random.default_rng(0)
# num of throws
N = 10**1
# max number of trials
total_n_trials = 10
# generate the throws' distributions of "performace" - overall scores
# mu_throws = rng.normal(0.7, 0.1, N)
mu_throws = np.linspace(0,1,N)

all_trials = np.zeros(N*total_n_trials)
for i in range(N):
    # generate trials per throw as Bernoulli trials with given mean
    all_trials[i*total_n_trials:(i+1)*total_n_trials] = rng.choice([0,1], size=total_n_trials, p=[1-mu_throws[i],mu_throws[i]])
    

更多解释-我想生成N个伯努利试验序列(即0和1，称为throws)，其中每个序列都有一个由另一个数组(mu_throws)中的值给出的平均值(概率p).这可以从正态分布中采样，或者在本例中，为了简单起见，我将其看作是一个从0到1的N=10个数字的序列.上述方法可行，但速度较慢.我希望N至少是10^4美元，然后total_n_trials可以是数百到(数万)千的任何东西(这会运行多次).我判断了以下post个，但没有找到答案.我还知道numpy random choice可以生成多维数组，但我没有找到为不同的行设置不同的p集的方法.基本上与我在上面所做的一样，只是reshape 了:

all_trials.reshape(N,-1)

array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 1., 1., 0., 0.],
       [1., 0., 0., 1., 0., 0., 0., 1., 1., 0.],
       [1., 0., 1., 0., 0., 1., 0., 1., 0., 1.],
       [1., 0., 1., 0., 0., 0., 1., 1., 0., 0.],
       [1., 0., 0., 1., 0., 1., 0., 1., 1., 0.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

我也想过这trick个，但还没有想到如何将其用于伯努利试验.谢谢你的提示.