Scipy：加速采样

问题描述

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

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)

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

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

解决方法

假设用 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

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