# pyengr / examples / num_trans / transform.py

  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 #!/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: