python and problems with passing array to a function

Question

I need to solve some ordinary differential equation dy/dx = f(x) = x^2 ln(x) and to proceed I create the array xpt between the limits 0. <= xpt <= 2. Because I have to be careful at xpt = 0, I have defined the function as follows:

def f(x):
    if x <= 1.e-6:
        return 0.
    else:
        return np.square(x)*np.log(x)

My calling programme reads:

Npt = 200

xpt = np.linspace(0.,2.,Npt)

fpt = np.zeros(Npt)

However, when I make the call fpt = f(xpt), I get the error:

"ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"

I can get around this by introducing a for loop, writing

for ip in range(Npt):
    fpt[ip] = f(xpt[ip])

But this seems a hack and unsatisfactory.

I have tried to look at the suggestion of using a.any() and redefined the function as

def Newf(x):

    if ((x <= 1.e-6).all()):
        return 0.
    else:
        return np.square(x*np.log(x))

But this seems to give f(0.) as NaN.

Any help on how to proceed gratefully received.

Answer 1

Try this code

import numpy as np

def f(x):
    return np.where(x <= 1.e-6, 0., np.square(x) * np.log(x))

Npt = 200
xpt = np.linspace(0., 2., Npt)
fpt = f(xpt)

print(fpt)

np.where(condition, value_if_true, value_if_false) applies the condition element-wise. If the condition is true, it uses value_if_true, otherwise it uses value_if_false.