ex2.4_dft_from_sine.rb

# -*- coding: utf-8 -*-
# http://rolz1.spaces.live.com/blog/cns!44B7739809D018D7!2807.entry

class SinWave
  attr_accessor :f0, :fs, :a, :n_max
  def initialize(args)
    @f0 = args[:f0]
    @fs = args[:fs]
    @a = args[:a]
    @n_max = args[:n_max]
  end

  def wav(n)
    if n > @n_max
      puts "overflow n:(#{n}) > n_max(#{@n_max})"
    end
    @a * Math.sin((2.0 * Math::PI * @f0 * n) / @fs)
  end
end

class SawWave
  attr_accessor :f0, :fs, :a, :n_max
  def initialize(args)
    @f0 = args[:f0]
    @fs = args[:fs]
    @a = args[:a]
    @n_max = args[:n_max]
  end

  def wav(n)
    if n > @n_max
      puts "overflow n:(#{n}) > n_max(#{@n_max})"
    end
    @a * Math.sin((2.0 * Math::PI * @f0 * n) / @fs) + 
      (@a/2) * Math.sin((2.0 * Math::PI * @f0 * n * 2) / @fs) + 
      (@a/3) * Math.sin((2.0 * Math::PI * @f0 * n * 3) / @fs) + 
      (@a/4) * Math.sin((2.0 * Math::PI * @f0 * n * 4) / @fs) + 
      (@a/5) * Math.sin((2.0 * Math::PI * @f0 * n * 5) / @fs)
  end
end

class SquareWave
  attr_accessor :f0, :fs, :a, :n_max
  def initialize(args)
    @f0 = args[:f0]
    @fs = args[:fs]
    @a = args[:a]
    @n_max = args[:n_max]
  end

  def wav(n)
    if n > @n_max
      puts "overflow n:(#{n}) > n_max(#{@n_max})"
    end
    @a * Math.sin((2.0 * Math::PI * @f0 * n) / @fs) + 
      (@a/3) * Math.sin((2.0 * Math::PI * @f0 * n * 3) / @fs) + 
      (@a/5) * Math.sin((2.0 * Math::PI * @f0 * n * 5) / @fs)
  end
end

#N_MAX = 1024
N_MAX = 64
#N_MAX = 8
A = 0.25
F0 = 250
FS = 8000
input = SinWave.new(:f0 => F0,
#input = SawWave.new(:f0 => F0,
#input = SquareWave.new(:f0 => F0,
                    :fs => FS,
                    :a => A,
                    :n_max => N_MAX)
a_max = -A
a_min = A
a_list = []
(0..N_MAX).each do |i|
  a = input.wav(i)
  a_list << a
  a = Math.sqrt(a*a)
  (a > a_max) && a_max = a
  (a_min > a) && a_min = a
end

puts "a_max = #{a_max}"
display_width = 40
offset = 40

puts '=' * display_width + '[input.wav]'
(0..N_MAX).each do |i|
  pos = (1.0 * (A + a_list[i]) / (2 * A)) * display_width + offset
  puts ' '*pos + '*' + '    ' + [i,pos].map{|v|v.to_s}.join(':')
end

module DSP
  class << self
    def _log2(x)
      y = 0
      while x > 1 do
        x = x >> 1
        y = y + 1
      end
      y
    end

    def _pow2(x)
      (x == 0) ? 1 : 2 << (x-1)
    end

    def fft(input_wav)
      xt_real = []
      xt_img = []
      n_max = input_wav.n_max
      (0..n_max).each do |i|
        xt_real[i] = input_wav.wav(i)
        xt_img[i] = 0.0
      end

      num_of_stage = _log2(n_max)

      (1..num_of_stage).each do |stage|
        (0..._pow2(stage - 1)).each do |i|
          (0..._pow2(num_of_stage - stage)).each do |j|
            n = _pow2(num_of_stage - stage + 1) * i + j
            m = _pow2(num_of_stage - stage) + n
            r = _pow2(stage - 1) * j
            a_real = xt_real[n]
            a_img  = xt_img[n]
            b_real = xt_real[m]
            b_img  = xt_img[m]
            c_real =  Math.cos((2.0 * Math::PI * r) / n_max)
            c_img  = -Math.sin((2.0 * Math::PI * r) / n_max)
            if stage < num_of_stage
              p a_real
              p b_real
              xt_real[n] = a_real + b_real
              xt_img[n]  = a_img + b_img
              xt_real[m] = (a_real - b_real) * c_real - (a_img  - b_img)  * c_img
              xt_img[m]  = (a_img  - b_img)  * c_real + (a_real - b_real) * c_img
            else
              xt_real[n] = a_real + b_real
              xt_img[n]  = a_img + b_img
              xt_real[m] = a_real - b_real
              xt_img[m]  = a_img - b_img
            end
          end
        end
      end

      # make index
      index = [0.0] * n_max
      (1..num_of_stage).each do |stage|
        (0..._pow2(stage-1)).each do |i|
          index[_pow2(stage-1)+i] = index[i] + _pow2(num_of_stage - stage)
        end
      end

      (0...n_max).each do |k|
        idx = index[k]
        if idx > k
          real = xt_real[idx]
          img = xt_img[idx]
          xt_real[idx] = xt_real[k]
          xt_img[idx] = xt_img[k]
          xt_real[k] = real
          xt_img[k] = img
        end
      end
      {
        # キー名はdftに合わせてるだけ
        :xf_real => xt_real,
        :xf_img => xt_img
      }
    end

    def dft(input_wav)
      xt_real = []
      xt_img = []
      xf_real = []
      xf_img = []
      n_max = input_wav.n_max

      (0..n_max).each do |i|
        xt_real[i] = input_wav.wav(i)
        xt_img[i] = 0.0
      end

      (0..n_max).each do |k|
        xf_real[k] = 0
        xf_img[k] = 0
        (0..n_max).each do |n|
          x = 2.0 * Math::PI * k * n / n_max
          w_real = Math.cos(x)
          w_img = -Math.sin(x)
          xf_real[k] += w_real * xt_real[n] - w_img * xt_img[n]
          xf_img[k]  += w_real * xt_img[n]  + w_img * xt_real[n]
        end
      end
      {
        :xf_real => xf_real,
        :xf_img => xf_img
      }
    end

    def amp_from_xf(xf)
      amps = []
      (0..xf[:xf_real].length-1).each do |i|
        amps[i] = Math.sqrt(xf[:xf_real][i] * xf[:xf_real][i] + xf[:xf_img][i] * xf[:xf_img][i])
      end
      amps
    end

    def theta_from_xf(xf)
      thetas = []
      (0..xf[:xf_real].length-1).each do |i|
        thetas[i] = Math.atan(xf[:xf_img][i] / xf[:xf_real][i])
      end
      thetas
    end
  end
end

_before_dft = Time.now
xf = DSP.dft(input)

_before_fft = Time.now
xf = DSP.fft(input)

_before_amp = Time.now
amps = DSP.amp_from_xf(xf)

_before_theta = Time.now
thetas = DSP.theta_from_xf(xf)

_after_theta = Time.now

puts "[Time]----"
puts "  total: #{_after_theta - _before_dft}"
puts "     dft: #{_before_fft - _before_dft}"
puts "     fft: #{_before_amp - _before_fft}"
puts "     amp: #{_before_theta - _before_amp}"
puts "    theta: #{_after_theta - _before_theta}"
puts "[Time]----"

puts "amps_min = #{amps.min}"
puts "amps_max = #{amps.max}"
puts "thetas_min = #{thetas.min}"
puts "thetas_max = #{thetas.max}"

puts '=' * display_width + '[xf_real]'
amp_width = amps.max - amps.min
amps.each_with_index do |a,i|
  pos = ((a - amps.min) / amp_width) * display_width + offset
  puts ' ' * pos + '*' + "[#{i}] #{a}"
end

#puts '=' * display_width + '[xf_img]'
#theta_width = thetas.max - thetas.min
#thetas.each_with_index do |t,i|
#  pos = ((t - thetas.min) / theta_width) * display_width + offset
#  puts ' ' * pos + '*' + "[#{i}] #{t}"
#end