tools.py

import numpy as np
import scipy.interpolate
import matplotlib.pyplot as plt
import math
import cmath
from matplotlib import animation
from dataclasses import dataclass

plt.rcParams[
    "animation.convert_path"
] = "C:\Program Files\ImageMagick-7.0.8-Q16\magick.exe"
from PIL import Image


def interpolate(points, output_count):
    base = points[0]
    ts = []
    for point in points:
        ts.append(abs(point - base))
        base = point

    ts = np.cumsum(ts)
    ts = ts / ts[-1]

    fn = scipy.interpolate.interp1d(ts, points, kind="linear")
    t = np.linspace(0, 1, output_count)
    return fn(t)


def plot_complex(points):
    plt.plot(points.real, points.imag)
    plt.show()


def _make_frequency_componet(freq, val):
    return lambda t: np.exp(complex(0, 2 * math.pi * freq * t)) * val


def make_idft_funs(data):
    data = np.fft.fft(data)
    half_length = len(data) // 2
    frequencies = list(range(half_length + len(data) % 2)) + list(
        range(-half_length, 0)
    )
    return [
        _make_frequency_componet(freq, val / len(data))
        for freq, val in zip(frequencies, data)
    ]


@dataclass
class DrawContainer:
    ifft_funs: list
    circle_color: str
    arrow_color: str
    fig_color: str
    points_list: list
    hide_offset: bool = True


def calculate_limts(*arrays, border_factor=0.2):
    all = np.concatenate(arrays)
    border = abs(np.max(all) - np.min(all)) * border_factor
    return (
        (np.min(all.real) - border, np.max(all.real) + border),
        (np.min(all.imag) - border, np.max(all.imag) + border),
    )


def draw(containers: DrawContainer, limits, backround_color="white"):
    fig = plt.figure(facecolor=backround_color)
    plt.axis("off")
    ax = plt.axes(xlim=limits[0], ylim=limits[1])
    lines = [ax.plot([], [], lw=2, color=c.fig_color)[0] for c in containers]
    delete_run = True
    t_last = 0

    def draw(center, val, container):
        if not (not center and container.hide_offset):
            circle = plt.Circle(
                (center.real, center.imag),
                abs(val),
                fill=False,
                color=container.circle_color,
            )
            ax.add_artist(circle)
            ax.arrow(
                center.real,
                center.imag,
                val.real,
                val.imag,
                color=container.arrow_color,
            )

    def animate(t):
        nonlocal delete_run
        if t == 0:
            delete_run = not delete_run
        arts = []
        ax.clear()
        ax.set_aspect("equal")
        ax.set_xlim(*limits[0])
        ax.set_ylim(*limits[1])
        ax.axis("off")
        for container in containers:
            origin = 0 + 0j
            vals = [fun(t) for fun in container.ifft_funs]
            for val in vals:
                draw(origin, val, container)
                origin += val
            if delete_run:
                points = np.array(container.points_list)
                container.points_list = container.points_list[1:]
            else:
                container.points_list.append(origin)
                points = np.array(container.points_list)
            ax.plot(points.real, points.imag, lw=3.5, color=container.fig_color)


    frames = np.concatenate((np.linspace(0, 1, 250), np.linspace(0, 1, 250)))
    anim = animation.FuncAnimation(
        fig, animate, init_func=lambda: None, frames=frames, interval=20
    )
    writer = animation.ImageMagickFileWriter(fps=30)
    anim.save('animation.gif', writer=writer,savefig_kwargs={'facecolor':backround_color})
    plt.show()