Square root of all values in numpy array, preserving sign

Question

I'd like to take the square root of every value in a numpy array, while preserving the sign of the value (and not returning complex numbers when negative) - a signed square root.

The code below demonstrates the desired functionality w/ lists, but is not taking advantage of numpy's optimized array manipulating superpowers.

def signed_sqrt(list):
    new_list = []
    for v in arr:
        sign = 1
        if v < 0:
            sign = -1
        sqrt = cmath.sqrt(abs(v))
        new_v = sqrt * sign
        new_list.append(new_v)


list = [1., 81., -7., 4., -16.]
list = signed_sqrt(list)
# [1., 9., -2.6457, 2. -4.]

For some context, I'm computing the Hellinger Kernel for [thousands of] image comparisons.

Any smooth way to do this with numpy? Thanks.

Answer 1

You can try using the numpy.sign function to capture the sign, and just take the square root of the absolute value.

import numpy as np
x = np.array([-1, 1, 100, 16, -100, -16])
y = np.sqrt(np.abs(x)) * np.sign(x)
# [-1, 1, 10, 4, -10, -4]