python - 优化使用 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

除了 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

其他想法

1. 这个实现可以进一步加速。只有 u 的最后一个/最大值用于创建 C[i, j] - 由于 u 的值较低，较早的条目被覆盖.您能想出一种方法来在构建 R 时立即确定最大的 u 吗？

2. 可能值得为每个 j 和 u 提前计算 Q[j] * gamma[:, u + 1] ，并且只在设置 C[i, j] 的值时执行查找。