numpy vectorize if elif elif else

Question

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!

You'll have to think what returning None means. Is it indicative of an error? Then raise an exception (ValueError probably). Should it perhaps also be False instead? Or is it indicative of "don't know, don't care"? Then use a masked array, and mask the None elements. — 9769953
– 9769953, Commented Sep 10, 2022 at 14:53
@9769953 returning None doesn't mean an error, nor does it mean False. There is a clear distinction between False and None in my app, as it says in the function docs "None means the angle is nor good nor bad" — astroboy
– astroboy, Commented Sep 10, 2022 at 14:57
A bool dtype array can only have True/False values. An integer one could have more. Object dtype could contain the None values (in fact np.empty(3, object) returns an array full of None). But you can't do further logic on an array containing None. A np.ma masked array could mark some values as "not valid". That in effect uses 2 bool arrays, the data and the mask (valid/invalid). — hpaulj
– hpaulj, Commented Sep 10, 2022 at 15:15

Mad Physicist · Accepted Answer · 2022-09-10 15:32:01Z

2

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

edited Sep 10, 2022 at 15:32

answered Sep 10, 2022 at 15:21

Mad Physicist

116k29 gold badges202 silver badges292 bronze badges

Sign up to request clarification or add additional context in comments.

18 Comments

astroboy Over a year ago

Wow! Thank you so much! You are very smart. I am very lucky to have such a good solution from you. Everything broken down and presented well, and you even found some ways to improve my conditions.

astroboy Over a year ago

Regarding the range conditions, I have updated my question, you were right.

astroboy Over a year ago

Is there a way I could thank you? I can probably award you a bounty, but I'm not sure if you even care about reputation at this point...

Mad Physicist Over a year ago

@acmpo6ou I'm glad to be of help. There's no need to go beyond an upvote and selecting the answer. I'm not quite done caring about reputation just yet: ~500pts to go until I get swag :) That being said, just knowing that I've helped to show you how to think numpythonically is enough. I'll sleep well knowing that there's one less person trying to use np.vectorize for everything

astroboy Over a year ago

500pts for you? That's as much as I have! I can do it! But I can award you a bounty only after 2 days, and by that time you might not need me anymore (because you would get there yourself). And I want to contribute the final drop that will overflow your reputation! What do you think? How many reputation a day do you get?

|

Collectives™ on Stack Overflow

numpy vectorize if elif elif else

1 Answer 1

18 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

18 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related