Python3中如何实现允许列符号差异的Numpy矩阵对比断言？

核心思路 #

针对每个内层n×n方阵，验证其每一列要么与对应方阵的同列完全（近似）相等，要么完全（近似）互为相反数，所有内层矩阵都满足条件则断言通过。

具体实现 #

可以封装一个类似numpy.testing.assert_almost_equal的自定义断言函数，步骤如下：

1. 单矩阵列检查逻辑
   对两个n阶方阵，遍历每一列，用np.allclose验证列的匹配关系（支持浮点近似）：

   • 列元素直接相等，或
   • 列元素互为相反数（即A_col ≈ -B_col）

2. 遍历三维数组的所有内层矩阵
   循环处理三维数组中m个内层矩阵，逐个验证，若任意矩阵不满足条件则抛出AssertionError。

完整代码示例 #

import numpy as np

def assert_columns_sign_invariant(arr1, arr2, rtol=1e-07, atol=0):
    # 先检查数组形状是否一致
    if arr1.shape != arr2.shape:
        raise AssertionError(f"Array shapes mismatch: {arr1.shape} vs {arr2.shape}")
    
    m, n, _ = arr1.shape
    for i in range(m):
        mat1 = arr1[i]
        mat2 = arr2[i]
        # 遍历每一列
        for col_idx in range(n):
            col1 = mat1[:, col_idx]
            col2 = mat2[:, col_idx]
            # 检查列是否相等或互为相反数（近似）
            if not (np.allclose(col1, col2, rtol=rtol, atol=atol) or 
                    np.allclose(col1, -col2, rtol=rtol, atol=atol)):
                raise AssertionError(
                    f"Matrix at index {i}, column {col_idx} does not match.\n"
                    f"Expected either {col1} or {-col1}, got {col2}"
                )

# 测试示例中的数组
arr1 = np.array([
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]
    ],
    [
        [10, 20, 30],
        [40, 50, 60],
        [70, 80, 90]
    ],
    [
        [100, 200, 300],
        [400, 500, 600],
        [700, 800, 900]
    ]
])

arr2 = np.array([
    [
        [1, -2, 3],
        [4, -5, 6],
        [7, -8, 9]
    ],
    [
        [10, 20, -30],
        [40, 50, -60],
        [70, 80, -90]
    ],
    [
        [100, 200, 300],
        [400, 500, 600],
        [700, 800, 900]
    ]
])

# 执行断言，不会抛出错误
assert_columns_sign_invariant(arr1, arr2)
print("断言通过")

补充说明 #

• 函数参数rtol和atol与np.allclose一致，用于控制浮点近似的精度，对应np.testing.assert_almost_equal的精度要求
• 如果需要严格匹配整数，可将rtol设为0，atol设为0，等价于用==判断

内容的提问来源于stack exchange，提问作者aura