#!/usr/bin/env python
"""
simulate a coupled system of masses being shaken....

start off with a coupled set of two 2nd-order ODEs

m1 \ddot{x1} + 2 b1 \dot{x1} + D1 x1 + D12(x1-x2) = k \sin \omega t
m2 \ddot{x2} + 2 b2 \dot{x2} + D2 x2 + D12(x2-x1) = k \sin \omega t

which can be translated to a system of two coupled 1st-order ODEs

vp1 = - 2 b1 /m1 v1 - D1/m1 x1 - D12(x1-x2) + k/m1 \sin omega t
v1 = xp1
vp2 = - 2 b2 /m2 v2 - D2/m2 x2 - D12(x2-x1) + k/m2 \sin omega t
v2 = xp2

(where p means a derivative)

Bob Wimmer
June 27, 2013

"""

from numpy import *
from math import *
from scipy.integrate import odeint
import Gnuplot, Gnuplot.funcutils
b1 = 0.05     #damping constant
D1 = 1.      #spring constant
b2 = 0.05
D2 = 2.
D12 = 2.    #coupling constant between mass 1 and mass 2
m1 = 0.3     #mass
m2 = .2
k = 2.     #driving amplitude is the same for both masses
om1 = sqrt(D1/m1) #resonance frequency of mass 1
om2 = sqrt(D2/m2) #resonance frequency of mass 2


def func(y,dt,omega):
    x1,v1,x2,v2 = hsplit(y,4)
    xp1 = v1
    vp1 = - 2*b1/m1*v1 - D1/m1*x1 - D12*(x1-x2) + k/m1*sin(omega*dt)
    xp2 = v2
    vp2 = - 2*b2/m2*v2 - D2/m2*x2 - D12*(x2-x1) + k/m2*sin(omega*dt)
    z = hstack([xp1, vp1, xp2, vp2])
    return z

def plot_results(t,y1,y2,y,xaxis,yaxis,title,outfile):
    g = Gnuplot.Gnuplot()
    g.reset()
    g("set term png")
    string = 'set output "' + outfile +'.png"'
    g(string)
    g('set size square')
    string = 'set xlabel "' + xaxis + '"'
    g(string)
    string = 'set ylabel "' + yaxis + '"'
    g(string)
    string = 'set title "' + title + '"'
#    g('set ylabel "y-axis [arb. units]"')   
#    g('set xlabel "x-axis [arb. units]"')
    data  = Gnuplot.Data(t,y,with_ ='line lc 1 lw 2', title = 'x1(t)+x2(t)')
    data1 = Gnuplot.Data(t,y1,with_ ='line lc 3', title = 'x1(t)')
    data2 = Gnuplot.Data(t,y2,with_ ='line lc 4', title = 'x2(t)')
    g.plot(data,data1,data2)
    raw_input('Please press the return key to continue...\n')

def main():
    time = []
    xval = []
    xval1 = []
    xval2 = []
    ampl = []
    ampl1 = []
    ampl2 = []
    vel1 = 0.
    x1   = 0.
    vel2 = 0.
    x2   = 0.
    om = []
    t = [0.]
    omega = 0.5*min(om1,om2)
    while omega < 2.*max(om1,om2):
        print 'omega = ', omega
        dt = 5.e-2*2*pi/omega
        xv = []
        x1v = []
        x2v = []
        for i in xrange(0,100):
            t = t[-1]
            t = append(t,t+dt)
            y = hstack([x1,vel1,x2,vel2])
            z = odeint(func,y,t,args=(omega,))
            x1,vel1,x2,vel2  = hsplit(z[-1],4)
            time.append(t[-1])
            xval.append(x1+x2)
            xval1.append(x1)
            xval2.append(x2)
            xv.append(x1+x2)
            x1v.append(x1)
            x2v.append(x2)
        om.append(omega)
        omega *= 1.02
        len_ = len(xv)
        len1 = len(x1v)
        len2 = len(x2v)
        ampl.append((amax(xv[len_/4:3*len_/4])-amin(xv[len_/4:3*len_/4]))/2)
        ampl1.append((amax(x1v[len1/4:3*len1/4])-amin(x1v[len1/4:3*len1/4]))/2)
        ampl2.append((amax(x2v[len2/4:3*len2/4])-amin(x2v[len2/4:3*len2/4]))/2)
    plot_results(time,xval1,xval2,xval,'time','x','temporal behavior','temporal_behavior')
    plot_results(om,ampl1,ampl2,ampl,'omega','amplitude','resonance behavior','resonance_behavior')

main()