Fast way to calculate conditional function

Question

What is the most fast way to calculate function like

# here x is just a number
def f(x):
    if x >= 0:
        return np.log(x+1)
    else:
        return -np.log(-x+1)

One possible way is:

# here x is an array
def loga(x)
    cond = [x >= 0, x < 0]
    choice = [np.log(x+1), -np.log(-x+1)
    return np.select(cond, choice)

But seems numpy goes through array element by element. Is there any way to use something conceptually similar to np.exp(x) to achieve better performance?

Answer 1

    def f(x):
        return (x/abs(x)) * np.log(1+abs(x))

Answer 2

In cases like these, masking helps -

def mask_vectorized_app(x):
    out = np.empty_like(x)
    mask = x>=0
    mask_rev = ~mask
    out[mask] = np.log(x[mask]+1)
    out[mask_rev] = -np.log(-x[mask_rev]+1)
    return out

Introducing numexpr module helps us further.

import numexpr as ne

def mask_vectorized_numexpr_app(x):
    out = np.empty_like(x)
    mask = x>=0
    mask_rev = ~mask

    x_masked = x[mask]
    x_rev_masked = x[mask_rev]
    out[mask] = ne.evaluate('log(x_masked+1)')
    out[mask_rev] = ne.evaluate('-log(-x_rev_masked+1)')
    return out

Inspired by @user2685079's post and then using the logarithmetic property : log(A**B) = B*log(A), we can push in the sign into the log computations and this allows us to do more work with numexpr's evaluate expression, like so -

s = (-2*(x<0))+1 # np.sign(x)
out = ne.evaluate('log( (abs(x)+1)**s)')

Computing sign using comparison gives us s in another way -

s = (-2*(x<0))+1

Finally, we can push this into the numexpr evaluate expression -

def mask_vectorized_numexpr_app2(x):
    return ne.evaluate('log( (abs(x)+1)**((-2*(x<0))+1))')

Runtime test

Loopy approach for comparison -

def loopy_app(x):
    out = np.empty_like(x)
    for i in range(len(out)):
        out[i] = f(x[i])
    return out

Timings and verification -

In [141]: x = np.random.randn(100000)
     ...: print np.allclose(loopy_app(x), mask_vectorized_app(x))
     ...: print np.allclose(loopy_app(x), mask_vectorized_numexpr_app(x))
     ...: print np.allclose(loopy_app(x), mask_vectorized_numexpr_app2(x))
     ...: 
True
True
True

In [142]: %timeit loopy_app(x)
     ...: %timeit mask_vectorized_numexpr_app(x)
     ...: %timeit mask_vectorized_numexpr_app2(x)
     ...: 
10 loops, best of 3: 108 ms per loop
100 loops, best of 3: 3.6 ms per loop
1000 loops, best of 3: 942 µs per loop

Using @user2685079's solution using np.sign to replace the first part and then with and without numexpr evaluation -

In [143]: %timeit np.sign(x) * np.log(1+abs(x))
100 loops, best of 3: 3.26 ms per loop

In [144]: %timeit np.sign(x) * ne.evaluate('log(1+abs(x))')
1000 loops, best of 3: 1.66 ms per loop