CODEWITHSUNDEEP
pythonnumpy
offeltoffel
offeltoffel

Reputation: 2801

Dynamically add to condition in numpy.where

Depending on the user input, I have different conditions to check for in my array. Assume there are a maximum of 3 conditions:

a must be 1

b must be positive

c must be True

Only where these three conditions evaluate as True, the array shall be processed:

myArr = np.random.rand(5)
a = np.array([1, 1, 1, 0, 1])
b = np.array([4, 3, -8, 7, 6])
c = np.array([True, False, True, True, True])
valid_indices = np.where((a == 1) & (b > 0) & (c == True))

>> valid_indices
>> Out: (array([0, 4], dtype=int64),)

Not knowing beforehand which of these conditions will be provided, I would have to check like this:

if a and not b and not c:
    valid_indices = np.where(a == 1)

elif a and not b and c:
    valid_indices = np.where((a == 1) & (c == True))

elif a and b and not c:
    valid_indices = np.where((a == 1) & (b > 0))

elif not a and b and c:
    valid_indices = np.where((b > 0) & (c == True))

elif not a and not b and c:
    valid_indices = np.where(c == True)

elif not a and b and not c:
    valid_indices = np.where((b > 0))

God forbid I add another condition. Things are getting really messy. I am looking for a way to dynamically add to the condition as if it were just a regular string or formatter. Is that possible?

Upvotes: 1

Views: 835

Answers (2)

daveg
daveg

Reputation: 535

If you set defaults to the true condition:

import numpy as np


myArr = np.random.rand(5)

a = np.array([1, 1, 1, 0, 1])
b = np.array([4, 3, -8, 7, 6])
c = np.array([True, False, True, True, True])


def get_id(**kwargs):
    """ a must be 1, b must be positive, c must be True """
    a = kwargs.get("a", np.ones(5))
    b = kwargs.get("b", np.ones(5))
    c = kwargs.get("c", np.ones(5) == 1)
    return np.where((a == 1) & (b > 0) & c)


print(get_id(a=a))
print(get_id(a=a, c=c))
print(get_id(a=a, b=b))
print(get_id(a=a, b=b, c=c))


(array([0, 1, 2, 4]),)
(array([0, 2, 4]),)
(array([0, 1, 4]),)
(array([0, 4]),)

Upvotes: 2

FBruzzesi
FBruzzesi

Reputation: 6495

Maybe something like the following can help:

myArr = np.random.rand(5)
size = myArr.size
all_true = np.repeat(True, size)

a = np.array([1, 1, 1, 0, 1])
b = np.array([4, 3, -8, 7, 6])
c = np.array([True, False, True, True, True])

valid_indices = np.where((a == 1 if 'a' in locals() else all_true) & 
                         (b > 0 if 'b' in locals() else all_true) & 
                         (c == True if 'c' in locals() else all_true))

In this way you can write all conditions and only check if the variable exists in the local varibales by 'a' in locals(). If you are defining a, b and c somewhere and refering them in a function you can check if they are defined in the global environment using 'a' in globals().

Upvotes: 1

Related Questions