numpy 向量化 if elif elif else

Question

我有一個 function：

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

我想對其進行矢量化處理，這樣我就可以傳入 numpy arrays 而不是為每個參數傳遞一個值。 簽名將如下所示：

def aspect_good(
    angles: np.ndarray[float],
    planet1_good: np.ndarray[bool],
    planet2_good: np.ndarray[bool],
) -> np.array[bool | None]:

我不知道該怎么做，我可以轉換每個 if, elif 語句：

((112 <= angles) & (angles <= 128)) | ((52 <= angles) & (angles <= 68))
((174 <= angles) & (angles <= 186)) | ((84 <= angles) & (angles <= 96))
((0 <= angles) & (angles <= 8)) & planets1_good & planets2_good
((0 <= angles) & (angles <= 6))
# how to convert the 'else' statement?

但我現在不確定如何連接它們。 有人可以幫忙嗎？ 我對 numpy 沒有太多經驗，也許它有一些有用的功能可以做到這一點。

更新

非常感謝大家，尤其是@Mad Physicist 。

所以，我可以使用這個：

def aspect_good(angles: np.typing.ArrayLike, planets1_good: np.typing.ArrayLike, planets2_good: np.typing.ArrayLike) -> np.typing.NDArray:
    """
    Decides if the angle represents a good aspect.
    """
    result = np.full_like(angle, -1, dtype=np.int8)

    false_mask = np.abs(angle % 90) <= 6
    result[false_mask] = 0

    true_mask = np.abs(angle % 60) <= 8
    result[true_mask] = 1

    return result

如果是的話，這太棒了，向瘋狂的物理學家致敬，解決方案是如此美麗和簡單，甚至比我以前的更簡單。 祝您生活愉快，好先生！

Answer 1

要在 numpy 中對此進行編碼，您需要通過多種方式調整您的想法。

最大的是矢量化。 您不能讓if語句單獨處理每個元素並且仍然高效，因此您需要將邏輯轉換為更簡化的邏輯。

還有一點就是不能有三值boolean。 這意味着您要么必須重新定義條件以適應真假二分法，要么使用不同的數據類型。 我將展示后一種方法，整數值0表示 False， 1表示 True， -1表示 None。

每種情況都着眼於角度，而我只看到一個范圍碰撞。 我會推薦以下方法：

• 制作一個 output 數組，其中填充了 None 值 (-1)
• 找出要設置為 False 的值，然后執行此操作
• 找出要設置為 True 的值，然后執行此操作

這里的順序很重要，因為當angle <= 6但兩個行星都很好時，您希望 True 值取代 False 值。

def aspect_good(angle: np.typing.ArrayLike, planet1_good: np.typing.ArrayLike, planet2_good: np.typing.ArrayLike) -> np.typing.NDArray:
    """
    Decides if the angle represents a good aspect.
    NOTE: returns None if the angle doesn't represent an aspect.
    """
    result = np.full_like(angle, -1, dtype=np.int8)

    false_mask = ((174 <= angle) & (angle <= 186)) | ((84 <= angle) & (angle <= 96)) | ((0 <= angle) & (angle <= 6))
    result[false_mask] = 0

    true_mask = ((112 <= angle) & (angle <= 128)) | ((52 <= angle) & (angle <= 68)) | ((0 <= angle) & (angle <= 8) & planet1_good & planet2_good)
    result[true_mask] = 1

    return result

這些額外的括號很重要：與邏輯運算符不同，位運算符比比較運算符具有更緊密的綁定。

這個 function 不關心您是否傳入多維數組或planet*_good的標量。 重要的是三個輸入相互廣播。

您可以將范圍條件重鑄為距中點的距離：

false_mask = (np.abs(angle - 180) <= 6) | (np.abs(angle - 90) <= 6) | (np.abs(angle - 3) <= 3)

出於好奇，您的意思是使用以下簡化嗎？

false_mask = np.abs(angle % 90) <= 6