PeteyCoco
PeteyCoco

Reputation: 149

Storing roots of a complex function in an array in SciPy

I have the complex function below:

 import numpy as np
 import scipy as sp
 from scipy.special import jv, hankel1, jvp, h1vp, h2vp, gamma

 nr = 2
 m = 20

 def Dm(x):
     return nr * jvp(m,nr*x,1) * hankel1(m,x) - jv(m,nr*x) * h1vp(m,x,1)

I'm looking to find as many complex roots of Dm(x) that lie in the 4th quadrant of the complex plane as I can using newton() from scipy.optimize and then store them into a one dimensional array. The best method I can come up with is to brute force it by using newton() at regularly spaced intervals across a finite portion of the 4th quadrant, check if the root is a duplicate of a previous root, check if the root is indeed a root, and then store it into an array. Once that algorithm finishes I want to sort the array by increasing real components. My questions are:

(i) Can I create an array of undefined length that I can keep adding values to as I find them?

(ii) Can I plot the function in such a way that I can visualize the roots? The math says that they are all on one sheet of the complex plane.

(iii) Is there a better method for finding the roots? I feel that I'll miss a lot of roots in the domain with my method.

Upvotes: 2

Views: 868

Answers (2)

Warren Weckesser
Warren Weckesser

Reputation: 114911

@DrV's answer looks good, so here I'll only suggest another approach for part (ii) of your question. A useful way to visualize the roots of a complex function is to plot the 0 contours of the real and imaginary parts. That is, compute z = Dm(...) on a reasonably dense grid, and then use matplotlib's contour function to plot the contours where z.real is 0 and where z.imag is zero. The roots of the function are the points where these contours intersect.

For example, here's some code that generates a plot of the contours.

import numpy as np
from scipy.special import jv, hankel1, jvp, h1vp
import matplotlib.pyplot as plt


nr = 2
m = 20

def Dm(x):
    return nr * jvp(m,nr*x,1) * hankel1(m,x) - jv(m,nr*x) * h1vp(m,x,1)


x = np.linspace(0, 40, 2500)
y = np.linspace(-15, 5, 2000)
X, Y = np.meshgrid(x, y)
z = Dm(X + 1j*Y)

plt.contour(x, y, z.real, [0], colors='b', lw=0.25)
plt.contour(x, y, z.imag, [0], colors='r', lw=0.25)
plt.savefig('contours.png')

Here's the plot. Each intersection of a red and blue line is a root.

contours

The plot suggests that you shouldn't expect to find roots with large negative imaginary parts. You could probably look into the asymptotic behavior of the Bessel and Hankel functions with large arguments for ideas to prove this.

Upvotes: 2

DrV
DrV

Reputation: 23520

Some answers:

(i) Use a list. Arrays have fixed size. Appending to a list is a very cheap option. When you append a new root to the list, check that the previous root is not in the list by, e.g., calculating np.amin(np.abs(np.array(a)-b)) where a is the list of existing roots and b is the new root. If this value is very small, you have arrived at an existing root. (How small depends on the function. It cannot be 0.0, as then you are usually not recognizing the same root due to floating point and iteration inaccuracies.)

If you have a very large number of roots (thousands), you may want to sort them as soon as you receive them. This makes searching the matching roots faster. On the other hand, most probably >90% of the time is spent in iterating the roots, and you do not need to worry about other performance issues. Then you just compile the list, sort it (list sorting is easy and fast) and convert into an array, if you need.

(ii) Yes. Two examples below: (For countour stuff, thank you belongs to Warren Weckesser and his very good answer!)

import numpy as np
from scipy.special import jv, hankel1, jvp, h1vp
import matplotlib.pyplot as plt


nr = 2
m = 20

# create a 2000 x 2000 sample complex plane between -2-2i .. +2+2i
x = np.linspace(-2, 2, 2000)
y = np.linspace(-2, 2, 2000)
X, Y = np.meshgrid(x, y)
C = X + 1j * Y 
z = 1-C**2

# draw a contour image of the imaginary part (red) and real part (blue)
csr = plt.contour(x, y, z.real, 5, colors='b')
plt.clabel(csr)
csi = plt.contour(x, y, z.imag, 5, colors='r')
plt.clabel(csi)
plt.axis('equal')
plt.savefig('contours.png')

# draw an image of the absolute value of the function, black representing zeros
plt.figure()
plt.imshow(abs(z), extent=[-2,2,-2,2], cmap=plt.cm.gray)
plt.axis('equal')
plt.savefig('absval.png')

This gives countours.png:

Countours of 1-C^2

and absval.png:

Absolute value of 1-C^2

Note that if you want to zoom the images, it is usually necessary to change the limits and recalculate the complex values z to avoid missing details. The images can of course be plotted on top of each other, the image colour palette can be changed, countour has a million options. If you only want to draw the zeros, replace the number 5 (number of contours) with [0] (plot only the listed contour) in the countour calls.

Of course, you will replace my (1-C^2) by your own function. The only thing to take care is that if the function receives an array of complex numbers, it returns an array of results in the same shape calculated pointwise. Imshow needs to receive an array of scalars. For more information, please see the imshow documentation.

(iii) There may be, but there are no general methods for finding all minima/maxima/zeros of arbitrary functions. (The function may even have an infinite number of roots.) Your idea of first plotting the function is a good one. Then you'll understand its behaviour easier.

Upvotes: 2

Related Questions