NumPy实现上三角矩阵移位性能能否匹敌甚至超越等效Cython实现？

TLDR：我正在实现一项无数学运算的数组操作，实测Cython版本性能显著更快，请问是否有方法可以优化NumPy或Cython版本的运行速度？

上下文 #

我需要编写函数实现如下操作：取NxN数组从index位置开始、左上角位于对角线上的子集，沿对角线向上平移一位；其次将首行从index开始的内容向左平移一位；最后将操作后数组的最后一列置零。

该数组为严格上三角矩阵，即对角线及以下所有元素均为0，是我设计的用于存储对象间历史碰撞数据的方案（对象由矩阵下标索引），等效于生成长度为n的索引列表的有序对嵌套列表，规模为n!/(2(n-2)!)。我希望通过该算法实现“从碰撞配对矩阵中移除单个对象”的功能。

该实现的优势在于：从矩阵中“移除碰撞对”的计算开销，远低于从嵌套列表中删除碰撞对、再调整待删除索引之后所有配对索引的开销。

本项目整体面向粉末床熔融增材制造的3D模型构建体积自动“排版”场景，算法基于模拟退火实现，因此碰撞集裁剪、历史信息存储、几何元素增删等能力至关重要，需要做深度优化。

示例 #

假设数组形式如下（不代表实际业务数据）：

arr = 
[[0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
 [0. 0. 2. 3. 4. 5. 6. 7. 8. 9.]
 [0. 0. 0. 3. 4. 5. 6. 7. 8. 9.]
 [0. 0. 0. 0. 4. 5. 6. 7. 8. 9.]
 [0. 0. 0. 0. 0. 5. 6. 7. 8. 9.]
 [0. 0. 0. 0. 0. 0. 6. 7. 8. 9.]
 [0. 0. 0. 0. 0. 0. 0. 7. 8. 9.]
 [0. 0. 0. 0. 0. 0. 0. 0. 8. 9.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 9.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]

当传入index = 3时，需要将子集index+1:n, index+1:n的内容赋值给index:n-1, index:n-1，再将首行从index开始的内容左移一位，最后将最后一列置零，输出结果如下：

fun(3, arr)
[[0. 1. 2. 4. 5. 6. 7. 8. 9. 0.]
 [0. 0. 2. 3. 4. 5. 6. 7. 8. 0.]
 [0. 0. 0. 3. 4. 5. 6. 7. 8. 0.]
 [0. 0. 0. 0. 5. 6. 7. 8. 9. 0.]
 [0. 0. 0. 0. 0. 6. 7. 8. 9. 0.]
 [0. 0. 0. 0. 0. 0. 7. 8. 9. 0.]
 [0. 0. 0. 0. 0. 0. 0. 8. 9. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 9. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]

实现1：纯NumPy版本 #

假设arr为NxN矩阵，实现代码如下：

def fun(index, n, arr):
     arr[index:-1, index:-1] = arr[index + 1:, index + 1:]
     arr[0, index:-1] = arr[0, index + 1:]
     arr[:, n-1:] = 0
     return arr

实现2：Cython版本 #

这是我首次尝试编写Cython代码，实现如下：

@cython.boundscheck(False)
def remove_from_collision_array(int index, int n, double[:,:] arr):

    cdef int i, j, x_shape, y_shape
    x_shape = arr.shape[0]

    for i in range(index, x_shape):
        for j in range(index, x_shape):
            if j <= i:
                # 位于对角线以下，无需操作
                continue
            elif i >= n-1 or j >= n-1:
                arr[i, j] = 0
            else:
                arr[i, j] = arr[i+1, j+1]

    arr[0, index:-1] = arr[0, index+1:]
    arr[:, n-1:] = 0

    return np.asarray(arr)

讨论 #

我对Cython的使用并不熟练，关闭了bounds_checking因为它能大幅提升运行速度，同时我在循环的elif分支中做了边界检查。

我原本认为循环实现的性能不可能超过NumPy，预先分配了5000x5000的NumPy数组避免运行时动态扩容，甚至测试过使用和NumPy版本完全相同的三行代码的Cython实现，性能表现同样不佳。

index=0是计算量最大的场景，我以此为基准循环测试，发现Cython实现比NumPy版本快50%以上，请问这是不是因为我没有正确使用NumPy的相关能力？

我并非计算机专业从业者，也不确定当前方案是否最优，我是正在做系统原型的设计师，如果有进一步优化性能的思路，欢迎告知。

答案更新 #

感谢Jerome的指导，该优化对提升本工具包的运行速度至关重要。我将他的思路落地到代码中后，性能得到了大幅提升，主要优化点有两个：

• 将j循环的起始位置设置为对角线以上，减少了n*(n-1)/2次循环迭代
• 移除了所有条件判断语句

更新后的Cython代码如下：

@cython.boundscheck(False)
@cython.wraparound(False)
def remove_from_collision_arrayV2(int index, int n, double[:,:] arr):

    cdef int i, j

    # 平移对角矩阵
    for i in range(index, n-1):
        for j in range(i, n-1):
                arr[i, j] = arr[i+1, j+1]

    # 平移首行
    for j in range(index, n-1):
        arr[0, j] = arr[0, j+1]

    # 将第n-1列置零
    for i in range(n):
        arr[i, n-1] = 0

    return np.asarray(arr)

基准测试结果：在500x500矩阵上以index=0执行500次操作，耗时如下：

• 原始NumPy代码：52.8s
• 原始Cython代码：16.47s，性能提升3.2倍
• 优化后Cython代码：0.014s，性能提升3550倍

内容的提问来源于stack exchange，提问作者ddm-j