MatplotlibとNumPyで周波数応答を複素平面上に描く，ついでに振幅/位相特性も描く

ディジタルフィルタの周波数応答を可視化する機会があったのでメモしておきます．

import

import numpy as np
import matplotlib.pyplot as plt

##複素平面上に描く
複素関数を高校数学の教科書みたいな複素平面に描きたい！

def plot_cp(func): # 複素平面上に描く
    # FigureとAxesを描画
    fig, ax = plt.subplots(figsize = (5, 5))
    ax.grid()

    # 表示範囲を設定 適宜弄る
    lim = [-2.5, 2.5]
    ax.set_xlim(lim)
    ax.set_ylim(lim)

    # 実軸
    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')
    # 虚軸
    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')
    # 原点
    plt.text(0.1*lim[0],0.1*lim[0], '$O$')

    # 余分な目盛りを削除
    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)

    # 虚軸の目盛りを"n j"に変更
    imlabel=[]
    for ticks in ax.get_yticks():
        imlabel.append(str(ticks)+" j")
    ax.set_yticklabels(imlabel) 

    #軸を中央に移動
    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))

    # 出力
    plt.show()

###例）
####入力
$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) # 複素平面上に描く

####出力

イイ感じ．

##振幅特性と位相特性を並べて描く

def plot_ap(func,theta): # 振幅特性と位相特性を描く
    # Figure
    fig=plt.figure(figsize = (10, 5))

    # ax1 振幅特性
    ax1=fig.add_subplot(121)
    ax1.grid()

    gain=abs(func)

    # 表示範囲を設定
    lim=[-np.pi, np.pi]
    ax1.set_xlim(lim)
    ax1.set_ylim(0,1.05*np.max(gain))

    # x軸のラベルを設定
    ax1.set_xlabel("$\omega$")
    # タイトルを設定
    ax1.set_title("amplitude characteristic")

    # x軸の目盛りを[-pi/T,pi/T]仕様に
    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 位相特性
    ax2=fig.add_subplot(122)
    ax2.grid()

    # 表示範囲を設定
    ax2.set_xlim(lim)
    ax2.set_ylim(lim)

    # x軸のラベルを設定
    ax2.set_xlabel("$\omega$")
    # タイトルを設定
    ax2.set_title("phase characteristic")

    # x軸の目盛りを[-pi/T,pi/T]仕様に
    ax2.set_xticks(xt)
    ax2.set_xticklabels(xl)

    # y軸の目盛りを[-pi,pi]仕様に
    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)))

    # 出力
    plt.show()

###例）
####入力
$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) # 振幅特性と位相特性を並べて描く

####出力

それっぽい．