専攻実験をPythonでやってみた（その３・終）

過去２回の続きです。

smon.hatenablog.com

smon.hatenablog.com

２次元フーリエ変換

最後のプログラミングの課題で，（数値データで）円を作る→フーリエ変換する→フーリエ逆変換（して元に戻ることを確認）する，という課題があるので，これをPythonでやってみます。

実験ではC言語でひたすら書いてもらっているのですが，Pythonだと拍子抜けするほど短くなります。

円を作る

import numpy as np
import matplotlib.pyplot as plt

nx = 256
ny = 256

i0 = 132
j0 = 124
r  = 84.0

re = np.zeros((nx, ny))
im = np.zeros((nx, ny))

for j in range(ny):
    for i in range(nx):
        if (i-i0)**2 + (j-j0)**2 < r**2:
            re[j][i] = 100.0
            im[j][i] = 0.0
        else:
            re[j][i] = 0.0
            im[j][i] = 0.0           
            
re.tofile('redata.flt')
im.tofile('imdata.flt')

これで円ができあがり，実部データと虚部データが保存されます。

２次元フーリエ変換

続いて保存した実部データと虚部データを使ってフーリエ変換していきます。

import numpy as np
import matplotlib.pyplot as plt

nx = 256
ny = 256

#fltファイルの読み込み
ispace_re = np.fromfile('redata.flt').reshape((nx, ny))
ispace_im = np.fromfile('imdata.flt').reshape((nx, ny))

ispace_data = np.vectorize(complex)(ispace_re, ispace_im)  #実部と虚部を複素数型として統合
ispace_data = ispace_data.reshape((nx, ny)) #1次元データを2次元データに並べ替え

kspace_data = np.fft.fft2(ispace_data)      #2次元フーリエ変換
kspace_data = np.fft.fftshift(kspace_data)  #直流成分が左上の端に来るので，中央に来るようにシフト

k_re = kspace_data.real  #実部のみ抽出
k_im = kspace_data.imag  #虚部のみ抽出

plt.imshow(k_re)

k_re.tofile('redata_fft.flt')
k_im.tofile('imdata_fft.flt')

なんとフーリエ変換が１行で書けてしまいます。恐るべしnumpy！
波数空間（k空間）の実部はこのようになります。

２次元フーリエ逆変換

上記と逆のことをすればOK。

import numpy as np
import matplotlib.pyplot as plt

nx = 256
ny = 256

#fltファイルの読み込み
kspace_re = np.fromfile('redata_fft.flt').reshape((nx, ny))
kspace_im = np.fromfile('imdata_fft.flt').reshape((nx, ny))

kspace_data = np.vectorize(complex)(kspace_re, kspace_im)  #実部と虚部を複素数型として統合
kspace_data = kspace_data.reshape((nx, ny)) #1次元データを2次元データに並べ替え

ispace_data = np.fft.ifft2(kspace_data)      #2次元フーリエ変換
#ispace_data = np.fft.fftshift(ispace_data)  #直流成分が左上の端に来るので，中央に来るようにシフト

ispace_re = ispace_data.real  #実部のみ抽出
ispace_im = ispace_data.imag  #虚部のみ抽出
ispace_abs = np.absolute(ispace_re, ispace_im)

plt.imshow(ispace_abs)

ispace_re.tofile('redata_ifft.flt')
ispace_im.tofile('imdata_ifft.flt')

絶対値画像を出してみると，きちんともとの円に戻っていることがわかります。

まとめ

もちろんC言語も勉強しておく必要がありますが，Pythonの便利さも（特に専攻実験を受けた方には）伝わったかなと思います。

Python初心者にとっては専攻実験の内容を焼き直すのは結構いい練習になるので，もし興味のある方は遊んでみると面白いかも。