SHARE
TWEET

tools.py

VirtualVoid Jun 20th, 2019 75 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import numpy as np
  2. import scipy.interpolate
  3. import matplotlib.pyplot as plt
  4. import math
  5. import cmath
  6. from matplotlib import animation
  7. from dataclasses import dataclass
  8.  
  9. plt.rcParams[
  10.     "animation.convert_path"
  11. ] = "C:\Program Files\ImageMagick-7.0.8-Q16\magick.exe"
  12. from PIL import Image
  13.  
  14.  
  15. def interpolate(points, output_count):
  16.     base = points[0]
  17.     ts = []
  18.     for point in points:
  19.         ts.append(abs(point - base))
  20.         base = point
  21.  
  22.     ts = np.cumsum(ts)
  23.     ts = ts / ts[-1]
  24.  
  25.     fn = scipy.interpolate.interp1d(ts, points, kind="linear")
  26.     t = np.linspace(0, 1, output_count)
  27.     return fn(t)
  28.  
  29.  
  30. def plot_complex(points):
  31.     plt.plot(points.real, points.imag)
  32.     plt.show()
  33.  
  34.  
  35. def _make_frequency_componet(freq, val):
  36.     return lambda t: np.exp(complex(0, 2 * math.pi * freq * t)) * val
  37.  
  38.  
  39. def make_idft_funs(data):
  40.     data = np.fft.fft(data)
  41.     half_length = len(data) // 2
  42.     frequencies = list(range(half_length + len(data) % 2)) + list(
  43.         range(-half_length, 0)
  44.     )
  45.     return [
  46.         _make_frequency_componet(freq, val / len(data))
  47.         for freq, val in zip(frequencies, data)
  48.     ]
  49.  
  50.  
  51. @dataclass
  52. class DrawContainer:
  53.     ifft_funs: list
  54.     circle_color: str
  55.     arrow_color: str
  56.     fig_color: str
  57.     points_list: list
  58.     hide_offset: bool = True
  59.  
  60.  
  61. def calculate_limts(*arrays, border_factor=0.2):
  62.     all = np.concatenate(arrays)
  63.     border = abs(np.max(all) - np.min(all)) * border_factor
  64.     return (
  65.         (np.min(all.real) - border, np.max(all.real) + border),
  66.         (np.min(all.imag) - border, np.max(all.imag) + border),
  67.     )
  68.  
  69.  
  70. def draw(containers: DrawContainer, limits, backround_color="white"):
  71.     fig = plt.figure(facecolor=backround_color)
  72.     plt.axis("off")
  73.     ax = plt.axes(xlim=limits[0], ylim=limits[1])
  74.     lines = [ax.plot([], [], lw=2, color=c.fig_color)[0] for c in containers]
  75.     delete_run = True
  76.     t_last = 0
  77.  
  78.     def draw(center, val, container):
  79.         if not (not center and container.hide_offset):
  80.             circle = plt.Circle(
  81.                 (center.real, center.imag),
  82.                 abs(val),
  83.                 fill=False,
  84.                 color=container.circle_color,
  85.             )
  86.             ax.add_artist(circle)
  87.             ax.arrow(
  88.                 center.real,
  89.                 center.imag,
  90.                 val.real,
  91.                 val.imag,
  92.                 color=container.arrow_color,
  93.             )
  94.  
  95.     def animate(t):
  96.         nonlocal delete_run
  97.         if t == 0:
  98.             delete_run = not delete_run
  99.         arts = []
  100.         ax.clear()
  101.         ax.set_aspect("equal")
  102.         ax.set_xlim(*limits[0])
  103.         ax.set_ylim(*limits[1])
  104.         ax.axis("off")
  105.         for container in containers:
  106.             origin = 0 + 0j
  107.             vals = [fun(t) for fun in container.ifft_funs]
  108.             for val in vals:
  109.                 draw(origin, val, container)
  110.                 origin += val
  111.             if delete_run:
  112.                 points = np.array(container.points_list)
  113.                 container.points_list = container.points_list[1:]
  114.             else:
  115.                 container.points_list.append(origin)
  116.                 points = np.array(container.points_list)
  117.             ax.plot(points.real, points.imag, lw=3.5, color=container.fig_color)
  118.  
  119.  
  120.     frames = np.concatenate((np.linspace(0, 1, 250), np.linspace(0, 1, 250)))
  121.     anim = animation.FuncAnimation(
  122.         fig, animate, init_func=lambda: None, frames=frames, interval=20
  123.     )
  124.     writer = animation.ImageMagickFileWriter(fps=30)
  125.     anim.save('animation.gif', writer=writer,savefig_kwargs={'facecolor':backround_color})
  126.     plt.show()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top