#!/usr/bin/env python
# coding: utf-8

# # Shallow water waves breaking on a beach
# 
# In this example we examine the water waves breaking on a beach.  We solve the shallow water equations:
# 
# \begin{align}
# (h)_t + (hu)_x & = 0 \\
# (hu)_t + (hu^2 + \frac{1}{2}gh^2)_x & = -ghb_x
# \end{align}
# 
# Here $h$ is the water depth, $u$ is the water velocity, $g$ is a constant representing the force of gravity, and $b$ is the height of the bottom (referred to as "bathymetry").
# This problem is surprisingly challenging, for two reasons.  First, if the surface height $h+b$ and the velocity $u$ are constant in space, then they should remain so for all time.  But if $b$ is not constant, then this requires that terms on the left and right side of the momentum equation cancel exactly.  Numerical discretizations that achieve this are said to be "well-balanced".  Second, on the dry beach the depth $h$ is zero, and the location of the wet-dry interface moves as waves approach.  Even the smallest numerical errors can lead to a negative depth near that interface.
# 
# For this example we make use of [a special Riemann solver developed by David George](http://www.bu.edu/pasi-tsunami/files/2012/11/George2008.pdf) that is well-balanced and handles dry states without generating negative depths.  Although the problem we're interested in is 1D, the solver is written for 2D flows.  We therefore use a very narrow 2D domain.

# In[ ]:


import numpy as np
import matplotlib.pyplot as plt


# In[ ]:


def wave_maker_bc(state,dim,t,qbc,auxbc,num_ghost):
    "Generate waves at left boundary as if there were a moving wall there."
    if dim.on_lower_boundary:
        qbc[0,:num_ghost,:]=qbc[0,num_ghost,:] 
        t=state.t;
        amp = state.problem_data['amp'];
        if t<=state.problem_data['t1']: 
            vwall = amp*(np.sin(t*np.pi/1.5))
        else: 
            vwall=0.
        for ibc in range(num_ghost-1):
            qbc[1,num_ghost-ibc-1,:] = 2*vwall - qbc[1,num_ghost+ibc,:]

def qinit(state):
    "Gaussian surface perturbation"
    b = state.aux[0,:,:] # Bathymetry
    X,Y = state.grid.p_centers
    xleft = X.min()
    surface = ambient_surface_height + pulse_amplitude * np.exp(-(X-(xleft+2.))**2/pulse_width)
    state.q[0,:,:] = np.maximum(0.0, surface - b)
    state.q[1,:,:] = 0.0
    state.q[2,:,:] = 0.0
    
def bathymetry(x):
    "Flat bottom for x<3; then a steep slope to x=5, followed by a gentle slope."
    return (x>3)*(x<5)*((x-3)*0.4) + (x>=5)*(0.8+(x-5)/40.)
        
def setup(num_cells=500,tfinal=30,solver_type='classic',num_output_times=150):

    from clawpack import riemann
    from clawpack import pyclaw

    if solver_type == 'classic':
        solver = pyclaw.ClawSolver2D(riemann.sw_aug_2D)
        solver.dimensional_split=True
        solver.limiters = pyclaw.limiters.tvd.minmod
        solver.cfl_max     = 0.45
        solver.cfl_desired = 0.4
    elif solver_type == 'sharpclaw':
        solver = pyclaw.SharpClawSolver2D(riemann.sw_aug_2D)

    solver.bc_lower[0] = pyclaw.BC.custom 
    solver.user_bc_lower = wave_maker_bc
    solver.bc_upper[0] = pyclaw.BC.extrap
    solver.bc_lower[1] = pyclaw.BC.periodic
    solver.bc_upper[1] = pyclaw.BC.periodic

    solver.aux_bc_lower[0] = pyclaw.BC.extrap
    solver.aux_bc_upper[0] = pyclaw.BC.extrap
    solver.aux_bc_lower[1] = pyclaw.BC.periodic
    solver.aux_bc_upper[1] = pyclaw.BC.periodic

    solver.fwave = True

    # Domain:
    xlower, xupper = -15.0,  15.0
    ylower, yupper = -0.5,  0.5

    mx = num_cells
    my = 2

    x = pyclaw.Dimension(xlower,xupper,mx,name='x')
    y = pyclaw.Dimension(ylower,yupper,my,name='y')
    domain = pyclaw.Domain([x,y])

    num_aux = 1
    state = pyclaw.State(domain,solver.num_eqn,num_aux)
    state.aux[:,:,:] = bathymetry(state.p_centers[0])

    state.problem_data['grav'] = 10.0  # Gravitational force
    state.problem_data['t1']   = 50.0  # Stop generating waves after this time
    state.problem_data['amp']  = 0.1   # Amplitude of incoming waves
    qinit(state)

    #===========================================================================
    # Set up controller and controller parameters
    #===========================================================================
    claw = pyclaw.Controller()
    claw.tfinal = tfinal
    claw.solution = pyclaw.Solution(state,domain)
    claw.solver = solver
    claw.num_output_times = num_output_times
    claw.keep_copy = True
    claw.output_format = None

    return claw


# Next we set the key parameters and call the problem setup function.  Running the code should take no more than a couple of minutes.  With only 250 cells, there is significant numerical dissipation; if you wish you can run it with more cells to get a more accurate solution.

# In[ ]:


# These parameters are used in qinit.
ambient_surface_height  = 1.0
pulse_amplitude         = 0.0 # Use this to add an initial Gaussian wave
pulse_width             = 1.0

claw = setup(num_cells=250,tfinal=30.0)
claw.verbosity = 0 # Use this to suppress output during the run
claw.run()


# # Plotting the results
# 
# The code below is used to plot individual frames of the solution.

# In[ ]:


from matplotlib import animation
from IPython.display import HTML

def plot_waves(claw,ylim=(0,1.2),save_plots=False):
    fig = plt.figure(figsize=[12,4])
    ax1 = fig.add_subplot(111)
    fills = []
    frame = claw.frames[0]
    b = frame.aux[0,:,:]
    h = frame.q[0,:,:]
    surface = np.maximum(b,h+b)

    x, y = frame.state.grid.p_centers    
    slice = 1
    #line, = ax1.plot(x[:,0],surface[:,slice],'-k',linewidth=3)
    fill = ax1.fill_between(x[:,0],b[:,slice],surface[:,slice],facecolor='blue')
    fill2 = ax1.fill_between(x[:,0],0*b[:,slice],b[:,slice],facecolor='brown')
    fills = [fill,fill2]
    ax1.set_xlim(-15,15)
    if ylim: ax1.set_ylim(ylim)

    def fplot(frame_number):
        fills[-2].remove()
        fills[-1].remove()
        frame = claw.frames[frame_number]
        b = frame.aux[0,:,:]
        h = frame.q[0,:,:]
        surface = np.maximum(b,h+b)
        #line.set_data(x[:,0],surface[:,slice])
        fill = ax1.fill_between(x[:,0],b[:,slice],surface[:,slice],facecolor='blue',
                                where=b[:,slice]<surface[:,slice])
        fill2 = ax1.fill_between(x[:,0],0*b[:,slice],b[:,slice],facecolor='brown')
        fills.append(fill)
        fills.append(fill2)
        if save_plots:
            fname = 'frame'+str(frame_number).zfill(4)+'.png'
            fig.savefig(fname)   
        return fill,

    anim = animation.FuncAnimation(fig, fplot, frames=len(claw.frames), interval=100, repeat=False)
    plt.close()
    return HTML(anim.to_jshtml())


# ## Plot full view

# In[ ]:


plot_waves(claw,ylim=(0,1.5))


# ## Zoomed-in view
# The next animation is zoomed in in the vertical dimension, to better show the structure of the waves and the runup on the beach.

# In[ ]:


plot_waves(claw,ylim=(0.9,1.25))


# The characteristic speeds for the shallow water equations are $u \pm \sqrt{gh}$.  Thus waves travel slower in shallower water.
# 
# In deep water, we see that the waves travel almost without changing shape.  As the waves reach the shallow sloping beach, two things happen.  First, their wavelength decreases because their speed decreases.  Second, they steepen because the difference in depth between the crest and trough becomes significant and the crest catches up to the trough.
# 
# Unlike real water waves, our waves cannot break (overturn) because $h$ is required to be a single-valued function of $x$.  Instead they form discontinuous shock waves.