Roots of Functions

In Q9 of the problem set, you're asked to solve the following equations $$x^2 + (a-4)x - 4a = 0,$$ where $a$ is a constant and $$4x^4 - 5x^2 - 6 = 0,$$ using the substitution $y = x^2$. Let's plot these out and see where the roots are!

There is an intresting article about roots here: https://www.quantamagazine.org/the-simple-math-behind-the-mighty-roots-of-unity-20210923/

import numpy as np
import matplotlib.pyplot as plt

# I took this from a forum on the internet (http://stackoverflow.com/questions/17044052/mathplotlib-imshow-complex-2d-array).
# It's very useful for plotting functions in the complex plane
from colorsys import hls_to_rgb

def colorize(z):
    r = np.abs(z)
    arg = np.angle(z) 

    h = (arg + np.pi)  / (2 * np.pi) + 0.5
    l = 1.0 - 1.0/(1.0 + r**0.5)
    s = 0.8

    c = np.vectorize(hls_to_rgb) (h,l,s) # --> tuple
    c = np.array(c)  # -->  array of (3,n,m) shape, but need (n,m,3)
    c = c.swapaxes(0,2)
    c = c.swapaxes(0,1)
    return c

lim = 6
x,y = lim*np.linspace(-1,1,500), lim*np.linspace(-1,1,500)
X,Y = np.meshgrid(x,y)
Z = X + 1j*Y

# Define the function we want to plot the roots of, and the parameter 'a'.
# Try changing 'a' to see how the root moves.  What happens if you make 'a' complex (i.e. a = 1.0*1j)?
a = 1.0
def f(x):
    return x**2 + (a-4.0)*x - 4.0*a 

# If you know the roots, you can find the minima becuase the roots of a quadratic must be symmetric about the minima
minima = (-a + 4) / 2

plt.figure(dpi=150)
plt.plot(x, f(x), 'k-') # Plot the function
plt.plot(4,0,'r*') # Plot the x = 4 root
plt.plot(-a,0,'b*') # Plot the x = -a root
plt.plot(minima, f(minima), 'mx') # Plot the minima
plt.grid(linestyle="--", color="k", alpha=0.5)
plt.xlabel("$x$")
plt.ylabel("$f(x)$")
plt.title("$f(x) = x^2 + (a-4)x - 4a$")

# Plot the function in the complex plane.  Note that for this function, all of the roots are along the real line: they have no imaginary part.
plt.figure(dpi=150)
plt.imshow(colorize(f(Z)), origin="lower", extent=(-lim,lim,-lim,lim))
plt.grid(linestyle="--", color="k", alpha=0.5)
plt.xlabel("Re[$z$]")
plt.ylabel("Im[$z$]")

plt.show()

lim = 2
x,y = lim*np.linspace(-1,1,500), lim*np.linspace(-1,1,500)
X,Y = np.meshgrid(x,y)
Z = X + 1j*Y

# Define the function we want to find the roots of
def f(z):
    return 4.0*z**4 - 5.0*z**2 - 6

# The roots
real_root = np.sqrt(2)
imag_root = np.sqrt(3/4)

plt.figure(dpi=150)
plt.imshow(colorize(f(Z)), origin="lower", extent=(-lim,lim,-lim,lim))
plt.grid(linestyle="--", color="k", alpha=0.5)

# Plot the REAL roots
plt.plot(-real_root,0,'r+')
plt.plot(real_root,0,'r+')
# Plot the COMPLEX roots
plt.plot(0.0, imag_root, 'bx')
plt.plot(0.0, -imag_root, 'bx')

plt.xlabel("Re[$z$]")
plt.ylabel("Im[$z$]")
plt.title("$f(z) = 4x^4 - 5x^2 - 6$")

plt.show()