Python loop doesn't return complex numbers (instead, returns nan)

Question

I've created a function that returns the roots of a 4th degree polynomial (I'm not using a ready library function for various reasons). The function returns up to 4 arguments, some can be real and some can be complex. When using the function outside of a loop it works exactly as expected. However, looping it causes the complex arguments to be returned as 'nan'.

I also tried putting the function code directly inside the loop, but I still got 'nan'.

Here is the function and loop. Only complex numbers should be returned when we set the parameters to mu = 0.5 and x >= -0.22 & x <= 0.22

def poly4(a, b, c, d, e):
    disc_0 = c**2 - 3*b*d + 12*a*e
    disc_1 = 2*c**3 - 9*b*c*d + 27*(b**2)*e + 27*a*d**2 - 72*a*c*e
    p = (8*a*c - 3*b**2) / (8*a**2)
    q = (b**3 - 4*a*b*c + 8*(a**2)*d) / (8*a**3)
    Q = ((disc_1 + (disc_1**2 - 4*disc_0**3)**0.5) / 2)**(1/3)
    S = 0.5 * (-(2/3)*p + (3*a)**(-1) * (Q + disc_0 / Q))**0.5

    x1 = -b/(4*a) - S + 0.5 * (-4*S**2 - 2*p + q/S)**0.5
    x2 = -b/(4*a) - S - 0.5 * (-4*S**2 - 2*p + q/S)**0.5
    x3 = -b/(4*a) + S + 0.5 * (-4*S**2 - 2*p - q/S)**0.5
    x4 = -b/(4*a) + S - 0.5 * (-4*S**2 - 2*p - q/S)**0.5
    return x1, x2, x3, x4

x = 0.1
mu = 0.5

a = 1
b = -2*x
c = x**2+mu**2-1
d = 2*x
e = -x**2

print(poly4(a, b, c, d, e))

x_vec = np.linspace(-0.2, 0.2, 5)
mu = 0.5

for i in x_vec:
    x = i
    a = 1
    b = -2*x
    c = x**2+mu**2-1
    d = 2*x
    e = -x**2
    result = poly4(a, b, c, d, e)
    print(result)

Answer 1

The problem is that inside your for i in x_vec: loop, i is a NumPy (rank-0) array -- a scalar. These have a specific type (likely numpy.float64) that will not be automatically cast to a complex-capable type (like e.g. numpy.complex128), contrarily to what happens with Python floats being silently casted to complex as needed.

A simple fix would just be to replace x = i with x = i.astype(complex) or other similar casting (e.g. x_vec = np.linspace(-0.2, 0.2, 5).astype(complex)), if you need i or x_vec to be a NumPy array.

Alternatively, you may want to convert x_vec to a Python list (e.g. x_vec = np.linspace(-0.2, 0.2, 5).tolist()) which will then ensure i/x to fall-back to Python automatic casting.