Scipy：加快采样速度

Question

我开发了以下程序，用于为特定任务以块的形式生成子样本列表。但是，它非常慢，因为我在循环中进行随机采样。

import scipy.stats as stats
import numpy as np

#GENERATE SOME RANDOMLY CHUNKED COUNT DATA
N_chunks=250
idx_chunks = np.random.randint(20, size=N_chunks)
idx_cumsum = np.cumsum(idx_chunks)
data_sample = stats.poisson(mu=5).rvs(size=np.sum(idx_chunks))
data_sample_split = np.split(data_sample, idx_cumsum)[:-1]

#GENERATE SUBSAMPLES OF THE LENGTH GIVEN BY EACH ELEMENT OF THE LIST
f = stats.poisson(mu=2)
output = []
total = 0
for _i1 in data_sample_split:
    temp = []
    for _ii1 in _i1:
        temp.append(f.rvs(_ii1))
    output.append(temp)

有什么方法可以加快程序的速度，同时获得完全相同的输出？

我特别希望在重塑为列表列表之前对我需要的所有样本进行预采样。但是，我不知道该怎么做。

Answer 1

假设用f.rvs生成多个值相当于用f.rvs生成部分，那么你可以生成一个大数组并自己拆分。这是一个例子：

# [...] -- Same code than above

def advancedSplit(rawData, data_sample_split):
    output = []
    partialSum = 0
    for _i1 in data_sample_split:
        temp = []
        for _ii1 in _i1:
            temp.append(rawData[partialSum:partialSum+_ii1])
            partialSum += _ii1
        output.append(temp)
    return output

def generateSubsamples(data_sample_split):
    total = sum(map(sum, data_sample_split))
    rawData = f.rvs(total) # Generate a big array to be split
    return advancedSplit(rawData, data_sample_split) # Split the array efficiently

# GENERATE SUBSAMPLES OF THE LENGTH GIVEN BY EACH ELEMENT OF THE LIST
output = generateSubsamples(data_sample_split)

这比我机器上的原始代码快 50 倍。

---

您还可以使用与原始代码中相同的逻辑将源数组拆分为np.split。但在实践中，这种方法明显较慢（由于执行许多小数组分配，而上述代码使用廉价切片）。

def advancedSplitVect(rawData, data_sample_split):
    output = []
    partial_sum = 0
    for _i1 in data_sample_split:
        if len(_i1) > 0:
            idx_cumsum = np.cumsum(_i1)
            total = idx_cumsum[-1]
            temp = np.split(rawData[partial_sum:partial_sum+total], idx_cumsum[:-1])
            partial_sum += total
            output.append(temp)
        else:
            output.append([])
    return output

请注意，advancedSplit 可以使用 Numba JIT 进行优化，以稍微加快代码速度。但是，这有点乏味，因为必须正确键入列表。因此，它可能不值得。