Vectorize a function with a condition

user_5053483 0
Viewed 213

I would like to vectorize a function with a condition, meaning to calculate its values with array arithmetic. np.vectorize handles vectorization, but it does not work with array arithmetic, so it is not a complete solution

An answer was given as the solution in the question "How to vectorize a function which contains an if statement?" but did not prevent errors here; see the MWE below.

import numpy as np

def myfx(x):
    return np.where(x < 1.1, 1, np.arcsin(1 / x))

y = myfx(x)

This runs but raises the following warnings:

<stdin>:2: RuntimeWarning: divide by zero encountered in true_divide
<stdin>:2: RuntimeWarning: invalid value encountered in arcsin

What is the problem, or is there a better way to do this?

I think this could be done by

  1. Getting the indices ks of x for which x[k] > 1.1 for each k in ks.
  2. Applying np.arcsin(1 / x[ks]) to the slice x[ks], and using 1 for the rest of the elements.
  3. Recombining the arrays.

I am not sure about the efficiency, though.

vectorization conditional-statements numpy
2 Answers

The statement np.where(x < 1.1, 1, np.arcsin(1 / x)) is equivalent to

mask = x < 1.1
a = 1
b = np.arcsin(1 / x)
np.where(mask, a, b)

Notice that you're calling np.arcsin on all the elements of x, regardless of whether 1 / x <= 1 or not. Your basic plan is correct. You can do the operations in-place on an output array using the where keyword of np.arcsin and np.reciprocal, without having to recombine anything:

def myfx(x):
    mask = (x >= 1.1)
    out = np.ones(x.shape)
    np.reciprocal(x, where=mask, out=out) # >= 1.1 implies != 0
    return np.arcsin(out, where=mask, out=out)

Using np.ones ensures that the unmasked elements of out are initialized correctly. An equivalent method would be

out = np.empty(x.shape)
out[~mask] = 1
user_29887300

You can always find an arithmetic expression that prevents the "divide by zero".

Example:

def myfx(x):
    return np.where( x < 1.1, 1, np.arcsin(1/np.maximum(x, 1.1)) )

The values where x<1.1 in the right wing are not used, so it's not an issue computing np.arcsin(1/1.1) where x < 1.1.

user_49267570
Related