如何调试 numpy 掩码

这个问题与this one https://stackoverflow.com/q/73672739/11004423.

我有一个正在尝试矢量化的函数。这是原来的函数：

def aspect_good(angle: float, planet1_good: bool, planet2_good: bool):
    """
    Decides if the angle represents a good aspect.
    NOTE: returns None if the angle doesn't represent an aspect.
    """

    if 112 <= angle <= 128 or 52 <= angle <= 68:
        return True
    elif 174 <= angle <= 186 or 84 <= angle <= 96:
        return False
    elif 0 <= angle <= 8 and planet1_good and planet2_good:
        return True
    elif 0 <= angle <= 6:
        return False
    else:
        return None

这就是我到目前为止所拥有的：

def aspects_good(
    angles: npt.ArrayLike,
    planets1_good: npt.ArrayLike,
    planets2_good: npt.ArrayLike,
) -> npt.NDArray:
    """
    Decides if the angles represent good aspects.

    Note: this function was contributed by Mad Physicist. Thank you.
    https://stackoverflow.com/q/73672739/11004423

    :returns: an array with values as follows:
        1 – the angle is a good aspect
        0 – the angle is a bad aspect
       -1 – the angle doesn't represent an aspect
    """
    result = np.full_like(angles, -1, dtype=np.int8)

    bad_mask = np.abs(angles % 90) <= 6
    result[bad_mask] = 0

    good_mask = (np.abs(angles - 120) <= 8) |\
                (np.abs(angles - 60) <= 8) |\
                ((np.abs(angles - 4) <= 4) & planets1_good & planets2_good)
    result[good_mask] = 1

    return result

它没有按预期工作，但是，我用 pytest 编写了一个测试：

def test_aspects_good():
    tests = np.array([
        [120, True, False, True],
        [60, True, False, True],
        [180, True, False, False],
        [90, True, False, False],

        [129, True, False, -1],
        [111, True, False, -1],
        [69, True, False, -1],
        [51, True, False, -1],
        [187, True, False, -1],
        [173, True, False, -1],
        [97, True, False, -1],
        [83, True, False, -1],

        [0, True, True, True],
        [0, True, False, False],
        [0, False, True, False],
        [0, False, False, False],

        [7, False, False, -1],
        [7, True, True, True],
        [9, True, True, -1],
    ])

    angles = tests[:, 0]
    planets1_good = tests[:, 1]
    planets2_good = tests[:, 2]
    expected = tests[:, 3]

    result = aspects_good(angles, planets1_good, planets2_good)
    assert np.array_equal(result, expected)

它失败了，说 False，数组不同。

这里我有result and expected数组并排组合：

array([[ 1,  1],
│      [ 1,  1],
│      [ 0,  0],
│      [ 0,  0],
│      [-1, -1],
│      [-1, -1],
│      [-1, -1],
│      [-1, -1],
│      [-1, -1],
│      [-1, -1],
│      [-1, -1],
│      [-1, -1],
│      [ 0,  1],
│      [ 0,  0],
│      [ 0,  0],
│      [ 0,  0],
│      [-1, -1],
│      [-1,  1],
│      [-1, -1]])

注意：第一列是result数组，第二个是expected。正如您所看到的，它们有两个地方不同。现在问题来了“如何调试这个？”通常我会使用调试器，并逐步执行每个 if/elif/else 条件。但我不知道如何调试 numpy 掩码。

这个问题似乎是三件事的结合：

1. Numpy 在整个数组中使用同构类型。

   你会发现tests.dtype is dtype('int64') or dtype('int32')取决于您的架构。这意味着包含的列planet1_good and planet2_good也是整数，而不是布尔值。

2. 按位与 (&) is not逻辑运算符。

   按位 AND 运算将返回具有最大输入类型的结果。具体来说，对于结果<=, which is一个布尔值，和一个int数组，结果将是int。这意味着你可以做类似的事情np.array([1, 2]) & np.array([True, True]) to get array([1, 0]), not array([True, False]).

3. Numpy 通过 dtype 区分布尔掩码和花式索引，即使花式索引仅包含零和一。如果你有一个 2 元素数组，x, then x[[True, True]] = 1分配1的两个元素x。然而，x[[1, 1]] = 1分配1只到第二个元素x.

这基本上就是这里发生的事情。bad_mask是一个布尔掩码，并且完全按照您的预期工作。然而，good_mask与几个整数数组进行 AND 运算，因此成为包含 0 和 1 的整数数组。表达方式result[good_mask] = 1实际上是分配第一个和第二个元素result to be 1，它偶然对应于您的两个测试。其余True结果不能也不会被分配1.

有几种方法可以解决这个问题，按优先顺序降序列出（我最喜欢的在顶部）：

1. 将所有数组转换为正确类型的 numpy 数组。现在你的函数不符合它接受任何类似数组的约定。如果您传递一个列表angles， 你会得到TypeError: unsupported operand type(s) for %: 'list' and 'int'。这是一个相当惯用的方法：

   angles = np.asanyarray(angles)
   planets1_good = np.asanyarray(planets1_good, dtype=bool)
   planets2_good = np.asanyarray(planets2_good, dtype=bool)
   
   result = np.full_like(angles, -1, dtype=np.int8)
   
   bad_mask = np.abs(angles % 90) <= 6
   result[bad_mask] = 0
   
   good_mask = (np.abs(angles - 120) <= 8) |\
               (np.abs(angles - 60) <= 8) |\
               ((np.abs(angles - 4) <= 4) & planets1_good & planets2_good)
   result[good_mask] = 1
   return result
   

2. 确保这件事good_mask敷之前其实是个面膜。你还是应该转换angles，但其他数组将被自动转换&操作员：

   good_mask = ((np.abs(angles - 120) <= 8) |\
                (np.abs(angles - 60) <= 8) |\
                ((np.abs(angles - 4) <= 4) & planets1_good & planets2_good)).astype(bool)
   

   您也可以执行与您所做的类似的操作bad_mask:

   good_mask = (np.abs(angles % 60) <= 8) & (angles >= -8) & (angles <= 128)
   

3. 将掩码转换为索引，这不会关心原始数据类型：

   result[np.flatnonzero(good_mask)] = 1