# Plot the stability region of a method

import matplotlib
import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt

def stab(z, g):
    # The stability function of an 2-stage SDIRK method 
    # of order >= 2 (Choose g).
    P = (g**2-2*g+0.5)*z**2-(2*g-1)*z+1
    Q = (1-g*z)**2
    return P/Q
    
g = 1/2+np.sqrt(3)/6   # gamma
x = np.linspace(-10,12,101)
y = np.linspace(-10,10,101)
X, Y = np.meshgrid(x, y)
Z = X+1j*Y
R = np.abs(stab(Z,g))
# R = np.abs(stab(Z,g)/np.exp(Z))

plt.clf()
plt.contour(X, Y, R,[1])
plt.axis('square')
plt.grid('on')
plt.contourf(X, Y, R,[-1,1])
plt.xlabel('Re(z)')
plt.ylabel('Im(z)')
plt.show()