From 5e25727877da020b0b23132fec8c0ea70288a18b Mon Sep 17 00:00:00 2001 From: samuel niederer Date: Sun, 24 Jul 2022 12:03:25 +0200 Subject: add presentation files --- buch/papers/kra/scripts/animation.py | 243 +++++++++++++++++++++++++++++++++++ 1 file changed, 243 insertions(+) create mode 100644 buch/papers/kra/scripts/animation.py (limited to 'buch/papers/kra/scripts/animation.py') diff --git a/buch/papers/kra/scripts/animation.py b/buch/papers/kra/scripts/animation.py new file mode 100644 index 0000000..5e805ae --- /dev/null +++ b/buch/papers/kra/scripts/animation.py @@ -0,0 +1,243 @@ +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches +import matplotlib.transforms +import matplotlib.text +from matplotlib.animation import FuncAnimation +import imageio + +from simulation import Simulation + + +class Mass: + def __init__(self, x_0, width, height, **kwargs): + self._x_0 = x_0 + xy = (x_0, 0) + self._rect = matplotlib.patches.Rectangle(xy, width, height, **kwargs) + + @property + def patch(self): + return self._rect + + @property + def x(self): + return self._rect.get_x() + + @property + def width(self): + return self._rect.get_width() + + def move(self, x): + self._rect.set_x(self._x_0 + x) + + +class Spring: + def __init__(self, n, height, ax, resolution=1000, **kwargs): + self._n = n + self._height = height + self._N = resolution + (self._line,) = ax.plot([], [], "-", **kwargs) + + def set(self, x_0, x_1): + T = (x_1 - x_0) / self._n + x = np.linspace(x_0, x_1, self._N, endpoint=True) + t = np.linspace(0, x_1 - x_0, self._N) + y = (np.sin(2 * np.pi * t / T) + 1.5) * self._height / 2 + self.line.set_data(x, y) + + @property + def line(self): + return self._line + + +class LinePlot: + def __init__(self, ax, **kwargs): + (self._line,) = ax.plot([], [], "-", **kwargs) + self._x = [] + self._y = [] + + @property + def line(self): + return self._line + + def update(self, x, y): + self._x.append(x) + self._y.append(y) + self._line.set_data(self._x, self._y) + + +class ScatterPlot: + def __init__(self, ax, **kwargs): + self._color = kwargs.get("color", "tab:green") + self._line = ax.scatter([], [], **kwargs) + self._ax = ax + self._x = [] + self._y = [] + + @property + def line(self): + return self._line + + def update(self, x, y, **kwargs): + self._x.append(x) + self._y.append(y) + self._line.remove() + self._line = self._ax.scatter(self._x, self._y, color=self._color, **kwargs) + + +class QuiverPlot: + def __init__(self, ax, **kwargs): + self.x = [] + self.y = [] + self.u = [] + self.v = [] + self.ax = ax + self.ln = self.ax.quiver([], [], [], []) + + def update(self, x, y, u, v): + self.x.append(x) + self.y.append(y) + self.u.append(u) + self.v.append(v) + self.ln.remove() + self.ln = self.ax.quiver(self.x, self.y, self.u, self.v) + + @property + def line(self): + return self.ln + + +anim_folder = "anim_0" +img_counter = 0 + +sim = Simulation() +params = { + "x_0": [2, -2], + "k_1": 1, + "k_c": 2, + "k_2": 1, + "m_1": 0.5, + "m_2": 0.5, +} + +time = 2.1 + + +# create axis +fig = plt.figure(figsize=(20, 15), constrained_layout=True) +fig.suptitle( + " ,".join([f"${key} = {val}$" for (key, val) in params.items()]), fontsize=20 +) +spec = fig.add_gridspec(3, 4) +ax0 = fig.add_subplot(spec[-1, :]) +ax1 = fig.add_subplot(spec[:-1, :2]) +ax2 = fig.add_subplot(spec[:-1, 2:]) + +ax0.set_yticks([]) + +mass_height = 0.5 +spring_height = 0.6 * mass_height +x_max = 21 +y_max = 2 * mass_height + +mass_1 = Mass( + 7, + 2, + mass_height, + color="tab:red", +) +mass_2 = Mass(14, 2, mass_height, color="tab:blue") +masses = [mass_1, mass_2] +patches = [mass.patch for mass in masses] + +spring_1 = Spring(4, spring_height, ax0, color="tab:red", linewidth=10) +spring_2 = Spring(4, spring_height, ax0, color="tab:gray", linewidth=10) +spring_3 = Spring(4, spring_height, ax0, color="tab:blue", linewidth=10) +springs = [spring_1, spring_2, spring_3] + +linePlot_1 = LinePlot(ax1, color="tab:red", label="$m_1$", alpha=1) +linePlot_2 = LinePlot(ax1, color="tab:blue", label="$m_2$", alpha=1) +linePlots = [linePlot_1, linePlot_2] + +# quiverPlot = QuiverPlot(ax2) +scatterPlot = ScatterPlot(ax2) + +lines = [spring.line for spring in springs] +lines.extend([plot.line for plot in linePlots]) +# lines.append(quiverPlot.line) +lines.append(scatterPlot.line) + +objects = lines + patches + +ax0.plot( + np.repeat(mass_1.x, 2), + [0, y_max], + "--", + color="tab:red", + label="Ruhezustand $m_1$", +) +ax0.plot( + np.repeat(mass_2.x, 2), + [0, y_max], + "--", + color="tab:blue", + label="Ruhezustand $m_2$", +) + + +def init(): + ax0.set_xlim(0, x_max) + ax0.set_ylim(0, y_max) + + ax1.set_xlim(0, time) + ax1.set_ylim(-4, 4) + ax1.set_xlabel("time", fontsize=20) + ax1.set_ylabel("$q$", fontsize=20) + + ax2.set_xlim(-4, 4) + ax2.set_ylim(-4, 4) + ax2.set_xlabel("$q_1$", fontsize=20) + ax2.set_ylabel("$q_2$", fontsize=20) + + for patch in patches: + ax0.add_patch(patch) + + spring_1.set(0, mass_1.x) + spring_2.set(mass_1.x + mass_1.width, mass_2.x) + spring_2.set(mass_2.x + mass_2.width, x_max) + + return objects + + +def update(frame): + global img_counter + x_1, x_2 = sim(frame, **params) + + mass_1.move(x_1) + mass_2.move(x_2) + + spring_1.set(0, mass_1.x) + spring_2.set(mass_1.x + mass_1.width, mass_2.x) + spring_3.set(mass_2.x + mass_2.width, x_max) + + linePlot_1.update(frame, x_1) + linePlot_2.update(frame, x_2) + + scatterPlot.update(x_1, x_2, alpha=0.25) + + img_counter += 1 + return objects + + +anim = FuncAnimation( + fig, + update, + frames=np.linspace(0, time, int(time * 30)), + init_func=init, + blit=False, +) + +ax0.legend(fontsize=20) +ax1.legend(fontsize=20) +FFwriter = matplotlib.animation.FFMpegWriter(fps=30) +anim.save("animation.mp4", writer=FFwriter) -- cgit v1.2.1