Draw the frequency response on the complex plane with Matplotlib and NumPy, and also draw the amplitude / phase characteristics

I had the opportunity to visualize the frequency response of the digital filter, so make a note of it.

import

import numpy as np
import matplotlib.pyplot as plt

Draw on the complex plane

I want to draw a complex function on a complex plane like a high school math textbook!

def plot_cp(func,theta): #Draw on the complex plane
    #Draw figures and axes
    fig, ax = plt.subplots(figsize = (5, 5))
    ax.grid()

    #Set the display range and play with it as appropriate
    lim = [-2.5, 2.5]
    ax.set_xlim(lim)
    ax.set_ylim(lim)

    #Real axis
    plt.quiver(lim[0],0,lim[1]-lim[0],0,angles='xy',scale_units='xy',width=0.005,headwidth=10,headlength=10,headaxislength=5,scale=1)
    plt.text(1.05*lim[1],0.02*lim[0], 'Re')
    #Imaginary axis
    plt.quiver(0,lim[0],0,lim[1]-lim[0],angles='xy',scale_units='xy',width=0.005,headwidth=10,headlength=10,headaxislength=5,scale=1)
    plt.text(0.03*lim[0],1.05*lim[1], 'Im')
    #origin
    plt.text(0.1*lim[0],0.1*lim[0], '$O$')

    #Remove extra scale
    xt=list(ax.get_xticks())
    for i in [0,np.floor(lim[0]),np.ceil(lim[1])]:
        xt.remove(i)
    ax.set_xticks(xt)
    ax.set_yticks(xt)

    #The scale of the imaginary axis"n j"change to
    imlabel=[]
    for ticks in ax.get_yticks():
        imlabel.append(str(ticks)+" j")
    ax.set_yticklabels(imlabel) 

    #Move axis to center
    ax.spines['bottom'].set_position('center')
    ax.spines['left'].set_position('center')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    plt.plot(np.real(func),np.imag(func))

    #output
    plt.show()

Example)

input

f(\omega)=1+\exp(-2\mathrm{j}\omega T)\quad \left(\omega \in \left[-\frac{\pi}{T},\frac{\pi}{T}\right]\right)

theta=np.linspace(-np.pi,np.pi,1000) # theta=omega*T
func=1+np.exp(-2j*theta)
plot_cp(func,theta) #Draw on the complex plane

output

It feels good.

Draw the amplitude characteristics and phase characteristics side by side

def plot_ap(func,theta): #Draw amplitude and phase characteristics
    # Figure
    fig=plt.figure(figsize = (10, 5))

    #ax1 Amplitude characteristics
    ax1=fig.add_subplot(121)
    ax1.grid()

    gain=abs(func)

    #Set display range
    lim=[-np.pi, np.pi]
    ax1.set_xlim(lim)
    ax1.set_ylim(0,1.05*np.max(gain))

    #Set x-axis label
    ax1.set_xlabel("$\omega$")
    #Set title
    ax1.set_title("amplitude characteristic")

    #x-axis scale[-pi/T,pi/T]To specifications
    xt=[-np.pi,-np.pi/2,0,np.pi/2,np.pi]
    ax1.set_xticks(xt)
    xl=["$-\pi/T$","$-\pi/2T$",0,"$\pi/2T$","$\pi/T$"]
    ax1.set_xticklabels(xl)

    ax1.plot(theta,gain)

    #ax2 phase characteristics
    ax2=fig.add_subplot(122)
    ax2.grid()

    #Set display range
    ax2.set_xlim(lim)
    ax2.set_ylim(lim)

    #Set x-axis label
    ax2.set_xlabel("$\omega$")
    #Set title
    ax2.set_title("phase characteristic")

    #x-axis scale[-pi/T,pi/T]To specifications
    ax2.set_xticks(xt)
    ax2.set_xticklabels(xl)

    #y-axis scale[-pi,pi]To specifications
    ax2.set_yticks(xt)
    yl=["$-\pi$","$-\pi/2$",0,"$\pi/2$","$\pi$"]
    ax2.set_yticklabels(yl)

    ax2.plot(theta,np.arctan2(np.imag(func),np.real(func)))

    #output
    plt.show()

Example)

input

f(\omega)=1+\exp(-2\mathrm{j}\omega T)\quad \left(\omega \in \left[-\frac{\pi}{T},\frac{\pi}{T}\right]\right)

theta=np.linspace(-np.pi,np.pi,1000) # theta=omega*T
func=1+np.exp(-2j*theta)
plot_ap(func,theta) #Draw the amplitude characteristics and phase characteristics side by side

output

It looks like that.