优化使用 numpy 创建 3d 矩阵

Question

我一直在尝试使用 NumPy 优化矩阵 C 的构造（见下文）。如何进一步优化我的代码，以便更快地构建矩阵 C？

给定以下矩阵：

Q:   array([[78.66  , 47.196 , 31.464 ],
           [40.3875, 24.2325, 16.155 ],
           [40.4775, 24.2865, 16.191 ],
           ...,
           [55.62  , 33.372 , 22.248 ],
           [76.7475, 46.0485, 30.699 ],
           [77.3325, 46.3995, 30.933 ]])

S:  [[[1,2,3],[],[],[1,...,1125]],
     [[],[1,...,200],[300,301][]],
     ...,
     [[1,1125],[],[12],[345,453]]]

gamma:   array([[0. , 1.4, 2.5, 3. , 3. ],
               [0. , 1.6, 3. , 3.7, 4. ],
               [0. , 1.8, 3.5, 4.4, 5. ]])

我有以下代码来构建三维矩阵C

# # Matrix C_ijk 
C = np.zeros((n,o,p))
for i in range(n):
    for j in range(o):
        for k in range(p):
            for u in range(m-1):
                if np.isin(i,S[j][u]):
                    C[i,j,k] = Q[j,k] * gamma[k,u+1]

编辑：m、n、o 和 p 是定义矩阵维度长度的整数。完整模型中分别为 5、1126、797 和 3。

Q 是尺码 (o,p) ; S 是尺寸 (o,m-1) ； gamma 是大小 (p,m-1) ； C 是尺码 (n,o,p) ;

这是一个小示例输出：

>>> Q
array([[10., 10.],
       [20., 20.],
       [30., 30.],
       [30., 30.]])
>>> S
[[[0, 1], [], [], [2]], [[2], [0], [1], []], [[], [1], [0, 2], []], [[], [2], [], [0, 1]]]
>>> gamma
array([[0.   , 0.575, 1.2  , 1.75 , 2.   ],
       [0.   , 0.625, 1.4  , 2.25 , 3.   ]])
>>> C
array([[[ 5.75,  6.25],
        [24.  , 28.  ],
        [52.5 , 67.5 ],
        [60.  , 90.  ]],

       [[ 5.75,  6.25],
        [35.  , 45.  ],
        [36.  , 42.  ],
        [60.  , 90.  ]],

       [[20.  , 30.  ],
        [11.5 , 12.5 ],
        [52.5 , 67.5 ],
        [36.  , 42.  ]]])

按照@AhmedMohamedAEK 的建议，以下列方式实施 numba 是否正确？

from numba import njit, prange

@njit(parallel=True)
def matrix_C(Q,S,gamma,n,o,p,m):
    C = np.zeros((n,o,p))
    for i in prange(n):
        for j in prange(o):
            for u in prange(m-1):
                if np.isin(i,S[j][u]):
                    C[i,j,:] = Q[j,:] * gamma[:,u+1]
    return C
C = matrix_C(Q,S,gamma,n,o,p,m)

Answer 1

您可以删除 k 中的循环，因为它被所有数组使用，如下所示：

# # Matrix C_ijk 
C = np.zeros((n,o,p))
for i in range(n):
    for j in range(o):
        for u in range(m-1):
            if np.isin(i,S[j][u]):
                C[i,j,:] = Q[j,:] * gamma[:,u+1]

然而，在 python 中使用嵌套循环是非常不鼓励的，应该使用外部模块将循环移动到 C 端，可以使用 cython 创建或 numba.

编辑： 对于上面的 numba 实现，如果数组非常大（几 MB 的大小），那么您可以使用如下并行实现：

from numba import njit, prange

@njit(parallel=True)
def matrix_C(Q,S,gamma,n,o,p,m):
    C = np.zeros((n,o,p))
    for i in prange(n):
        for j in range(o):
            for k in range(p):
                for u in range(m-1):
                    if np.isin(i,S[j][u]):
                        C[i,j,k] = Q[j,k] * gamma[k,u+1]
    return C
C = matrix_C(Q,S,gamma,n,o,p,m)

然而，如果数组相对较小，那么我将删除 parallel 和 prange 并使用以下内容：

@njit()
def matrix_C(Q,S,gamma,n,o,p,m):
    C = np.zeros((n,o,p))
    for i in range(n):
        for j in range(o):
            for k in range(p):
                for u in range(m-1):
                    if np.isin(i,S[j][u]):
                        C[i,j,k] = Q[j,k] * gamma[k,u+1]
    return C
C = matrix_C(Q,S,gamma,n,o,p,m)

并记住第一次调用该函数是在编译时，所以它会有点慢，但以后的调用会很快。

Answer 2

除了 Python 循环之外，特别伤害您的是最内层循环中的线性时间 np.isin 查找。这很容易补救。我们可以只创建那些索引 i, j, u，其中 i 在 S[j][u] 中，因此我们以后不需要搜索它们。

这是通过以下代码实现的，创建索引生成器R。表达式有点长，但不难理解。

R = ((i, j, u)
     for j in range(o)
     for u in range(m - 1)
     for i in S[j][u]
     if i < n)

这个索引生成器大大简化了 C 的计算：

for i, j, u in R:
    C[i, j] = Q[j] * gamma[:, u + 1]

由于现在完全避免了很多工作，因此这应该比您最初的实施要快得多。

完整代码：

def matrix_C(Q, S, gamma, n, o, p, m):
    C = np.zeros((n, o, p))
    R = ((i, j, u)
         for j in range(o)
         for u in range(m - 1)
         for i in S[j][u]
         if i < n) 

    for i, j, u in R:
        C[i, j] = Q[j] * gamma[:, u + 1]

    return C

其他想法

这个实现可以进一步加速。只有 u 的 last/largest 值用于创建 C[i, j] - 由于 u 的值较低而导致的较早条目将被覆盖。你能想出一种方法可以在构建 R 时立即确定最大的 u 吗？
可能值得为每个 j 和 u 提前计算 Q[j] * gamma[:, u + 1]，并且仅在设置 [= 的值时执行查找20=].

Answer 3

你可以使用np.fromfunction

C = np.fromfunction(function=f, shape=(n, o, p))

其中 f 是

def f(i, j, k):
    ...
    # your logic here
    ...
    return value

python for 循环很慢（参见 here），所以这更有效（至少对于外部 3 个循环）

祝你好运

优化使用 numpy 创建 3d 矩阵

Optimize the use of numpy for creating a 3d-matrix

python

numpy

matrix

numba

其他想法