Numpy/Scipy: Генерация случайных чисел и разделение

#list #numpy #scipy #numpy-slicing

Вопрос:

Предположим, я сгенерировал два набора чисел: один-случайно сгенерированный набор целых чисел, I0 а другой-набор целых I1 чисел одинаковой длины.

У меня есть еще один набор целых I1 чисел, которые я хочу разделить на списки длины , заданные общей суммой количества элементов, в которых я вижу I0 .

Результатом должен быть список списков , составленных из элементов I1 , каждый из которых имеет длину, заданную соответствующим индексом I0

Я могу добиться этого с помощью следующего:

 import scipy.stats as stats
import numpy as np

N = 1000 #some initial number of 'events'

I0 = stats.poisson(mu=5).rvs(size=N)

N0 = np.sum(I0)

N0_cumsum = np.cumsum(I0)

I1 = stats.binom(n=1, p = 0.5).rvs(size=N0)

I1_split = np.split(I1, N0_cumsum)[:-1]

print(N == len(I1_split))
print(I0[0]==len(I1_split[0]))
 

Выход:

 True

True
 

Теперь я хочу выполнить вложенное разделение, чтобы в качестве вывода у меня был список списков списков или трехслойный список списков.

Например, теперь определение I2 :

 import scipy.stats as stats
import numpy as np

N = 1000 #some initial number of 'events'

I0 = stats.poisson(mu=5).rvs(size=N)

N0 = np.sum(I0)

N0_cumsum = np.cumsum(I0)

I1 = stats.poisson(mu=5).rvs(size=N0)

I1_split = np.split(I1, N0_cumsum)[:-1]

print(N == len(I1_split))
print(I0[0]==len(I1_split[0]))

I2 = []
for _i1 in I1_split:
    temp = []
    for _ii1 in _i1:
        temp.append(list(stats.poisson(mu=5).rvs(size=_ii1)))
    I2.append(temp)


print(len(output)==N)
print(len(output[0]) == I0[0])
print(len(output[0][0]) == I1[0])
 

Выход:

 True
True
True
True
True
 

This gives me the result I am after, but it is painfully slow. I would like to try and pre-generate all the samples and split them.

Has anyone got an idea as to how I would do this efficiently?