scipy.integrate.quad() を使ってフーリエ級数展開してみるテスト

fourier.py

# -*- coding: utf-8 -*-
from pylab import *
from scipy.integrate import quad

x_min = -4.0
x_max = 8.0
xs = linspace(x_min, x_max, 256)
space = (x_max - x_min) / 80

def fourier(fun, n_max):
  a = []
  b = []
  for n in xrange(0, 1+n_max):
    res, err = quad(lambda x:fun(x)*cos(n*x), -pi, pi)
    a.append(res/pi)
    res, err = quad(lambda x:fun(x)*sin(n*x), -pi, pi)
    b.append(res/pi)
  def fn(x):
    sum = a[0] / 2
    for n in xrange(1, 1+n_max):
      sum += a[n]*cos(n*x) + b[n]*sin(n*x)
    return sum
  return fn

def yrange(fn):
  y_min = 100000.0
  y_max = -100000.0
  for x in xs:
    y = float(fn(x))
    if y < y_min:
      y_min = y
    if y > y_max:
      y_max = y
  return (y_min, y_max)

def draw(fn, txt, prefix):
  y_min, y_max = yrange(fn)
  baselineskip = (y_max - y_min) / 18
  y_max += baselineskip*2.75
  y_min -= baselineskip*2.75
  for til in xrange(0,20):
    axis([x_min, x_max, y_min, y_max])
    text(x_max-space, y_max-baselineskip*1.25, '$f(x)='+txt+'$', fontsize=13, horizontalalignment='right')
    txt_n = '$n\leq %d$' % til
    text(x_max-space, y_max-baselineskip*2.25, txt_n, fontsize=13, horizontalalignment='right')
    f_fn = fourier(fn, til)
    plot(xs, amap(fn,xs), 'b:', lw=1)
    plot(xs, amap(f_fn,xs), 'r-', lw=1)
    savefig("%s%02d" % (prefix,til))
    clf()

def regularize(x):
  return (x + pi) % (pi * 2) - pi

def f1(x):
  # 1 | 0
  x = regularize(x)
  if x >= 0:
    return 1
  else:
    return 0

def f2(x):
  # x
  x = regularize(x)
  return x

def f3(x):
  # |x|
  x = regularize(x)
  return abs(x)

def f4(x):
  # x^2
  x = regularize(x)
  return x*x

def f5(x):
  # x | 0
  x = regularize(x)
  if x >= 0:
    return x
  else:
    return 0

def f6(x):
  # {3,0,4,1,2}
  x = regularize(x)
  ix = int(floor((pi + x) / (pi*2) * 5))
  ys = [3,0,4,1,2]
  return ys[ix]

draw(f1, '1 (0\leq x<\pi), 0 (-\pi\leq x<0)', 'f1_')
draw(f2, 'x (-\pi\leq x<\pi)', 'f2_')
draw(f3, '|x| (-\pi\leq x<\pi)', 'f3_')
draw(f4, 'x^2 (-\pi\leq x<\pi)', 'f4_')
draw(f5, 'x (0\leq x<\pi), 0 (-\pi\leq x<0)', 'f5_')
draw(f6, '{3,0,4,1,2}', 'f6_')