numpy vectorize if elif elif else-CodePudding

I have a 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

I want to vectorize it, such that instead of passing one value for each argument I could pass in numpy arrays. The signature would look like this:

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

I'm not sure how to do it though, I could convert each if, elif statement:

((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?

But I'm not really sure how to connect them now. Can somebody please help? I don't have a lot of experience with numpy, maybe it has some useful functions to do this.

UPDATE

Big thanks to everybody, and especially to @Mad Physicist.

So, I can use this:

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

This is awesome! Kudos to Mad Physicist, the solution is so beautiful and simple, even simpler than what I had before. Have a happy life, good sir!

CodePudding user response：

To code this in numpy, you will need to adjust your thinking in a couple of ways.

The biggest is vectorization. You can't have if statements processing each element individually and still be efficient, so you'll need to convert the logic to something more streamlined.

The other point is that you can't have a three-value boolean. That means that you'll either have to redefine your conditions to fit into a true-false dichotomy, or use a different datatype. I'm going to show the latter approach with integers valued 0 for False, 1 for True, and -1 for None.

Each condition looks at angle, and I see only one range collision. I would recommend the following approach:

Make an output array filled with None values (-1)
Figure out which values to set to False, and do it
Figure out which values to set to True, and do it

The order here is important, because you want True values to supersede the False ones when angle <= 6 but both planets are good.

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

Those extra parentheses are important: unlike logical operators, bitwise operators have tighter binding than comparison operators.

This function doesn't care if you pass in a multi-dimensional array or a scalar for planet*_good. All that matters is that the three inputs broadcast against each other.

You can recast the range conditions as a distance from the midpoint:

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

Out of curiosity, did you mean to use the following simplification?

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