fortran66のブログ

fortran について書きます。

2と3のべきの積の次数のFFT

昔書いたFFTOOP風にしてみました。とりあえず、べきは固定になってますが、イニシャライズのところでパラメータをとるようにすればいいような気もしなくも無いです。まぁ将来直すということで。

普通教科書に載っているFFTは2のべき乗のものですが、同じようにして2以外の数のべきのFFTも求められます。さらに我慢して考えると、因数分解したべきの形のFFTも導けます。Numerical RecipesFFTのところにバタフライ演算の図と、再帰の形での説明があります。(NR初版のFFTのコード自体にはバグがあるという話しです。NRの本のコードは色々腐っているので、パッチの当たった別売りコードを使うか、注意深く使用するのが良いようです。)漸化式の形かなんかできっちり示したいところですが、まぁ将来考えるということで。

ここのやり方では、バタフライ演算で順序が入れ替わるので、FFTデータをあらかじめ並べ替えておく必要があります。そのためのインデックスを初期化ルーチンで用意しています。

■実行結果

素数864(=2^5*3^3)のFFT。f(t) = exp(-2πi t/864) + exp(-2πi t/864) + exp(-2πi t/864)のFFTを計算してます。基本波、3倍波、5倍波を出しているつもりです。規格化因子・フェーズの符号の定義などは、よく覚えていないので、適宜解釈することにしますw

■ソース・コード

ここではサイズ固定になっていますが、2と3のべきの2つのパラメータをかえればいいはずです。

MODULE m_fft
   IMPLICIT NONE
   PRIVATE
   INTEGER, PARAMETER :: np2 = 5, np3 = 3 ! 2^5 * 3^3 = 864
   INTEGER, PARAMETER :: nn2 = 2**np2, nn3 = 3**np3, nn = nn2 * nn3, &
                           nn1_2 = nn / 2, nn1_3 = nn / 3, nn2_3 = 2 * nn1_3
   REAL(8), PARAMETER:: pi = 4 * ATAN(1.0d0), pi2 = 2 * pi , pi2_n = pi2 / nn
   COMPLEX(8), PARAMETER:: sqrt3_2 = CMPLX(0.0d0, 0.5d0 * SQRT(3.0d0), KIND = 8)
   
   INTEGER :: i_
   COMPLEX(8), PARAMETER :: omega(0:nn - 1) = [( EXP( CMPLX(0.0d0, 2 * pi / nn * i_, KIND = 8) ), i_ = 0, nn - 1 )]
   REAL(8)   , PARAMETER :: hann(0:nn - 1)  = [( (1.0d0 - COS(2 * pi * i_ / nn)) / 2,             i_ = 0, nn - 1 )]
   
   TYPE, PUBLIC :: t_fft
     COMPLEX(8) :: fft(0:nn - 1)
     INTEGER    :: indx(0:nn - 1)
    CONTAINS
     PROCEDURE :: init
     PROCEDURE :: fft_2_3
   END TYPE t_fft   
   
 CONTAINS
   !--------------------------------------------------------------------------
   SUBROUTINE init(this)
     CLASS(t_fft), INTENT(OUT) :: this
     INTEGER :: i, j2, j3, k, n
     
     DO i = 0, nn - 1
      n = 0
      k = i
      DO j3 = np3 - 1, 0, -1
       n = n + MOD(k, 3) * nn2 * 3**j3
       k = k / 3
      END DO
      DO j2 = np2 - 1, 0, -1
       n = n + MOD(k, 2) * 2**j2
       k = k / 2
      END DO
      this%indx(i) = n  ! indx = 0..N - 1
     END DO
     
     RETURN
   END SUBROUTINE init
   !--------------------------------------------------------------------------
   SUBROUTINE fft_2_3(this)
     CLASS(t_fft), INTENT(IN OUT) :: this
     COMPLEX(8):: tmp1, tmp2, tmp3, c1, c2, c3, c4, c5, c6
     INTEGER :: i, j, k, k2, k3, m1, m2, m3, iphase1, iphase2, kn2, kn3
     
     this%fft = this%fft(this%indx)
     ! FFT_3
     DO k3 = 0, np3 - 1
      kn3 = 3**k3
      DO i = 0, nn - 1, 3 * kn3 
       DO j = 0, kn3 - 1
        iphase1 = 3**(np3 - k3 - 1) * j * nn2
        iphase2 = 2 * iphase1
        c1 = omega( MOD(iphase1, nn) )
        c2 = omega( MOD(iphase2, nn) )
        m1 = i + j 
        m2 = m1 + kn3
        m3 = m2 + kn3
        tmp1 =      this%fft(m1)
        tmp2 = c1 * this%fft(m2)
        tmp3 = c2 * this%fft(m3)
        this%fft(m1) = tmp1 + tmp2 + tmp3
        c3 = tmp1 - 0.5d0 * ( tmp2 + tmp3 )
        c5 =     sqrt3_2  * ( tmp2 - tmp3 ) ! sqrt3_2 = i sqrt(3) / 2
        this%fft(m2) = c3 + c5
        this%fft(m3) = c3 - c5
       END DO
      END DO
     END DO
     ! FFT_2
     DO k2 = 0, np2 - 1
      kn2 = 2**k2
      DO i = 0, nn - 1, 2 * kn2 * nn3
       DO j = 0, kn2 * nn3 - 1
        iphase1 = 2**(np2 - k2 - 1) * j
        c1 = omega( MOD(iphase1, nn) )
        m1 = i + j
        m2 = m1 + kn2 * nn3
        tmp1 =      this%fft(m1)
        tmp2 = c1 * this%fft(m2)
        this%fft(m1) = tmp1 + tmp2 ! 1 = x^2 ; x = 1 or -1
        this%fft(m2) = tmp1 - tmp2
       END DO
      END DO
     END DO
     
     RETURN
   END SUBROUTINE fft_2_3
   !------------------------------------------------------------------------
END MODULE m_fft
!==========================================================================
PROGRAM test
   USE m_fft
   IMPLICIT NONE
   TYPE (t_fft) :: fft
   INTEGER :: i, j, nn
   REAL(8), PARAMETER :: pi = 4 * ATAN(1.0d0)
   
   nn = SIZE(fft%fft)
   CALL fft%init()
   DO j = 0, nn - 1
    fft%fft(j) = EXP( CMPLX(0.0d0, -     j * 2 * pi / nn, KIND = 8) ) &
               + EXP( CMPLX(0.0d0, - 3 * j * 2 * pi / nn, KIND = 8) ) &
               + EXP( CMPLX(0.0d0, - 5 * j * 2 * pi / nn, KIND = 8) ) 
   END DO
   CALL fft%fft_2_3()
   DO i = 0, nn - 1
    PRINT '(i10, 3f20.12)', i, fft%fft(i) 
   END DO

   STOP
END PROGRAM test