fortran66のブログ

fortran について書きます。

2^n3^m型のFFT

MPEG1/Layer2,3では1フレームあたりのデータ数が変わって、FFTにかけるデータ数も変わるので一般形にしておきます。

Fortran2003には本来、パラメトリックTYPEというのがあって、TYPEの中の要素のKINDや定数をTYPEインスタンス生成のときに指定できるんですが、Intel Fortranではまだ実装されていないので、適当にやります。チラシ裏のメモということで。

■ソース・コード

MODULE m_fft
  IMPLICIT NONE
  PRIVATE
  REAL(8), PARAMETER:: pi = 4 * ATAN(1.0d0), pi2 = 2 * pi
  COMPLEX(8), PARAMETER:: sqrt3_2 = CMPLX(0.0d0, 0.5d0 * SQRT(3.0d0), KIND = 8)

  TYPE, PUBLIC :: t_fft
     INTEGER :: np2, np3
     INTEGER   , ALLOCATABLE :: indx(:)
     REAL(8)   , ALLOCATABLE :: hann(:)
     COMPLEX(8), ALLOCATABLE :: fft(:)
     COMPLEX(8), ALLOCATABLE :: omega(:)
   CONTAINS
     PROCEDURE :: init
     PROCEDURE :: fft_2_3
  END TYPE t_fft

 CONTAINS
  !--------------------------------------------------------------------------
  SUBROUTINE init(this, np2, np3)
    CLASS(t_fft), INTENT(OUT) :: this
    INTEGER, INTENT(IN) :: np2, np3 ! nn = 2^np2 * 3^np3
    INTEGER :: i, j2, j3, k, n
    INTEGER :: nn2, nn3, nn, nn1

    this%np2 = np2
    this%np3 = np3
    nn2 = 2**np2
    nn3 = 3**np3
    nn   = nn2 * nn3
    nn1  = nn - 1
    ALLOCATE( this%fft(0:nn1), this%indx(0:nn1), this%omega(0:nn1), this%hann(0:nn1) )
    this%omega = [( EXP( CMPLX(0.0d0, 2 * pi / nn * i, KIND = 8) ), i = 0, nn1 )]
    this%hann  = [( (1.0d0 - COS(2 * pi * i / nn)) / 2            , i = 0, nn1 )]

    DO i = 0, nn1
     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
    INTEGER :: np2, np3, nn2, nn3, nn, nn1

    np2 = this%np2
    np3 = this%np3
    nn2 = 2**np2
    nn3 = 3**np3
    nn   = nn2 * nn3
    nn1  = nn - 1
    this%fft = this%fft(this%indx)
    ! FFT_3
    DO k3 = 0, np3 - 1
     kn3 = 3**k3
     DO i = 0, nn1, 3 * kn3
      DO j = 0, kn3 - 1
       iphase1 = 3**(np3 - k3 - 1) * j * nn2
       iphase2 = 2 * iphase1
       c1 = this%omega( MOD(iphase1, nn) )
       c2 = this%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, nn1, 2 * kn2 * nn3
      DO j = 0, kn2 * nn3 - 1
       iphase1 = 2**(np2 - k2 - 1) * j
       c1 = this%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)

  CALL fft%init(7, 2)
  nn = SIZE(fft%fft)
  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