cms.smoothing.smoothing_fft のソースコード

import sys
import csv
import numpy as np
from numpy import sin, cos, tan, pi
import ast

# mode:  'file': read data from csvfile, other: make from func(x)
mode = 'file'

# for mode != 'file'
x0 = 0.0
x1 = 1.0
ndata = 1024

# low and high index limits to remove from FT data
#kcut0 = 5
#kcut1 = 1020
kcut0 =  0
kcut1 = 10

csvfile = 'xrd.csv'
outfile  = 'smoothing-fft.csv'




def getarg(iarg, defval = None):
    try:
        return argv[iarg]
    except:
        return defval




def getintarg(iarg, defval = None):
    try:
        return int(argv[iarg])
    except:
        return defval

        


def getfloatarg(iarg, defval = None):
    try:
        return float(argv[iarg])
    except:
        return defval



if __name__ == "__main__":
    argv = sys.argv
    narg = len(argv)
    if narg <= 2:
        print("")
        print("Usage: python smmothing-fft.py csvfile kcut0 kcut1")
        print("")
        exit()
    else:
        csvfile = getarg(1,  csvfile)
        kcut0   = getfloatarg(2, kcut0)
        kcut1   = getfloatarg(3, kcut1)


# three trigonometric functions with randomized frequencies
# A: amplitude, f: frequency, p1: phase    


def func(x):
    x += 0.03 * np.random.rand();
    f1, p1, A1 = ( 1.5, pi/4.0, 1.0);
    f2, p2, A2 = ( 3.0, pi/3.0, 0.3);
    f3, p3, A3 = (10.0, pi/6.0, 0.5);
    return A1 * sin(2.0*pi * f1 * x + p1) \
         + A2 * sin(2.0*pi * f2 * x + p2) \
         + A3 * sin(2.0*pi * f3 * x + p3);




def main():
    global x0, x1, ndata
    global csvfile
    global outfile

    print("")
    print("========================================")
    print("mode:", mode)
    print("csv file:", csvfile)
    print("k range: {} - {}".format(kcut0, kcut1))
    
    if mode == 'file':
        x = []
        y = []
        print("Read data from [{}]".format(csvfile))
        with open(csvfile) as f:
            fin = csv.reader(f)
            next(fin)
            c = 0
            for row in fin:
                if c == 0:
                    label = row
                    print("{}\t{}".format(label[0], label[1]))
                else:
                     print("{}\t{}".format(row[0], row[1]))
                     x.append(float(row[0]))
                     y.append(float(row[1]))
                c += 1
        ndata = len(x)
        xstep = x[1] - x[0]
        x0 = x[1]
        x1 = x0 + xstep * (ndata-1)
    else:
        x = np.zeros(ndata)
        y = np.zeros(ndata)
        print("Make data by func(x) with random scattering")
        xstep = (x1 - x0) / ndata
        print("x = ({}, {}, {})".format(x0, x1, xstep))
        print("ndata={}".format(ndata))
        for i in range(0, ndata):
            x[i] = x0 + i * xstep
            y[i] = func(x[i])
        ndata = len(x)

    print("")
    print("ndata:", ndata)
# FFT
    F = np.fft.fft(y)
    
# calculate FFTed x axis
    xfrange = 1.0 / xstep
    dxf = xfrange / (ndata-1)
    xf = np.zeros(ndata)
    for i in range(0, ndata):
        xf[i] = i * dxf

    print("")
    print("x range     : {} - {}".format(min(x), max(x)))
    print("x(fft) range: {} - {}".format(min(xf), max(xf)))

# cut FT data outside [lf, hf] for smoothing
# note the FFTed data is symmetric for real part and anti-symmetric for imagenary part with respect to ndata/2
# LF data are near ix = 0 and ix = ndata - 1 HF data are around ndata/2
    Fs = np.zeros(ndata, dtype = 'complex')
    for i in range(0, int(ndata/2)):
        if kcut0 <= xf[i] <= kcut1:
            Fs[i] = F[i]
            Fs[ndata - i - 1] = F[ndata - i - 1]
        else:
#            print("cut: {} {}".format(i, xf[i]))
            Fs[i] = 0.0
            Fs[ndata - i - 1] = 0.0

# iFFT
    print("")
    
    fs = np.fft.ifft(Fs)

    f = open(outfile, 'w')
    fout = csv.writer(f, lineterminator='\n')
    fout.writerow(['x', 'y', 'x(fft)', 'y(fft).r', 'y(fft).i', 
                   'ys(fft,cut).r', 'ys(fft,cut).i',
                   'ys(ifft).r', 'ys(ifft).i', '|ys(ifft)|'])
    for i in range(0, ndata):
        fout.writerow([x[i], y[i], xf[i], F[i].real, F[i].imag, 
                       Fs[i].real, Fs[i].imag, 
                       fs[i].real, fs[i].imag, abs(fs[i])])

    print("")

    
    
if __name__ == "__main__":
    main()