%matplotlib inline
import pandas as pd
import matplotlib.pyplot as pl
import numpy as np
import sympy as sp
from numpy import linalg as LA
k = 10
m = 1.0
Mat = np.array([[-(2*k)/m, k/m, 0],[k/m, -(2*k)/m, k/m],[0, k/m, -(2*k)/m]])
vals, vecs = LA.eig(Mat)
print("eigenvalues:")
print(np.around(vals,3))
print("eigenvector:")
print(np.around(vecs, 3))
L = 10
k = 10
m = 1.0
dt = 0.03
t = 0.0
x1i =0.5*L/10 # initial displacement from equil, m1
v1i = 0.0 # initial vel, m1
x2i = 0.71*L/10 # initial displacement from equil, m2
v2i = 0.0 # initial vel, m2
x3i = 0.5*L/10 # initial displacement from equil, m3
v3i = 0.0 # initial vel, m3
s = np.array([x1i, v1i, x2i, v2i, x3i, v3i]) # initial state
def derivs_3m(s, t):
x1=s[0] # get the variables from the state
v1=s[1]
x2=s[2]
v2=s[3]
x3=s[4]
v3=s[5]
a1 = (-(2*k*x1) + k*x2)/m
a2 = (-(2*k)*x2 + k*x1 + k*x3)/m
a3 = (k*x2 - (2*k)*x3)/m
return np.array([v1, a1, v2, a2, v3, a3])
def RK4Step(s, dt, t, derivs):
"""
Taking a single RK4 step.
"""
f1 = derivs(s, t)
f2 = derivs(s+f1*dt/2.0, t+dt/2.0)
f3 = derivs(s+f2*dt/2.0, t+dt/2.0)
f4 = derivs(s+f3*dt, t+dt)
return s + (f1+2*f2+2*f3+f4)*dt/6.0
x1list=[s[0]]
x2list=[s[2]]
x3list=[s[4]]
tlist=[0.0]
t=0.0
while t<6:
s = RK4Step(s, dt, t, derivs_3m)
t += dt
x1list.append(s[0])
x2list.append(s[2])
x3list.append(s[4])
tlist.append(t)
pl.subplot(221)
pl.ylabel("x1")
pl.title("Motion of Coupled Masses")
pl.plot(tlist,x1list,label="x1")
pl.subplot(222)
pl.ylabel("x2")
pl.xlabel("t")
pl.plot(tlist,x2list,label="x2")
pl.subplot(223)
pl.ylabel("x3")
pl.xlabel("t")
pl.plot(tlist,x3list,label="x3")
from scipy.optimize import curve_fit
def cosFit(t, A, omega, phi):
"""
Function def for a cosine fit
"""
return A*np.cos(omega*t+phi)
x1a=np.array(x1list)
ta=np.array(tlist)
popt, pcov = curve_fit(cosFit, ta, x1a, p0=(0.5, np.sqrt(5.858), 0.0))
A=popt[0]
omega=popt[1]
phi=popt[2]
print("A =>", A)
print("omega**2 =>", omega**2)
print("phi =>", phi)
pl.title('Fit to find frequency.')
pl.xlabel('t')
pl.ylabel('x1')
pl.plot(ta, cosFit(ta, A, omega, phi), 'b-', label="fit")
pl.plot(ta, x1a, 'r.', label='data')
pl.legend()
def cosFit(t, A, omega, phi):
"""
Function def for a cosine fit
"""
return A*np.cos(omega*t+phi)
x2a=np.array(x2list)
ta=np.array(tlist)
popt, pcov = curve_fit(cosFit, ta, x2a, p0=(0.7071, np.sqrt(5.858), 0.0))
A=popt[0]
omega=popt[1]
phi=popt[2]
print("A =>", A)
print("omega**2 =>", omega**2)
print("phi =>", phi)
pl.title('Fit to find frequency.')
pl.xlabel('t')
pl.ylabel('x2')
pl.plot(ta, cosFit(ta, A, omega, phi), 'b-', label="fit")
pl.plot(ta, x2a, 'r.', label='data')
pl.legend()
from scipy.optimize import curve_fit
def cosFit(t, A, omega, phi):
"""
Function def for a cosine fit
"""
return A*np.cos(omega*t+phi)
x3a=np.array(x3list)
ta=np.array(tlist)
popt, pcov = curve_fit(cosFit, ta, x3a, p0=(0.7071, np.sqrt(5.858), 0.0))
A=popt[0]
omega=popt[1]
phi=popt[2]
print("A =>", A)
print("omega**2 =>", omega**2)
print("phi =>", phi)
pl.title('Fit to find frequency.')
pl.xlabel('t')
pl.ylabel('x3')
pl.plot(ta, cosFit(ta, A, omega, phi), 'b-', label="fit")
pl.plot(ta, x3a, 'r.', label='data')
pl.legend()
