python - 具有不同概率的二项式分布(伯努利试验)

Question

我想加快下面的代码 - 即 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)，其中每个序列都有一个平均值(概率 p) 由另一个数组 (mu_throws) 中的值给出。这可以从正态分布中采样，或者在这种情况下，为了简单起见，我将其作为从 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但还没有想出如何将它用于伯努利试验。感谢您提供任何提示。

Answer 1

同时 np.random.Generator.binomial() (或遗留 np.random.binomial() )可能是计算您所追求的内容的最简洁方法(如 @MechanicPig's answer 中所建议)，看起来它不是最快的。

实际上，使用您指定的与您相关的输入大小，OP 的方法甚至可以更快。

相反，基于对均匀分布的随机数组进行阈值处理的方法(本质上是 @SalvatoreDanieleBianco's answer 的精良版本)似乎是最快的。

此外，较新的随机生成器似乎速度更快。

最终，阈值方法可以在 Numba 中实现，但这并没有带来显着的加速。

以下是一些等效的方法:

import numpy as np


def bin_samples_OP(n, mu, seed=0):
    m = len(mu)
    rng = np.random.default_rng(seed)
    result = np.zeros((m, n), dtype=np.uint8)
    for i in range(m):
        result[i, :] = rng.choice([0, 1], size=n, p=[1 - mu[i], mu[i]])
    return result

import numpy as np


def bin_samples_binom(n, mu, seed=0):
    np.random.seed(seed)
    return np.random.binomial(1, mu[:, None], (len(mu), n))

import numpy as np


def bin_samples_binom2(n, mu, seed=0):
    rng = np.random.default_rng(seed)
    return rng.binomial(1, mu[:, None], (len(mu), n))

import numpy as np


def bin_samples_rand(n, mu, seed=0):
    np.random.seed(seed)
    return (np.random.random_sample((len(mu), n)) < mu[:, None]).astype(np.uint8)

import numpy as np


def bin_samples_rand2(n, mu, seed=0):
    rng = np.random.default_rng(seed)
    return (rng.random(size=(len(mu), n)) < mu[:, None]).astype(np.uint8)

import numpy as np


def bin_samples_rand3(n, mu, seed=0):
    rng = np.random.default_rng(seed)
    return (rng.random(size=(len(mu), n)).T < mu).astype(np.uint8).T

import numpy as np
import numba as nb
import random


@nb.njit
def bin_samples_nb(n, mu, seed=0):
    m = len(mu)
    random.seed(seed)
    result = np.empty((m, n), dtype=np.uint8)
    for i in range(m):
        for j in range(n):
            result[i, j] = random.random() < mu[i]
    return result

下面是一些测试，看看它们是如何等价的:

funcs = (
    bin_samples_OP, bin_samples_binom, bin_samples_binom2,
    bin_samples_rand, bin_samples_rand2, bin_samples_rand3,
    bin_samples_nb)


n = 10
m = 5
mu = np.linspace(0, 1, m)

print(f"{'Target':>24}  {mu}")
for i, func in enumerate(funcs, 1):
    res = func(n, mu)
    mu_exp = np.mean(res, -1)
    is_good = np.isclose(mu, mu_exp, atol=0.1, rtol=0.1).astype(np.uint8)
    print(f"{func.__name__:>24}  {mu_exp}  {is_good}")

                  Target  [0.   0.25 0.5  0.75 1.  ]
          bin_samples_OP  [0.   0.31 0.53 0.73 1.  ]  [1 1 1 1 1]
       bin_samples_binom  [0.   0.22 0.5  0.77 1.  ]  [1 1 1 1 1]
      bin_samples_binom2  [0.   0.31 0.56 0.68 1.  ]  [1 1 1 1 1]
        bin_samples_rand  [0.   0.22 0.5  0.77 1.  ]  [1 1 1 1 1]
       bin_samples_rand2  [0.   0.22 0.47 0.77 1.  ]  [1 1 1 1 1]
       bin_samples_rand3  [0.   0.22 0.47 0.77 1.  ]  [1 1 1 1 1]
          bin_samples_nb  [0.   0.35 0.53 0.78 1.  ]  [1 1 1 1 1]
         bin_samples_pnb  [0.   0.22 0.5  0.78 1.  ]  [1 1 1 1 1]

(bin_samples_pnb() 与 bin_sampls_nb() 相同，但带有 @nb.njit(parallel=True) 和 range() 替换为 nb.prange())

以及它们各自的输入大小更接近您指定的时间:

timings = {}
for p in range(12, 14):
    for q in range(12, 15):
        n = 2 ** p
        m = 2 ** q
        mu = np.linspace(0, 1, m)
        print(f"n={n}, m={m};  (p={p}, q={q})")
        timings[n, m] = []
        for func in funcs:
            res = func(n, mu)
            mu_exp = np.mean(res, axis=-1)
            
            is_good = np.allclose(mu, mu_exp, atol=0.1, rtol=0.1)
            timed = %timeit -r 1 -n 1 -o -q func(n, mu)
            timings[n, m].append(timed.best)
            print(f"{func.__name__:>24}  {is_good}  {float(timed.best * 10 ** 3):>10.4f} ms")

n=4096, m=4096;  (p=12, q=12)
          bin_samples_OP  True    533.5690 ms
       bin_samples_binom  True    517.4861 ms
      bin_samples_binom2  True    487.3975 ms
        bin_samples_rand  True    202.7566 ms
       bin_samples_rand2  True    135.5819 ms
       bin_samples_rand3  True    139.0170 ms
          bin_samples_nb  True    110.7469 ms
         bin_samples_pnb  True    141.4636 ms
n=4096, m=8192;  (p=12, q=13)
          bin_samples_OP  True   1031.1586 ms
       bin_samples_binom  True   1035.9100 ms
      bin_samples_binom2  True    957.9257 ms
        bin_samples_rand  True    392.7915 ms
       bin_samples_rand2  True    246.8801 ms
       bin_samples_rand3  True    250.9648 ms
          bin_samples_nb  True    222.7287 ms
         bin_samples_pnb  True    270.5297 ms
n=4096, m=16384;  (p=12, q=14)
          bin_samples_OP  True   2055.6572 ms
       bin_samples_binom  True   2036.8609 ms
      bin_samples_binom2  True   1865.3614 ms
        bin_samples_rand  True    744.9770 ms
       bin_samples_rand2  True    493.8508 ms
       bin_samples_rand3  True    476.5701 ms
          bin_samples_nb  True    449.1700 ms
         bin_samples_pnb  True    523.8716 ms
n=8192, m=4096;  (p=13, q=12)
          bin_samples_OP  True    886.8198 ms
       bin_samples_binom  True   1021.3093 ms
      bin_samples_binom2  True    946.2922 ms
        bin_samples_rand  True    372.9902 ms
       bin_samples_rand2  True    221.8423 ms
       bin_samples_rand3  True    234.7383 ms
          bin_samples_nb  True    218.5655 ms
         bin_samples_pnb  True    276.3884 ms
n=8192, m=8192;  (p=13, q=13)
          bin_samples_OP  True   1744.0382 ms
       bin_samples_binom  True   2101.7884 ms
      bin_samples_binom2  True   1985.6555 ms
        bin_samples_rand  True    720.7352 ms
       bin_samples_rand2  True    462.9820 ms
       bin_samples_rand3  True    453.3408 ms
          bin_samples_nb  True    455.7062 ms
         bin_samples_pnb  True    541.0242 ms
n=8192, m=16384;  (p=13, q=14)
          bin_samples_OP  True   3490.9539 ms
       bin_samples_binom  True   4106.5607 ms
      bin_samples_binom2  True   3719.5692 ms
        bin_samples_rand  True   1363.6868 ms
       bin_samples_rand2  True    884.4729 ms
       bin_samples_rand3  True    868.3783 ms
          bin_samples_nb  True    888.8390 ms
         bin_samples_pnb  True   1030.0389 ms

这些表明 bin_samples_rand2()/bin_samples_rand3()(它们的性能基本相同)和 bin_samples_nb() 处于相同的范围内。

在 Numba 中强制并行化不会有助于提高速度(因为 LLVM 可以自行产生更优化的速度，而不是依赖于 Numba 的并行化模型)并且它还会使可预测性变得松散。