# Untitled

a guest
Sep 23rd, 2020
168
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. import numpy as np
2. import matplotlib.pyplot as plt
3.
4. #%%###########################################################################
5. def dft_k(x, k):
6.     N = len(x)
7.     t = np.linspace(0, 1, N, False)
8.     exp = np.exp(1j * 2 * np.pi * k * t)
9.     return (x * exp).sum()
10.
11. def allfft(x):
12.     coef = np.fft.fft(x)
13.     return coef, np.angle(coef), np.abs(coef)
14.
15. def plot(x, y=None, size=0, show=1, ax_equal=False, **kw):
16.     if y is None:
17.         plt.plot(x, **kw)
18.     else:
19.         plt.plot(x, y, **kw)
20.     _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)
21.
22. def scat(x, y=None, size=0, show=1, ax_equal=False, s=18, **kw):
23.     if y is None:
24.         plt.scatter(np.arange(len(x)), x, s=s, **kw)
25.     else:
26.         plt.scatter(x, y, s=s, **kw)
27.     _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)
28.
29. def _scale_plot(fig, ax, size=True, show=False, ax_equal=False):
30.     xmin, xmax = ax.get_xlim()
31.     rng = xmax - xmin
32.     ax.set_xlim(xmin + .02 * rng, xmax - .02 * rng)
33.     if size:
34.         fig.set_size_inches(15, 7)
35.     if ax_equal:
36.         yabsmax = max(np.abs([*ax.get_ylim()]))
37.         mx = max(yabsmax, max(np.abs([xmin, xmax])))
38.         ax.set_xlim(-mx, mx)
39.         ax.set_ylim(-mx, mx)
40.         fig.set_size_inches(8, 8)
41.     if show:
42.         plt.show()
43. #%%###########################################################################
44. N = 128
45. f_signal, f_winding = 1, 0.5
46.
47. t = np.linspace(0, 1, N, False)
48. x = np.cos(2 * np.pi * f_signal * t)
49. exp = np.exp(1j * 2 * np.pi * f_winding * t)
50. coef, theta, r = allfft(x)
51.
52. dftk = (x * exp)[:128]
53. scat(dftk.real, dftk.imag, size=0, ax_equal=True, show=0)
54.
55. re, im = dftk.sum().real, dftk.sum().imag
56. plt.quiver([0], [0], [re], [im], color='r', scale=8500 / max(np.abs([re, im])))
57. plt.show()
58. #%%##################################################################
59. x_p = np.pad(x, [0, len(x)])
60. coef_p, theta_p, r_p = allfft(x_p)
61. scat(r_p, show=1)
62. #%%##################################################################
63. coef_pk = [dft_k(x, k) for k in np.arange(0, len(x), 0.5)]
64. theta_pk, r_pk = np.angle(coef_pk), np.abs(coef_pk)
65. print(np.mean(abs(r_pk - r_p)))
66. #%%###########################################################################
67. t_p = np.linspace(0, 1, 2 * N, False)
68. f_winding_p = 1
69. exp_p = np.exp(1j * 2 * np.pi * f_winding_p * t_p)
70.
71. dftk_p = (x_p * exp_p)
72. scat(dftk_p.real, dftk_p.imag, size=0, ax_equal=True, show=0)
73.
74. rp, ip = dftk_p.sum().real, dftk_p.sum().imag
75. plt.quiver([0], [0], [rp], [ip], color='r', scale=8500 / max(np.abs([rp, ip])))
76. plt.show()
77. #%%##################################################################
78. scat(x_p, show=0)
79. plot(exp_p.real, color='orange')
80.