Source

pyengr / examples / num_trans / transform.py

Full commit
#!/usr/bin/env python
# -*- coding: utf8 -*-

from numpy import empty, arange, absolute, sinc, pi, sin
from numpy.fft import fft, fftshift
from fftw3 import Plan
from matplotlib import pyplot as plt

ngrid = 256
extent = (-3.*pi, 3.*pi)
#extent = (-16., 16.)    # power of 2 will match.

xgrid = arange(ngrid, dtype='float64')
xgrid /= ngrid - 1
xgrid *= extent[1] - extent[0]
xgrid += extent[0]
bw = 2*pi / (xgrid[1]-xgrid[0])
kgrid = arange(ngrid, dtype='float64')
kgrid /= ngrid 
kgrid *= bw
kgrid -= bw/2
#kscale = (extent[1]-extent[0]) / ngrid
kscale = 1.0 / (ngrid/2)
xarr = empty(ngrid, dtype='complex128')
karr = empty(ngrid, dtype='complex128')
karrw = empty(ngrid, dtype='complex128')
plan = Plan(xarr, karrw, direction='forward', flags=['estimate'])

#xarr.fill(0)
#xarr[absolute(xgrid) <= 0.5] = 1
xarr[:] = sin(xgrid)
# use numpy.fft.fft.
karr = fft(xarr)
karrs = fftshift(karr)
karrs *= kscale
# use fftw3.Plan.
plan()
karrws = fftshift(karrw)
karrws *= kscale
# analytical solution.
kana = sinc(kgrid/(2*pi))

fig = plt.figure(figsize=(20, 8))

xax = fig.add_subplot(1, 2, 1)
xax.plot(xgrid, xarr.real)
xax.set_xlim(xgrid[0], xgrid[-1])
xax.set_ylim(-1.1, 1.1)
xax.set_xlabel('$x$')
xax.grid()

kax = fig.add_subplot(1, 2, 2)
kax.plot(kgrid/pi, absolute(karrs), label='numpy.fft.fft')
kax.plot(kgrid/pi, absolute(karrws), label='fftw3.Plan')
#kax.plot(kgrid/pi, absolute(kana), label='analytical')
kax.set_xlim(kgrid[0]/pi, kgrid[-1]/pi)
kax.set_xlabel('$k$ (radius in $\pi$)')
kax.grid()
kax.legend()

plt.show()

# vim: set ft=python ff=unix fenc=utf8 ai et sw=4 ts=4 tw=79: