#!/usr/bin/env python
# coding: utf-8

# ## ThinkDSP
# 
# This notebook contains code examples from Chapter 7: Discrete Fourier Transform
# 
# Copyright 2015 Allen Downey
# 
# License: [Creative Commons Attribution 4.0 International](http://creativecommons.org/licenses/by/4.0/)

# In[1]:


# Get thinkdsp.py

import os

if not os.path.exists('thinkdsp.py'):
    get_ipython().system('wget https://github.com/AllenDowney/ThinkDSP/raw/master/code/thinkdsp.py')


# In[2]:


import numpy as np
import matplotlib.pyplot as plt

from thinkdsp import decorate
PI2 = 2 * np.pi


# In[3]:


# suppress scientific notation for small numbers
np.set_printoptions(precision=3, suppress=True)


# ## Complex sinusoid
# 
# Here's the definition of ComplexSinusoid, with print statements to display intermediate results.

# In[4]:


from thinkdsp import Sinusoid

class ComplexSinusoid(Sinusoid):
    """Represents a complex exponential signal."""
    
    def evaluate(self, ts):
        """Evaluates the signal at the given times.

        ts: float array of times
        
        returns: float wave array
        """
        print(ts)
        phases = PI2 * self.freq * ts + self.offset
        print(phases)
        ys = self.amp * np.exp(1j * phases)
        return ys


# Here's an example:

# In[5]:


signal = ComplexSinusoid(freq=1, amp=0.6, offset=1)
wave = signal.make_wave(duration=1, framerate=4)
print(wave.ys)


# The simplest way to synthesize a mixture of signals is to evaluate the signals and add them up.

# In[6]:


from thinkdsp import SumSignal

def synthesize1(amps, freqs, ts):
    components = [ComplexSinusoid(freq, amp)
                  for amp, freq in zip(amps, freqs)]
    signal = SumSignal(*components)
    ys = signal.evaluate(ts)
    return ys


# Here's an example that's a mixture of 4 components.

# In[7]:


amps = np.array([0.6, 0.25, 0.1, 0.05])
freqs = [100, 200, 300, 400]
framerate = 11025

ts = np.linspace(0, 1, framerate, endpoint=False)
ys = synthesize1(amps, freqs, ts)
print(ys)


# Now we can plot the real and imaginary parts:

# In[8]:


n = 500
plt.plot(ts[:n], ys[:n].real)
plt.plot(ts[:n], ys[:n].imag)
decorate(xlabel='Time')


# The real part is a mixture of cosines; the imaginary part is a mixture of sines.  They contain the same frequency components with the same amplitudes, so they sound the same to us:

# In[9]:


from thinkdsp import Wave

wave = Wave(ys.real, framerate)
wave.apodize()
wave.make_audio()


# In[10]:


wave = Wave(ys.imag, framerate)
wave.apodize()
wave.make_audio()


# We can express the same process using matrix multiplication.

# In[11]:


def synthesize2(amps, freqs, ts):
    args = np.outer(ts, freqs)
    M = np.exp(1j * PI2 * args)
    ys = np.dot(M, amps)
    return ys


# And it should sound the same.

# In[12]:


amps = np.array([0.6, 0.25, 0.1, 0.05])
ys = synthesize2(amps, freqs, ts)
print(ys)


# In[13]:


wave = Wave(ys.real, framerate)
wave.apodize()
wave.make_audio()


# To see the effect of a complex amplitude, we can rotate the amplitudes by 1.5 radian:

# In[14]:


phi = 1.5
amps2 = amps * np.exp(1j * phi)
ys2 = synthesize2(amps2, freqs, ts)

n = 500
plt.plot(ts[:n], ys.real[:n], label=r'$\phi_0 = 0$')
plt.plot(ts[:n], ys2.real[:n], label=r'$\phi_0 = 1.5$')
decorate(xlabel='Time')


# Rotating all components by the same phase offset changes the shape of the waveform because the components have different periods, so the same offset has a different effect on each component.

# ### Analysis
# 
# The simplest way to analyze a signal---that is, find the amplitude for each component---is to create the same matrix we used for synthesis and then solve the system of linear equations.

# In[15]:


def analyze1(ys, freqs, ts):
    args = np.outer(ts, freqs)
    M = np.exp(1j * PI2 * args)
    amps = np.linalg.solve(M, ys)
    return amps


# Using the first 4 values from the wave array, we can recover the amplitudes.

# In[16]:


n = len(freqs)
amps2 = analyze1(ys[:n], freqs, ts[:n])
print(amps2)


# If we define the `freqs` from 0 to N-1 and `ts` from 0 to (N-1)/N, we get a unitary matrix. 

# In[17]:


N = 4
ts = np.arange(N) / N
freqs = np.arange(N)
args = np.outer(ts, freqs)
M = np.exp(1j * PI2 * args)
print(M)


# To check whether a matrix is unitary, we can compute $M^* M$, which should be the identity matrix:

# In[18]:


MstarM = M.conj().transpose().dot(M)
print(MstarM.real)


# The result is actually $4 I$, so in general we have an extra factor of $N$ to deal with, but that's a minor problem.
# 
# We can use this result to write a faster version of `analyze1`:
# 

# In[19]:


def analyze2(ys, freqs, ts):
    args = np.outer(ts, freqs)
    M = np.exp(1j * PI2 * args)
    amps = M.conj().transpose().dot(ys) / N
    return amps


# In[20]:


N = 4
amps = np.array([0.6, 0.25, 0.1, 0.05])
freqs = np.arange(N)
ts = np.arange(N) / N
ys = synthesize2(amps, freqs, ts)

amps3 = analyze2(ys, freqs, ts)
print(amps3)


# Now we can write our own version of DFT:

# In[21]:


def synthesis_matrix(N):
    ts = np.arange(N) / N
    freqs = np.arange(N)
    args = np.outer(ts, freqs)
    M = np.exp(1j * PI2 * args)
    return M


# In[22]:


def dft(ys):
    N = len(ys)
    M = synthesis_matrix(N)
    amps = M.conj().transpose().dot(ys)
    return amps


# And compare it to analyze2:

# In[23]:


print(dft(ys))


# The result is close to `amps * 4`.
# 
# We can also compare it to `np.fft.fft`.  FFT stands for Fast Fourier Transform, which is an even faster implementation of DFT.

# In[24]:


print(np.fft.fft(ys))


# The inverse DFT is almost the same, except we don't have to transpose $M$ and we have to divide through by $N$.

# In[25]:


def idft(amps):
    N = len(amps)
    M = synthesis_matrix(N)
    ys = M.dot(amps) / N
    return ys


# We can confirm that `dft(idft(amps))` yields `amps`:

# In[26]:


ys = idft(amps)
print(dft(ys))


# ### Real signals
# 
# Let's see what happens when we apply DFT to a real-valued signal.

# In[27]:


from thinkdsp import SawtoothSignal

framerate = 10000
signal = SawtoothSignal(freq=500)
wave = signal.make_wave(duration=0.1, framerate=framerate)
wave.make_audio()


# `wave` is a 500 Hz sawtooth signal sampled at 10 kHz.

# In[28]:


hs = dft(wave.ys)
len(wave.ys), len(hs)


# `hs` is the DFT of this wave, and `amps` contains the amplitudes.

# In[29]:


amps = np.abs(hs)
plt.plot(amps)
decorate(xlabel='Frequency (unspecified units)', ylabel='DFT')


# The DFT assumes that the sampling rate is N per time unit, for an arbitrary time unit.  We have to convert to actual units -- seconds -- like this:

# In[30]:


N = len(hs)
fs = np.arange(N) * framerate / N


# Also, the DFT of a real signal is symmetric, so the right side is redundant.  Normally, we only compute and plot the first half:

# In[31]:


plt.plot(fs[:N//2+1], amps[:N//2+1])
decorate(xlabel='Frequency (Hz)', ylabel='DFT')


# Let's get a better sense for why the DFT of a real signal is symmetric.  I'll start by making the inverse DFT matrix for $N=8$.

# In[32]:


M = synthesis_matrix(N=8)


# And the DFT matrix:

# In[33]:


Mstar = M.conj().transpose()


# And a triangle wave with 8 elements:

# In[34]:


from thinkdsp import TriangleSignal

wave = TriangleSignal(freq=1).make_wave(duration=1, framerate=8)
wave.ys


# Here's what the wave looks like.

# In[35]:


wave.plot()


# Now let's look at rows 3 and 5 of the DFT matrix:

# In[36]:


row3 = Mstar[3, :]
print(row3)


# In[37]:


row5 = Mstar[5, :]
row5


# They are almost the same, but row5 is the complex conjugate of row3.

# In[38]:


def approx_equal(a, b, tol=1e-10):
    return np.sum(np.abs(a-b)) < tol


# In[39]:


approx_equal(row3, row5.conj())


# When we multiply the DFT matrix and the wave array, the element with index 3 is:

# In[40]:


X3 = row3.dot(wave.ys)
X3


# And the element with index 5 is:

# In[41]:


X5 = row5.dot(wave.ys)
X5


# And they are the same, within floating point error.

# In[42]:


abs(X3 - X5)


# Let's try the same thing with a complex signal:

# In[43]:


wave2 = ComplexSinusoid(freq=1).make_wave(duration=1, framerate=8)
plt.plot(wave2.ts, wave2.ys.real)
plt.plot(wave2.ts, wave2.ys.imag)


# Now the elements with indices 3 and 5 are different:

# In[44]:


X3 = row3.dot(wave2.ys)
X3


# In[45]:


X5 = row5.dot(wave2.ys)
X5


# Visually we can confirm that the FFT of the real signal is symmetric:

# In[46]:


hs = np.fft.fft(wave.ys)
plt.plot(abs(hs))


# And the FFT of the complex signal is not.

# In[47]:


hs = np.fft.fft(wave2.ys)
plt.plot(abs(hs))


# Another way to think about all of this is to evaluate the DFT matrix for different frequencies.  Instead of $0$ through $N-1$, let's try $0, 1, 2, 3, 4, -3, -2, -1$.

# In[48]:


N = 8
ts = np.arange(N) / N
freqs = np.arange(N)
freqs = [0, 1, 2, 3, 4, -3, -2, -1]
args = np.outer(ts, freqs)
M2 = np.exp(1j * PI2 * args)


# In[49]:


approx_equal(M, M2)


# So you can think of the second half of the DFT as positive frequencies that get aliased (which is how I explained them), or as negative frequencies (which is the more conventional way to explain them).  But the DFT doesn't care either way.
# 
# The `thinkdsp` library provides support for computing the "full" FFT instead of the real FFT.

# In[50]:


framerate = 10000
signal = SawtoothSignal(freq=500)
wave = signal.make_wave(duration=0.1, framerate=framerate)


# In[51]:


spectrum = wave.make_spectrum(full=True)


# In[52]:


spectrum.plot()
decorate(xlabel='Frequency (Hz)', ylabel='DFT')


# In[ ]: