pythonsympympmath

How to find roots of a complex function with python?

I'm working with a complex function whose coordinate is omega. I formed it as follows:

from sympy import *

def alpha(n, rho, l):
     return n**2 + (2*rho + 2)*n + 2*rho + 1

def beta(n, rho, l):
     return -(2*n**2 + (8*rho + 2)*n + 8*rho**2 + 4*rho + l*(l+1) -3)

def gamma(n, rho, l):
     return n**2 + 4*rho*n + 4*rho**2 - 4

def ContinuedFrac(n, rho, l):
     cf = 0
     for i in reversed(range(1,n+1)):
         cf = cf + beta(i, rho, l)/alpha(i, rho, l)
         cf = (-gamma(i, rho, l)/alpha(i, rho, l)) / cf
     return cf

And I would like to find the roots for which the function ContinuedFrac is equal to beta(0, i*omega, 3)/alpha(0, i*omega, 3) and admits the values n=3, rho = i*omega and l=3. I tried to find the roots using the findroot function from the mpmath library:

from sympy import *
import mpmath as mp

omega = symbols("ω")

f = ContinuedFrac(1, i*omega, 3) + beta(0, i*omega, 3)/alpha(0, I*omega, 3)
findroot(f, 2-0.5j)

where 2-0.5j is an initial values that is needed for convergence. I wanted an answer like (1.198 - 0.185 i) but the result was an error (photo below). What could I be doing wrong?

enter image description here

Solution

  • Note that you used i to indicate the imaginary number, but in SymPy it is I.

    The reason you see that error is because findroot requires a numerical function, but you passed in a symbolic expression. You could use lambdify to convert f to a numerical function, but we can get away with a different procedure instead.

    You could try to use solve or solveset, but it is not guaranteed that they find all roots. For example:

    sol = solve(f, omega)
sol = [s.n() for s in sol]
for s in sol:
    print(s)
# out:
# -0.92071676505822 + 0.74526156471082*I
# 0.92071676505822 + 0.74526156471082*I
# 1.13988532510002 + 0.19223843528918*I
# -1.13988532510002 + 0.19223843528918*I

    In this case, we can create a 3D plot with the magnitude of the function and verify that all zeros have been found. To do that, I'm going to use this plotting module:

    from spb import *
l = 2
plot_complex(log(f), (omega, -l-l*1j, l+l*1j), n=500,
             zlabel="log(abs(f))", threed=True, backend=KB)

    enter image description here

    The downward spikes are the zeros. The upward spikes are the poles. We can count 4 zeros, which match what solve found.

    Now, let's talk about SymPy's nsolve (numerical solver), which is going to use mpmath's findroot: by providing appropriate initial guesses, we can use nsolve to find the roots. To find appropriate initial guesses, I'm going to plot the magnitude of the complex function:

    p1 = plot_complex(f, (omega, -l-l*1j, l+l*1j), {"interpolation": "spline36"},
    coloring="e", grid=False, n=1000)

    enter image description here

    These zebra-stripes represents the magnitude: they are circling around the zeros and poles. In this kind of plot, we can't distinguish between them, however we can obtain a initial guess for nsolve. For example:

    s1 = nsolve(f, omega, -1+0.25j)
s2 = nsolve(f, omega, 1+0.25j)
s3 = nsolve(f, omega, -0.9+0.75j)
s4 = nsolve(f, omega, 0.9+0.75j)
p2 = plot_complex_list([s1, s2, s3, s4], {"marker": "^", "color": "r"}, show=False)
(p1 + p2).show()

    enter image description here

    This last picture overlays the solutions to the previous plot. If you try to use nsolve, you'll soon discover that it is very sensitive to the initial guess, in particular for complex functions where, as you can see from the 3D plot, the spatial distance between a zero and a pole might be very small. This means that saddle points might force the algorithm to fall into the same root even with very different initial guesses.