"""
Sam Dolan. Registration number: xxxxxxxxx
"""
# (1) Import modules
import numpy as np
import matplotlib.pyplot as plt
# (2) Define all the functions we will use here
def f1(x, A=3):
"""
Given a real number x, this function returns the value of f1(x) defined in the assignment.
Here A is an optional argument whose default value is 3.
"""
return A*x*np.exp(-x) - 1
def f1prime(x, A=3):
"""
Given a real number x, this function returns the value of the derivative of f1(x).
"""
return A*(1-x)*np.exp(-x)
def NRmethod(f, fprime, x0, max_steps=20, tol=1e-12, A=3):
"""
Seeks a root of a function by applying the Newton-Raphson step repeatedly to generate a sequence of estimates.
The first two arguments are functions that return f(x) and f’(x).
’x0’ is the initial guess.
The method should stop once either of two conditions is met:
(i) the absolute difference between successive estimates is
less than ’tol’ (indicating that the sequence has
converged sufficiently), or,
(ii) the number of steps exceeds ’max_steps’.
The function should return a list of the sequence of estimates,
[x0, x1, x2, ...]
"""
xs = [x0] # Create a list whose first entry is initial guess
for i in range(max_steps):
x = xs[-1] # take the last value in the list
xnew = x - f(x,A)/fprime(x,A) # one iteration of NR method
xs.append(xnew) # add the new value
if (np.abs(xs[-1] - xs[-2]) < tol):
# if the difference between the last two values in the list xs
# is less than 'tol', then end the loop now
break
return xs
# These are the functions for part 2
def f2(z, A=1):
return z**3 - 1
def f2prime(z, A=1):
return 3*z**2
# (3) Produce nicely-formatted output for each task on the sheet.
print("Part 1(a)")
print("See code: I have added a docstring to each function.")
print("-----")
print("Part 1(b)")
x0 = 0.0
xseq = NRmethod(f1, f1prime, x0)
print("With a starting value of x0=%.2f the NR method converged after %i iterations to: " % (x0, len(xseq)-1))
print(" %.10f" % xseq[-1])
x0 = 2.0 # Look for a second root
xseq = NRmethod(f1, f1prime, x0)
print("With a starting value of x0=%.2f the NR method converged after %i iterations to: " % (x0, len(xseq)-1))
print(" %.10f" % xseq[-1])
print("-----")
print("Part 1(c)")
print("The method fails for x = 1 because the the derivative f1'(x) is zero there.")
print("As this quantity appears in the denominator in the Newton-Raphson step,")
print("the right-hand side is not defined, and the method fails.")
print("The method also fails for values of x larger than approximately 9.9, ")
print("because the next value is negative and larger in magnitude - too large")
print("for the np.exp() function to handle - it returns a nan ('not a number')")
print()
#print(NRmethod(f1, f1prime, 1.0)) # uncomment these lines to see the failure
#print(NRmethod(f1, f1prime, 10.0))
print("-----")
print("Part 1(d)")
xseq = NRmethod(f1, f1prime, 0.0, max_steps=21, A=np.exp(1))
# max_steps=20 or 21 is acceptable here.
# xexact = 1.0
print("The estimate after 20 iterations is %.10f." % xseq[-1])
print("It differs from the exact root 1 by %.4e." % (1-xseq[-1]))
print("The NR method converges more slowly for A=e because the root is a repeated root, not a simple root,")
print("meaning that the derivative of the function is zero at the root, so the")
print("function is tangent to the axis there. ")
print("It can be shown that if x is a simple root then the convergence of the NR method is quadratic,")
print("whereas if x is a repeated root then the convergence of the NR method is only linear.")
print("")
print("------")
print("Part 2")
print("------")
print("Part 2(a)")
print("To find the three values, we solve z1 = 0 = z - (z^3 - 1) / (3 z^2)..")
print("which is equivalent to 2 z^3 + 1 = 0, which has three roots:")
r = (-0.5)**(1/3)
zs = [r, r*np.exp(2*np.pi/3*1j),r*np.exp(4*np.pi/3*1j)]
print(zs)
print("(In fact, if we start with one of these three values, z1 is not quite zero")
print("due to round-off error. z2 is then extremely large in magnitude.)")
print("------")
print("Part 2(b)")
zs = [0.02,0.02*1j, -0.01*(1+1j), -2**(-1/3) + 0.001*(1-1j)]
labels = ["(i)", "(ii)", "(iii)", "(iv)"]
for i in range(4):
zseq = NRmethod(f2, f2prime, zs[i], max_steps=40)
print(labels[i] + " z = ", zseq[-1])
print("In cases (i) and (ii) it converges to 1; in case (iii) to the cube root of unity with negative imaginary part, and in case (iv) to the cube root of unity with positive imaginary part.")
print("------")
print("Part 2(c)")
npts = 500 # I recommended n=300, but as the code is vectorized, it runs fast, so larger values can be used.
xs = np.linspace(-1.5,1.5,npts)
ys = np.linspace(-1.5,1.5,npts)
# Use broadcasting to make a 2D array of initial values for the complex number
zs = xs.reshape((npts,1)) + ys.reshape((1,npts))*1j
maxits = 20
for i in range(maxits):
zs = zs - (zs**3 - 1)/(3*zs**2)
roots = [1.0, np.exp(2*np.pi/3*1j), np.exp(4*np.pi/3*1j)]
basins = np.zeros((npts,npts), dtype=np.int8) # 2D array to contain values 0,1,2,3 only
eps = 1e-6 # A small number
for i in range(len(roots)):
diff = np.abs(zs - roots[i]) # How close to root i ?
basins += i*(diff < eps) # If close enough, assign the value i
plt.imshow(basins.transpose())
# Now make a labelled plot to save to file.
plt.figure()
plt.rcParams.update({'font.size': 14})
plt.gcf().set_size_inches(6,6)
plt.xlabel("Re(z)")
plt.ylabel("Im(z)")
plt.title("Basins of attraction")
plt.imshow(basins.transpose(), aspect='auto', extent=(-1.5,1.5,-1.5,1.5))
plt.tight_layout()
plt.savefig("basins.pdf")