高速フーリエ変換、畳み込み和、自己相関を求めるプログラム

fft_conv_corr.cpp

/*
fft_conv_corr.cpp : 高速フーリエ変換、畳み込み和、自己相関を求めるプログラム
author            : Koyo Miyamura 
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <climits>
#include <iostream>

#define N_S 1024
#define SIM_S 10000

typedef struct{
    double r, i;
} complex;

complex complex_add(complex z1, complex z2){
    complex tmp;
    tmp.r = z1.r + z2.r;
    tmp.i = z1.i + z2.i;
    return tmp;
}

complex complex_sub(complex z1, complex z2){
    complex tmp;
    tmp.r = z1.r - z2.r;
    tmp.i = z1.i - z2.i;
    return tmp;
}


complex complex_mul(complex z1, complex z2){
    complex tmp;
    tmp.r = z1.r * z2.r - z1.i * z2.i;
    tmp.i = z1.r * z2.i + z1.i * z2.r;
    return tmp;
}

complex complex_div(complex z1, complex z2){
    complex tmp;
    double  den;
    den = z2.r * z2.r + z2.i * z2.i;

    if(den == 0){
        printf("Fatal Error: divided by zero!\n");
        exit(1);
    }
    tmp.r = (z1.r * z2.r + z1.i * z2.i)/ den;
    tmp.i = (z2.r * z1.i - z1.r * z2.i)/ den;
    return tmp;
}

complex complex_jexp(double x){
    complex tmp;
    tmp.r = cos(x);
    tmp.i = sin(x);
    return tmp;
}

// fft_mode: 0->FFT, 1->IFFT
void fft(complex *x, long N, int fft_mode){
    int j;
    complex fe[N/2],fo[N/2];
    if(fft_mode == 0){//FFT
        if(N > 2){
            for(j = 0; j < N/2; j++){
                fe[j].r = x[2*j].r;
                fe[j].i = x[2*j].i;

                fo[j].r = x[2*j+1].r;
                fo[j].i = x[2*j+1].i;
            }
            //再起
            fft(fe,N/2,0);
            fft(fo,N/2,0);

            for(j = 0; j < N/2; j++){
                x[j].r = ( (fe[j].r + (complex_mul(complex_jexp(-(2*M_PI/N) * j), fo[j])).r ) ) / sqrt(2);
                x[j].i = ( (fe[j].i + (complex_mul(complex_jexp(-(2*M_PI/N) * j), fo[j])).i ) ) / sqrt(2);
            }
            for(j = 0; j < N/2; j++){
                x[j+N/2].r = ( (fe[j].r + (complex_mul(complex_jexp(-(2*M_PI/N) * (j+N/2)), fo[j])).r) ) / sqrt(2);
                x[j+N/2].i = ( (fe[j].i + (complex_mul(complex_jexp(-(2*M_PI/N) * (j+N/2)), fo[j])).i) ) / sqrt(2);
            }
        }
        else{
            complex f[2];
            f[0].r = x[0].r;
            f[0].i = x[0].i;
            f[1].r = x[1].r;
            f[1].i = x[1].i;

            x[0].r = (f[0].r + f[1].r) / sqrt(2);
            x[0].i = (f[0].i + f[1].i) / sqrt(2);

            x[1].r = (f[0].r - f[1].r) / sqrt(2);
            x[1].i = (f[0].i - f[1].i) / sqrt(2);
        }
    }
    else{//IFFT
        if(N > 2){
            for(j = 0; j < N/2; j++){
                fe[j].r = x[2*j].r;
                fe[j].i = x[2*j].i;

                fo[j].r = x[2*j+1].r;
                fo[j].i = x[2*j+1].i;
            }
            //再起
            fft(fe,N/2,1);
            fft(fo,N/2,1);

            for(j = 0; j < N/2; j++){
                x[j].r = ( (fe[j].r + (complex_mul(complex_jexp((2*M_PI/N) * j), fo[j])).r ) ) / sqrt(2);
                x[j].i = ( (fe[j].i + (complex_mul(complex_jexp((2*M_PI/N) * j), fo[j])).i ) ) / sqrt(2);
            }
            for(j = 0; j < N/2; j++){
                x[j+N/2].r = ( (fe[j].r + (complex_mul(complex_jexp((2*M_PI/N) * (j+N/2)), fo[j])).r) ) / sqrt(2);
                x[j+N/2].i = ( (fe[j].i + (complex_mul(complex_jexp((2*M_PI/N) * (j+N/2)), fo[j])).i) ) / sqrt(2);
            }
        }
        else{
            complex f[2];
            f[0].r = x[0].r;
            f[0].i = x[0].i;
            f[1].r = x[1].r;
            f[1].i = x[1].i;

            x[0].r = (f[0].r + f[1].r) / sqrt(2);
            x[0].i = (f[0].i + f[1].i) / sqrt(2);

            x[1].r = (f[0].r - f[1].r) / sqrt(2);
            x[1].i = (f[0].i - f[1].i) / sqrt(2);
        }
    }    
}

// 畳み込み和
void conv(complex* x, complex* y, complex* z, long N){
    int i;
    complex cmp_sq, tmp_x[N], tmp_y[N];
    cmp_sq.r = sqrt(N);
    cmp_sq.i = 0;

    for(i = 0; i < N; i++){
        tmp_x[i].r = x[i].r;
        tmp_x[i].i = x[i].i;
        tmp_y[i].r = y[i].r;
        tmp_y[i].i = y[i].i;
    }

    fft(tmp_x, N, 1);   
    fft(tmp_y, N, 1);

    for(i = 0; i < N; i++){
        z[i] = complex_mul(cmp_sq, complex_mul(tmp_x[i],tmp_y[i]));
    }

    fft(z, N, 0);
}

// 自己相関
void corr(complex* x, complex* y, complex* z, long N){
    int i;
    complex t[N];
    complex cmp_n;
    cmp_n.r = N;
    cmp_n.i = 0;

    complex cmp_1_n;
    cmp_1_n.r = 1.0/(N);
    cmp_1_n.i = 0;
    
    t[0].r = y[0].r;
    t[0].i = y[0].i;
    t[0] = complex_div(t[0], cmp_n);
    for(i = 1; i < N ; i++){
        t[i].r = y[N-i].r;
        t[i].i = y[N-i].i;
        t[i] = complex_div(t[i], cmp_n);
    }

    conv(x, t, z, N);
}