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

Question

我有：

• 猫，由10个类别组成的形状数组（10，）
• probs，一个形状为（10，50）的概率数组，代表每个类别被选为50个不同变量的概率
• n\u choices，一个具有形状（num\u sims，50）的数组，包含整数，表示要选择的类别数，并替换每个变量。例如，变量1可以是0个选项，变量2可以是33个选项，等等
• sims，一个填充有形状为0的数组（num_sims，50，10），稍后将填充结果

我想做的是：

• 对于数组中的每一行（表示一个模拟）以及该行中的每个变量，从“cats”中选择N个选项，其中N等于“N\u choices”中的相应值
• 一旦做出了选择，每选择一个类别，就在“模拟人生”中加1。换句话说，我想根据“probs”将“n\u choices”中的值分配给10个类别，并将结果保存到“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解决方案不是最快的，我会感到惊讶。你知道吗