#!/usr/bin/python
#
"""
Quasiteilchen und Waermetransport
WS 17/18, C. Henkel und M. Bargheer

Solve the equations of motion for the Fermi-Pasta-Ulam-Tsingou 
chain (also known as nonlinear Rubin model). Two generic scenarios:
-- excite first oscillator with a short pulse (Q_pulsed = True)
-- start with one (linear) mode excited (Q_pulsed = False)
"""

from pylab import *
from scipy.integrate import odeint
from scipy.optimize import curve_fit
from time import time

tsize = 18
xsize = 7.5
ysize = xsize*0.5*(sqrt(5)-1)
plt.rcParams['figure.figsize'] = (xsize, ysize)
plt.rcParams['figure.autolayout'] = True
plt.rcParams['font.size'] = tsize-2
plt.rcParams['axes.labelsize'] = tsize
plt.rcParams['xtick.labelsize'] = tsize-4
plt.rcParams['ytick.labelsize'] = tsize-2


plot_in_t = True
Q_pulsed = False # all oscillators at rest, if false: excite mode
excite_mode = 0
ampl_mode = -30. # 3, -3 gives nice "2HG"
# 1, -10. works with alpha = 0.2 
# -20. works with beta = 0.2
Q_period = True # extract period of initially excited mode
Q_carpet = False
Q_mode_energy = False # analyse mode energy vs amplitude


small = 1. # overall scale parameter to enter linear regime
nMax = 15 # number of points - 1 (points labelled from 0 to nMax)
omega0 = 1. # order of magnitude of bond (vibrational) frequency
m = 1. # mass
a = 1 # lattice spacing (must be integer, not float!)

# discrete mode dispersion, n is integer, with asymmetric boundary conditions:
# anti-node to the left of the open end
# node to the right of the fixed end 
# effective length
L = a*(nMax + 1.5)
# dispersion relation
def disp(n):
    q = (n + 0.5)*pi/L
    return omega0*sqrt(2*(1 - cos(q*a)))
period = 2*pi/omega0
c0 = omega0*a # speed of sound in linear (long-wavelength) regime
alpha = small*0.20 # nonlinear parameter for cubic force
beta = small*0.0 # nonlinear parameter for quartic force
# --- parameters for pulsed force
#
f0, omega1, tau = 1.03*omega0, 0.0*omega0, 0.25*period 
#
# ---
tMax = 8*period/max(alpha,beta)
# actually, the equations of motion are solved 'internally' with an
# adaptive time step. The stepping chosen here can be used for
# 'stroboscopic images'
if Q_pulsed:
    dtmin = 0.05/omega0 # minimum time step to resolve bond vibration
    F0 = f0
else:
    dtmin = 2*pi/disp(excite_mode)
Nb_t = max(100, int(tMax/dtmin))
ti, dt = linspace(0., tMax, Nb_t, retstep = True)
x = array(range(nMax+1))*a

if Q_pulsed:
    print( 'FPUT with %4i sites, force applied to on 0\'th site' % (nMax+1) )
else:
    print( 'FPUT with %4i sites, mode %i excited' % (nMax+1, excite_mode) )

if alpha:
    print('alpha = %3.4g' %(alpha))
if beta:
    print('beta = %3.4g' %(beta))


def force(u, alpha, beta):
    f = zeros_like(u)
    for xi in range(a,len(u)-a):
        f[xi] = omega0**2*( u[xi + a] + u[xi - a] - 2*u[xi] ) + \
                  (-alpha)*( (u[xi + a] - u[xi])**2 - (u[xi] - u[xi-1])**2 ) + \
                  (beta)*( (u[xi + a] - u[xi])**3 - (u[xi] - u[xi-1])**3 )
    # left (open) end: no molecular bond to xi - a
    xi = 0
    f[xi] = omega0**2*( u[xi + a] - u[xi] ) + \
                  (-alpha)*( (u[xi + a] - u[xi])**2 ) + \
                  (beta)*( (u[xi + a] - u[xi])**3 )
    # right (fixed) end: site fixed to zero at xi + a
    xi = a*nMax
    f[xi] = omega0**2*( u[xi - a] - 2*u[xi] ) + \
                  (-alpha)*( (0. - u[xi])**2 - (u[xi] - u[xi-1])**2 ) + \
                  (beta)*( (0. - u[xi])**3 - (u[xi] - u[xi-1])**3 )
    return f

# Newton equations of motion, written for phase-space coordinates 
# q = (u, p)
if Q_pulsed:
    def newton(q, t):
        # extract velocities, conventional choice of mass
        v = q[nMax+1:2*(nMax+1)]/m
        # extract displacements
        u = q[0:nMax+1]
        dq = zeros_like(q)
        dq[0:nMax+1] = v
        dq[nMax+1:2*(nMax+1)] = force(u, alpha, beta)
        # add pulsed force 
        dq[0] += F0*exp(-0.5*(t/tau - 1)**2)*cos(omega1*(t - tau))
        return dq
else:
# do not add pulsed force (and speed up calculation)
    def newton(q, t):
        # extract velocities, conventional choice of mass
        v = q[nMax+1:2*(nMax+1)]/m
        # extract displacements
        u = q[0:nMax+1]
        dq = zeros_like(q)
        dq[0:nMax+1] = v
        dq[nMax+1:2*(nMax+1)] = force(u, alpha, beta)
        return dq


# discrete modes, both x and n are integers, orthogonal
# with respect to scalar product 
# sum( mode(n1, :)*mode(n2, :) ) = delta(n1, n2)
# anti-node to the left of the open end
# node to the right of the fixed end 
# effective length
L = a*(nMax + 1.5)
def mode(n, x):
    q = (n + 0.5)*pi/L
    return cos(q*(x + 0.5*a))/sqrt(0.5*(nMax + 1.5))

def sketch_modes(scale_f = 0.1):
    figure('phonon modes', figsize=(xsize,xsize))
    # extend spatial grid to include boundary positions
    x_ex = pad(x, (1,), mode='linear_ramp', end_values=(x[0]-a, x[-1]+a))
    for k in range(nMax,-1,-1):
        fk = zeros_like(x, dtype=float)
        for ii, xi in enumerate(x):
            y = mode(k, xi)
            fk[ii] = y
            # print fk[ii],
        # print
        # extend spatial mode to include boundary values (mixed BC)
        fk = pad(fk, (1,), mode='constant', constant_values=(fk[0], 0.))
        l, = plot( x_ex, disp(k) + scale_f*fk, 'o-', label = str(k) )
        plot( [x_ex[0], x_ex[-1]], [disp(k), disp(k)], '-', c = getp(l, 'c') )
    xlabel('position')
    ylabel('phonon modes')
    legend(fontsize = tsize-4)
    show(block=False)
    return

if Q_pulsed:
    # initial conditions
    q0 = zeros(2*(nMax+1))
else:
    # excite one mode, `a la FPUT
    u0 = ampl_mode*mode(excite_mode, x)
    q0 = zeros(2*(nMax+1))
    q0[:nMax+1] = u0 # displacement
    q0[nMax+1:2*(nMax+1)] = zeros(nMax+1) # zero momentum

Q_solved = True
try:
    u00 = rubin[0,0]
except NameError:
    # solve equations of motion
    print( 'Solve for %4i time steps (omega0 dt / 2pi = %4.4g)' \
       % (Nb_t, omega0*dt/(2*pi)) )
    Q_solved = False

if Q_solved:
    Q_reuse = raw_input('Use computed solution "rubin" [y/n]? ')
    if (Q_reuse == 'n') or (Q_reuse == 'N'):
        Q_solved = False
    elif (Q_reuse == 'y') or (Q_reuse == 'Y'):
        Q_solved = True

if Q_solved:
    Nb_t, Nb_x = shape(rubin)
    print('%4i time steps (omega0 dt / 2pi = %4.4g)' %(Nb_t, omega0*dt/(2*pi)))
else:
    clock_0 = time()
    rubin = odeint(newton, q0, ti, full_output = 1)
    cpu_time = time() - clock_0
    rubin = rubin[0]
    print('... FPUT solved (%3.4g sec).' %(cpu_time))



# make insightful plots
scale_u = 0.3/f0
if plot_in_t:
    figure(1)
    xi = 0
    ut = rubin[:,xi]
    plot( ti/period, -xi - scale_u*ut )
    for xi in range(1,nMax+1):
    #    ut = rubin[:,xi] - rubin[:,xi-1] # plot difference
        ut = rubin[:,xi] # plot displacement
        plot( ti/period, -xi - scale_u*ut )

    ylim((-1-(nMax+1), 2))

    # add straight line with speed of sound

    t0 = 2.1 # experimentally found first zero of u(0, t)
    for xi in range(nMax+1):
        plot( [(t0 + xi/c0)/period], [-xi], 'o' )
    plot( [t0/period, (t0 + a*(nMax)/c0)/period], [0, -(nMax)], '-' ) 
    grid()
    xlabel(r'$\mathsf{time\ } t / T$')
    ylabel(r'$\mathsf{displacement\ } u(x, t)$')

#
# plot global variables like energy, total expansion
# actually, this is *not* the expansion (wrong word), that would
# be given by the displacement u[0], simply. But the quantity
# computed here has a physical meaning: it corresponds to the
# x-ray diffraction amplitude when the scattering vector G is
# such that G a = 0 mod 2 pi (Bragg condition) and G u[x] << 1.
#
def plot_energy():
    energy = zeros_like(ti)
    expand = zeros_like(ti)
    for ii, t in enumerate(ti):
        xi = 0 # free (left) end
        energy[ii] = (0.5/m)*rubin[ii,nMax + 1 + xi]**2 # kinetic energy only
        expand[ii] = rubin[ii,xi]
        for xi in range(1,nMax+1):
            bond = rubin[ii,xi] - rubin[ii,xi - a]
            energy[ii] += (0.5/m)*rubin[ii,nMax + 1 + xi]**2 + \
                0.5*omega0**2*bond**2 - (alpha/3.)*bond**3 + (0.25*beta)*bond**4
            expand[ii] += rubin[ii,xi]
        # fixed (right) end, last bond to fixed substrate
        xi = nMax
        bond = - rubin[ii,xi]
        energy[ii] += 0.5*omega0**2*bond**2 - (alpha/3.)*bond**3 + (0.25*beta)*bond**4
    #
    figure(20)
    plot( ti / period, energy, 'r', label = 'total energy' )
    xlabel(r'$\mathsf{time\ } t / T$')
    legend()

    figure(7)
    plot( ti / period, 1.*expand, 'b--', label = 'x-ray amplitude' )
    grid()
    legend()
    xlabel(r'$\mathsf{time\ } t / T$')
    xlabel(r'$\mathsf{average\ displacement}$')
    show(block=False)
    return energy

#
# analyse energy of an arbitrary configuration, used for 'energy per mode'
# exercise
#
def tot_energy(k = None, uk = 0., pk = 0., it = 0):
    if k != None:
        q = zeros(2*(nMax+1))
        q[:nMax+1] = uk*mode(k, x) # displacement
        q[nMax+1:2*(nMax+1)] = pk*mode(k, x) # momentum
    else:
        q = rubin[it,:]
    energy = 0.
    x_ray = 0.
    xi = 0 # free (left) end
    energy = (0.5/m)*q[nMax + 1 + xi]**2 # kinetic energy only
    x_ray = q[xi]
    for xi in range(1,nMax+1):
        bond = q[xi] - q[xi - a]
        energy += (0.5/m)*q[nMax + 1 + xi]**2 + \
            0.5*omega0**2*bond**2 - (alpha/3.)*bond**3 + (0.25*beta)*bond**4
        x_ray += q[xi]
    # fixed (right) end, last bond to fixed substrate
    xi = nMax
    bond = - q[xi]
    energy += 0.5*omega0**2*bond**2 - (alpha/3.)*bond**3 + (0.25*beta)*bond**4
    #
    return energy, x_ray

# analyse mode energy to check amplitude-dependent mode frequency
#
if Q_mode_energy:
    k = 7
    omega = disp(k)
    p_max = 10.
    Nb_q = 10
    # vary separately momentum and position
    pk = linspace(0., p_max, Nb_q)
    Ekin = zeros_like(pk)
    for ii, pk_i in enumerate(pk):
        Ekin[ii], dum = tot_energy(k = k, pk = pk_i)
    uk = linspace(0., p_max/(m*omega), Nb_q)
    Epot = zeros_like(uk)
    for ii, uk_i in enumerate(uk):
        Epot[ii], dum = tot_energy(k = k, uk = uk_i)
    # find cubic fit
    pot3 = polyfit(uk, Epot, 3)
    print('polynomial fit for E_pot(uk):')
    print('%3.4g + %3.4g u_k + %3.4g u_k^2 + %3.4g u_k^3' 
           %(pot3[3], pot3[2], pot3[1], pot3[0]))
    print('linear theory: %3.4g u_k^2' %((0.5/m)*(omega/omega0)**2))
    # 
    figure(23)
    plot( pk, Ekin, 'bo', label = r'$%i: E_{\rm kin}$' %(k) )
    plot( pk, (0.5/m)*pk**2, 'b--' )
    #
    plot( uk, Epot, 'rx', label = r'$%i: E_{\rm pot}$' %(k) )
    plot( uk, (0.5/m)*(omega/omega0*uk)**2, 'r--', label = 'linear' )
    #
    xlabel('$\mathsf{displacement\ } u_{%i}\ \mathsf{or\ } p_{%i}$' %(k, k))
    ylabel('$\mathsf{energy\ }$')
    legend()


# xt plot('carpet')
#
if Q_carpet:
    wmax = 0.
    data = zeros_like(rubin[:,:nMax])
    for ii in range(Nb_t):
         data[ii,:] = diff(rubin[ii,:nMax+1])
         wmax = max(wmax, max(data[ii,:]))
    figure(19, figsize=(ysize,ysize))
    contourf(range(nMax), ti/period, data/wmax)
    xlabel('space')
    ylabel('time')

# animated movie, is getting really slow if all instants in ti are shown,
# perhaps better to switch to 'animation' command
#
def make_movie(tpause = 0.001, di = 10):
    figure(17)
    ymax = scale_u
    for ii in range(Nb_t):
        ymax = max(ymax, max(rubin[ii,:nMax+1]))
    ymax = 1.1*ymax
    lh, = plot(rubin[0,:nMax+1],'b')
    ylim(-ymax,ymax)
    th = text(0.9*nMax, 0.9*ymax, str(0)+'/'+str(Nb_t))
    xlabel('position')
    for ii in range(1, Nb_t, di):
        setp(lh, ydata = rubin[ii,:nMax+1])
        pause(tpause)
        setp(th, text = str(ii)+'/'+str(Nb_t))
    return

# plot dynamics for modes (reciprocal space)
#
def mode_analysis():
    figure(31)
    ek = zeros((len(ti), nMax+1))
    scale_m = 0.1/f0
    for k in range(nMax+1):
        uk = zeros_like(ti)
        pk = zeros_like(ti)
        ek[:,k] = zeros_like(ti)
        for xi in x:
            uk += mode(k, xi)*rubin[:,xi]
            pk += mode(k, xi)*rubin[:,nMax + 1 + xi]
        ek[:,k] += 0.5 * disp(k)**2 * uk**2
        ek[:,k] += 0.5 * pk**2
        l, = plot( ti / period, disp(k) + scale_m*uk )
        if k == excite_mode:
            uk0 = uk
            c0 = getp(l,'c')
        if True:
            # add a cosine plot with the (linear) mode frequency to check 
            # oscillations
            amp_k = max(uk)
            amp_k2 = min(uk)
            if -amp_k2 > amp_k: # negative value has largest amplitude
                amp_k = amp_k2
                t_k = ti[argmin(uk)] # instant in time where minimum is reached
            else:
                t_k = ti[argmax(uk)] # instant in time where maximum is reached
            plot( ti / period, disp(k) + scale_m*amp_k*cos(disp(k)*(ti - t_k)), '--', 
                  color = c0 )
    xlabel(r'$\mathsf{time\ } \omega_0 t / 2 \pi$')
    ylabel(r'$\mathsf{mode\ amplitude\ } u_k\, [\mathsf{a.u.}]$')
    if Q_period:
        # extract (shifted) period of initially excited mode
        print('extract shift in n = %i mode frequency' %(excite_mode))
        omega = disp(excite_mode)
        # define simple cos model with DC shift
        def cos_model(t, A0, A1, omega1, phi1):
            return A0 + A1*cos(omega1*t+phi1)
        params, param_cov = curve_fit(cos_model, ti, uk0, 
               p0 = (0., ampl_mode, omega, 0.))
        A0, A1, omega1, phi1 = params
        plot(ti / period, omega + scale_m*cos_model(ti, *params), '.',
              color = c0)
        print('fit parameters:')
        print('DC offset %3.4g' %(A0))
        print('frequency %3.4g (lin: %3.4g)' %(omega1, omega))
        print('AC amplitude %3.4g (input: %3.4g)' %(A1, ampl_mode))
        print('phase %3.4g' %(phi1))
    #
    figure(41)
    for k in range(nMax,-1,-1):
        plot( [ti[0] / period, ti[-1] / period], [disp(k), disp(k)], '-', c='gray' )
        plot( ti / period, disp(k) + scale_m*ek[:,k], label = str(k) )
    xlabel(r'$\mathsf{time\ } \omega_0 t / 2 \pi$')
    ylabel(r'$\mathsf{mode\ energy\ } E_k\, [\mathsf{a.u.}]$')
    legend(fontsize=tsize-4)
    show(block=False)
    return ek

if True:
    ek = mode_analysis()
    energy = plot_energy()
    figure(21)
    plot( ti / period, energy, 'r', label = 'total energy' )
    eModeTot = zeros_like(ti)
    for k in range(nMax,-1,-1):
        eModeTot += ek[:,k] 
        plot( ti / period, eModeTot, 'k--' )
    # xlabel(r'$\mathsf{time\ } \omega_0 t / 2 \pi$')
    # ylabel(r'$\mathsf{mode\ energy\ } E_k\, [\mathsf{a.u.}]$')


show(block=False)





Displacements of 16 nonlinearly coupled oscillators, initially excited in the lowest mode. Mixed Neumann/Dirichlet boundary conditions, cubic nonlinearity in the potential energy. The tilted line illustrates the speed of sound. Time is scaled to the characteristic oscillator period of one bond.