给定一个 numpy 矩阵,生成结果联合分布的最有效方法是什么?

What is the most efficient way to generate the joint distribution of outcomes given a numpy matrix?

假设有 i 个客户 (i = 1,...,5),他们对商品 D_i 的需求可能是低、中或高。我有以下带有需求数字的 NumPy 矩阵:

import numpy as np

# Client demand
demand = np.array([[150, 160, 170],
                   [100, 120, 135],
                   [250, 270, 300],
                   [300, 325, 350],
                   [600, 700, 800]])

现在,我想获得所有客户的结果的联合分布。由于每个客户有 3 个(统计上独立的)事件,总共有 5 个客户,因此有 3^5 = 243 种可能的组合。获得新矩阵的最有效方法是什么,该矩阵给出每个客户 i 和场景 j (j=1,...243) 的所有需求数据?

编辑:

我发现 np.meshgrid 可以满足我的要求,但它似乎只需要表示网格坐标的一维数组,因此无法提供 NumPy 矩阵:

import numpy as np

# Client demand
demand = np.array([[150, 160, 170],
                   [100, 120, 135],
                   [250, 270, 300],
                   [300, 325, 350],
                   [600, 700, 800]])

# working
scenarios = np.array(np.meshgrid([150, 160, 170],
                     [100, 120, 135],
                     [250, 270, 300],
                     [300, 325, 350],
                     [600, 700, 800]
                     )).T.reshape(-1,5)

print(scenarios.shape)

# not working
np.array(np.meshgrid(demand)).T.reshape(-1,5)

这里有两个选项。第二个比第一个快得多:

import numpy as np

# Client demand
demand = np.array([[150, 160, 170],
                   [100, 120, 135],
                   [250, 270, 300],
                   [300, 325, 350],
                   [600, 700, 800]])

import itertools
idxs = list(itertools.product(np.arange(3),repeat=5))

# 803 µs ± 17.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
res = np.array([demand[np.arange(5),idx] for idx in idxs])
# 155 µs ± 3.3 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
res2 = demand[np.arange(5), idxs]