-
Notifications
You must be signed in to change notification settings - Fork 0
/
fourier.py
76 lines (66 loc) · 2.77 KB
/
fourier.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from symfit import parameters, variables, sin, cos, Fit
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib import animation
import os
class FourierApp:
sns.set()
def __init__(self, func, low_lim=-np.pi, high_lim=np.pi, n=3, figsize=(15, 12), k=5, save=False):
if n > 10:
print('Calculating fourier series for more than 10 items is computationally inefficient')
self.func_name = func.__name__
self.n = n
self.k = k
self.func = func
self.low_lim = low_lim
self.high_lim = high_lim
self.fig = plt.figure(figsize=figsize)
self.xdata, self.ydata = self.make_data()
plt.plot(self.xdata, self.ydata)
self.initial, = plt.plot([], [], color="blue")
self.initial.set_data(self.xdata, self.ydata)
self.fourier, = plt.plot([], [], color="red", ls=':')
self.save = save
if self.save:
Writer = animation.writers['pillow']
self.writer = Writer(fps=1, metadata=dict(artist='Me'), bitrate=1800)
self.make_dir()
def make_dir(self):
if not os.path.exists('generated_gifs'):
os.mkdir('generated_gifs')
def animate(self, i):
plt.title('Fourier series for {0} with {1} elements'.format(self.func_name, i + 1))
y = self.app(i)
self.fourier.set_data(self.xdata, y)
return self.fourier,
def init(self):
self.fourier.set_data([], [])
return self.fourier,
def fourier_series(self, x, f, n=2):
a0, *cos_a = parameters(','.join(['a{}'.format(i) for i in range(0, n + 1)]))
sin_b = parameters(','.join(['b{}'.format(i) for i in range(1, n + 1)]))
series = a0 + sum(ai * cos(i * f * x) + bi * sin(i * f * x)
for i, (ai, bi) in enumerate(zip(cos_a, sin_b), start=1))
return series
def make_data(self):
xdata = np.linspace(self.low_lim, self.high_lim)
if self.func_name == 'power_func':
ydata = self.func(xdata, k=self.k)
else:
ydata = self.func(xdata)
return xdata, ydata
def app(self, n):
n = n + 1
x, y = variables('x, y')
w, = parameters('w')
model_dict = {y: self.fourier_series(x, f=w, n=n)}
fit = Fit(model_dict, x=self.xdata, y=self.ydata)
fit_result = fit.execute()
return fit.model(x=self.xdata, **fit_result.params).y
def run(self):
anim = animation.FuncAnimation(self.fig, self.animate, init_func=self.init,
frames=self.n, blit=False, repeat=True, save_count=50)
if self.save:
anim.save('generated_gifs/' + self.func_name + '_fourier_fit_{}.gif'.format(self.n), writer=self.writer)
plt.show()