Merge branch 'master' of https://github.com/AndreasFMueller/SeminarSpezielleFunktionen

2 files changed, 283 insertions, 0 deletions

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)

diff --git a/buch/papers/kra/scripts/simulation.py b/buch/papers/kra/scripts/simulation.py
new file mode 100644
index 0000000..8bccb6a
--- /dev/null
+++ b/buch/papers/kra/scripts/simulation.py
@@ -0,0 +1,40 @@
+import sympy as sp

+

+

+class Simulation:

+    def __init__(self):

+        self.k_1, self.k_2, self.k_c = sp.symbols("k_1 k_2 k_c")

+        self.m_1, self.m_2 = sp.symbols("m_1 m_2")

+        self.t = sp.symbols("t")

+        K = sp.Matrix(

+            [[-(self.k_1 + self.k_c), self.k_c], [self.k_c, -(self.k_2 + self.k_c)]]

+        )

+        M = sp.Matrix([[1 / self.m_1, 0], [0, 1 / self.m_2]])

+        A = M * K

+

+        self.eigenvecs = []

+        self.eigenvals = []

+        for ev, mult, vecs in A.eigenvects():

+            self.eigenvecs.append(sp.Matrix(vecs))

+            self.eigenvals.extend([ev] * mult)

+

+    def __call__(self, t, x_0, k_1, k_c, k_2, m_1, m_2):

+        params = {

+            self.k_1: k_1,

+            self.k_c: k_c,

+            self.k_2: k_2,

+            self.m_1: m_1,

+            self.m_2: m_2,

+        }

+        x_0 = sp.Matrix(x_0)

+        eig_mat = sp.Matrix.hstack(*self.eigenvecs).subs(params)

+        g = eig_mat.inv() * x_0

+        L = sp.Matrix(

+            [

+                g[0] * sp.cos(self.eigenvals[0].subs(params) * self.t),

+                g[1] * sp.cos(self.eigenvals[1].subs(params) * self.t),

+            ]

+        )

+        x = eig_mat * L

+        f = sp.lambdify(self.t, x, "numpy")

+        return f(t).squeeze()