##!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
 KdV_UnSolitonZK.py 
 Résolution de l'équation de Korteweg de Vries par le schéma de Zabusky et Kruskal
 ut + mu u ux + nu uxxx = 0 
"""

__author__      = "Dominique Lefebvre"
__copyright__   = "Copyright 2020 - TangenteX.com"
__version__     = "1.0"
__date__        = "10 novembre 2020" 
__email__       = "dominique.lefebvre@tangentex.com"

from scipy import cosh, sqrt, zeros, linspace
from matplotlib.pyplot import figure,grid,xlabel,ylabel,tight_layout,show, title
import matplotlib.animation as animation

# Fonction sécante hyperbolique
def sech(x):
    y = 1./cosh(x)
    return y

# Solution analytique de l'équation KdV
def KdVAnalytique(x,t,c,x0,mu,nu): 
    u = (3.*c/mu)*sech(0.5*sqrt(c/nu)*(x - c*t - x0))**2
    return u

# Définition des coefficients non-linéaire et dispersif de l'équation kdV
mu = 6.       # coefficient du terme non linéaire uux
nu = 1.       # coefficient du terme dispersif uxxx

# Domaine spatial
dx = 0.1
x0 = -8.
xmax = 8.
Nx =  int((xmax - x0)/dx)
x = linspace(x0,xmax,Nx)

# Paramètres du soliton initial
c1 = 8.; shift1 = -1.
u0 = c1/2.           # amplitude max du soliton

# Domaine temporel - la valeur de dt est calculée pour remplir le critère
# de stabilité du schéma de Zabusky et Kruskal 
dt = dx/(mu*u0 + 4.*nu/(dx)**2)

# Coefficients de calcul pour le schéma ZK
A = (mu/3.)*dt/dx
B = nu*dt/(dx**3)

# Définition du buffer de travail
U = zeros((Nx,3),float)

# Construction du soliton initial
U[:,0] = KdVAnalytique(x,0,c1,shift1,mu,nu)
U[0,1] = 0. ; U[0,2] = 0. 
U[Nx-1,1] = 0. ; U[Nx-1,2] = 0. 

# tracé de l'évolution
fig1 = figure(figsize=(8,6))
ax = fig1.add_subplot(111,autoscale_on=False, xlim=(x0,xmax), ylim=(0,u0))
# affichage du soliton initial
line, = ax.plot(x,U[:,0])

# Initialisation du calcul - calcul du terme  U[:, 1]
for  i in range (1, Nx-1):                              
    if i > 1 and  i < Nx-2: 
       a2 = U[i+2,0] + 2*U[i-1,0] - 2*U[i+1,0] - U[i-2,0]
    else: 
       a2 = U[i-1, 0] - U[i+1, 0]
    U[i, 1] = U[i, 0] - \
              (mu/6.)*(dt/dx)*(U[i + 1, 0] + U[i, 0] + U[i - 1, 0])*(U[i+1, 0] - U[i-1, 0]) - \
              (nu/2.)*(dt/(dx**3))*a2

# fonction de calcul de chaque pas du schéma
def KdV(num):
    for  i in range (1, Nx-2): 
         if i > 1 and i < Nx-2: 
            a2 = U[i+2,1] + 2*U[i-1,1] - 2*U[i+1,1] - U[i-2,1]
         else: 
            a2 = U[i-1, 1] - U[i+1, 1]
         U[i, 2] = U[i, 0] - \
                   A*(U[i + 1, 1] + U[i, 1] + U[i - 1, 1])*(U[i+1, 1] - U[i-1, 1]) - \
                   B*a2
    # sauvegarde des résultats
    line.set_data(x,U[:,2])
    # permutation des buffers     
    U[:, 0] = U[:, 1]
    U[:, 1] = U[:, 2]
    return line,
    
# préparation de l'affichage
title(u"Soliton de Korteweg de Vries - Schéma de Zabusky et Kruskal")
grid(True)
xlabel("X")
ylabel("U(x,t)")
tight_layout()

# lancement animation
ani = animation.FuncAnimation(fig1,KdV,frames=3000,interval=10,repeat=False)

show()