Search code examples
pythonnumpycomplex-numbers

Calculate only real roots from fractional power in numpy


I have a numpy array with negative and positive values, and I'm trying to raise it to the power of 1/3. I get Nan for all the negative numbers (which I understand because it involves complex roots), however, I want to print out the real roots instead of Nan.

Is there a fast pythonic way to do this? Beacuse for the third root, I know there will always be a real root.

import numpy as np

x = np.linspace(-5,5,10)
z = x**(1/3)

And the result for z is:

array([       nan,        nan,        nan,        nan,        nan,
       0.82207069, 1.1856311 , 1.40572111, 1.57256466, 1.70997595])

I want z to contain all the real roots of the operation without nan.


Solution

  • You can use trick like this:

    z = np.where(x<0, -np.abs(x)**(1/3), x**(1/3))