import sys import csv import numpy as np from numpy import sin, cos, tan, pi from numpy import exp, log, sin, cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, sqrt, abs from random import random as rand import pandas as pd import matplotlib.pyplot as plt from tklib.tkutils import replace_path, getarg, getintarg, getfloatarg, pint, pfloat from tklib.tksci.tksci import kB, e from tklib.tkvariousdata import tkVariousData from tklib.tkapplication import tkApplication from tklib.tkparams import tkParams from tklib.tksci.tksci import asin, acos, atan, degcos, degsin, degtan, degacos, degasin, degatan, arcsin, arccos, arctan from tklib.tksci.tksci import factorial, log10, gamma, combination, eVTonm, nmToeV, Bn from tklib.tksci.tksci import h, h_bar, hbar, e, kB, NA, c, pi, pi2, torad, todeg, basee from tklib.tksci.tksci import me, mp, mn from tklib.tksci.tksci import u0, e0, a0, R, F, g, G from tklib.tksci.tksci import HartreeToeV, RyToeV, KToeV, eVToK, JToeV, eVToJ, Debye app = tkApplication() cparams = tkParams() cparams.func_str = 'sin(10.0 * pi * x) + 0.5 * rand()' cparams.mode = 'FFT' cparams.xmin = -1.0 cparams.xmax = 1.0 cparams.nx = 1024 cparams.figsize = [12, 8] cparams.fontsize = 12 cparams.legend_fontsize = 10 logfile = 'fft_func-out.txt' #app.replace_path(cparams.infile) print(f"Open logfile [{logfile}]") app.redirect(targets = ["stdout", logfile], mode = 'w') def usage(app): print(f"Usage: python {sys.argv[0]} func_str x0 x1 nx") def update_vars(app): cparams.mode = getarg(1, cparams.mode) cparams.func_str = getarg(2, cparams.func_str) cparams.xmin = getfloatarg(3, cparams.xmin) cparams.xmax = getfloatarg(4, cparams.xmax) cparams.nx = getintarg (5, cparams.nx) def func(x): y = eval(cparams.func_str, globals(), {"x": x}) return y def common_func(): print("") print("========================================") print(f"mode: {cparams.mode}") print(f"func: {cparams.func_str}") print(f"xrange : {cparams.xmin} - {cparams.xmax} for {cparams.nx} points") print("") print(f"Build function [{cparams.func_str}]") cparams.xstep = (cparams.xmax - cparams.xmin) / (cparams.nx - 1) x = [cparams.xmin + i * cparams.xstep for i in range(cparams.nx)] y = [func(_x) for _x in x] return x, y def IFFT(): x, y = common_func() F = np.fft.ifft(y) # calculate FFTed x axis xfrange = 1.0 / cparams.xstep dxf = xfrange / (cparams.nx - 1) xf = [i * dxf for i in range(cparams.nx)] print("") print("x range : {} - {}".format(min(x), max(x))) print("x(fft) range: {} - {}".format(min(xf), max(xf))) Fr = [F[i].real for i in range(cparams.nx)] Fi = [F[i].imag for i in range(cparams.nx)] Fr = [F[i].real for i in range(cparams.nx)] df = pd.DataFrame(np.array([x, y, xf, Fr, Fi, Fr]).T, columns = ['x', 'y', 'x(fft)', 'y(fft).r', 'y(fft).i', '|y(fftr)|']) outfile = 'fft_func-ifft.xlsx' #app.replace_path(cparams.infile, template = ["{dirname}", "{filebody}-ifft.xlsx"]) print("") print(f"Save to [{outfile}]") df.to_excel(outfile, index = False) print("") print("Plot") fig, axes = plt.subplots(1, 2, figsize = cparams.figsize) axes = axes.flatten() # ax2 = axes[1].twinx() axes[0].tick_params(labelsize = cparams.fontsize) axes[1].tick_params(labelsize = cparams.fontsize) axes[0].plot(x, y, label = 'input data', color = 'black', linewidth = 0.5) axes[0].set_xlabel('x', fontsize = cparams.fontsize) axes[0].set_ylabel(f'y = {cparams.func_str}', fontsize = cparams.fontsize) axes[0].legend(fontsize = cparams.legend_fontsize) ins1 = axes[1].plot(xf, Fr, label = 'real', color = 'black', linewidth = 0.5) ins2 = axes[1].plot(xf, Fi, label = 'imag', color = 'red', linewidth = 0.5) ins3 = axes[1].plot(xf, Fr, label = 'abs', color = 'blue', linewidth = 0.5) axes[1].set_xlabel('k', fontsize = cparams.fontsize) axes[1].set_ylabel('Inverse $FT^*$', fontsize = cparams.fontsize) axes[1].legend(fontsize = cparams.legend_fontsize) plt.tight_layout() plt.pause(0.001) app.terminate("\n", usage = usage, pause = True) def FFT(): x, y = common_func() F = np.fft.fft(y) # calculate FFTed x axis xfrange = 1.0 / cparams.xstep dxf = xfrange / (cparams.nx - 1) xf = [i * dxf for i in range(cparams.nx)] print("") print("x range : {} - {}".format(min(x), max(x))) print("x(fft) range: {} - {}".format(min(xf), max(xf))) Fr = [F[i].real for i in range(cparams.nx)] Fi = [F[i].imag for i in range(cparams.nx)] Fr = [F[i].real for i in range(cparams.nx)] df = pd.DataFrame(np.array([x, y, xf, Fr, Fi, Fr]).T, columns = ['x', 'y', 'x(fft)', 'y(fft).r', 'y(fft).i', '|y(fftr)|']) outfile = 'fft_func-fft.xlsx' #app.replace_path(cparams.infile, template = ["{dirname}", "{filebody}-fft.xlsx"]) print("") print(f"Save to [{outfile}]") df.to_excel(outfile, index = False) print("") print("Plot") fig, axes = plt.subplots(1, 2, figsize = cparams.figsize) axes = axes.flatten() # ax2 = axes[1].twinx() axes[0].tick_params(labelsize = cparams.fontsize) axes[1].tick_params(labelsize = cparams.fontsize) axes[0].plot(x, y, label = 'input data', color = 'black', linewidth = 0.5) axes[0].set_xlabel('x', fontsize = cparams.fontsize) axes[0].set_ylabel(f'y = {cparams.func_str}', fontsize = cparams.fontsize) axes[0].legend(fontsize = cparams.legend_fontsize) ins1 = axes[1].plot(xf, Fr, label = 'real', color = 'black', linewidth = 0.5) ins2 = axes[1].plot(xf, Fi, label = 'imag', color = 'red', linewidth = 0.5) ins3 = axes[1].plot(xf, Fr, label = 'abs', color = 'blue', linewidth = 0.5) axes[1].set_xlabel('k', fontsize = cparams.fontsize) axes[1].set_ylabel('$FT^*$', fontsize = cparams.fontsize) axes[1].legend(fontsize = cparams.legend_fontsize) plt.tight_layout() plt.pause(0.001) app.terminate("\n", usage = usage, pause = True) def main(): print("") print( "#######################################") print(f"# {sys.argv[0]}") print( "#######################################") update_vars(app) cparams.print_parameters() if cparams.mode == 'FFT': FFT() elif cparams.mode == 'IFFT': IFFT() else: app.terminate(f"\nmode [{cparams.mode}] is not implemented.\n", usage = usage, pause = True) if __name__ == "__main__": main()