#!/usr/bin/env python2.7
"""
Make an animation of the linear shallow-water equations in 2D

Based on the exact solution for axisymmetrical waves in:
G. F. Carrier and H. Yeh (2005)
Tsunami propagation from a finite source.
Computer Modelling in Engineering & Sciences,
vol. 10, no. 2, pp. 113-122
doi: 10.3970/cmes.2005.010.113

The dimensionless and linear shallow-water equations
in polar coordinates are:
  d/dt^2 eta(r,t) - 1/r * d/dr( r * d/dr eta(r,t) ) = 0.

The initial conditions are a series of Gaussian humps, on
an equidistant grid in the (x,y)-plane. One Gaussian hump
released at t=t0 is given by the initial conditions:
  eta(r,t0) = 2 * exp( -r^2 ),      d/dt eta(r,t0) = 0.

"""

#%% import libs
from __future__ import print_function
from future.builtins import input
import numpy as np
import scipy.special as sp
import matplotlib.pyplot as plt
import multiprocessing
import time
from scipy.integrate import quad
from scipy.interpolate import RectBivariateSpline

#%% input
n   = int( 181 )       # number of points in r-direction
m   = int( 151 )       # number of time steps
dr  = np.float64(0.2)  # step in r-direction
dt  = np.float64(0.2)  # time step
res = int( 10 )        # resample factor for 2D interpolation
mpl = int( 11 )        # no. of elevation snapshots to plot

#%% integrand
r   = np.float64()
t   = np.float64()
rho = np.float64()

def integrand(rho,r,t):
    return rho * sp.jn(0,rho*r) * np.cos(rho*t) * np.exp( -(rho**2)/4 )

def integrate(args):
    r_now = args[0]
    t_now = args[1]
    return quad( integrand, 0, np.inf,
                 args=(r_now,t_now), epsabs=1.e-5, epsrel=1.e-5, limit=100)

#%% compute time integral of surface evolution as a function of r and t

print( '=== compute time integral of surface evolution as a function of r and t' )

start_time = time.time()

eta      = np.empty( [n,m], dtype=np.float64 ) # surface elevation
err      = np.empty( [n,m], dtype=np.float64 ) # surface elevation integrated in time
r        = np.linspace( 0, (n-1)*dr, num=n, dtype=np.float64 ) # radial coordinate
t        = np.linspace( 0, (m-1)*dt, num=m, dtype=np.float64 ) # time coordinate
eta0     = 2 * np.exp( - r**2 )
eta[:,0] = eta0

num_cores= int( multiprocessing.cpu_count() )
pool     = multiprocessing.Pool(num_cores)
print( '--- number of cores : {0:d}'.format(num_cores) )

for jt in range(1,m):
    print( '\r--- jt = {0:4d}'.format(jt), end='' )
    t_now  = t[jt]
    params = [ (r[jr],t_now) for jr in range(n) ]
    data = pool.map( integrate, params )
    for jr in range(0,n):
        [ eta[jr,jt], err[jr,jt] ] = data[jr]
pool.close()
pool.join()
print( ' ' )

no_nans = np.count_nonzero(np.isnan(eta))
no_infs = np.count_nonzero(np.isinf(eta))
print( '--- number of NaN\'s in eta(r,t): {0:d}'.format(no_nans) )
print( '--- number of inf\'s in eta(r,t): {0:d}'.format(no_infs) )

#%% bivariate spline interpolate to fine grid
print( '=== interpolate to fine grid' )
n_f      = res * (n-1) + 1
m_f      = res * (m-1) + 1
eta_f    = np.empty( [n_f,m_f], dtype=np.float64 ) # surface elevation
r_f      = np.linspace( 0, (n-1)*dr, num=n_f, dtype=np.float64 ) # radial coordinate
t_f      = np.linspace( 0, (m-1)*dt, num=m_f, dtype=np.float64 ) # time coordinate

spl_eta  = RectBivariateSpline( r, t, eta )
eta_f    = spl_eta( r_f, t_f )

#%% save to file
print( '=== save results to file' )
with open( 'mat.npz', 'wb' ) as outfile:
#    dr_f = dr / res
#    dt_f = dt / res
    np.savez( outfile, n=n, m=m, dr=dr, dt=dt, eta=eta, res=res )

#%% asymptotic behaviour in terms of parabolic cylinder function D
print( '=== compute asymptotic behaviour' )
eta_asym = np.zeros( [n,m], dtype=np.float64 ) # surface elevation
for jt in range(0,m):
    for jr in range(1,n):
        print( '\r--- (jt,jr) = ({0:d},{1:d})    '.format(jt,jr), end='' )
        rnow = r[jr]
        tnow = t[jt]
        eta_asym[jr,jt] = 1. / np.sqrt( np.sqrt(2.) * rnow ) \
                        * sp.pbdv( np.float64(0.5), np.sqrt(2.) * (rnow-tnow) )[0] \
                        * np.exp( -0.5 * (rnow-tnow)**2 )
print( ' ' )
eta_asym[0,:] = 3*eta_asym[1,:] - 3*eta_asym[2,:] + eta_asym[3,:]
spl_asym = RectBivariateSpline( r, t, eta_asym )
eta_asym = spl_asym( r_f, t_f )

#%% execution time

end_time = time.time()
print( '--- execution time : {0:.3f} s'.format( end_time - start_time ) )

#%% plot quadrature error

print( '=== make plots' )

tp, rp   = np.meshgrid( t, r )

fig, ax  = plt.subplots(1)
pc       = ax.pcolormesh( rp, tp, err, cmap='viridis',
                          vmin=0, vmax=5*np.std(err) )
ax.set_title( 'quadrature error' )
ax.axis( 'equal' )
ax.axis( 'image' )
fig.colorbar( pc, format='%.1e' )
ax.set_xlabel( 'r' )
ax.set_ylabel( 't' )
plt.show( block=False )
plt.savefig( 'fig.sol/swe_ani_2D_quad_error.png', dpi=300, bbox_inches='tight' )

#%% plot surface elevation on coarse grid

fig, ax  = plt.subplots(1)
pc       = ax.pcolormesh( rp, tp, eta, cmap='RdBu_r',
                          vmin=-2, vmax=2 )
ax.set_title( 'surface elevation eta(r,t)' )
ax.axis( 'equal' )
ax.axis( 'image' )
fig.colorbar( pc )
ax.set_xlabel( 'r' )
ax.set_ylabel( 't' )
plt.show( block=False )
plt.savefig( 'fig.sol/swe_ani_2D_eta_coarse.png', dpi=300, bbox_inches='tight' )

#%% plot surface elevation on fine grid
tpf, rpf = np.meshgrid( t_f, r_f )

fig, ax  = plt.subplots(1)
pc       = ax.pcolormesh( rpf, tpf, eta_f, cmap='RdBu_r',
                          vmin=-2, vmax=2 )
ax.set_title( 'surface elevation eta(r,t)' )
ax.axis( 'equal' )
ax.axis( 'image' )
fig.colorbar( pc )
ax.set_xlabel( 'r' )
ax.set_ylabel( 't' )
plt.show( block=False )
plt.savefig( 'fig.sol/swe_ani_2D_eta_fine.png', dpi=300, bbox_inches='tight' )

#%% difference with asymptotic solution
fig, ax  = plt.subplots(1)
pc       = ax.pcolormesh( rpf, tpf, np.sqrt(rpf)*(eta_asym-eta_f), cmap='RdBu_r',
                          vmin=-0.02, vmax=0.02 )
ax.set_title( 'difference of exact and asymptotic eta(r,t)' )
ax.axis( 'equal' )
ax.axis( 'image' )
fig.colorbar( pc )
ax.set_xlabel( 'r' )
ax.set_ylabel( 't' )
plt.show( block=False )
plt.savefig( 'fig.sol/swe_ani_2D_eta_diff_asym.png', dpi=300, bbox_inches='tight' )

#%% some snapshots of the surface elevation at several instants

fig, ax  = plt.subplots(1)
for j in range(0,mpl):
    jt = int( np.rint( np.float64(j) / np.float64(mpl-1) * np.float64(m_f-1) ) )
    ax.plot( r_f, eta_f[:,jt], label=r'$t$ = {0:.3f}'.format(t_f[jt]) )
plt.gca().set_prop_cycle(None)
for j in range(0,mpl):
    jt = int( np.rint( np.float64(j) / np.float64(mpl-1) * np.float64(m_f-1) ) )
    if j==0:
        ax.plot( r_f, eta_asym[:,jt], '--', label=r'approximation' )
    else:
        ax.plot( r_f, eta_asym[:,jt], '--' )
ax.grid
ax.axis( [ 0, np.max(r_f), -0.6, 2.2 ] )
ax.set_xlabel( r'$r$' )
ax.set_ylabel( r'$\eta(r,t)$' )
ax.set_title( r'surface elevation at several instants' )
lgd = ax.legend( loc='center left', bbox_to_anchor=(1.0, 0.5) )
plt.show(block=False)
fig.set_size_inches( 12, 5, forward=True )
plt.savefig( 'fig.sol/swe_ani_2D_eta_snap.pdf',
             bbox_extra_artists=(lgd,), bbox_inches='tight' )

#%% snapshots to check self-similar behaviour

fig, ax  = plt.subplots(1)
ax.plot( [-6, 3], [0, 0], 'k-', linewidth=0.5 )
plt.gca().set_prop_cycle(None)
for j in range(0,mpl):
    jt = int( np.rint( np.float64(j) / np.float64(mpl-1) * np.float64(m_f-1) ) )
    if t_f[jt] >= 4 :
        ax.plot( r_f-t_f[jt], np.sqrt(r_f)*eta_f[:,jt],
                 label=r'$t$ = {0:.3f}'.format(t_f[jt]) )
ax.plot( r_f-t_f[jt], np.sqrt(r_f)*eta_asym[:,jt], 'k--',
         label=r'approximation', linewidth=1.0 )
ax.grid
ax.axis( [ -6, +3, -0.35, 0.7 ] )
ax.set_xlabel( r'$r-t$' )
ax.set_ylabel( r'$\sqrt{r}$ $\eta(r-t,t)$' )
ax.set_title( r'self-similar behaviour of surface elevation' )
lgd = ax.legend( loc='center left', bbox_to_anchor=(1.0, 0.5) )
plt.show(block=False)
fig.set_size_inches( 12, 5, forward=True )
plt.savefig( 'fig.sol/swe_ani_2D_eta_similarity.pdf',
             bbox_extra_artists=(lgd,), bbox_inches='tight' )

#%% ready

print( '=== ready' )

plt.subplots(1) # needed to show the last figure in spyder
plt.show( block=False )
plt.close()

input( '--- hit \'Enter\' to close ...' )
plt.close( 'all' )

#!/usr/bin/env python2.7
"""
Make an animation of the linear shallow-water equations in 2D

Based on the exact solution for axisymmetrical waves in:
G. F. Carrier and H. Yeh (2005)
Tsunami propagation from a finite source.
Computer Modelling in Engineering & Sciences,
vol. 10, no. 2, pp. 113-122
doi: 10.3970/cmes.2005.010.113

The dimensionless and linear shallow-water equations
in polar coordinates are:
  d/dt^2 eta(r,t) - 1/r * d/dr( r * d/dr eta(r,t) ) = 0.

The initial conditions are a series of Gaussian humps, on
an equidistant grid in the (x,y)-plane. One Gaussian hump
released at t=t0 is given by the initial conditions:
  eta(r,t0) = 2 * exp( -r^2 ),      d/dt eta(r,t0) = 0.

"""

from __future__ import print_function
import matplotlib.pyplot as plt, mpl_toolkits.mplot3d
import numpy as np
import os
import subprocess
import time
import sys

from scipy.interpolate import RectBivariateSpline, UnivariateSpline
from matplotlib.colors import LightSource
from mayavi import mlab

# animation parameters
Lx    = np.float64( 30.   ) # domain size in x-direction
Ly    = np.float64( 30.   ) # domain size in y-direction
x0    = np.float64(  6.0  ) # x-coordinate of centre of splashes
y0    = np.float64(  6.0  ) # y-coordinate of centre of splashes
t0    = np.float64(  5.0  ) # time of splash
dta   = np.float64(  0.2  ) # time step
nta   = int( 751 )          # number of time steps
nx    = int( 301 )          # spatial step in x-direction
ny    = int( 301 )          # spatial step in y-direction
mkani = True                # save animation to file
nte   = int(  48 )          # no. of duplicate frames at end of animation

#%% load surface elevation data from file
print( '=== load results from file' )
data = np.load( 'mat.npz' ) #, n=n, m=m, dr=dr, dt=dt, eta=eta, res=res )
n    = data['n']
m    = data['m']
dr   = data['dr']
dt   = data['dt']
eta  = data['eta']
res  = data['res']

r    = np.linspace( 0, (n-1)*dr, n, dtype=np.float64 )
t    = np.linspace( 0, (m-1)*dt, m, dtype=np.float64 )

#%% create function to interpolate surface elevation for any r and t
rnow = np.float64()
tnow = np.float64()
rmax = np.max(r)
tmax = np.max(t)
msel = int( np.floor( ( rmax - 3.0 ) / dt ) )
if msel > m : 
    msel = m
tsel = (msel-1) * dt
print( '--- tsel = {0:.3f}'.format(tsel) )

# define spline interpolations
sim_ksi  = r - tsel
sim_eta  = np.sqrt(r) * eta[:,msel-1]
spl2_eta = RectBivariateSpline( r, t, eta )
spl1_eta = UnivariateSpline( sim_ksi, sim_eta, s=0 )

# plot self-similar form
fig, ax  = plt.subplots(1)
ax.plot( [-tsel, +4], [0, 0], 'k-', linewidth=0.5 )
plt.gca().set_prop_cycle(None)
ax.plot( r-tsel, np.sqrt(r)*eta[:,msel-1] )
ax.grid
ax.axis( [ -tsel, +4, -0.35, 0.7 ] )
ax.set_xlabel( r'$r-t$' )
ax.set_ylabel( r'$\sqrt{r}$ $\eta(r-t,t)$' )
ax.set_title( r'self-similar behaviour of surface elevation' )
plt.show(block=False)

def eta_polar( args ) :
    rnow    = args[0]
    tnow    = args[1]
    rshape  = np.shape( rnow )
    rr      = np.ravel( rnow )
    eta_now = np.zeros( np.shape( rr ), dtype=np.float64 )
    # print( '--- shape of rr : {}'.format(np.shape(rr)) )
    if tnow >= 0:
        if tnow <= tsel : # interpolation
            jj          = np.where( rr <= rmax )
            tt          = tnow * np.ones( rr[jj].shape )
            eta_now[jj] = spl2_eta( rr[jj], tt, grid=False )
        else:             # self-similar solution based on t=tsel
            ksi         = rr - tnow
            jj          = np.where( ( ksi > -tsel ) & ( ksi < (rmax-tsel) ) & ( rr > 0 ) )        
            # print( '--- length of jj : {0:d}'.format(np.size(jj)) )
            eta_now[jj] = spl1_eta( ksi[jj] ) / np.sqrt(rr[jj]) 
    eta_now = np.reshape( eta_now, rshape )
    return eta_now

#%% create an animation with mayavi

print( '=== create mayavi animation' )

x    = np.linspace( 0, Lx, nx )
y    = np.linspace( 0, Ly, ny )
x, y = np.meshgrid( x, y )
z    = np.zeros( ( ny, nx, nta ), dtype=np.float64 )

# compute number of copies of the domain needed
tmax = nta * dta - t0
repx = int( np.ceil( tmax / Lx ) )
repy = int( np.ceil( tmax / Ly ) )
print( '--- repetitions in x-dir : {0:d}'.format(repx) )
print( '--- repetitions in y-dir : {0:d}'.format(repy) )

# create the surface fields needed in the animation
print( '--- compute surface elevation as a function of time' )

start_time = time.time()

for jx in range(-repx,+repx+1) :
    if jx % 2 == 0 :
        xc = +x0   + jx*Lx # even
    else :
        xc = Lx-x0 + jx*Lx # odd
    for jy in range(-repy,+repy+1) :
        if jy % 2 == 0 :
            yc = +y0   + jy*Ly # even
        else :
            yc = Ly-y0 + jy*Ly # odd
        sys.stdout.write( '\r--- jx = {0:3d} ; jy = {1:3d}'.format(jx,jy) )
        sys.stdout.flush()
        rnow   = np.sqrt( (x-xc)**2 + (y-yc)**2 )
        for i in range(nta) :
            tnow      = np.float64(i)*dta - t0
            params    = ( rnow, tnow )
            z[:,:,i] += eta_polar( params )
print( ' ' )

end_time = time.time()
print( '--- elapsed time for free-surface construction : {0:.3f} s'.format( end_time - start_time ) )

no_nans = np.count_nonzero(np.isnan(z))
no_infs = np.count_nonzero(np.isinf(z))
print( '--- number of NaN\'s in z(x,y,t): {0:d}'.format(no_nans) )
print( '--- number of inf\'s in z(x,y,t): {0:d}'.format(no_infs) )

time.sleep(5) # pause 5 seconds

# mayavi plot
fig  = mlab.figure( size=(528,337), bgcolor=(1,1,0.9) )
Ld   = np.sqrt( Lx**2 + Ly**2 )
surf = mlab.surf( x.T, y.T, z[:,:,0].T, colormap='Blues', warp_scale=4,
                  vmin=-2, vmax=+2 )
# Change the visualization parameters.
surf.actor.property.interpolation = 'phong'
surf.actor.property.specular = 0.1
surf.actor.property.specular_power = 100
surf.actor.property.ambient=0
mlab.view( azimuth=235, elevation=70, distance=0.9*Ld,
           focalpoint=(0.35*Lx,0.25*Ly,0.0) )

# animation
print( '--- show animation (and plot to png-files)' )
@mlab.show
@mlab.animate( delay=10 )
def anim() :
    cnt = True
    while cnt :
        for i in range(nta):
            # print( '\r--- i = {0:5d}'.format(i), end='' )
            tnow  = np.float64(i)*dta - t0
            znow  = z[:,:,i] # eta_polar( rnow, tnow )
            surf.mlab_source.scalars = znow.T
            if mkani :
                # create png-files
                fname = os.path.join( 'fig.ani', 'ani_{0:04d}.png'.format(i) )
                mlab.savefig( filename=fname )
                if i == (nta-1) :
                    for j in range(nte) :
                        fname = os.path.join( 'fig.ani', 'ani_{0:04d}.png'.format(nta+j) )
                        mlab.savefig( filename=fname )
                cnt = False
            yield
        
anim()    
# print( ' ' )

if mkani :

    print( '--- create animation with ffmpeg' )
    fps          = int(24)
    ffmpeg_fname = os.path.join( 'fig.ani', 'ani_%04d.png' )
    prefix       = 'ani'
    cmd          = 'ffmpeg -y -f image2 -r {} -i {} -c:v libvpx -crf 4 -b:v 1M {}.webm'.format(fps,ffmpeg_fname,prefix)
    # cmd          = 'ffmpeg -f image2 -r {} -i {} -vcodec mpeg4 -y {}.mp4'.format(fps,ffmpeg_fname,prefix)
    print( '{0:s}'.format(cmd) )
    subprocess.check_output(['bash','-c', cmd])

    # Remove temp image files with extension
    print( '--- delete png-files' )
    [os.remove( os.path.join( 'fig.ani', f ) ) for f in os.listdir('fig.ani') if f.endswith('.png')]