Can cmath.sqrt be applied to a NumPy array?-CodePudding

I want to calculate the square root of a numpy array of negative numbers.

I tried with np.sqrt() but it gives error beacuse the domain.

Then, I found that for complex numbers you can use cmath.sqrt(x) but it also gives me an error.

Here's my code

import numpy as np
import cmath
from cmath import sqrt

x = np.arange(-10, 10, 0.01)
E = 1

p1 = cmath.sqrt(E - x**2)

And got this error

Traceback (most recent call last):
  File "C:\Users\os.py", line 49, in <module>
    p1 = cmath.sqrt(E - x**2)
TypeError: only length-1 arrays can be converted to Python scalars

Later I tried to use a for loop and it's not possible either. Here's the code:

import numpy as np
import cmath
from cmath import sqrt

x = np.arange(-10, 10, 0.01)
E = 1

for i in range(0, len(x)):
    p1 = cmath.sqrt(E - x(i)**2)

and the message error

Traceback (most recent call last):
  File "C:\Users\os.py", line 48, in <module>
    p1 = cmath.sqrt(E - x(i)**2)
TypeError: 'numpy.ndarray' object is not callable

I don't know what am I dpoing wrong, can anyone help me?, please. I need to calculate the square root of an numpy array of negative numbers, does anyone know how to do this?

CodePudding user response：

Yes, you can iterate to perform cmath on individual elements of x:

In [254]: np.array([cmath.sqrt(E-i**2) for i in x])
Out[254]: 
array([0. 9.94987437j, 0. 9.93982394j, 0. 9.92977341j, ...,
       0. 9.91972278j, 0. 9.92977341j, 0. 9.93982394j])

but np.sqrt will work if you give it a complex dtype array:

In [255]: np.sqrt(E-x.astype(complex)**2)
Out[255]: 
array([0. 9.94987437j, 0. 9.93982394j, 0. 9.92977341j, ...,
       0. 9.91972278j, 0. 9.92977341j, 0. 9.93982394j])

some comparative timings

In [259]: timeit np.emath.sqrt(E-x**2)
166 µs ± 336 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [260]: timeit np.sqrt(E-x.astype(complex)**2)
129 µs ± 82.1 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [261]: timeit np.array([cmath.sqrt(E-i**2) for i in x])
2.54 ms ± 4.36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

CodePudding user response：

You could use numpy.emath.sqrt() which will work with negative numbers and return the complex square roots:

import numpy as np

x = np.array([-4, -16, -25, 100, 64])

np.emath.sqrt(x)
# array([ 0. 2.j,  0. 4.j,  0. 5.j, 10. 0.j,  8. 0.j])