%pylab inline

x = arange(0, 1, 0.01)
sin_n = lambda n: sin(2 * pi * n * x)

plot(x, sin_n(1))

for i in [1, 2, 5, 8, 9, 12]:
    plot(x, sin_n(i), label='n = %i' % i)
legend()

sum_i_j = lambda i,j: sum([sin_n(k) for k in range(i, j+1)], axis=0)

for i in [1, 2, 3, 4, 5, 10, 20, 40]:
    plot(x, sum_i_j(1, i))

plot(x, x<1/2.)
ylim(-1.1, 2.1)

F_n = lambda n: (1/2. + sum([2./pi/(2*i+1) * sin(2 * pi * (2*i+1) * x) for i in range(n+1)], axis=0))

plot(x, x<1/2.)
for i in [0, 1, 2, 3, 5, 10]:
    plot(x, F_n(i))
ylim(-1.1, 2.1)

n = 20
plot(x, x<1/2.)
plot(x, F_n(n))
ylim(-1.1, 2.1)
title('square wave approximation with %i terms' % n)

rect_fft = fft.fft(x<1/2.)

sample_freq = x.size / 1.
plot(linspace(0, 1, rect_fft.size) * sample_freq, abs(rect_fft))
xlabel('frequency (Hz)')
ylabel('amplitude')
xlim(0, sample_freq / 2)

zero_padded_fft = fft.fft(hstack((x<1/2., zeros(x.shape[0] * 5))))

plot(linspace(0, 1, zero_padded_fft.size) * sample_freq, abs(zero_padded_fft))
xlabel('frequency (Hz)')
ylabel('amplitude')
xlim(0, sample_freq / 2)