##
## This is a python+matplotlib visualisation of Drill Exercise 4.2a
##
## Run it with python:
## >> python 3470_kap4_z-transform_ex_42a.py

from pylab import * 

# Create z-plane
z_min,z_max = -2,2
x = linspace(z_min,z_max,500)
X,Y = meshgrid(x,x)
z = X + 1j*Y

# Compute impulse reponse H(z) for all z
H = z / (0.5-z)

# Create figure
fig = figure()
ax = fig.add_subplot(111)

# Plot |H(z)|, and specify extent from -2 to 2 in both directions
cax = ax.imshow(abs(H),extent=(z_min,z_max,z_min,z_max))

# Peak value is very high. Clip at 5
cax.set_clim([0,3])

# Add colorbar and text
ccax = colorbar(cax)
ccax.set_label('$|$H(z)$|$')
ax.set_xlabel('Re\{z\}')
ax.set_ylabel('Im\{z\}')

# Draw a ROC circle
theta = linspace(0,2*pi,500)
ax.plot(0.5*sin(theta),0.5*cos(theta),'r')
ax.grid(True)
ax.text(-1.2,.7,'ROC: $|$z$|>$0.5',size=16,color='r')

# Mark pole and zero
ax.plot(0.5,0,'x',markersize=10,markeredgewidth=2,markeredgecolor=[0.5,1,0.5])
ax.plot(0,0,'o',markersize=10,markeredgewidth=2,markeredgecolor=[0.5,1,0.5])

# Save figure
savefig( 'pole_plot.pdf', bbox_inches='tight', pad_inches=0, format='pdf')