Numpy: Boolean operation on np.array passed into function

Question

Numpy allows to pass a numpy.array as argument into a function and evaluate the function for every element of the array.:

def f(n):
    return n**2

arr = np.array([1,2,3])
f(arr)

outputs:

>>>[1 4 9]

This works fine, as long as f(n)doesn't perform boolean operations on n like this:

def f(n):
    if n%2 == 0:
        print(n)

The above code throws following exception: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

This makes sense since the debugger showed, that the function f(n) received the entire numpy.arrayas input. Now, I'm looking for a way to change the code behaviour, so that I can perform boolean operations on the inputs. Is this possible while passing the entire numpy.array, or do I have to call the function by manually iterating over each element of the array?

---Edit:---

def get_far_field_directivity(k,a,theta):
    temp = k*a*np.sin(theta)
    if temp == 0: return 1
    else: return (2*sp.jn(1,(temp)))/(temp)

The function returns to other functions, which have to use its value on further calculation which is why the indexing approach by @Chrysophylaxs won't work.

Answer 1

do you want to try something like this? https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html

import numpy as np

arr = np.array([1,2,3])

def f(n):
    if n%2 == 0:
        print(n)
    return n**2

vfunc = np.vectorize(f)      
vfunc(arr)          

outputs:

2
array([1, 4, 9])

whereas this

f(arr)     

outputs:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in f
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()