python - 如何在不使用循环的情况下为 3D numpy 数组中的每个值做出 N 个随机选择

Question

我有:

• 猫，一个包含 10 个类别的数组，形状为 (10,)
• probs，形状为 (10, 50) 的概率数组，代表每个类别被 50 个不同变量选中的几率
• n_choices，一个形状为 (num_sims, 50) 的数组，包含整数，表示要为每个变量选择替换的类别数。例如，这可能是变量 1 的 0 个选择，变量 2 的 33 个选择等
• sims，一个用形状为 (num_sims, 50, 10) 的零填充的数组，稍后将填充结果

我想做的是如下:

• 对于数组中的每一行(代表一个模拟)，以及该行中的每个变量，从“cats”中做出 N 个选择，其中 N 等于“n_choices”中的相应值
• 做出选择后，每次选择类别时，“sims”加 1。换句话说，我想根据'probs'从'n_choices'中分配出10个类别的值，并将结果保存到'sims'

目前，我已设法使用循环实现此功能，如下所示。这对于少量 sims 来说很好，但实际上 num_sims 将达到数千，这意味着我的代码太慢了。

def allocate_N(N, var_index):
  """Make N choices from cats for a given variable, and return
  the incides of each category
  var_index is the position of the variable in n_choices"""
  allocation = np.random.choice(cats, size=N, p=probs[:, var_index])
  allocation_sorted = np.argsort(cats)
  ypos = np.searchsorted(cats[allocation_sorted], allocation)
  cat_indices = allocation_sorted[ypos]
  return cat_indices

def add_to_sim(sims, cat_indices, var_index):
  """Takes the category indices from allocate_n and adds 1 to
  sims at the corresponding location for each occurrence of
  the category in cat_indices"""
  from collections import Counter
  a = Counter(list(cat_indices))
  vals = [1*a[j] for j in cat_indices]
  pos = [(var_index, x) for x in cat_indices]
  sims[tuple(np.transpose(pos))] = vals

# For each variable and each row in sims, make N allocations
# and add results to 'sims'
for var_index in range(len(n_choices.T)):
  sim_count = 0
  # slice is (vars x cats), a single row of 'sims'
  for slice in sims:      
    N = n_choices[sim_count, var_index]
    if N > 0:
      cat_indices = allocate_N(N, var_index)
      add_to_sim(slice, cat_indices, var_index)
    sim_count += 1

我确定一定有办法对其进行矢量化？我能够使用 here 方法同时为每个变量做出一个随机选择。 ，但我不确定如何将其应用于我的特定问题。

感谢您的帮助!

Answer 1

您似乎描述的是 multinomial distribution 的样本.您可以直接从分布中获取样本。不幸的是，每个模拟和变量的分布参数(试验次数和概率)都会发生变化，np.random.multinomial 和 scipy.stats.multinomial 都不允许具有多组参数的矢量化采样。这意味着，如果你想这样做，你仍然必须用循环来做。至少，您的代码可以简化为以下内容:

import numpy as np

np.random.seed(0)
# Problem size
n_cats = 10
n_vars = 50
n_sims = 100
n_maxchoices = 50
# Make example problem
probs = np.random.rand(n_cats, n_vars)
probs /= probs.sum(0)
n_choices = np.random.randint(n_maxchoices, size=(n_sims, n_vars))
sims = np.zeros((n_sims, n_vars, n_cats), np.int32)
# Sample multinomial distribution for each simulation and variable
for i_sim in range(n_sims):
    for i_var in range(n_vars):
        sims[i_sim, i_var] = np.random.multinomial(n_choices[i_sim, i_var],
                                                   probs[:, i_var])
# Check number of choices per simulation and variable is correct
print(np.all(sims.sum(2) == n_choices))
# True

请注意，如果您愿意使用 Numba，您仍然可以使它更快，具有如下功能:

import numpy as np
import numba as nb

@nb.njit(parallel=True)
def make_simulations(probs, n_choices, sims):
    for i_sim in nb.prange(n_sims):
        for i_var in nb.prange(n_vars):
            sims[i_sim, i_var] = np.random.multinomial(n_choices[i_sim, i_var],
                                                       probs[:, i_var])

---

编辑:不使用仅一个循环的多项式采样的可能替代解决方案可能是这样的:

import numpy as np

np.random.seed(0)
# Problem size
n_cats = 10
n_vars = 50
n_sims = 100
n_maxchoices = 50
# Make example problem
probs = np.random.rand(n_cats, n_vars)
probs /= probs.sum(0)
n_choices = np.random.randint(n_maxchoices, size=(n_sims, n_vars))
sims = np.zeros((n_sims, n_vars, n_cats), np.int32)
# Fill simulations array
n_choices_var = n_choices.sum(0)
sims_r = np.arange(n_sims)
# For each variable
for i_var in range(n_vars):
    # Take choices for all simulations
    choices_var = np.random.choice(n_cats, n_choices_var[i_var], p=probs[:, i_var])
    # Increment choices counts in simulations array
    i_sim = np.repeat(sims_r, n_choices[:, i_var])
    np.add.at(sims, (i_sim, i_var, choices_var), 1)
# Check result
print(np.all(sims.sum(2) == n_choices))
# True

我不确定这是否真的会更快，因为它会生成许多中间数组。我想这取决于问题的特定参数，但如果 Numba 解决方案不是最快的，我会感到惊讶。