#!/usr/bin/env python
# -*- coding: utf-8 -*-
#Die zweite Zeile braucht es, damit wir Umlaute verwenden können

"""
   Pythonskript zur Simulation von fünf gekoppelten Pendeln mit Massen $m_i$ 
   und Reibungskoeffizienten b_i.
   Die Federkonstanten sind hier noch alle gleich. 
"""


from scipy.integrate import odeint
from numpy import loadtxt
import matplotlib.pyplot as plt

#Definition von Konstanten des Problems (kann man ändern und damit spielen!)
#Massen:
m1 = 1.; m2 = 1.; m3 = 1.; m4 = 1.; m5 = 1.
#Reibungskoeffizienten
b1 = 0.25; b2 = 0.25; b3 = 0.25; b4 = 0.25; b5 = 0.25
#Federkonstante
D = 1.

#Verpacken der Parameter in einen "Vektor"
p = [m1, m2, m3, m4, m5, b1, b2, b3, b4, b5, D]


def eqs(x,t,p):
    """
    Hier wird die zu integrierende Funktion definiert.
    x:  Vektor mit den Variablen x_i, die den aktuellen Zustand beschreiben
    t:  Zeit
    p:  Parameter (m_i, b_i, D)
    """

    x1, y1, x2, y2, x3, y3, x4, y4, x5, y5  = x
    m1, m2, m3, m4, m5, b1, b2, b3, b4, b5, D = p

    #Nun beschreiben wir die Ableitungen
    f = [y1,
         (-b1*y1 - 2*D*x1 + D*x2         )/m1,
         y2,
         (-b2*y2 + D*x1 - 2*D*x2 + D*x3  )/m3,
         y3,
         (-b3*y3          + D*x2 - 2*D*x3)/m3,
         y4,
         (-b4*y4 + D*x3 - 2*D*x4 + D*x5  )/m4,
         y5,
         (-b5*y5          + D*x4 - 2*D*x5)/m5]
    return f


#Definition der Anfangsbedingunen
x1 = 0.   #Alle Pendel sind in ihrer Ruheposition
y1 = 0.1  #Nur Pendel 1 hat eine Anfangsgeschwindigkeit
x2 = 0.
y2 = 0.
x3 = 0.
y3 = 0.
x4 = 0.
y4 = 0.
x5 = 0.
y5 = 0.

#Verpacke nun die Anfangsbedingungen in einen "Vektor" z0 (Zustand 0)
z0 = [x1,y1,x2,y2,x3,y3,x4,y4,x5,y5]

# Definiere Parameter für die Integrationsroutine
abserr = 1.0e-8
relerr = 1.0e-6
stoptime = 25.0
numpoints = 750

# Erzeuge die Zeitpunkte, für die die Lösungen berechnet werden sollen.
t = [stoptime * float(i) / (numpoints - 1) for i in range(numpoints)]

#und rufe die Integrationsroutine auf

loesung = odeint(eqs, z0, t, args = (p,), atol = abserr, rtol = relerr)

#Speichere nun die Resultate ab

with open('fünf-oszillatoren.dat','w') as f:
    for zeit, l in zip(t, loesung):
        print >> f, zeit, l[0], l[1], l[2], l[3], l[4], l[5], l[6], l[7], l[8], l[9]


t, x1, y1, x2, y2, x3, y3, x4, y4, x5, y5 = loadtxt('fünf-oszillatoren.dat',unpack=True)

fig, ax = plt.subplots()
ax.plot(t,x1,'k',lw=2,label=r'$x_1$')
ax.plot(t,x2,'r',lw=2,label=r'$x_2$')
ax.plot(t,x3,'b',lw=2,label=r'$x_3$')
ax.plot(t,x4,'c',lw=2,label=r'$x_4$')
ax.plot(t,x5,'m',lw=2,label=r'$x_5$')
ax.set_xlabel('Zeit',fontsize=16)
ax.set_ylabel('Auslenkung',fontsize=16)
ax.tick_params(labelsize=14)
ax.legend(loc='upper right')
ax.set_title(u'Fünf gekoppelte Pendel',fontsize=16)
plt.savefig(u'fünf-oszillatoren.png',dpi=100)
plt.savefig(u'fünf-oszillatoren.eps',dpi=100)
plt.show()